Skip to content

[ROCm] Fix the Tensor Core (Matrix Core) classification for AMD GPUs in Kernel Stats tool - #3230

Open
clarkechong wants to merge 4 commits into
openxla:masterfrom
clarkechong:clchong/xprof-amd-matrix-core-classification
Open

[ROCm] Fix the Tensor Core (Matrix Core) classification for AMD GPUs in Kernel Stats tool#3230
clarkechong wants to merge 4 commits into
openxla:masterfrom
clarkechong:clchong/xprof-amd-matrix-core-classification

Conversation

@clarkechong

Copy link
Copy Markdown
Contributor

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.

"Potentially" as 1 op -> multiple kernels, so the kernel itself is not necessarily tensor core eligible even if the op it originates from is.

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::IsKernelUsingMatrixCore to rocm_type_utils.cc:

    • Tensile's _MI{M}x{N}x{B} MatrixInstruction token and ck_tile's five tile-shape types.
    • A literal pattern list for the remaining producers: gemm_fusion, Custom_Cijk_, _xdl, gtcx, miopenSp3AsmConvRage, moe_gemm, moe_mxgemm, wmma, aiter, FMHA_FWD, attn_fwd, bwd_kernel_, mfma.
    • An exclusion list checked first (_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.cc into cuda_type_utils.cc as cuda::IsKernelUsingTensorCore.

    • The public IsKernelUsingTensorCore now takes a device_vendor argument and dispatches on it. An unrecognised/unimplemented vendor returns false.
    • ConvertDeviceTraceXPlaneToKernelReports reads the vendor once per device plane through GetDeviceCaps.
  • IsOpTensorCoreEligible now also matches the XLA/HLO op names dot_general and conv_general_dilated.

  • REMOVE a DCHECK_EQ that asserts every kernel sharing an op name has the same is_op_tensor_core_eligible value.

    • WHY: Eligibility is computed per kernel in ConvertDeviceTraceXPlaneToKernelReports (xplane_to_kernel_stats_db.cc:75-83)

          bool tensor_core_eligible =
              IsEinsumTensorCoreEligible(stats.equation) ||
              IsOpTensorCoreEligible(kernel.op_name());
          if (!tensor_core_eligible && kernel.is_kernel_using_tensor_core()) {
            VLOG(1) << "Detected new Op using TensorCores: " << kernel.op_name()
                    << std::endl;
            tensor_core_eligible = true;
          }
          kernel.set_is_op_tensor_core_eligible(tensor_core_eligible);
      
    • 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_eligible flag.

    • However, is_op_tensor_core_eligible is 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_eligible flags, causing the assertion to fire.

    • Specific example of this given below in 'Evidence' section (NVIDIA CASE 3)

  • GroupKernelReportsByOpName aggregates is_op_tensor_core_eligible with OR

  • Five 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

  • On AMD traces, "Is Kernel using TensorCore" is now correctly populated on Kernel Stats
  • A trace with no recognised vendor classifies nothing (previously, would match with NVIDIA kernel name patterns).
  • Ops lowered from XLA dot_general and conv_general_dilated now report as TensorCore-eligible.
    • This affects NVIDIA traces. See evidence.

Evidence (XProf v2.23.1, profiled via JAX)

AMD (8x MI300X gfx942)

(BEFORE PR CHANGES)
image

  • No correct Tensor Core (Matrix Core on AMD) classification on kernels which used the matrix core (verified with rocprof --pmc hardware counter traces).
  • No correct Tensor Core (Matrix Core on AMD) eligibility classification on known-eligible ops (e.g. fusion GEMM).

(AFTER PR CHANGES)
image

  • Correct Tensor Core (Matrix Core on AMD) classification on kernels which used the matrix core.
  • Correct Tensor Core (Matrix Core on AMD) eligibility classification on known-eligible ops.

NVIDIA (4x H100) CASE 1

(BEFORE CHANGES)
image

  • dot_general ops are incorrectly classified as NOT Tensor Core eligible

(AFTER PR CHANGES)
image

  • dot_general ops are now correctly classified as Tensor Core eligible
  • Remaining classification is unchanged

NVIDIA CASE 2

(CONTEXT: A dot_general op 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)
image

  • Tensor Core eligibility is incorrectly classified differently for the same op (jit(train_step)/transpose(jvp(Model))/block_7/{qkv, out_proj}/dot_general) depending on which kernel row it appears on.
  • The is_op_tensor_core_eligible flag is set true for the xmma_gemm kernel, as xmma matches with the existing kernel name patterns.
  • However, it is set on a PER KERNEL basis, despite the flag referring to an op.
  • Hence why, for the exact same op, the nccl kernel has the is_op_tensor_core_eligible flag as false

(AFTER PR CHANGES)
image

  • Tensor Core eligibility is now correctly classified for the same op (jit(train_step)/transpose(jvp(Model))/block_7/{qkv, out_proj}/dot_general).
  • is_op_tensor_core_eligible flag is still being set on a per kernel basis (here, by the xmma kernel again), however due to this change:

GroupKernelReportsByOpName aggregates is_op_tensor_core_eligible with OR

xprof/utils/kernel_stats_utils.cc:324-326

KernelStatsByOpName GroupKernelReportsByOpName(

...

  stats.is_op_tensor_core_eligible =
          stats.is_op_tensor_core_eligible ||
          kernel_report.is_op_tensor_core_eligible();

The nccl kernel now also sees the same is_op_tensor_core_eligible value.

NVIDIA CASE 3:

(Again, this bug pattern would have appeared on AMD side if the required kernel name patterns existed)

(BEFORE CHANGES)
image

XProf crashes when loading the Overview page or the Framework Op Stats page, for the H100 trace shown in CASE 2.

Crash log:

I0000 00:00:1787583433.536114 1488923 overview_page_processor.cc:90] OverviewPageProcessor::ProcessSession: Starting ConvertOpStatsToOverviewPage
F0000 00:00:1787583433.537244 1488923 kernel_stats_utils.cc:342] Check failed: stats.is_op_tensor_core_eligible == kernel_report.is_op_tensor_core_eligible() (1 vs. 0) 

(AFTER CHANGES)
image

  • No crash occurs.

ROCm Kernel Name Pattern Coverage

Details

Methodology

  1. Enumerate the producers.
  • XLA's ROCm codegen backends are a closed set of six (xla/backends/gpu/autotuner/factory_rocm.cc:87-104).
  • Custom calls are an open set chosen by the framework and cannot be enumerated, however their kernel names are still visible as the collector reports whatever was dispatched.
  1. Locate the code that builds the name and the code that selects the matrix instruction.

  2. Verify by disassembly. Unbundle the shipped code objects with clang-offload-bundler --unbundle --targets=hipv4-amdgcn-amd-amdhsa--gfx942, disassemble with llvm-objdump -d --show-all-symbols, attribute instructions to the enclosing ELF FUNC symbol, and count v_mfma, v_smfmac and v_wmma. Score each candidate token as precision and recall against those counts. --show-all-symbols is required because Tensile places a local label_ASM_Start at the kernel address and plain objdump prints only that.

  3. 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} and Custom_Cijk_

Tensile backs hipBLASLt and rocBLAS, and therefore most GEMMs. One flag, EnableMatrixInstruction, gates both halves. It sets the MatrixInstM key that builds the name token (SolutionStructs.py:1918-1937, emitted at :4820-4823), and it selects mfmaIter() over macIter() 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 the Cijk_ 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 writes MI{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 _MI occurrence instead of stopping at the first, because the same name carries other _MI tokens such as _MIAV0 and _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, gtcx and miopenSp3AsmConvRage

Xdlops names a solver class and reaches no kernel symbol. Solver names come from ComputeSolverDbId (solver.hpp:112-122), while kernels are looked up by literal name through hipModuleGetFunction (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. _xdl is how Composable Kernel spells the matrix pipeline in the instances MIOpen dispatches (implicitgemm_ck_util.hpp:341-344). gtcx marks the GTC dynamic implicit-GEMM assembly, suffixed per architecture as gtcx2, gtcx3 or gtcx35 and constructed at conv_asm_implicit_gemm_gtc_perf_config.cpp:261-272, where the x is the matrix-instruction marker and the families without it, such as igemm_bwd_gtc_*, issue none. miopenSp3AsmConvRage is the Winograd shader, gated to gfx942 (conv_wino_rage_RxS.cpp:84) and hand-written with 96 v_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 the WarpGemm type is computed inside a constexpr policy instead of being a template parameter. WarpGemm, Mfma and mfma appear 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, TileFlatmmShape and TileSageAttnShape. 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. And FmhaBwdDQDKDVKernel contains a nested DqAccPrezeroKernel that zeroes a buffer while inheriting the entire enclosing type name including its shape, so it has to be excluded by name.

AITER: namespace aiter, plus FMHA_FWD

AITER is hand-written assembly whose names are pre-mangled strings in checked-in CSV manifests (aiter/hsa/gfx942/*/**.csv), looked up at csrc/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 and topksoftmax, 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_FWD catches the one family carrying no namespace. The match is on bare aiter rather than aiter::, because the mangled form is _ZN5aiter34fmha_fwd_... with a length digit between the two, and XProf keeps these mangled since the .kd suffix 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-dot triton_* 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_fwd and the bwd_kernel_* kernels issue MFMA, while bwd_preprocess and bwd_postprocess do not, which is why the token is bwd_kernel_ and not bwd_.

Op eligibility

IsOpTensorCoreEligible answers 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 reads xla::OpMetadata through HloInstructionWrapper::TfOpName(), which is the framework-level annotation rather than the HLO opcode. XLA preserves that metadata when it rewrites a dot into custom-call(__cublas$gemm) (hlo_computation.cc:1814-1823) or a convolution into __cudnn$convForward (conv_rewriter.cc:735), so dot_general and conv_general_dilated still identify the op after lowering. They are matched with StrContains because the string carries both a scope prefix and an :op_type suffix.

Validation

Classifier run over the disassembly ground truth:

set kernels false positives missed
Tensile gfx942 (hipBLASLt + rocBLAS) 191,905 0 0
MIOpen gfx942 15,871 0 0
AITER, all architectures 1,284 0 0
ck_tile, 13 families 84 0 0
Tensile gfx1100 (WMMA) 11,227 0 0

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, Xdlops matches nothing, a bare ck_tile::sequence warp tile takes 27 non-matmul kernels, and MT{M}x{N}x{K} takes 791.

CAVEATS

  • the aiter pattern match assumes ALL AITER kernels utilize the matrix core.
  • NOT COVERED: rocWMMA is header-only and declares no kernels of its own, i.e. kernel name is determined by the caller.
  • NOT COVERED: Hand-written Triton kernels take the Python function name.
  • On RDNA, Tensile's token and ck_tile's shape types indicate the WMMA unit rather than MFMA. Reporting those as matrix-core use is intended as both are the matrix pipeline.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant