[TLE] Support cross-dtype Hopper WGMMA accumulator reuse

[July-h5kf3]

Background

TLE user-promise pipelines can chain Hopper WGMMA operations whose input
dtypes use different instruction K shapes while accumulating into the same
FP32 C fragment. One concrete case is an FP8 m64n64k32 operation followed
by a BF16 m64n64k16 operation.

The physical ownership of the FP32 accumulator registers is determined by
the instruction M/N shape and the warp/CTA distribution. K determines the
current WGMMA instruction form, but does not change that C-register
ownership. Therefore, an accumulator can be reused across the FP8 and BF16
operations when their physical C layouts are compatible.

Problem

convertDot currently selects instrMNK only from the result accumulator's
NvidiaMmaEncodingAttr. In a cross-dtype async chain, that result may retain
the preceding FP8/K32 encoding even though the current operands are BF16 and
require K16.

This has two forms:

• With register A, the current DotOperandEncodingAttr carries the BF16/K16
  parent MMA encoding, but lowering ignores it and uses the stale
  accumulator encoding.
• With shared A, there is no dot-operand parent encoding from which to recover
  the current instruction K. The current operand dtypes are the authoritative
  source for K, while M/N and C-register ownership must remain those of the
  accumulator.

As a result, lowering may select the wrong WGMMA instruction shape for the
current operands instead of preserving the compatible accumulator layout and
emitting the BF16/K16 instruction.

Minimal reproducer

The following input models a BF16/K16 WGMMA consuming an FP32 accumulator
whose result encoding came from a preceding FP8/K32 WGMMA:

#mma_fp8 = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 32]}>
#mma_bf16 = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [4, 1], instrShape = [16, 64, 16]}>
#dot_bf16 = #ttg.dot_op<{opIdx = 0, parent = #mma_bf16, kWidth = 2}>
#shared_b = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}>
#smem = #ttg.shared_memory

module attributes {
  "ttg.num-ctas" = 1 : i32,
  "ttg.num-warps" = 4 : i32,
  "ttg.target" = "cuda:90",
  "ttg.threads-per-warp" = 32 : i32
} {
  tt.func @fp8_acc_to_bf16_wgmma(
      %a: tensor<64x64xbf16, #dot_bf16>,
      %b: !ttg.memdesc<64x64xbf16, #shared_b, #smem>,
      %acc: tensor<64x64xf32, #mma_fp8>) {
    %res = ttng.warp_group_dot %a, %b, %acc {
      inputPrecision = 0 : i32,
      isAsync = true,
      tle.wgmma_accumulator_chain_c
    } : tensor<64x64xbf16, #dot_bf16>
      * !ttg.memdesc<64x64xbf16, #shared_b, #smem>
      -> tensor<64x64xf32, #mma_fp8>
    tt.return
  }
}

Run it with:

triton-opt repro.mlir \
  --allocate-shared-memory-nv='compute-capability=90 ptx-version=81' \
  --convert-triton-gpu-to-llvm='compute-capability=90 ptx-version=81' \
  --convert-nv-gpu-to-llvm

Before this change, lowering aborts in WGMMAOpPattern::getPtxAsm with:

Assertion `supported && "WGMMA type or shape is not supported"' failed.

After this change, lowering succeeds and emits:

wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16

Fix

• Keep the result MMA encoding as the accumulator layout and ownership
  contract.
• For register-A WGMMA, select the current instruction shape from A's parent
  MMA encoding only when both encodings are Hopper encodings with matching
  M/N, warpsPerCTA, and CTALayout.
• For shared/shared WGMMA, infer only the instruction K from the current
  operand dtypes (f16/bf16 -> 16, TF32 -> 8, FP8/int8 -> 32), while
  preserving the accumulator's M/N layout.
• Reject an explicitly marked accumulator chain when the physical C layouts
  are incompatible or a valid shared/shared Hopper instruction K cannot be
  selected.
• Guard the behavior with __TLE__; ordinary non-TLE lowering is unchanged.

Same-dtype WGMMA is a no-op under this selection because the inferred/current
instruction shape matches the accumulator encoding.

Tests

Added MLIR regression coverage for:

• FP8/K32 accumulator -> BF16/K16 register-A WGMMA.
• FP8/K32 accumulator -> BF16/K16 shared/shared WGMMA.
• Rejection of a marked chain with incompatible accumulator N layout.

Local validation on H800/SM90:

pre-commit (changed files): PASS
incremental FlagTree build: PASS
cross-dtype WGMMA lit tests: 2/2 PASS
python/test/tle/unit/test_tle_wgmma_pipeline_routing.py: 6/6 PASS

The complete local TLE lit suite reports 65/72 PASS. The seven failures are
the same pre-existing failures observed in the environment's unmodified
compiler control; neither new cross-dtype test is among them.

A downstream FP8/BF16 MLA reproducer was also rebuilt with the patched
compiler:

• 28 correctness and boundary cases passed.
• Candidate and control outputs remained bit-exact.
• A 24-shape paired H800 benchmark showed no systematic performance
  regression; normalized geomean latency versus CUDA changed from 0.92844
  to 0.92780 (approximately -0.07%, noise-level).

Risk and scope

The change is limited to Hopper WGMMA instruction-shape selection under
__TLE__. It does not relax accumulator compatibility: M/N, warp
distribution, and CTA layout must still describe the same physical C
register ownership. Explicitly marked incompatible chains fail with a
diagnostic rather than silently selecting an instruction.

No public API is changed.

Related issue

None.