fix(hip): share MatMul and Gemm shapes with dynamic batch validation

CLAUDE.md

@@ -146,6 +146,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co
 - Each category splits into a pure `infer*` function of static shapes and a `reify*` function that may emit index SSA. A `reify*` helper must validate through its `infer*` counterpart **before** touching the builder: a rewrite or reification that reports failure must leave the IR unchanged.
 - Keep the shape machinery internal: only the `infer*`/`reify*` rules belong in `HipShapeUtils.h`. Dimension maps and static folds stay file-static in the `.cpp`, and non-template helper bodies belong in `OnnxToHipUtils.cpp` rather than inline in the header every converter includes.
 - Express "not known at compile time" as `std::optional`, not a parallel `bool` flag, so a caller cannot pass a value that contradicts the flag.
+- MatMul uses the reified output batch count plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Dynamic contraction K is carried independently from A and B and checked for equality in the runtime before cache/descriptor creation or dispatch. Multi-axis dynamic batches are accepted; before flattening, lowering validates every right-aligned axis with `A == B || A == 1 || B == 1` and requires the output extent to equal `select(A == 1, B, A)`. It also multiplies all leading extents into output/A/B matrix counts, and the checked runtime requires the axis bit plus each operand count being 1 or the output count before BLAS. Static partial per-axis broadcasts remain compile-time errors. HIP lowering calls `wrap_hipblasLtMatmul_checked`; the unversioned 11-argument `wrap_hipblasLtMatmul` remains the Hipsr compatibility ABI and forwards a true axis bit. Invalidate HIP LLVM-IR/native artifacts compiled against the former checked ABI.
 - See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md).
 
 ### Allocation and memory planning

build.py

@@ -326,7 +326,8 @@ def build_targets(args, build_dir):
 def run_tests(args, build_dir):
     """Run the GPU-free test suites (no device needed, so they run on the build
     machine in every CI job): the MLIR LIT pass-verification suite plus the
-    compiler-plugin registrar and output-allocator ctest unit tests."""
+    compiler-plugin registrar, output-allocator, constant-externalizer, and
+    MatMul/Gemm ctest unit tests."""
     step("Test (check-hip-mlir-lit)")
     run_subprocess(
         [
@@ -351,7 +352,7 @@ def run_tests(args, build_dir):
             "-C",
             args.config,
             "-R",
-            "StaticPlugins|OutputAllocator",
+            "StaticPlugins|OutputAllocator|MatmulGemmContractUnitTest",
             "--output-on-failure",
         ]
     )

docs/design/compiler-runtime-contract.md

@@ -182,6 +182,21 @@ The generated LLVM IR plus `RuntimeState` therefore carry pool behavior. See
 [pool-allocs-memory-planning.md](pool-allocs-memory-planning.md) for the
 attribute, lowering, and grow-on-demand runtime contract.
 
+HIP MatMul and Gemm carry two contraction extents in generated calls. MatMul
+calls `wrap_hipblasLtMatmul_checked` with `A[-1]`, `B[-2]`, and one i1 proving
+all right-aligned runtime batch axes and output extents agree with ONNX
+broadcast. Gemm passes the transpose-aware K extent from each operand. The real
+and mock wrappers validate these contracts before descriptor/cache creation or
+dispatch, then use the equal K value in cache keys. Runtime-invalid dimensions
+and checked output-size overflow set the shared error flag and skip BLAS work;
+a known, valid nonempty output allocation is zero-filled on failure.
+
+The unversioned `wrap_hipblasLtMatmul` remains the exact 11-argument Hipsr ABI
+and forwards once to the checked implementation with equal K and legacy
+batch/stride assumptions plus a true axis-validity bit. HIP artifacts built
+against an earlier checked MatMul signature and artifacts built against the old
+one-K Gemm signature must be recompiled; existing Hipsr MatMul artifacts remain
+valid.
 ---
 
 ## Consumers

docs/design/hip-shape-inference.md

@@ -182,15 +182,69 @@ materialize index SSA only after validation succeeds.
 A reifier must validate every precondition before touching the builder.
 Failure must leave the IR unchanged, including when the valid result shape is
 rank zero; `FailureOr` distinguishes that empty success from failure.
-Conversion-side destination builders in `OnnxToHipUtils.cpp` consume the same
-reified shape and validate imported static result metadata before creating
-`tensor.empty`.
+Conversion first runs the pure shape rule and checks imported static result
+metadata. Only after every fallible static check succeeds may reification emit
+dimension SSA and destination builders create `tensor.empty`.
 
 Common DPS verification is similarly centralized in `verifyDpsComputeOp`. It
 checks ranked tensor/memref uniformity, destination count, result count, and
 tensor result/init type equality before a category-specific verifier examines
 shape semantics.
 
+MatMul and Gemm accept a dynamic contraction K as unknown-compatible. Static
+equal K remains valid, while static unequal K is rejected by the pure shape
+rule before reification or conversion emits IR. HIP-to-LLVM passes both runtime
+extents independently (`MatMul`: A[-1]/B[-2]; `Gemm`: transpose-aware A/B K)
+to the wrappers. The wrappers compare them before descriptor creation, cache
+lookup, or dispatch and key caches only with the equal value.
+
+Runtime rejection is failure-contained. A K mismatch, negative/overflowing
+dimension, dynamically concealed partial batch broadcast, invalid output
+pointer, or checked output element/byte overflow records the shared recoverable
+error flag and skips BLAS work. When the exact nonempty output byte count is
+known and storage is valid, the wrapper queues an exact zero-fill so downstream
+consumers never observe uninitialized output. A zero-element output may have
+null storage and dispatches no BLAS work.
+
+### MatMul strided-batch representability
+
+The hipBLASLt MatMul lowering takes the batch count from the reified output
+shape and carries independent A/B batch strides, so either whole matrix may
+broadcast across the other's batches. One constant stride per operand can
+express exactly two layouts: stride 0 reuses a single matrix across every output
+batch, and a stride of the matrix size walks one matrix per output batch. An
+operand's matrix count must therefore be either 1 or the output's.
+
+A partial per-axis broadcast falls strictly between the two — batch `[2, 1]`
+against an output batch of `[2, 3]` holds 2 matrices where the output needs 6.
+`verifyStridedBatchMatmul` rejects partial layouts visible in the static types.
+Dynamic extents can conceal the same layout, including across multiple batch
+axes, so those layouts are accepted statically and validated at runtime.
+Before flattening, lowering right-aligns A/B batch axes (using 1 for implicit
+leading axes) and requires per axis `A == B || A == 1 || B == 1`. It computes
+the exact reification broadcast choice `select(A == 1, B, A)` and requires the
+runtime output descriptor extent to equal it. The axis predicates are ANDed
+into `batch_axes_valid`.
+
+Lowering then multiplies every leading extent to form the output and
+per-operand matrix counts. It selects matrix-size stride only when an operand
+count equals the output count, otherwise stride zero. The checked wrapper first
+requires `batch_axes_valid`, then requires each count to be exactly 1 or the
+output count, all before descriptor/cache creation or BLAS dispatch. The
+per-axis bit is necessary because incompatible batches such as `[2,3]` and
+`[3,2]` have the same flattened matrix count.
+
+`wrap_hipblasLtMatmul_checked` carries `batch_axes_valid`, both contraction
+extents, both operand batch counts, and both strides. B's matrix stride is
+formed from B's own contraction extent, not A's. The exact 11-argument
+`wrap_hipblasLtMatmul` symbol remains the Hipsr compatibility ABI and forwards
+once with `batch_axes_valid=true` using its legacy assumptions. A false axis
+bit or invalid matrix count records the shared recoverable error before any
+BLAS work and zeroes a known nonempty output.
+HIP LLVM-IR and native model artifacts compiled against the former unversioned
+full-batch MatMul ABI or the previous one-K Gemm ABI must be invalidated. Hipsr
+artifacts using the legacy 11-argument MatMul ABI remain compatible.
+
 ## `--hip-infer-shapes`
 
 `--hip-infer-shapes` is a module pass that runs after ONNX-to-HIP conversion and before One-Shot Bufferize. It is restricted to HIP dialect operations.

docs/design/output-allocator-design.md

@@ -265,7 +265,7 @@ MlirCustomOp::Compute(context)
    │     ├─ main_graph(state, inputs)            thin wrapper: unpacks input descriptors,
    │     │  └─ main_graph_internal(...)          calls body, DISCARDS its returned descriptor
    │     │     ├─ hipdnn_ep_get_pool_base(state, domain_id, size)   grow-on-demand GPU pool
-   │     │     ├─ <compute ops> (e.g. wrap_hipblasLtMatmul, wrap_miopen*) -> pool slots
+   │     │     ├─ <compute ops> (e.g. wrap_hipblasLtMatmul_checked, wrap_miopen*) -> pool slots
    │     │     ├─ hipdnn_ep_alloc_output(state, out_idx, shape, rank, elem)   <-- OUTPUT ALLOC
    │     │     │  └─ output_allocator.cpp: forwards to alloc.allocate(self, ...)
    │     │     │     └─ output_allocate_cb(...)   noexcept

docs/design/per-op-profiling.md

@@ -15,10 +15,14 @@ Per-operator GPU profiling measures GPU and CPU time for each runtime wrapper fu
 Each operator wrapper contains one macro call:
 
 ```cpp
-int wrap_hipblasLtMatmul(RuntimeState *state, ..., int64_t M, int64_t N, int64_t K, ...) {
+int wrap_hipblasLtMatmul_checked(RuntimeState *state, ..., int64_t M,
+                                 int64_t N, int64_t K_a, int64_t K_b, ...) {
+  // The wrapper has already established K_a == K_b.
+  int64_t K_equal = K_a;
   OP_PROFILE("matmul", [&] {
     char b[64];
-    snprintf(b, sizeof(b), "%lldx%lldx%lld", (long long)M, (long long)N, (long long)K);
+    snprintf(b, sizeof(b), "%lldx%lldx%lld", (long long)M, (long long)N,
+             (long long)K_equal);
     return std::string(b);
   }, state);
   // ... operator implementation ...

docs/hip_dialect_intro.md

@@ -44,11 +44,14 @@ Matrix multiplication backed by the hipBLASLt library (`hipblasLtMatmul`).
 
 | Op | DPS Syntax | Runtime | Status |
 |---|---|---|---|
-| `hip.hipblaslt.matmul` | `(%ctx) ins(%A, %B : ...) outs(%C : ...)` | `hip_hipblaslt_matmul(handle, A, B, C, rankA, rankB, batch, M, K, N)` | Full impl |
-
-Rank-generic: batch is determined from A's rank (3D -> batched, 2D -> single).
-If B has fewer dims than A (e.g. `X[B,S,D] @ W[D,D]`), B is broadcast across
-batches (`stride_B = 0`). Supports strided batched GEMM via hipBLASLt.
+| `hip.matmul` | `(%ctx) ins(%A, %B : ...) outs(%C : ...)` | `wrap_hipblasLtMatmul_checked(state, slot, A, B, C, batch_axes_valid, M, N, K_a, K_b, batch, elem, a_batches, b_batches, a_stride, b_stride)` | Full impl |
+
+The checked generated-code ABI carries both contraction extents and validates
+their equality plus every right-aligned batch axis before dispatch. Either
+operand may be a single matrix broadcast across all output batches or provide
+one matrix per output batch. The exact 11-argument `wrap_hipblasLtMatmul`
+symbol remains the legacy Hipsr ABI and forwards to the checked entry point
+with its original A/B batch assumptions.
 
 ---
 

include/hip/Dialect/IR/HipOps.td

@@ -831,11 +831,12 @@ def Hip_ConvTransposeOp : Hip_DpsOp<"conv_transpose"> {
 // reason to write that boilerplate by hand. Matmul is the worked
 // example for the orthogonality of these two layers on #260; #262
 // extends the same split to gemm / qmoe / matmul_nbits.
-def Hip_MatmulOp : Hip_DpsOp<"matmul", /*traits=*/[OpStateOpInterface],
-                             /*outsAccessor=*/"Output",
-                             /*autoReify=*/0,
-                             /*autoInfer=*/1,
-                             /*declareInfer=*/1> {
+def Hip_MatmulOp :
+    Hip_DpsOp_Semantic<"matmul", /*traits=*/[OpStateOpInterface],
+                       /*outsAccessor=*/"Output",
+                       /*autoReify=*/0,
+                       /*autoInfer=*/1,
+                       /*declareInfer=*/1> {
   let summary = "Matrix multiplication (supports batched N-D x N-D)";
   let description = [{
     Performs matrix multiplication: output = A @ B
@@ -849,9 +850,10 @@ def Hip_MatmulOp : Hip_DpsOp<"matmul", /*traits=*/[OpStateOpInterface],
       B:      [..., K, N]
       output: [broadcast(A.batch, B.batch), M, N]
 
-    `?` (kDynamic) is permitted on any dim and behaves as a wildcard for
-    verification; `reifyResultShapes` lifts dynamic dims into `tensor.dim`
-    / `memref.dim` of the corresponding operand. See
+    `?` (kDynamic) is supported for M, N, and structurally representable batch
+    dimensions; dynamic contraction K is unsupported. `reifyResultShapes`
+    lifts supported dynamic dims into `tensor.dim` / `memref.dim` of the
+    corresponding operand. See
     [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md)
     for the dim-source contract and the consuming pipeline.
 
@@ -907,13 +909,14 @@ def Hip_MatmulOp : Hip_DpsOp<"matmul", /*traits=*/[OpStateOpInterface],
 // Quantized Operations
 //===----------------------------------------------------------------------===//
 
-def Hip_MatMulNBitsOp : Hip_DpsOp<"matmul_nbits",
-                                  /*traits=*/[AttrSizedOperandSegments,
-                                              OpStateOpInterface],
-                                  /*outsAccessor=*/"Output",
-                                  /*autoReify=*/0,
-                                  /*autoInfer=*/1,
-                                  /*declareInfer=*/1> {
+def Hip_MatMulNBitsOp :
+    Hip_DpsOp_Semantic<"matmul_nbits",
+                       /*traits=*/[AttrSizedOperandSegments,
+                                   OpStateOpInterface],
+                       /*outsAccessor=*/"Output",
+                       /*autoReify=*/0,
+                       /*autoInfer=*/1,
+                       /*declareInfer=*/1> {
   let summary = "Quantized N-bit matrix multiplication (com.microsoft MatMulNBits)";
   let description = [{
     Performs matrix multiplication with N-bit quantized weights.
@@ -986,6 +989,7 @@ def Hip_MatMulNBitsOp : Hip_DpsOp<"matmul_nbits",
           "hipdnn_ep_op_state_construct_matmul_nbits", {});
     }
   }];
+  let hasVerifier = 1;
 }
 
 def Hip_QMoEOp : Hip_DpsOp<"qmoe",
@@ -3629,11 +3633,12 @@ def Hip_LinearAttentionOp : Hip_DpsOp<"linear_attention",
   }];
 }
 
-def Hip_GemmOp : Hip_DpsOp<"gemm", /*traits=*/[OpStateOpInterface],
-                           /*outsAccessor=*/"Output",
-                           /*autoReify=*/0,
-                           /*autoInfer=*/1,
-                           /*declareInfer=*/1> {
+def Hip_GemmOp :
+    Hip_DpsOp_Semantic<"gemm", /*traits=*/[OpStateOpInterface],
+                       /*outsAccessor=*/"Output",
+                       /*autoReify=*/0,
+                       /*autoInfer=*/1,
+                       /*declareInfer=*/1> {
   let summary = "General matrix multiplication (ONNX Gemm)";
   let description = [{
     Computes Y = alpha * A' * B' + beta * C, where:
@@ -3688,6 +3693,11 @@ def Hip_GemmOp : Hip_DpsOp<"gemm", /*traits=*/[OpStateOpInterface],
           {});
     }
   }];
+
+  // Verifies the DPS contract plus the shared `inferGemmShape` rule, so the
+  // `outs` shape, the converter destination, and `reifyResultShapes` are all
+  // held to one ONNX Gemm shape function.
+  let hasVerifier = 1;
 }
 
 def Hip_CumSumOp : Hip_DpsOp<"cumsum", /*traits=*/[],