[ROCm] Fix the Tensor Core (Matrix Core) classification for AMD GPUs in Kernel Stats tool - #3230
Open
clarkechong wants to merge 4 commits into
Open
Conversation
GroupKernelReportsByOpName asserts that every kernel report sharing an op name carries the same is_op_tensor_core_eligible. That assumption does not hold. Eligibility is computed per kernel in ConvertDeviceTraceXPlaneToKernelReports: it starts from the einsum equation and the op name, then is promoted to true for any kernel that actually used TensorCores. An op launching a mix of TensorCore and non-TensorCore kernels, under a name the eligibility list does not match, therefore produces reports that disagree, and debug builds abort. Aggregate with OR instead, which matches how the flag is produced and how it is consumed: gpu_tensorcore_utilization divides TensorCore duration by total duration, so an op with any TensorCore kernel is eligible.
IsOpTensorCoreEligible matches TF-style op names only -- Conv2D, /MatMul, BatchMatMul, XlaDot. JAX and other XLA frontends produce names like "jit(train_step)/jit(main)/dot_general[dimension_numbers=...]", which match nothing, so the op is reported as not eligible for TensorCores even when its kernels demonstrably used them. This is not AMD-specific: it is why the "GPU TensorCore utilization" column also reads 0 on an H100 running JAX. Match dot_general and conv_general_dilated by substring, since the names carry both a scope prefix and a parameter suffix.
kernel_stats_utils.cc holds the list of Nvidia SASS and CUTLASS tokens that IsKernelUsingTensorCore matches against. That was fine while Nvidia was the only vendor the classifier knew about, but the next commit adds the AMD equivalents, and holding two vendors' naming conventions makes the file a vendor grab-bag rather than a kernel stats util. Move the list to cuda_type_utils.cc, alongside the Nvidia hardware model, and expose it as cuda::IsKernelUsingTensorCore. The public IsKernelUsingTensorCore keeps its signature and its VLOG, and delegates. The vendor term is kept rather than neutralised. The three PerCore signatures are neutral because the dispatcher treats them as interchangeable. The next commit dispatches this classifier too, but on an explicit branch naming each vendor's function rather than through a shared signature, so the convention that applies is the one cuda_type_utils already follows for cuda_core and tensor_core: inside a vendor namespace, name the vendor's own hardware. Pure move, no behaviour change, pinned by a new test on the Nvidia names.
IsKernelUsingTensorCore matches only Nvidia SASS and CUTLASS tokens (h884, hmma,
xmma_gemm and so on), so no AMD kernel has ever matched. Every kernel in an MI300
trace reports "Is Kernel using TensorCore: False", and because
gpu_tensorcore_utilization is derived from that flag, the Framework Op Stats
column reads 0.0 on every row -- worse than absent, since a visible zero implies
the matrix cores were idle.
Add the AMD equivalents to rocm_type_utils, mirroring the Nvidia list the
previous commit moved to cuda_type_utils. Rather than collect whichever tokens
looked plausible, each pattern is taken from the naming rule of the library that
produces the kernel. XLA's ROCm codegen backends are a closed set of six in
xla/backends/gpu/autotuner/factory_rocm.cc, so they can be worked down in turn.
hipBLASLt and the fission backend reach Tensile, which writes its
MatrixInstruction parameter into the name and omits the token when the
parameter is unset. One flag drives both halves: EnableMatrixInstruction gates
the key that builds the token and selects mfmaIter() over macIter(). Two
arities ship at once, MI{M}x{N}x{K}x{B} from classic Tensile in rocBLAS and
MI{M}x{N}x{B} from TensileLite in hipBLASLt, so three fields are matched.
Tensile's hand-written custom kernels return a name chosen by hand and reach
no token at all, so they are matched by their own prefix.
Triton is named by XLA, which calls a dot fusion "gemm_fusion_*". Being a dot
fusion is necessary rather than sufficient, since whether Triton emits MFMA
depends on the shape and dtype.
MIOpen runs Composable Kernel instances and its own GTC assembly. Its MFMA
solver classes are named "...Xdlops", but a solver name is not a kernel name
and that string reaches no kernel symbol at all; CK spells the matrix
pipeline "xdl" and the assembly spells it as the "x" of "gtcx". Its Winograd
Rage shader is hand-written MFMA carrying neither token. Classic CK's
mixture-of-experts GEMM is the one exception that carries no "xdl", because
it takes the instruction shape as index_t parameters that never materialise.
XLA's own emitters never reach the matrix cores, so there is nothing to match.
Custom calls have no such list, since the framework rather than XLA decides what
is linked, but their kernel names are visible all the same: the collector
reports whatever was dispatched, whoever launched it. Transformer Engine's fused
attention is matched on that basis across both backends it selects between, the
ck_tile and aiter path and AOTriton behind it.
ck_tile has a naming rule of its own, though not where one would first look for
it. Its single device entry point is templated on the kernel type, so the symbol
is that type's whole instantiation, and the matrix instruction never appears in
it because the WarpGemm type is computed inside a constexpr policy rather than
being a template parameter. What does appear is the policy's input, the tile
shape, and only five shape types feed a warp tile to WarpGemmDispatcher. Those
five separate the matmul kernels from their datamovement siblings, which take
scalar tile sizes and carry no shape type at all.
aiter is matched on its namespace rather than on family stems, because it is a
matmul kernel library throughout: of the 1,284 kernels it ships across
attention, mixture-of-experts, MLA, paged attention and its GEMMs, 1,214 issue
MFMA and the 70 that do not are its backward-attention datamovement helpers and
its top-k softmax. Matching the namespace covers families aiter adds later, and
accepts the reverse risk that a future kernel inside it doing no matmul is
claimed until excluded.
Each of the three attention libraries pairs its matmul kernels with helpers that
issue nothing and sit one token away, so all of them are excluded ahead of the
patterns: aiter's _odo, _dq_convert and _dq_shuffle, ck_tile's nested
DqAccPrezeroKernel which inherits its enclosing kernel's whole type name and so
its shape, and AOTriton's bwd_preprocess and bwd_postprocess, which is why that
token is bwd_kernel_ rather than bwd_.
"mfma" is carried as well, for kernels that name the instruction outright, which
rocSOLVER's mfma_gemm_kernel does.
The Tensile, MIOpen, ck_tile and aiter entries were established by disassembling
kernels and counting v_mfma per symbol, or for aiter by joining its shipped
kernel manifests to that disassembly. Over Tensile's 191,905 shipped gfx942
kernels the patterns select exactly the 181,159 containing v_mfma; over MIOpen's
15,871 with a body, exactly the 15,440; over aiter's 1,284, exactly the 1,214;
over 84 ck_tile kernels compiled across 13 families, exactly the 24. Nothing
missed, no false positive. "Cijk_" over the Tensile set would claim 3,357
kernels containing none, "Xdlops" over the MIOpen set matches nothing
whatsoever, and a bare ck_tile::sequence warp tile takes 27 non-matmul kernels,
which is why none of those can be what this keys on.
The remaining entries rest on a MaxText DeepSeek-V2-Lite training step on 8x
MI300X, joining all 457 kernel names to SQ_INSTS_VALU_MFMA_MOPS_BF16: 43 of 43
gemm_fusion kernels issued MFMA while the non-dot triton_* fusions did not.
The scan for the Tensile token steps over _MIAV and _MIWT, unrelated parameters
appearing in the same name.
Together these recognise 64.4% of that workload's kernel time, which is every
kernel the counters measured as issuing MFMA and nothing that issued none.
Coverage still cannot be claimed for custom calls, since a framework linking a
different library contributes kernels this list has never seen, and rocWMMA is
permanently beyond it: being header-only it declares no kernels of its own, so
anything built on it is named by its caller.
The name alone decides, with no architecture check, so on RDNA both the Tensile
token and ck_tile's shape types report the WMMA unit rather than MFMA: all 5,827
kernels carrying the token in the gfx1100 libraries execute v_wmma and none
executes v_mfma. Reporting those as matrix-core use is intended, since both are
the matrix pipeline, and it costs nothing on CDNA, where no shipped kernel
carrying either executes anything but v_mfma.
The vendors are selected on device_vendor, as the roofline models are, which is
why IsKernelUsingTensorCore now takes it and why the caller reads it from the
device plane. Unioning the two matchers would be simpler and would survive a
missing vendor stat, but it is wrong, because not every pattern here is an
AMD-only string: XLA builds gemm_fusion_<dot> on either vendor, Triton names
attention attn_fwd and bwd_kernel_* on either, and Nvidia has a WMMA API of its
own. A union would therefore reclassify Nvidia kernels on evidence gathered only
on ROCm. An unrecognised or unimplemented vendor classifies nothing, the same
choice GetSharedMemoryBandwidthPerCore makes for a vendor it has no model for.
The stat is dependable: the collectors write it unconditionally, and captured
traces carry it, eight reading "AMD" on 8x MI300X and four reading "Nvidia" on
4x H100.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation 1
In the 'Kernel Stats' page of XProf, there are two columns: "Kernel uses TensorCore" and "Op is TensorCore eligible". This helps identify kernels that do not use the matrix core despite being (potentially) being eligible to use it.
Within XProf, the internal "Kernel uses TensorCore" flag is set by matching the kernel name string to a set of string patterns for kernels that are known to use the matrix core. However, currently, these string patterns only cover NVIDIA kernel names. To fix this functionality for AMD GPU traces, we require equivalent string patterns for known AMD kernel names that use the matrix core.
Motivation 2
Summary of changes
Add
rocm::IsKernelUsingMatrixCoretorocm_type_utils.cc:_MI{M}x{N}x{B}MatrixInstruction token and ck_tile's five tile-shape types.gemm_fusion,Custom_Cijk_,_xdl,gtcx,miopenSp3AsmConvRage,moe_gemm,moe_mxgemm,wmma,aiter,FMHA_FWD,attn_fwd,bwd_kernel_,mfma._odo_,_dq_convert,_dq_shuffle,topksoftmax,DqAccPrezeroKernel), for non-matmul kernels that sit inside an otherwise matched family.Move the existing NVIDIA patterns AS WRITTEN, out of
kernel_stats_utils.ccintocuda_type_utils.ccascuda::IsKernelUsingTensorCore.IsKernelUsingTensorCorenow takes adevice_vendorargument and dispatches on it. An unrecognised/unimplemented vendor returns false.ConvertDeviceTraceXPlaneToKernelReportsreads the vendor once per device plane throughGetDeviceCaps.IsOpTensorCoreEligiblenow also matches the XLA/HLO op namesdot_generalandconv_general_dilated.REMOVE a
DCHECK_EQthat asserts every kernel sharing an op name has the sameis_op_tensor_core_eligiblevalue.WHY: Eligibility is computed per kernel in
ConvertDeviceTraceXPlaneToKernelReports(xplane_to_kernel_stats_db.cc:75-83)It starts from the einsum equation and the op name (op-level properties) and so initially, kernels from the same op share the same
is_op_tensor_core_eligibleflag.However,
is_op_tensor_core_eligibleis then promoted true PER KERNEL, for any kernel that matches against the kernel name patterns in this PR.This crashes XProf when, for example, an eligible GEMM (kernel name matches) and non-eligible helper (no kernel name match), produce different
is_op_tensor_core_eligibleflags, causing the assertion to fire.Specific example of this given below in 'Evidence' section (NVIDIA CASE 3)
GroupKernelReportsByOpNameaggregatesis_op_tensor_core_eligiblewith ORFive new unit tests in
kernel_stats_utils_test.cc. The AMD fixtures are kernel names captured verbatim from an MI300X trace, abbreviated in the middle where noted.Consequences
dot_generalandconv_general_dilatednow report as TensorCore-eligible.Evidence (XProf v2.23.1, profiled via JAX)
AMD (8x MI300X gfx942)
(BEFORE PR CHANGES)

rocprof --pmchardware counter traces).(AFTER PR CHANGES)

NVIDIA (4x H100) CASE 1
(BEFORE CHANGES)

dot_generalops are incorrectly classified as NOT Tensor Core eligible(AFTER PR CHANGES)

dot_generalops are now correctly classified as Tensor Core eligibleNVIDIA CASE 2
(CONTEXT: A
dot_generalop is lowered to a GEMM kernel and a NCCL collective)(This same bug pattern would have appeared on AMD side if the required kernel name patterns existed, but of course they are only being added here in this PR)
(BEFORE CHANGES)

jit(train_step)/transpose(jvp(Model))/block_7/{qkv, out_proj}/dot_general) depending on which kernel row it appears on.is_op_tensor_core_eligibleflag is set true for thexmma_gemmkernel, asxmmamatches with the existing kernel name patterns.ncclkernel has theis_op_tensor_core_eligibleflag as false(AFTER PR CHANGES)

jit(train_step)/transpose(jvp(Model))/block_7/{qkv, out_proj}/dot_general).is_op_tensor_core_eligibleflag is still being set on a per kernel basis (here, by thexmmakernel again), however due to this change:xprof/utils/kernel_stats_utils.cc:324-326The
ncclkernel now also sees the sameis_op_tensor_core_eligiblevalue.NVIDIA CASE 3:
(Again, this bug pattern would have appeared on AMD side if the required kernel name patterns existed)
(BEFORE CHANGES)

XProf crashes when loading the Overview page or the Framework Op Stats page, for the H100 trace shown in CASE 2.
Crash log:
(AFTER CHANGES)

ROCm Kernel Name Pattern Coverage
Details
Methodology
xla/backends/gpu/autotuner/factory_rocm.cc:87-104).Locate the code that builds the name and the code that selects the matrix instruction.
Verify by disassembly. Unbundle the shipped code objects with
clang-offload-bundler --unbundle --targets=hipv4-amdgcn-amd-amdhsa--gfx942, disassemble withllvm-objdump -d --show-all-symbols, attribute instructions to the enclosing ELFFUNCsymbol, and countv_mfma,v_smfmacandv_wmma. Score each candidate token as precision and recall against those counts.--show-all-symbolsis required because Tensile places a locallabel_ASM_Startat the kernel address and plainobjdumpprints only that.Cross-check on a real trace (DEEPSEEKV2-16B TRAIN STEP) with
SQ_INSTS_VALU_MFMA_MOPS_BF16, which is non-zero exactly when the matrix pipeline issued.Tensile:
_MI{M}x{N}x{B}andCustom_Cijk_Tensile backs hipBLASLt and rocBLAS, and therefore most GEMMs. One flag,
EnableMatrixInstruction, gates both halves. It sets theMatrixInstMkey that builds the name token (SolutionStructs.py:1918-1937, emitted at:4820-4823), and it selectsmfmaIter()overmacIter()for the inner loop (KernelWriter.py:2236-2243). The key defaults to empty (Common.py:1630), so the token is absent exactly when the parameter is unset. That is why the token decides, and not theCijk_problem-type prefix, which is present on source GEMMs too.Two arities ship at the same time, from different forks of the generator rather than from version drift. Classic Tensile in rocBLAS writes
MI{M}x{N}x{K}x{B}, and TensileLite in hipBLASLt writesMI{M}x{N}x{B}(Naming.py:131). Matching three digit fields accepts both; requiring four would miss every hipBLASLt kernel. The matcher scans every_MIoccurrence instead of stopping at the first, because the same name carries other_MItokens such as_MIAV0and_MIWT4.Custom kernels bypass the generator and return a hand-chosen name (
Naming.py:96-99,CustomKernels.py:32-33), so they never carry the token and need a prefix of their own. All 82 shipped for gfx942 issue MFMA.MIOpen:
_xdl,gtcxandmiopenSp3AsmConvRageXdlopsnames a solver class and reaches no kernel symbol. Solver names come fromComputeSolverDbId(solver.hpp:112-122), while kernels are looked up by literal name throughhipModuleGetFunction(hipoc_kernel.hpp:224-231). It appears in 0 of 23,816 gfx942 symbols, so keying on it would match nothing.Three tokens do reach a symbol.
_xdlis how Composable Kernel spells the matrix pipeline in the instances MIOpen dispatches (implicitgemm_ck_util.hpp:341-344).gtcxmarks the GTC dynamic implicit-GEMM assembly, suffixed per architecture asgtcx2,gtcx3orgtcx35and constructed atconv_asm_implicit_gemm_gtc_perf_config.cpp:261-272, where thexis the matrix-instruction marker and the families without it, such asigemm_bwd_gtc_*, issue none.miopenSp3AsmConvRageis the Winograd shader, gated to gfx942 (conv_wino_rage_RxS.cpp:84) and hand-written with 96v_mfma; being Winograd it carries no implicit-GEMM token, so neither of the other two reaches it.ck_tile: five tile-shape types
ck_tile is reached through Transformer Engine's fused attention. It builds one device entry point templated on the kernel type (
kernel_launch.hpp:81-99), so the symbol is the whole instantiation, but the matrix instruction still cannot appear in it, because theWarpGemmtype is computed inside aconstexprpolicy instead of being a template parameter.WarpGemm,Mfmaandmfmaappear in 0 of 979 symbols.What does reach the symbol is the policy's input, the tile shape. Only five shape types feed a warp tile to
WarpGemmDispatcher:TileGemmShape,TileFmhaShape,TileFmhaBwdShape,TileFlatmmShapeandTileSageAttnShape. Helper kernels take scalar tile sizes and carry no shape type at all, which is what separates them.Two things make this narrower than it first appears. A bare
sequence<M,N,K>warp tile is not sufficient, since non-GEMM ops have one too, which gives 27 false positives. AndFmhaBwdDQDKDVKernelcontains a nestedDqAccPrezeroKernelthat zeroes a buffer while inheriting the entire enclosing type name including its shape, so it has to be excluded by name.AITER: namespace
aiter, plusFMHA_FWDAITER is hand-written assembly whose names are pre-mangled strings in checked-in CSV manifests (
aiter/hsa/gfx942/*/**.csv), looked up atcsrc/cpp_itfs/mha_fwd.cu:249-255. Nothing constructs a name at runtime, so the shipped set is enumerable: 1,214 of 1,284 kernels issue MFMA, and the 70 that do not are the backward-attention datamovement helpers andtopksoftmax, both excluded explicitly.Matching the namespace rather than each family is deliberate. AITER is a matmul library throughout, so the namespace is exact today and also covers families added later. It accepts the reverse risk, noted below.
FMHA_FWDcatches the one family carrying no namespace. The match is on bareaiterrather thanaiter::, because the mangled form is_ZN5aiter34fmha_fwd_...with a length digit between the two, and XProf keeps these mangled since the.kdsuffix defeats the demangler.Triton and AOTriton
XLA names its own dot fusions
gemm_fusion_<dot>(gemm_fusion.cc:1596). Being a dot fusion is necessary but not sufficient, since whether Triton emits MFMA depends on shape and dtype. In the reference trace all 43 issued MFMA, while the non-dottriton_*fusions issued none.AOTriton ships in PyTorch-ROCm and is Transformer Engine's second-choice attention backend, so both naming styles can appear in one run.
attn_fwdand thebwd_kernel_*kernels issue MFMA, whilebwd_preprocessandbwd_postprocessdo not, which is why the token isbwd_kernel_and notbwd_.Op eligibility
IsOpTensorCoreEligibleanswers the other half of the question, whether the op should have used the matrix cores, and the gap between the two is what the utilisation metric reports. It readsxla::OpMetadatathroughHloInstructionWrapper::TfOpName(), which is the framework-level annotation rather than the HLO opcode. XLA preserves that metadata when it rewrites adotintocustom-call(__cublas$gemm)(hlo_computation.cc:1814-1823) or a convolution into__cudnn$convForward(conv_rewriter.cc:735), sodot_generalandconv_general_dilatedstill identify the op after lowering. They are matched withStrContainsbecause the string carries both a scope prefix and an:op_typesuffix.Validation
Classifier run over the disassembly ground truth:
Trace cross-check: a MaxText DeepSeek-V2-Lite step on 8× MI300X (ROCm 7.14.0, hipBLASLt 1.4.1, rocBLAS 5.5.0, MIOpen 3.5.2, JAX 0.10.0). 64.39% of kernel time was classified as matrix-core, against 64.39% measured by counters.
Alternatives rejected on the same data:
Cijk_takes 3,357 non-MFMA split-K epilogues,Xdlopsmatches nothing, a bareck_tile::sequencewarp tile takes 27 non-matmul kernels, andMT{M}x{N}x{K}takes 791.CAVEATS
aiterpattern match assumes ALL AITER kernels utilize the matrix core.