From a204fde40a858ff70407b44e1cf30152fbce6fe8 Mon Sep 17 00:00:00 2001 From: fhanuman Date: Tue, 28 Jul 2026 05:40:03 -0600 Subject: [PATCH 1/9] fix(hip): unify result-shape construction and reification Share broadcast, Gemm, and MatMul extent logic so converters allocate the same shapes exposed downstream, while preserving cached MatMul ABI compatibility and rejecting unsupported partial batch broadcasts. Co-Authored-By: GPT-5.6 Sol Made-with: Cursor --- CLAUDE.md | 9 + docs/design/hip-shape-inference.md | 51 +++- include/hip/Dialect/IR/HipOps.td | 8 +- include/hip/Dialect/IR/HipShapeUtils.h | 72 +++-- lib/Conversion/HipToLLVM/HipToLLVMUtils.h | 1 + lib/Conversion/HipToLLVM/MatmulLowering.cpp | 141 +++++---- lib/Conversion/OnnxToHip/GemmConversion.cpp | 20 +- lib/Conversion/OnnxToHip/MatMulConversion.cpp | 38 ++- lib/Conversion/OnnxToHip/MaxConversion.cpp | 30 +- lib/Conversion/OnnxToHip/MinConversion.cpp | 26 +- lib/Conversion/OnnxToHip/OnnxToHipUtils.h | 98 +++--- lib/Conversion/OnnxToHip/WhereConversion.cpp | 64 ---- lib/Dialect/IR/HipDialect.cpp | 22 +- lib/Dialect/IR/HipReifyResultShapesImpl.cpp | 85 +----- lib/Dialect/IR/HipShapeUtils.cpp | 288 +++++++++++++----- lib/Runtime/hipdnn_ep_runtime.h | 28 +- lib/Runtime/mock/mock_gpu.cpp | 10 + lib/Runtime/real/matmul.cpp | 62 ++-- .../Conversion/hip-to-llvm/test_matmul.mlir | 29 +- .../lit/Conversion/onnx-to-hip/test_gemm.mlir | 50 +++ .../onnx-to-hip/test_gemm_invalid_c.mlir | 15 + .../test_gemm_invalid_transpose.mlir | 15 + .../Conversion/onnx-to-hip/test_matmul.mlir | 34 +++ test/lit/Conversion/onnx-to-hip/test_max.mlir | 37 ++- test/lit/Conversion/onnx-to-hip/test_min.mlir | 37 ++- .../Dialect/hip-broadcast-reify-shapes.mlir | 74 +++++ test/lit/Dialect/hip-gemm-reify-shapes.mlir | 47 +++ test/lit/Dialect/hip-matmul-reify-shapes.mlir | 27 ++ .../Dialect/hip-matmul-shape-verifier.mlir | 13 + test/lit/Pipeline/pipeline-pool-lower.mlir | 4 +- test/numeric/tests/test_gemm.py | 54 ++++ test/numeric/tests/test_matmul.py | 44 +++ test/numeric/tests/test_max_min_broadcast.py | 55 ++++ test/numeric/tests/test_where.py | 28 ++ 34 files changed, 1162 insertions(+), 454 deletions(-) create mode 100644 test/lit/Conversion/onnx-to-hip/test_gemm_invalid_c.mlir create mode 100644 test/lit/Conversion/onnx-to-hip/test_gemm_invalid_transpose.mlir create mode 100644 test/lit/Dialect/hip-broadcast-reify-shapes.mlir create mode 100644 test/lit/Dialect/hip-gemm-reify-shapes.mlir create mode 100644 test/numeric/tests/test_gemm.py create mode 100644 test/numeric/tests/test_max_min_broadcast.py diff --git a/CLAUDE.md b/CLAUDE.md index 830667edd..fd731e270 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -138,6 +138,15 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co - Keep tiny host-written shape buffers out of the GPU pool; `hip-materialize-host-scalars` redirects them to host-mapped scratch. - See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md). +### Result-shape agreement + +- Converter destination construction and `reifyResultShapes` must use the same `OpFoldResult` shape helper for broadcast, Gemm, and MatMul. +- Fully dynamic broadcast uses `select(lhs == 1, rhs, lhs)`, not integer maximum: broadcasting extents 0 and 1 produces 0. +- Rank-zero success is an empty shape carried by `FailureOr`; never use an empty vector as both success and failure. +- Variadic Max/Min derive every pairwise intermediate rank from the shared broadcast shape. +- MatMul uses the reified output batch product plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Per-axis partial batch broadcasting is rejected because the single-stride runtime cannot represent it. New artifacts call `wrap_hipblasLtMatmul_v2`; retain the legacy wrapper for cached artifacts compiled with the old signature. +- See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md). + ### Allocation and memory planning - Every transient allocation must be pooled or rewritten as an output allocation. Leftover `hip.alloc`/`memref.alloc` paths are not a supported per-inference allocator strategy. diff --git a/docs/design/hip-shape-inference.md b/docs/design/hip-shape-inference.md index 6b84c20b5..cbde1cb46 100644 --- a/docs/design/hip-shape-inference.md +++ b/docs/design/hip-shape-inference.md @@ -4,7 +4,7 @@ Licensed under the MIT License. --> # HIP dialect shape inference -**Date:** 2026-07-24 +**Date:** 2026-07-28 **Document Type:** Design **Status:** Implemented **Related:** [unranked-tensor-handling.md](unranked-tensor-handling.md), [pool-allocs-memory-planning.md](pool-allocs-memory-planning.md), [output-allocator-design.md](output-allocator-design.md), [compiler-runtime-contract.md](compiler-runtime-contract.md), [pipeline_pass_menu.md](../pipeline_pass_menu.md) @@ -76,13 +76,14 @@ Choose the smallest mechanism that matches the operation's semantics: |---|---| | Result shape equals DPS init shape, including most multi-result DPS ops | Shared `HipDpsOpInterface` default | | Result shape equals a named input | `reifyElementwiseSameShape` or a small dedicated thunk | -| NumPy-style broadcast | `Hip_DpsOp_Broadcast` and broadcast helpers | +| NumPy-style broadcast | `Hip_DpsOp_Broadcast`, `reifyBroadcastResultShape`, and the shared converter bridge | | Reduction with constant axes/keepdims | `Hip_DpsOp_Reduction` and reduction helpers | | Permutation | `reifyTransposeByPerm` | | Gather/GatherND/GatherElements | Gather-specific helpers or thunks | | OneHot, Compress, TopK | Dedicated reification thunks | | Pad, Tile, Expand, Slice, Range | Fold-or-bail helpers with fallback to DPS-init shape | -| MatMul/Gemm/MatMulNBits | Dedicated shape logic based on operand dimensions and attributes | +| MatMul/Gemm | Shared operand-based helpers used by conversion and reification | +| MatMulNBits | Dedicated shape logic based on A and the N attribute | | Attention or normalization with multiple destinations | One shape vector per DPS init unless an op supplies a dedicated thunk | | Convolution, pooling, or resize with converter-computed destinations | DPS-init shape, with semantic validity handled by conversion or verification | | Runtime-dependent count, such as NonZero | DPS-init shape; unresolved dimensions remain dynamic | @@ -116,6 +117,47 @@ Reification is allowed to create IR at the caller's insertion point. Helpers the Reification is per result: `reifyResultShapes` returns one shape vector for every tensor result. The number and rank of those vectors must match the operation's tensor results even when the implementation derives them from DPS init operands. +### Shared converter/reification shape helpers + +Converter destination construction and operation reification must not +independently implement the same shape category. Broadcast, Gemm, and MatMul +use helpers in `HipShapeUtils` that return `FailureOr>`. +The `FailureOr` is required because a valid rank-zero result has a successful +empty shape. + +Broadcast dimensions are right-aligned. Static 1 yields to the other side; +equal non-unit static dimensions agree; dynamic/static-non-1 tightens to the +static extent under the ONNX input contract. Two dynamic dimensions emit: + +```mlir +%lhs_is_one = arith.cmpi eq, %lhs_dim, %c1 : index +%extent = arith.select %lhs_is_one, %rhs_dim, %lhs_dim : index +``` + +Do not replace this with an integer maximum: broadcasting dimensions 0 and 1 +produces 0, not 1. + +ONNX conversion keeps the imported ranked result type and uses the shared +`OpFoldResult`s only to populate `tensor.empty` dynamic-size operands. When a +reified dimension is constant but the imported dimension is dynamic, the +converter materializes a constant index size. `--hip-infer-shapes` remains the +single owner of later type narrowing, destination rebuilding, and cast +barriers. + +Variadic Max/Min derive every pairwise intermediate type from the shared +broadcast shape. Gemm derives M/N from A/B with transpose-aware indices and +checks optional C without using C as an extent source. MatMul broadcasts only +the leading batch slices, then appends M from A[-2] and N from B[-1]. + +The hipBLASLt MatMul lowering takes the batch product from the reified output +shape and carries independent A/B batch strides. A rank-2 or all-leading-one +operand uses stride 0, allowing the entire matrix to broadcast against a +batched operand on the other side. Per-axis partial batch broadcasting where +both operands contain fewer matrices than the output batch product is not yet +representable by one constant stride and is rejected by the MatMul verifier. +New lowering calls the versioned `wrap_hipblasLtMatmul_v2`; the legacy wrapper +remains available so cached artifacts with the previous signature still load. + ## `--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. @@ -224,6 +266,8 @@ Primary regression coverage: | `test/lit/Dialect/hip-infer-shapes.mlir` | Module-level static-dimension refinement and cast barriers | | `test/lit/Dialect/hip-infer-loop-body-shapes.mlir` | Pre-conversion rank establishment | | `test/lit/Dialect/hip-dps-op-interface.mlir` | Shared `HipDpsOpInterface` reification | +| `test/lit/Dialect/hip-broadcast-reify-shapes.mlir` | Broadcast dynamic SSA, zero extents, and rank-zero success | +| `test/lit/Dialect/hip-gemm-reify-shapes.mlir` | Transpose-aware Gemm M/N reification | | `test/lit/Dialect/hip-matmul-reify-shapes.mlir` | Per-op reification through `--resolve-shaped-type-result-dims` | | `test/lit/Dialect/hip-matmul-shape-verifier.mlir` | Static MatMul shape validation | | `test/lit/Dialect/hip-loop-verifier.mlir` | Loop-carried type contract | @@ -234,6 +278,7 @@ Complex operations may use dedicated files; common shape categories should exten ## Current limitations - ONNX MatMul rank-1 operands require promotion to rank 2 before constructing `hip.matmul`; the runtime and current verifier require rank at least 2. +- MatMul supports whole-matrix batch broadcast (one operand's leading product is 1) and equal flattened batch counts; per-axis partial batch broadcast remains unsupported by the strided-batch runtime. - Converter migration to inferred-type builders is incremental; explicit result-type builders remain supported. - A future multi-result operation that needs custom `InferTypeOpInterface` logic may require a dedicated result-type inference implementation file. - Runtime-dependent extents without pre-execution SSA remain dynamic. diff --git a/include/hip/Dialect/IR/HipOps.td b/include/hip/Dialect/IR/HipOps.td index 30926f48b..be6288342 100644 --- a/include/hip/Dialect/IR/HipOps.td +++ b/include/hip/Dialect/IR/HipOps.td @@ -1446,7 +1446,10 @@ def Hip_MinOp : Hip_DpsOp_Broadcast<"min", /*operandGetters=*/["Lhs", "Rhs"], }]; } -def Hip_MaxOp : Hip_DpsOp<"max", /*traits=*/[OpStateOpInterface]> { +def Hip_MaxOp : Hip_DpsOp_Broadcast<"max", + /*operandGetters=*/["Lhs", "Rhs"], + /*outsAccessor=*/"Output", + /*traits=*/[OpStateOpInterface]> { let summary = "Elementwise maximum"; // OpStateOpInterface: shared OpTensorState (see hip.mul). @@ -2552,8 +2555,7 @@ def Hip_SubOp : Hip_DpsOp_Broadcast<"sub", /*operandGetters=*/["Lhs", "Rhs"]> { let description = [{ Performs elementwise subtraction: output = lhs - rhs - Both inputs must have the same shape. Broadcasting is not currently supported - and must be handled before conversion to HIP dialect. + Supports multidirectional (NumPy-style) broadcasting. Uses destination-passing style: output buffer is provided as argument. diff --git a/include/hip/Dialect/IR/HipShapeUtils.h b/include/hip/Dialect/IR/HipShapeUtils.h index 753695d4b..1e678e9d8 100644 --- a/include/hip/Dialect/IR/HipShapeUtils.h +++ b/include/hip/Dialect/IR/HipShapeUtils.h @@ -41,6 +41,14 @@ SmallVector inferMatmulShape(ArrayRef aShape, ArrayRef bShape, function_ref emitError); +/// Verify that MatMul's broadcasted batches are representable by one constant +/// strided-batch offset per operand. Whole-matrix broadcast (rank 2 or all +/// leading extents 1) and equal static batch products are supported. Partial +/// per-axis broadcast and two nontrivial dynamic batch shapes are rejected. +LogicalResult +verifyStridedBatchMatmul(ArrayRef aShape, ArrayRef bShape, + function_ref emitError); + /// Verify that the actual `outs` operand shapes of a DPS HIP op match the /// shapes returned by `computeExpected`. `op` must implement /// `DestinationStyleOpInterface`. @@ -84,15 +92,29 @@ OpFoldResult reifyDimOrConstant(OpBuilder &b, Location loc, int64_t staticDim, SmallVector reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source); -/// Compute the NumPy-broadcast result shape over `operands` and lift each -/// output dim to an `OpFoldResult`. Static result dims become `IndexAttr` -/// (no IR emitted); dynamic result dims become `tensor.dim` against -/// whichever operand contributes the runtime extent — right-aligned, and -/// preferring the canonical side (in-range and != 1) when multiple -/// operands could contribute. The canonical-side preference matches the -/// batch-dim contract in `MatmulOp::reifyResultShapes` and ensures that -/// a future `tensor.dim` of the result folds back to the operand that -/// actually determines the size at runtime. +/// Compute a NumPy-broadcast result shape from already-reified operand shapes. +/// Static extents remain `IndexAttr`; dynamic/dynamic pairs materialize the +/// runtime broadcast rule as ordinary index SSA. +/// +/// Before (choosing either dynamic operand is incorrect when it is 1): +/// %lhs_dim = tensor.dim %lhs, %c0 +/// %init = tensor.empty(%lhs_dim) : tensor +/// After: +/// %lhs_dim = tensor.dim %lhs, %c0 +/// %rhs_dim = tensor.dim %rhs, %c0 +/// %lhs_is_one = arith.cmpi eq, %lhs_dim, %c1 : index +/// %extent = arith.select %lhs_is_one, %rhs_dim, %lhs_dim : index +/// %init = tensor.empty(%extent) : tensor +/// +/// The `FailureOr` distinguishes failure from a successful rank-zero shape. +FailureOr> +reifyBroadcastShape(OpBuilder &b, Location loc, + ArrayRef> inputShapes, + function_ref emitError); + +/// Compute the NumPy-broadcast result shape over ranked tensor `operands`. +/// This is the ValueRange convenience wrapper around the mixed-shape helper +/// above and uses `tensor::getMixedSizes` for each operand. /// /// Used by elementwise ops that take broadcast-shape operands and write /// the broadcast result into their `outs` (add, mul, sub, div, min, mod, @@ -102,12 +124,25 @@ SmallVector reifyElementwiseSameShape(OpBuilder &b, Location loc, /// and the helper handles both cases identically (it only looks at /// shapes). /// -/// All operands must be `RankedTensorType`-typed Values (the interface -/// contract for `reifyResultShapes` callers). Returns an empty vector -/// if broadcast fails — verifiers should already have caught this, but -/// reify bails defensively to avoid materializing nonsense IR. -SmallVector reifyBroadcastShape(OpBuilder &b, Location loc, - ValueRange operands); +/// All operands must be `RankedTensorType`-typed Values. +FailureOr> +reifyBroadcastResultShape(OpBuilder &b, Location loc, ValueRange operands, + function_ref emitError); + +/// Reify ONNX MatMul's result shape: broadcast the leading batch dimensions, +/// append M from `A[-2]`, and append N from `B[-1]`. Rank-1 MatMul is outside +/// the current HIP op contract. +FailureOr> +reifyMatmulResultShape(OpBuilder &b, Location loc, Value A, Value B, + function_ref emitError); + +/// Reify ONNX Gemm's rank-2 `{M, N}` result using transpose-aware dimensions. +/// Optional C is checked for static unidirectional broadcast compatibility but +/// never supplies M or N. +FailureOr> +reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, + Value optionalC, int64_t transA, int64_t transB, + function_ref emitError); /// Reify the result shape of a transpose op as `output[i] = input[perm[i]]`. /// `perm` must be a permutation of `[0, rank-1)` and have the same length @@ -200,15 +235,14 @@ LogicalResult reifyReductionShape(OpBuilder &b, Location loc, Value data, /// One-shot reify body for elementwise NumPy-broadcast ops (add, mul, /// sub, div, min, mod, equal, less, and, where, ...). Wraps -/// `reifyBroadcastShape` with the per-op guards (no-results bail, +/// `reifyBroadcastResultShape` with the per-op guards (no-results bail, /// every operand must be `RankedTensorType`) and writes the lifted /// dim list into `reified`. /// /// `operands` is the list of broadcast input operands in the order /// they should be aligned (right-aligned for NumPy broadcast). -/// Returns `failure()` on any defensive bail or when broadcast itself -/// fails (verifier should already have caught the latter; reify bails -/// to avoid materializing nonsense IR). +/// Returns `failure()` on any defensive bail or when broadcast itself fails. +/// A successful rank-zero result writes one empty shape into `reified`. /// /// Used as the body of `Hip_DpsOp_Broadcast`'s auto-emitted reify /// dispatcher; see `Hip_DpsOp_Broadcast` in `HipOps.td`. diff --git a/lib/Conversion/HipToLLVM/HipToLLVMUtils.h b/lib/Conversion/HipToLLVM/HipToLLVMUtils.h index a32c638cd..259fcc834 100644 --- a/lib/Conversion/HipToLLVM/HipToLLVMUtils.h +++ b/lib/Conversion/HipToLLVM/HipToLLVMUtils.h @@ -52,6 +52,7 @@ inline constexpr const char *kMiopenConvolutionForward = inline constexpr const char *kMiopenConvolutionTranspose = "wrap_miopenConvolutionTranspose"; inline constexpr const char *kWrapHipblasltMatmul = "wrap_hipblasLtMatmul"; +inline constexpr const char *kWrapHipblasltMatmulV2 = "wrap_hipblasLtMatmul_v2"; inline constexpr const char *kWrapMiopenT5LayerNormForward = "wrap_miopenT5LayerNormForward"; inline constexpr const char *kWrapSkipSimplifiedLayerNorm = diff --git a/lib/Conversion/HipToLLVM/MatmulLowering.cpp b/lib/Conversion/HipToLLVM/MatmulLowering.cpp index 9ce34eab6..bdac586f2 100644 --- a/lib/Conversion/HipToLLVM/MatmulLowering.cpp +++ b/lib/Conversion/HipToLLVM/MatmulLowering.cpp @@ -5,6 +5,8 @@ #include "HipToLLVMUtils.h" +#include "hip/Dialect/IR/HipShapeUtils.h" + namespace mlir { namespace hip { namespace { @@ -12,8 +14,10 @@ namespace { // ===== hipBLASLt ops ========================================================= // hip.hipblaslt.matmul(handle) ins(A, B) outs(C) -// -> hip_hipblaslt_matmul(handle, A, B, C, rankA, rankB, batch, M, K, N) -// Rank-generic: batch from A if 3D, B broadcast if rankB < rankA. +// -> wrap_hipblasLtMatmul(state, A, B, C, M, N, K, batch, A-stride, +// B-stride) +// Rank-generic for whole-tensor batch broadcast: a rank-2 (or leading-one) +// operand uses stride 0 while the other operand supplies the output batches. struct MatmulOpLowering : public ConvertOpToLLVMPattern { using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; @@ -42,13 +46,19 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { // Get memref types and shapes auto AType = cast(op.getA().getType()); auto BType = cast(op.getB().getType()); + auto outputType = cast(op.getOutput().getType()); int64_t ARank = AType.getRank(); int64_t BRank = BType.getRank(); + int64_t outputRank = outputType.getRank(); + if (failed(verifyStridedBatchMatmul(AType.getShape(), BType.getShape(), + [&]() { return op.emitOpError(); }))) + return failure(); // === DYNAMIC SHAPE SUPPORT === // For dynamic shapes, we compute dimensions at runtime MemRefDescriptor ADesc(adaptor.getA()); MemRefDescriptor BDesc(adaptor.getB()); + MemRefDescriptor outputDesc(adaptor.getOutput()); // Compute M, K, N from runtime dimensions // A: [..., M, K], B: [..., K, N] or B: [K, N] @@ -57,90 +67,71 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { Value K = ADesc.size(rewriter, loc, ARank - 1); Value N = BDesc.size(rewriter, loc, BRank - 1); - // Compute batch count from leading dimensions of A - Value batchCount; - if (ARank == 2) { - batchCount = createI64Const(1); - } else { - batchCount = ADesc.size(rewriter, loc, 0); - for (int64_t i = 1; i < ARank - 2; ++i) { - Value dim = ADesc.size(rewriter, loc, i); - batchCount = LLVM::MulOp::create(rewriter, loc, batchCount, dim); - } + // The output shape already contains the broadcasted leading dimensions. + Value batchCount = createI64Const(1); + for (int64_t i : llvm::seq(0, outputRank - 2)) { + Value dim = outputDesc.size(rewriter, loc, i); + batchCount = LLVM::MulOp::create(rewriter, loc, batchCount, dim); } // Compute element size in bytes unsigned elemBits = AType.getElementType().getIntOrFloatBitWidth(); Value elemSize = createI64Const(elemBits / 8); - // b_batch_stride: the per-batch stride (in elements) for hipBLASLt's - // STRIDED_BATCH_OFFSET on layA when batch_count > 1. Two distinct B - // memref shapes both reach this site with batch_count > 1 and need - // DIFFERENT strides: - // * Rank-2 broadcast weight `[K, N]` → one matrix shared across all - // batches → stride = 0. - // * Rank-N per-batch weight `[d_0, ..., d_{N-3}, K, N]` whose leading - // dims hold more than one matrix → stride = K * N. - // * Rank-N weight whose leading dims multiply to 1 (e.g. `[1, K, N]`, - // `[1, 1, K, N]`) → still ONE matrix; stride = 0. - // Previously encoded as a `b_batched` bool keyed on `BRank > 2`. That - // misclassified the `[1, K, N]` leading-one case as per-batch, causing - // OOB reads of `K*N` elements past the end of the weight on every batch - // beyond the first. - // - // Leading-dim product is computed at compile time when all leading dims - // are static (the common case for shipping models). When any leading dim - // is dynamic we emit a runtime `select` over the leading-product so the - // descriptor cache picks the right pre-built layout per call. + // Each operand independently uses stride 0 when it contains one matrix, + // or its matrix size when it contains one matrix per output batch. // // Before: - // hip.matmul ins(%A, %B : memref<2x128x4096xf16>, - // memref<1x4096x1024xf16>) + // hip.matmul ins(%A, %B : memref<128x4096xf16>, + // memref<2x4096x1024xf16>) // After: - // %stride = llvm.mlir.constant(0 : i64) : i64 // leading product is 1 - // llvm.call @wrap_hipblasLtMatmul(..., %stride) : (...) -> i32 - Value bBatchStride; - if (BRank == 2) { - bBatchStride = createI64Const(0); - } else { + // %a_stride = llvm.mlir.constant(0 : i64) : i64 + // %b_stride = llvm.mul %K, %N : i64 + // llvm.call @wrap_hipblasLtMatmul(..., %a_stride, %b_stride) + auto computeBatchStride = [&](MemRefType type, MemRefDescriptor desc, + Value matrixElements) -> Value { + int64_t rank = type.getRank(); + if (rank == 2) + return createI64Const(0); + bool allLeadingStatic = true; int64_t staticLeadingProduct = 1; - for (int64_t i : llvm::seq(0, BRank - 2)) { - if (BType.isDynamicDim(i)) { + for (int64_t i : llvm::seq(0, rank - 2)) { + if (type.isDynamicDim(i)) { allLeadingStatic = false; break; } - staticLeadingProduct *= BType.getDimSize(i); + staticLeadingProduct *= type.getDimSize(i); } - if (allLeadingStatic) { - bBatchStride = (staticLeadingProduct <= 1) - ? createI64Const(0) - : LLVM::MulOp::create(rewriter, loc, K, N).getRes(); - } else { - Value one = createI64Const(1); - Value zero = createI64Const(0); - Value leadingProduct = one; - for (int64_t i : llvm::seq(0, BRank - 2)) { - Value dim = BDesc.size(rewriter, loc, i); - leadingProduct = - LLVM::MulOp::create(rewriter, loc, leadingProduct, dim).getRes(); - } - Value isBroadcast = LLVM::ICmpOp::create( - rewriter, loc, LLVM::ICmpPredicate::sle, leadingProduct, one); - Value kn = LLVM::MulOp::create(rewriter, loc, K, N).getRes(); - bBatchStride = - LLVM::SelectOp::create(rewriter, loc, isBroadcast, zero, kn) - .getRes(); + if (allLeadingStatic) + return staticLeadingProduct <= 1 ? createI64Const(0) : matrixElements; + + Value one = createI64Const(1); + Value leadingProduct = one; + for (int64_t i : llvm::seq(0, rank - 2)) { + Value dim = desc.size(rewriter, loc, i); + leadingProduct = + LLVM::MulOp::create(rewriter, loc, leadingProduct, dim).getRes(); } - } + Value isBroadcast = LLVM::ICmpOp::create( + rewriter, loc, LLVM::ICmpPredicate::sle, leadingProduct, one); + return LLVM::SelectOp::create(rewriter, loc, isBroadcast, + createI64Const(0), matrixElements) + .getRes(); + }; + Value aMatrixElements = LLVM::MulOp::create(rewriter, loc, M, K); + Value bMatrixElements = LLVM::MulOp::create(rewriter, loc, K, N); + Value aBatchStride = computeBatchStride(AType, ADesc, aMatrixElements); + Value bBatchStride = computeBatchStride(BType, BDesc, bMatrixElements); // Runtime signature: // int wrap_hipblasLtMatmul(RuntimeState* state, // const void* A, const void* B, void* output, // int64_t M, int64_t N, int64_t K, // int64_t batch_count, int64_t elem_size, - // int64_t b_batch_stride, int op_state_slot) - SmallVector paramTypes = { + // int64_t a_batch_stride, + // int64_t b_batch_stride) + SmallVector paramTypes = { ptrType, // state i32Type, // op_state_slot ptrType, // A @@ -151,21 +142,27 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { i64Type, // K i64Type, // batch_count i64Type, // elem_size + i64Type, // a_batch_stride i64Type // b_batch_stride }; FailureOr funcOp = LLVM::lookupOrCreateFn( - rewriter, module, kWrapHipblasltMatmul, paramTypes, i32Type); + rewriter, module, kWrapHipblasltMatmulV2, paramTypes, i32Type); if (failed(funcOp)) return failure(); - SmallVector args = { - statePtr, getOpStateSlotValue(op, rewriter, loc), - APtr, BPtr, - outputPtr, M, - N, K, - batchCount, elemSize, - bBatchStride}; + SmallVector args = {statePtr, + getOpStateSlotValue(op, rewriter, loc), + APtr, + BPtr, + outputPtr, + M, + N, + K, + batchCount, + elemSize, + aBatchStride, + bBatchStride}; LLVM::CallOp::create(rewriter, loc, *funcOp, args); rewriter.eraseOp(op); diff --git a/lib/Conversion/OnnxToHip/GemmConversion.cpp b/lib/Conversion/OnnxToHip/GemmConversion.cpp index 834b808f7..a61351acb 100644 --- a/lib/Conversion/OnnxToHip/GemmConversion.cpp +++ b/lib/Conversion/OnnxToHip/GemmConversion.cpp @@ -11,6 +11,12 @@ namespace { //===----------------------------------------------------------------------===// // ONNX Gemm -> HIP Gemm +// +// Before: +// %init = tensor.empty(tensor.dim %A, 0, tensor.dim %A, 1) +// After: +// %shape = gemm_shape(%A, %B, transA, transB) +// %init = tensor.empty(%shape.M, %shape.N) //===----------------------------------------------------------------------===// struct GemmToHip : public mlir::RewritePattern { GemmToHip(mlir::MLIRContext *ctx) @@ -44,7 +50,17 @@ struct GemmToHip : public mlir::RewritePattern { mlir::Location loc = op->getLoc(); auto resultType = mlir::cast(op->getResult(0).getType()); - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, inputA); + mlir::FailureOr> resultShape = + mlir::hip::reifyGemmResultShape(rewriter, loc, inputA, inputB, inputC, + transA, transB, + [&]() { return op->emitError(); }); + if (mlir::failed(resultShape)) + return mlir::failure(); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "Gemm result type is incompatible with inferred shape"); // hip.gemm llvm::SmallVector attrs; @@ -61,7 +77,7 @@ struct GemmToHip : public mlir::RewritePattern { if (hasInputC) { operands.push_back(inputC); } - operands.push_back(init); + operands.push_back(*init); // Result type inferred from `init` via InferTypeOpInterface — DPS contract: // result type == outs operand type. auto hipOp = mlir::hip::GemmOp::create(rewriter, loc, operands, attrs); diff --git a/lib/Conversion/OnnxToHip/MatMulConversion.cpp b/lib/Conversion/OnnxToHip/MatMulConversion.cpp index 893fd829e..f4952f956 100644 --- a/lib/Conversion/OnnxToHip/MatMulConversion.cpp +++ b/lib/Conversion/OnnxToHip/MatMulConversion.cpp @@ -10,6 +10,13 @@ namespace hip { namespace { /// onnx.MatMul -> hip.hipblaslt.matmul +/// +/// Before: +/// %batch = tensor.dim %A, %c0 +/// %init = tensor.empty(%batch, ...) : tensor +/// After: +/// %shape = matmul_shape(%A, %B) +/// %init = tensor.empty(%shape...) : tensor struct MatMulToHip : public mlir::RewritePattern { MatMulToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.MatMul", /*benefit=*/1, ctx) {} @@ -33,32 +40,23 @@ MatMulToHip::matchAndRewrite(mlir::Operation *op, auto resultType = mlir::cast(op->getResult(0).getType()); - // MatMul: result[..., M, N] = A[..., M, K] @ B[..., K, N]. - // Batch and M dims come from A; N comes from B's last dim. - llvm::SmallVector dynSizes; - const int64_t rank = resultType.getRank(); - const auto bType = mlir::cast(b.getType()); - for (int64_t dimIdx : llvm::seq(rank)) { - if (!resultType.isDynamicDim(dimIdx)) - continue; - if (dimIdx == rank - 1) { - dynSizes.push_back( - mlir::tensor::DimOp::create(rewriter, loc, b, bType.getRank() - 1)); - } else { - dynSizes.push_back(mlir::tensor::DimOp::create(rewriter, loc, a, dimIdx)); - } - } - - mlir::Value init = - mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); + mlir::FailureOr> resultShape = + mlir::hip::reifyMatmulResultShape(rewriter, loc, a, b, + [&]() { return op->emitError(); }); + if (mlir::failed(resultShape)) + return mlir::failure(); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "MatMul result type is incompatible with inferred shape"); // Inferred-type Op::create overload: result type is read from the typed // outs operand via the auto-emitted MatmulOp::inferReturnTypes (HipOps.td // base, autoInfer=1). Equivalent to passing `resultType` explicitly -- // outs.getType() == resultType by construction here -- but keeps the DPS // contract `result_type == outs_operand_type` closed by ODS rather than // restated at the callsite. - auto hipOp = mlir::hip::MatmulOp::create(rewriter, loc, context, a, b, init); + auto hipOp = mlir::hip::MatmulOp::create(rewriter, loc, context, a, b, *init); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); } diff --git a/lib/Conversion/OnnxToHip/MaxConversion.cpp b/lib/Conversion/OnnxToHip/MaxConversion.cpp index 985aab5a1..03edf0a5b 100644 --- a/lib/Conversion/OnnxToHip/MaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/MaxConversion.cpp @@ -14,6 +14,12 @@ namespace { /// Handles variadic inputs by pairwise chaining: /// max(a, b, c) = max(max(a, b), c) /// Single input is identity (pass through). +/// +/// Before: +/// %tmp = hip.max ... outs(%same_rank_as_a) +/// After: +/// %tmp_shape = broadcast_shape(%a, %b) +/// %tmp = hip.max ... outs(%empty_for_tmp_shape) struct MaxToHip : public mlir::RewritePattern { MaxToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.Max", /*benefit=*/1, ctx) {} @@ -48,15 +54,25 @@ MaxToHip::matchAndRewrite(mlir::Operation *op, mlir::Value accumulate = op->getOperand(0); for (unsigned i = 1; i < numInputs; ++i) { mlir::Value rhs = op->getOperand(i); - auto accType = mlir::cast(accumulate.getType()); + mlir::FailureOr> stepShape = + mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, + [&]() { return op->emitError(); }); + if (mlir::failed(stepShape)) + return mlir::failure(); + + bool isFinal = i == numInputs - 1; mlir::RankedTensorType stepResultType = - (i == numInputs - 1) ? resultType : accType; + isFinal ? resultType + : getTensorTypeFromReifiedShape(*stepShape, + resultType.getElementType()); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, stepResultType, *stepShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "Max result type is incompatible with broadcast shape"); - mlir::Value source = - (accType.getRank() == stepResultType.getRank()) ? accumulate : rhs; - mlir::Value init = createEmptyTensor(rewriter, loc, stepResultType, source); - auto maxOp = mlir::hip::MaxOp::create(rewriter, loc, stepResultType, - context, accumulate, rhs, init); + auto maxOp = mlir::hip::MaxOp::create(rewriter, loc, context, accumulate, + rhs, *init); accumulate = maxOp->getResult(0); } diff --git a/lib/Conversion/OnnxToHip/MinConversion.cpp b/lib/Conversion/OnnxToHip/MinConversion.cpp index 18daa913a..9b8a0fc81 100644 --- a/lib/Conversion/OnnxToHip/MinConversion.cpp +++ b/lib/Conversion/OnnxToHip/MinConversion.cpp @@ -14,6 +14,12 @@ namespace { /// Handles variadic inputs by pairwise chaining: /// min(a, b, c) = min(min(a, b), c) /// Single input is identity (pass through). +/// +/// Before: +/// %tmp = hip.min ... outs(%same_rank_as_a) +/// After: +/// %tmp_shape = broadcast_shape(%a, %b) +/// %tmp = hip.min ... outs(%empty_for_tmp_shape) struct MinToHip : public mlir::RewritePattern { MinToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.Min", /*benefit=*/1, ctx) {} @@ -49,15 +55,23 @@ MinToHip::matchAndRewrite(mlir::Operation *op, mlir::Value accumulate = op->getOperand(0); for (unsigned i = 1; i < numInputs; ++i) { mlir::Value rhs = op->getOperand(i); - auto accType = mlir::cast(accumulate.getType()); - mlir::RankedTensorType stepResultType = - (i == numInputs - 1) ? resultType : accType; + mlir::FailureOr> stepShape = + mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, + [&]() { return op->emitError(); }); + if (mlir::failed(stepShape)) + return mlir::failure(); - mlir::FailureOr initOrFailure = createBroadcastEmptyTensor( - rewriter, loc, stepResultType, {accumulate, rhs}); + bool isFinal = i == numInputs - 1; + mlir::RankedTensorType stepResultType = + isFinal ? resultType + : getTensorTypeFromReifiedShape(*stepShape, + resultType.getElementType()); + mlir::FailureOr initOrFailure = + createEmptyTensorFromReifiedShape(rewriter, loc, stepResultType, + *stepShape); if (mlir::failed(initOrFailure)) return rewriter.notifyMatchFailure( - op, "Min: no ranked operand spans dynamic result dim"); + op, "Min result type is incompatible with broadcast shape"); auto minOp = mlir::hip::MinOp::create(rewriter, loc, context, accumulate, rhs, *initOrFailure); accumulate = minOp->getResult(0); diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h index 6ba6b329a..d12cea994 100644 --- a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h @@ -16,6 +16,7 @@ #include "hip/Conversion/OnnxToHip/Passes.h" #include "hip/Dialect/IR/HipDialect.h" +#include "hip/Dialect/IR/HipShapeUtils.h" #include "hip/Dialect/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -26,6 +27,7 @@ #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" @@ -144,13 +146,51 @@ inferReduceResultType(mlir::Operation *op, mlir::Value data, return mlir::RankedTensorType::get(outShape, inputType.getElementType()); } +/// Build a tensor.empty with the imported result type and the dynamic sizes +/// described by `reifiedShape`. Static reified dimensions are materialized as +/// constant index operands when the imported type keeps that dimension +/// dynamic; the existing `hip-infer-shapes` pass owns later type refinement. +inline mlir::FailureOr createEmptyTensorFromReifiedShape( + mlir::OpBuilder &builder, mlir::Location loc, + mlir::RankedTensorType resultType, + llvm::ArrayRef reifiedShape) { + if (static_cast(reifiedShape.size()) != resultType.getRank()) + return mlir::failure(); + + llvm::SmallVector dynSizes; + for (int64_t dimIdx : llvm::seq(resultType.getRank())) { + std::optional reifiedStatic = + mlir::getConstantIntValue(reifiedShape[dimIdx]); + if (!resultType.isDynamicDim(dimIdx)) { + if (reifiedStatic && *reifiedStatic != resultType.getDimSize(dimIdx)) + return mlir::failure(); + continue; + } + dynSizes.push_back(mlir::getValueOrCreateConstantIndexOp( + builder, loc, reifiedShape[dimIdx])); + } + return mlir::Value( + mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), + resultType.getElementType(), dynSizes)); +} + +/// Derive a tensor type for a synthesized intermediate from a reified shape. +/// Constant dimensions become static; all other dimensions stay dynamic. +inline mlir::RankedTensorType +getTensorTypeFromReifiedShape(llvm::ArrayRef reifiedShape, + mlir::Type elementType, + mlir::Attribute encoding = {}) { + llvm::SmallVector shape; + shape.reserve(reifiedShape.size()); + for (mlir::OpFoldResult dim : reifiedShape) + shape.push_back( + mlir::getConstantIntValue(dim).value_or(mlir::ShapedType::kDynamic)); + return mlir::RankedTensorType::get(shape, elementType, encoding); +} + /// Create a tensor.empty for a DPS init whose shape is the NumPy-style -/// broadcast of \p operands. Operand shapes are right-aligned with the -/// result. For each dynamic dimension of \p resultType, the size is taken -/// from the first operand that truly contributes at that axis -- i.e. whose -/// corresponding dim is not statically 1. Shorter-rank operands (left-padded -/// with 1) and statically-1 dims are skipped. If every spanning operand is -/// statically 1 at the axis, fall back to the first operand that spans it. +/// broadcast of \p operands. Converter destination construction delegates to +/// the same dialect helper used by ReifyRankedShapedTypeOpInterface. /// /// Use this for binary/multinary broadcast elementwise ops (Add, Mul, Where, /// ...). Do NOT use `createEmptyTensor(resultType, source)` when operands can @@ -160,46 +200,12 @@ inline mlir::FailureOr createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, mlir::RankedTensorType resultType, mlir::ValueRange operands) { - int64_t resultRank = resultType.getRank(); - llvm::SmallVector dynSizes; - for (int64_t dimIdx : llvm::seq(resultRank)) { - if (!resultType.isDynamicDim(dimIdx)) - continue; - - mlir::Value chosen; - int64_t chosenDim = -1; - mlir::Value fallback; - int64_t fallbackDim = -1; - for (mlir::Value operand : operands) { - auto t = mlir::dyn_cast(operand.getType()); - if (!t) - continue; - int64_t offset = resultRank - t.getRank(); - if (dimIdx < offset) - continue; - int64_t operandDim = dimIdx - offset; - if (!fallback) { - fallback = operand; - fallbackDim = operandDim; - } - if (!t.isDynamicDim(operandDim) && t.getDimSize(operandDim) == 1) - continue; - chosen = operand; - chosenDim = operandDim; - break; - } - if (!chosen) { - chosen = fallback; - chosenDim = fallbackDim; - } - if (!chosen) - return mlir::failure(); - dynSizes.push_back( - mlir::tensor::DimOp::create(builder, loc, chosen, chosenDim)); - } - return mlir::Value( - mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), - resultType.getElementType(), dynSizes)); + mlir::FailureOr> shape = + mlir::hip::reifyBroadcastResultShape( + builder, loc, operands, [&]() { return mlir::emitError(loc); }); + if (mlir::failed(shape)) + return mlir::failure(); + return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); } /// Get !hip.context from function argument 0. Returns failure if the diff --git a/lib/Conversion/OnnxToHip/WhereConversion.cpp b/lib/Conversion/OnnxToHip/WhereConversion.cpp index 65231f74b..1b6cd9f5a 100644 --- a/lib/Conversion/OnnxToHip/WhereConversion.cpp +++ b/lib/Conversion/OnnxToHip/WhereConversion.cpp @@ -9,70 +9,6 @@ namespace mlir { namespace hip { namespace { -/// Create a tensor.empty for a DPS init operand whose shape is the result of -/// ONNX-style multidirectional (NumPy) broadcasting over multiple inputs. -/// -/// Operand shapes are right-aligned with the result. For each dynamic -/// dimension of \p resultType, the size is taken from the first operand that -/// truly contributes to the broadcast extent at that axis -- i.e., an operand -/// whose corresponding dim is not statically 1. Operands whose rank does not -/// span the dimension (shorter rank, conceptually padded with 1 on the left) -/// are skipped, as are operands whose dim is statically 1. If every operand -/// is statically 1 at the axis (degenerate case for a dynamic result), we -/// fall back to the first operand spanning the dim. -/// -/// Returns failure if no ranked operand spans a dynamic result dim (e.g. -/// every operand is unranked while the result is ranked-and-dynamic). We -/// surface this through `FailureOr` rather than `assert` so that the check -/// remains active in Release builds (assertions are stripped under NDEBUG) -/// and the pattern fails cleanly via `notifyMatchFailure` instead of -/// dereferencing a null Value at the next builder call. -mlir::FailureOr -createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, - mlir::RankedTensorType resultType, - mlir::ValueRange operands) { - int64_t resultRank = resultType.getRank(); - llvm::SmallVector dynSizes; - for (int64_t dimIdx : llvm::seq(resultRank)) { - if (!resultType.isDynamicDim(dimIdx)) - continue; - - mlir::Value chosen; - int64_t chosenDim = -1; - mlir::Value fallback; - int64_t fallbackDim = -1; - for (mlir::Value operand : operands) { - auto t = mlir::dyn_cast(operand.getType()); - if (!t) - continue; - int64_t offset = resultRank - t.getRank(); - if (dimIdx < offset) - continue; // operand is broadcast (padded to 1) at this axis - int64_t operandDim = dimIdx - offset; - if (!fallback) { - fallback = operand; - fallbackDim = operandDim; - } - if (!t.isDynamicDim(operandDim) && t.getDimSize(operandDim) == 1) - continue; // statically broadcast, does not define the extent - chosen = operand; - chosenDim = operandDim; - break; - } - if (!chosen) { - chosen = fallback; - chosenDim = fallbackDim; - } - if (!chosen) - return mlir::failure(); - dynSizes.push_back( - mlir::tensor::DimOp::create(builder, loc, chosen, chosenDim)); - } - return mlir::Value( - mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), - resultType.getElementType(), dynSizes)); -} - /// onnx.Where -> hip.where /// ONNX Where: output[i] = condition[i] ? X[i] : Y[i]. /// Supports multidirectional (NumPy-style) broadcasting between condition, diff --git a/lib/Dialect/IR/HipDialect.cpp b/lib/Dialect/IR/HipDialect.cpp index 9aaf29994..442113818 100644 --- a/lib/Dialect/IR/HipDialect.cpp +++ b/lib/Dialect/IR/HipDialect.cpp @@ -567,18 +567,22 @@ LogicalResult MatmulOp::verify() { /*numInits=*/1))) return failure(); + ArrayRef aShape = getShapeOf(getA()); + ArrayRef bShape = getShapeOf(getB()); // Static shape check via the shared matmul helper. The lambda is // invoked once; it returns `{outputShape}` on success or `{}` on shape // mismatch (in which case it has already issued a diagnostic on `*this`). - return mlir::hip::verifyHipOpShape( - *this, [&]() -> SmallVector> { - SmallVector outShape = - mlir::hip::inferMatmulShape(getShapeOf(getA()), getShapeOf(getB()), - [&]() { return this->emitOpError(); }); - if (outShape.empty()) - return {}; - return {std::move(outShape)}; - }); + if (failed(mlir::hip::verifyHipOpShape( + *this, [&]() -> SmallVector> { + SmallVector outShape = mlir::hip::inferMatmulShape( + aShape, bShape, [&]() { return this->emitOpError(); }); + if (outShape.empty()) + return {}; + return {std::move(outShape)}; + }))) + return failure(); + return mlir::hip::verifyStridedBatchMatmul( + aShape, bShape, [&]() { return this->emitOpError(); }); } // `MatmulOp::reifyResultShapes` lives in diff --git a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp index cfaa289c9..2e86ba4b0 100644 --- a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp +++ b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp @@ -44,11 +44,9 @@ ArrayRef getShapeOf(Value v) { //===----------------------------------------------------------------------===// // MatmulOp // -// Reify recomputes the result shape via `inferMatmulShape`, then lifts -// each dim to an OpFoldResult: static dims become `IndexAttr`; dynamic -// dims become `tensor.dim` of whichever operand contributes the runtime -// size — M from A[-2], N from B[-1], batch from the broadcast-canonical -// side. +// Reify delegates to the shared MatMul helper used by converter destination +// construction: M comes from A[-2], N from B[-1], and leading dimensions use +// NumPy broadcast semantics. // // Before: // %m = hip.matmul ins(%a, %b : tensor, tensor<4096x4096xf16>) @@ -66,58 +64,11 @@ MatmulOp::reifyResultShapes(OpBuilder &b, if (getNumResults() == 0) return failure(); - ArrayRef aShape = getShapeOf(getA()); - ArrayRef bShape = getShapeOf(getB()); - if (aShape.empty() || bShape.empty()) - return failure(); - - // Re-run the matmul-shape helper. verify() has already passed by reify - // time, but bail on empty() in case a pre-verify call sneaks in. - SmallVector outShape = mlir::hip::inferMatmulShape( - aShape, bShape, [&]() { return this->emitOpError(); }); - if (outShape.empty()) + FailureOr> dims = mlir::hip::reifyMatmulResultShape( + b, getLoc(), getA(), getB(), [&]() { return this->emitOpError(); }); + if (failed(dims)) return failure(); - - Location loc = getLoc(); - Value A = getA(); - Value B = getB(); - size_t outRank = outShape.size(); - size_t aRank = aShape.size(); - size_t bRank = bShape.size(); - - // Loop-invariant: right-alignment padding for A's and B's batch dims. - size_t batchRank = outRank - 2; - size_t aPad = batchRank - (aRank >= 2 ? aRank - 2 : 0); - size_t bPad = batchRank - (bRank >= 2 ? bRank - 2 : 0); - - SmallVector dims; - dims.reserve(outRank); - for (size_t i : llvm::seq(0, outRank)) { - // M dim: A[-2]. - if (i + 2 == outRank) { - dims.push_back( - mlir::hip::reifyDimOrConstant(b, loc, outShape[i], A, aRank - 2)); - continue; - } - // N dim: B[-1]. - if (i + 1 == outRank) { - dims.push_back( - mlir::hip::reifyDimOrConstant(b, loc, outShape[i], B, bRank - 1)); - continue; - } - // Batch dim: prefer the side that contributes the size (in range, not 1). - // When neither contributes, prefer A in range so folds see a stable source. - int64_t aDim = i < aPad ? 1 : aShape[i - aPad]; - int64_t bDim = i < bPad ? 1 : bShape[i - bPad]; - bool aCanonical = i >= aPad && aDim != 1; - bool bCanonical = i >= bPad && bDim != 1; - bool pickA = aCanonical || (!bCanonical && i >= aPad); - Value src = pickA ? A : B; - size_t srcDim = pickA ? i - aPad : i - bPad; - dims.push_back( - mlir::hip::reifyDimOrConstant(b, loc, outShape[i], src, srcDim)); - } - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -276,24 +227,12 @@ GemmOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { if (getNumResults() == 0) return failure(); - ArrayRef aShape = getShapeOf(getInputA()); - ArrayRef bShape = getShapeOf(getInputB()); - if (aShape.size() != 2 || bShape.size() != 2) + FailureOr> dims = mlir::hip::reifyGemmResultShape( + b, getLoc(), getInputA(), getInputB(), getInputC(), getTransA(), + getTransB(), [&]() { return this->emitOpError(); }); + if (failed(dims)) return failure(); - - Location loc = getLoc(); - Value A = getInputA(); - Value B = getInputB(); - bool transA = getTransA() != 0; - bool transB = getTransB() != 0; - - size_t mDim = transA ? 1 : 0; - size_t nDim = transB ? 0 : 1; - SmallVector dims; - dims.reserve(2); - dims.push_back(mlir::hip::reifyDimOrConstant(b, loc, aShape[mDim], A, mDim)); - dims.push_back(mlir::hip::reifyDimOrConstant(b, loc, bShape[nDim], B, nDim)); - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index bb0d12d42..8719dbbc8 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -17,8 +17,10 @@ #include "hip/Dialect/IR/HipDialect.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Traits.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Matchers.h" @@ -29,6 +31,8 @@ #include "llvm/Support/raw_ostream.h" #include +#include +#include using namespace mlir; using namespace mlir::hip; @@ -50,6 +54,41 @@ std::string formatShape(ArrayRef shape) { return os.str(); } +FailureOr +broadcastDim(OpBuilder &b, Location loc, OpFoldResult lhs, OpFoldResult rhs, + function_ref emitError) { + if (lhs == rhs) + return lhs; + + std::optional lhsStatic = getConstantIntValue(lhs); + std::optional rhsStatic = getConstantIntValue(rhs); + + if (lhsStatic && rhsStatic) { + if (*lhsStatic == 1) + return rhs; + if (*rhsStatic == 1 || *lhsStatic == *rhsStatic) + return lhs; + emitError() << "incompatible broadcast dimensions " << *lhsStatic << " and " + << *rhsStatic; + return failure(); + } + + // Under the ONNX broadcastability precondition, a dynamic extent paired + // with a known non-unit extent must be either 1 or that known extent. + if (lhsStatic) + return *lhsStatic == 1 ? rhs : lhs; + if (rhsStatic) + return *rhsStatic == 1 ? lhs : rhs; + + Value lhsValue = getValueOrCreateConstantIndexOp(b, loc, lhs); + Value rhsValue = getValueOrCreateConstantIndexOp(b, loc, rhs); + Value one = arith::ConstantIndexOp::create(b, loc, 1); + Value lhsIsOne = + arith::CmpIOp::create(b, loc, arith::CmpIPredicate::eq, lhsValue, one); + return OpFoldResult( + arith::SelectOp::create(b, loc, lhsIsOne, rhsValue, lhsValue)); +} + } // namespace SmallVector @@ -92,6 +131,47 @@ mlir::hip::inferMatmulShape(ArrayRef aShape, ArrayRef bShape, return result; } +LogicalResult mlir::hip::verifyStridedBatchMatmul( + ArrayRef aShape, ArrayRef bShape, + function_ref emitError) { + if (aShape.size() < 2 || bShape.size() < 2) + return failure(); + + ArrayRef aBatch = aShape.drop_back(2); + ArrayRef bBatch = bShape.drop_back(2); + auto isSingleMatrix = [](ArrayRef batch) { + return llvm::all_of(batch, [](int64_t dim) { return dim == 1; }); + }; + if (isSingleMatrix(aBatch) || isSingleMatrix(bBatch)) + return success(); + // With at most one batch axis, every valid broadcast has operand matrix + // counts in {1, output_count}, which one constant stride can represent. + if (std::max(aBatch.size(), bBatch.size()) <= 1) + return success(); + + auto isDynamic = [](int64_t dim) { return ShapedType::isDynamic(dim); }; + if (llvm::any_of(aBatch, isDynamic) || llvm::any_of(bBatch, isDynamic)) { + emitError() << "matmul with two nontrivial dynamic batch shapes is not " + "supported by the strided-batch runtime"; + return failure(); + } + + SmallVector outputBatch; + if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, outputBatch)) + return failure(); + auto product = [](ArrayRef shape) { + return std::accumulate(shape.begin(), shape.end(), int64_t{1}, + std::multiplies()); + }; + int64_t outputCount = product(outputBatch); + if (product(aBatch) == outputCount && product(bBatch) == outputCount) + return success(); + + emitError() << "matmul partial per-axis batch broadcast is not supported by " + "the strided-batch runtime"; + return failure(); +} + LogicalResult mlir::hip::verifyHipOpShape( Operation *op, function_ref>()> computeExpected) { @@ -163,85 +243,150 @@ mlir::hip::reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source) { return dims; } -SmallVector mlir::hip::reifyBroadcastShape(OpBuilder &b, - Location loc, - ValueRange operands) { - if (operands.empty()) - return {}; +FailureOr> +mlir::hip::reifyBroadcastShape(OpBuilder &b, Location loc, + ArrayRef> inputShapes, + function_ref emitError) { + if (inputShapes.empty()) { + emitError() << "broadcast requires at least one input shape"; + return failure(); + } + + size_t resultRank = 0; + for (const SmallVector &shape : inputShapes) + resultRank = std::max(resultRank, shape.size()); + + SmallVector result(resultRank, b.getIndexAttr(1)); + for (const SmallVector &shape : inputShapes) { + size_t pad = resultRank - shape.size(); + for (size_t i : llvm::seq(0, resultRank)) { + OpFoldResult inputDim = + i < pad ? OpFoldResult(b.getIndexAttr(1)) : shape[i - pad]; + FailureOr merged = + broadcastDim(b, loc, result[i], inputDim, emitError); + if (failed(merged)) + return failure(); + result[i] = *merged; + } + } + return result; +} + +FailureOr> mlir::hip::reifyBroadcastResultShape( + OpBuilder &b, Location loc, ValueRange operands, + function_ref emitError) { + if (operands.empty()) { + emitError() << "broadcast requires at least one operand"; + return failure(); + } - // Collect operand shapes; bail if any is non-ranked (verifier should - // already have caught this on the op). - SmallVector> shapes; + SmallVector> shapes; shapes.reserve(operands.size()); - for (Value v : operands) { - auto t = dyn_cast(v.getType()); - if (!t) - return {}; - shapes.push_back(t.getShape()); + for (size_t i : llvm::seq(0, operands.size())) { + Value operand = operands[i]; + if (!isa(operand.getType())) { + emitError() << "broadcast operand must be a ranked tensor"; + return failure(); + } + bool reused = false; + for (size_t j : llvm::seq(0, i)) { + if (operand == operands[j]) { + shapes.push_back(shapes[j]); + reused = true; + break; + } + } + if (reused) + continue; + shapes.push_back(tensor::getMixedSizes(b, loc, operand)); } + return reifyBroadcastShape(b, loc, shapes, emitError); +} - // Compute the broadcast result shape via sequential pairwise reduction. - // `getBroadcastedShape` follows NumPy/ONNX semantics: - // - 1 broadcasts against any other dim - // - dynamic + static>1 -> static (the strictly-correct tightening, - // since the dynamic side must equal the static side at runtime) - // - dynamic + dynamic -> dynamic - // - equal static -> static; unequal non-1 static -> failure - SmallVector outShape(shapes[0].begin(), shapes[0].end()); - for (size_t k : llvm::seq(1, shapes.size())) { - SmallVector tmp; - if (!OpTrait::util::getBroadcastedShape(outShape, shapes[k], tmp)) - return {}; - outShape = std::move(tmp); +FailureOr> mlir::hip::reifyMatmulResultShape( + OpBuilder &b, Location loc, Value A, Value B, + function_ref emitError) { + auto aType = dyn_cast(A.getType()); + auto bType = dyn_cast(B.getType()); + if (!aType || !bType) { + emitError() << "matmul operands must be ranked tensors"; + return failure(); } + if (inferMatmulShape(aType.getShape(), bType.getShape(), emitError).empty()) + return failure(); + if (failed(verifyStridedBatchMatmul(aType.getShape(), bType.getShape(), + emitError))) + return failure(); - size_t outRank = outShape.size(); - // Right-alignment padding per operand (operand `k` doesn't reach output - // dims in `[0, pads[k])`; those positions are an implicit 1 contribution). - SmallVector pads(operands.size()); - for (size_t k : llvm::seq(0, operands.size())) - pads[k] = outRank - shapes[k].size(); + SmallVector aSizes = tensor::getMixedSizes(b, loc, A); + SmallVector bSizes = tensor::getMixedSizes(b, loc, B); + ArrayRef aBatch = ArrayRef(aSizes).drop_back(2); + ArrayRef bBatch = ArrayRef(bSizes).drop_back(2); + SmallVector> batchShapes = { + SmallVector(aBatch.begin(), aBatch.end()), + SmallVector(bBatch.begin(), bBatch.end())}; + FailureOr> result = + reifyBroadcastShape(b, loc, batchShapes, emitError); + if (failed(result)) + return failure(); + result->push_back(aSizes[aSizes.size() - 2]); + result->push_back(bSizes.back()); + return result; +} - SmallVector dims; - dims.reserve(outRank); - for (size_t i : llvm::seq(0, outRank)) { - // Pick the operand to reify this dim against: - // 1. earliest operand that is in-range AND has a non-1 dim - // (that operand actually determines the runtime extent; - // `tensor.dim %that, i` folds to the constant when that - // operand's dim is static) - // 2. else earliest operand that is in-range (all in-range - // operands have a 1 here, so reifying against any of them - // is correct; first wins for stability) - // 3. else operand 0 dim 0 — defensive fallback that should be - // unreachable when the broadcast result rank == max input rank. - Value bestSrc; - size_t bestSrcDim = 0; - bool foundCanonical = false; - for (size_t k : llvm::seq(0, operands.size())) { - if (i < pads[k]) +FailureOr> +mlir::hip::reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, + Value optionalC, int64_t transA, int64_t transB, + function_ref emitError) { + auto aType = dyn_cast(A.getType()); + auto bType = dyn_cast(B.getType()); + if (!aType || !bType || aType.getRank() != 2 || bType.getRank() != 2) { + emitError() << "gemm A and B must be rank-2 tensors"; + return failure(); + } + if ((transA != 0 && transA != 1) || (transB != 0 && transB != 1)) { + emitError() << "gemm transA and transB must be 0 or 1"; + return failure(); + } + + int64_t aKDim = transA ? 0 : 1; + int64_t bKDim = transB ? 1 : 0; + int64_t aK = aType.getDimSize(aKDim); + int64_t bK = bType.getDimSize(bKDim); + if (!ShapedType::isDynamic(aK) && !ShapedType::isDynamic(bK) && aK != bK) { + emitError() << "gemm contraction dim mismatch: A has " << aK + << " but B has " << bK; + return failure(); + } + + SmallVector aSizes = tensor::getMixedSizes(b, loc, A); + SmallVector bSizes = tensor::getMixedSizes(b, loc, B); + SmallVector result = {aSizes[transA ? 1 : 0], + bSizes[transB ? 0 : 1]}; + + if (optionalC) { + auto cType = dyn_cast(optionalC.getType()); + if (!cType || cType.getRank() > 2) { + emitError() << "gemm C must be a ranked tensor of rank at most 2"; + return failure(); + } + SmallVector cSizes = tensor::getMixedSizes(b, loc, optionalC); + size_t pad = result.size() - cSizes.size(); + for (size_t i : llvm::seq(0, result.size())) { + if (i < pad) continue; - size_t kDim = i - pads[k]; - if (!bestSrc) { - bestSrc = operands[k]; - bestSrcDim = kDim; - } - if (shapes[k][kDim] != 1) { - bestSrc = operands[k]; - bestSrcDim = kDim; - foundCanonical = true; - break; + std::optional cStatic = getConstantIntValue(cSizes[i - pad]); + std::optional resultStatic = getConstantIntValue(result[i]); + if (cStatic && resultStatic && *cStatic != 1 && + *cStatic != *resultStatic) { + emitError() << "gemm C dimension " << *cStatic + << " is not broadcastable to output dimension " + << *resultStatic; + return failure(); } } - (void)foundCanonical; - if (!bestSrc) { - bestSrc = operands[0]; - bestSrcDim = 0; - } - dims.push_back( - reifyDimOrConstant(b, loc, outShape[i], bestSrc, bestSrcDim)); } - return dims; + return result; } SmallVector @@ -434,10 +579,11 @@ mlir::hip::reifyBroadcastShapeFor(OpBuilder &b, Location loc, for (Value v : operands) if (!isa(v.getType())) return failure(); - SmallVector dims = reifyBroadcastShape(b, loc, operands); - if (dims.empty()) + FailureOr> dims = reifyBroadcastResultShape( + b, loc, operands, [&]() { return op->emitOpError(); }); + if (failed(dims)) return failure(); - reified.assign({std::move(dims)}); + reified.assign({std::move(*dims)}); return success(); } diff --git a/lib/Runtime/hipdnn_ep_runtime.h b/lib/Runtime/hipdnn_ep_runtime.h index 0bf4159f0..8a6a65674 100644 --- a/lib/Runtime/hipdnn_ep_runtime.h +++ b/lib/Runtime/hipdnn_ep_runtime.h @@ -804,23 +804,20 @@ int wrap_hipblasLtGemm(void *handle, // hipBLASLt handle // MatMul operation wrapper (batched matrix multiplication) // Called by generated IR for onnx.MatMul lowering // Computes output = A @ B for each batch -// A: [batch_count x M x K], B: [K x N] (broadcast) or [batch_count x K x N] +// A/B may each contain one broadcast matrix or one matrix per output batch. // output: [batch_count x M x N] // -// `b_batch_stride` is hipBLASLt's STRIDED_BATCH_OFFSET on layA when -// `batch_count > 1`: the per-batch advance in elements through B. It MUST be: -// * 0 when B is a broadcast weight — one matrix reused across all -// batches. Includes both rank-2 `[K, N]` and rank-N -// `[1, ..., 1, K, N]` (any leading-dim product == 1). -// * K*N when B is per-batch — leading-dim product > 1, so the buffer -// actually holds multiple `[K, N]` matrices laid out contiguously. -// Mis-setting this to K*N for a broadcast B causes hipBLASLt to step K*N -// elements past the end of the weight buffer on every batch beyond the -// first, reading uninitialised memory into the GEMM and producing wrong -// (often NaN) outputs for batch > 0. For batch_count == 1 the value is -// ignored. Always pass an exact stride; the compiler computes 0 vs K*N -// at compile time when B's leading dims are static, else at runtime. -int wrap_hipblasLtMatmul( +// `a_batch_stride` / `b_batch_stride` are hipBLASLt's per-batch advances in +// elements. A stride is 0 when one matrix is broadcast across all batches; +// otherwise it is M*K for A or K*N for B. +// Legacy wrapper retained for cached model artifacts compiled before the A +// batch-stride parameter was added. It assumes A is per-batch. +int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, + const void *B, void *output, int64_t M, int64_t N, + int64_t K, int64_t batch_count, int64_t elem_size, + int64_t b_batch_stride); + +int wrap_hipblasLtMatmul_v2( RuntimeState *state, int op_state_slot, // per-instance op-state slot (shared algo table) const void *A, // Matrix A GPU pointer @@ -831,6 +828,7 @@ int wrap_hipblasLtMatmul( int64_t K, // Columns of A / Rows of B int64_t batch_count, // Number of batches int64_t elem_size, // Element size in bytes (2=f16, 4=f32) + int64_t a_batch_stride, // 0 = broadcast; M*K = per-batch int64_t b_batch_stride); // 0 = broadcast (any rank); K*N = per-batch // GroupQueryAttention operation wrapper (Full MS spec) diff --git a/lib/Runtime/mock/mock_gpu.cpp b/lib/Runtime/mock/mock_gpu.cpp index 155c54bad..11bd83e3e 100644 --- a/lib/Runtime/mock/mock_gpu.cpp +++ b/lib/Runtime/mock/mock_gpu.cpp @@ -596,6 +596,16 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, int64_t b_batch_stride) { + return wrap_hipblasLtMatmul_v2(state, op_state_slot, A, B, output, M, N, K, + batch_count, elem_size, M * K, b_batch_stride); +} + +int wrap_hipblasLtMatmul_v2(RuntimeState *state, int op_state_slot, + const void *A, const void *B, void *output, + int64_t M, int64_t N, int64_t K, + int64_t batch_count, int64_t elem_size, + int64_t a_batch_stride, int64_t b_batch_stride) { + (void)a_batch_stride; (void)b_batch_stride; (void)op_state_slot; if (!state) { diff --git a/lib/Runtime/real/matmul.cpp b/lib/Runtime/real/matmul.cpp index c60a45c8c..85d30bbe1 100644 --- a/lib/Runtime/real/matmul.cpp +++ b/lib/Runtime/real/matmul.cpp @@ -45,22 +45,13 @@ static bool autotune_enabled() { struct MatmulCacheKey { int64_t M, N, K, batch_count, elem_size; - // hipBLASLt's STRIDED_BATCH_OFFSET on layA, in elements. Two distinct - // values reach this site at the same (M,N,K,batch,elem_size): - // * 0 — B is a broadcast weight (rank-2 [K,N], or rank-N - // [1,...,1,K,N] whose leading-dim product is 1). - // * K*N — B is per-batch (leading-dim product > 1; the buffer holds - // multiple [K,N] matrices laid out contiguously). - // Part of the cache key because the layout descriptor is parameterised - // by the stride: mixing the two would silently route one path through - // the other's stride and read past the end of a broadcast weight buffer. - // Keyed on the actual int stride (not a 0/1 bool) so any future site - // that legitimately uses a stride other than {0, K*N} also gets its - // own cache entry rather than aliasing one of these two. - int64_t b_batch_stride; + // hipBLASLt STRIDED_BATCH_OFFSET values for the user's A and B buffers. + // They are part of the key because both matrix layouts depend on them. + int64_t a_batch_stride, b_batch_stride; bool operator==(const MatmulCacheKey &o) const { return M == o.M && N == o.N && K == o.K && batch_count == o.batch_count && - elem_size == o.elem_size && b_batch_stride == o.b_batch_stride; + elem_size == o.elem_size && a_batch_stride == o.a_batch_stride && + b_batch_stride == o.b_batch_stride; } }; @@ -72,15 +63,17 @@ struct MatmulCacheKeyHash { hash_combine_val(h, k.K); hash_combine_val(h, k.batch_count); hash_combine_val(h, k.elem_size); + hash_combine_val(h, k.a_batch_stride); hash_combine_val(h, k.b_batch_stride); return h; } }; /// Cached hipBLASLt descriptors + multi-algorithm auto-tune state for a -/// single (M, N, K, batch, elem_size) shape. Descriptors are created in -/// queryOrCreateMatmul() and owned by the MatmulAlgoTable, which frees them -/// when the last session sharing it is destroyed. +/// single (M, N, K, batch, elem_size, A stride, B stride) configuration. +/// Descriptors are created in queryOrCreateMatmul() and owned by the +/// MatmulAlgoTable, which frees them when the last session sharing it is +/// destroyed. struct MatmulCacheEntry { hipblasLtMatmulDesc_t desc; hipblasLtMatrixLayout_t layA, layB, layC; @@ -198,19 +191,10 @@ static MatmulCacheEntry *queryOrCreateMatmul(MatmulAlgoTable &table, if (key.batch_count > 1) { int64_t bc = key.batch_count; - // layA → user's B. The stride is whatever the compiler computed — - // 0 for broadcast B (rank-2 [K,N] or rank-N [1,...,1,K,N]), K*N for - // per-batch B (rank-N with leading-dim product > 1). Setting sA = K*N - // for a broadcast weight reads K*N elements PAST the end of the - // buffer on batch 1+ and feeds garbage into the GEMM — typical symptom - // on vision models is image-0 correct, image-1+ NaN (the OOB read - // often lands in a fp16-NaN pattern from adjacent pool slots / - // constants). Mis-setting sA = 0 for a per-batch B does the opposite: - // every batch reads matrix 0 instead of its own. - // layB → user's A and layC → output are always per-batch (the BATCH - // partition comes from A's leading dim by construction). + // Row-major -> column-major swaps operands: layA describes user's B and + // layB describes user's A. int64_t sA = key.b_batch_stride; - int64_t sB = M * K, sC = M * N; + int64_t sB = key.a_batch_stride, sC = M * N; MATMUL_CACHE_CHECK(hipblasLtMatrixLayoutSetAttribute( entry.layA, HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &bc, sizeof(bc))); MATMUL_CACHE_CHECK(hipblasLtMatrixLayoutSetAttribute( @@ -436,6 +420,15 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, int64_t b_batch_stride) { + return wrap_hipblasLtMatmul_v2(state, op_state_slot, A, B, output, M, N, K, + batch_count, elem_size, M * K, b_batch_stride); +} + +int wrap_hipblasLtMatmul_v2(RuntimeState *state, int op_state_slot, + const void *A, const void *B, void *output, + int64_t M, int64_t N, int64_t K, + int64_t batch_count, int64_t elem_size, + int64_t a_batch_stride, int64_t b_batch_stride) { OP_PROFILE( "matmul", [&] { @@ -468,11 +461,11 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const char *type_name = (elem_size == 2) ? "f16" : "f32"; RUNTIME_DEBUG_LOG("[REAL] wrap_hipblasLtMatmul: M=%lld, N=%lld, K=%lld, " - "batch=%lld, b_batch_stride=%lld, elem_size=%lld (%s), " - "total_bytes=%lld\n", + "batch=%lld, a_batch_stride=%lld, b_batch_stride=%lld, " + "elem_size=%lld (%s), total_bytes=%lld\n", (long long)M, (long long)N, (long long)K, - (long long)batch_count, (long long)b_batch_stride, - (long long)elem_size, type_name, + (long long)batch_count, (long long)a_batch_stride, + (long long)b_batch_stride, (long long)elem_size, type_name, (long long)(batch_count * M * N * elem_size)); MatmulState *ms = MatmulState::get_op_state(state, op_state_slot); @@ -482,7 +475,8 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, return -1; } - MatmulCacheKey key{M, N, K, batch_count, elem_size, b_batch_stride}; + MatmulCacheKey key{ + M, N, K, batch_count, elem_size, a_batch_stride, b_batch_stride}; MatmulCacheEntry *cached = queryOrCreateMatmul(*ms->table, handle, key); if (!cached) { fprintf(stderr, diff --git a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir index d3494fb0a..830f29e66 100644 --- a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir +++ b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir @@ -24,10 +24,10 @@ module { // CHECK-LABEL: llvm.func @test_matmul_rank2_b // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64) -> i32 -// Verify 11 parameters: +// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// Verify 12 parameters: // - 4 pointers: state, A, B, output -// - 6 i64: M=128, N=1024, K=4096, batch_count=1, elem_size=2, b_batch_stride=0 +// - 7 i64: M, N, K, batch_count, elem_size, A stride, B stride // (B is rank-2 [K, N] = broadcast weight → stride = 0, compile-time const) // - 1 i32: op_state_slot (-1 here — --assign-op-state-slots not run in this RUN) @@ -54,4 +54,25 @@ module { // CHECK-LABEL: llvm.func @test_matmul_rank3_leading_one_b // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 + +// ----- Rank-2 broadcast A against rank-3 B ------------------------------- +// A contains one matrix and therefore uses A stride 0. B and the output carry +// two batches. + +module { + func.func @test_matmul_rank2_a(%ctx: !hip.context, + %A: memref<128x4096xf16, 1>, + %B: memref<2x4096x1024xf16, 1>, + %output: memref<2x128x1024xf16, 1>) { + hip.matmul(%ctx) + ins(%A, %B : memref<128x4096xf16, 1>, memref<2x4096x1024xf16, 1>) + outs(%output : memref<2x128x1024xf16, 1>) + return + } +} + +// CHECK-LABEL: llvm.func @test_matmul_rank2_a +// CHECK-NOT: llvm.icmp +// CHECK-NOT: llvm.select +// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 diff --git a/test/lit/Conversion/onnx-to-hip/test_gemm.mlir b/test/lit/Conversion/onnx-to-hip/test_gemm.mlir index 481a97bf2..236a5409f 100644 --- a/test/lit/Conversion/onnx-to-hip/test_gemm.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_gemm.mlir @@ -68,6 +68,56 @@ module { return %output : tensor<1x5120xf16> } + // --- Cases 4-7: dynamic M/N source mapping for every transpose pair --- + func.func @test_gemm_dynamic_00(%a: tensor, %b: tensor<4x?xf16>) -> tensor { + %none = "onnx.NoValue"() {value} : () -> none + %output = "onnx.Gemm"(%a, %b, %none) : (tensor, tensor<4x?xf16>, none) -> tensor + return %output : tensor + } + // CHECK-LABEL: func.func @test_gemm_dynamic_00 + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor<4x?xf16>) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] + // CHECK: tensor.empty(%[[M]], %[[N]]) : tensor + + func.func @test_gemm_dynamic_10(%a: tensor<4x?xf16>, %b: tensor<4x?xf16>) -> tensor { + %none = "onnx.NoValue"() {value} : () -> none + %output = "onnx.Gemm"(%a, %b, %none) {transA = 1 : si64} : (tensor<4x?xf16>, tensor<4x?xf16>, none) -> tensor + return %output : tensor + } + // CHECK-LABEL: func.func @test_gemm_dynamic_10 + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor<4x?xf16>, %[[B:[A-Za-z0-9_]+]]: tensor<4x?xf16>) + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C1]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] + // CHECK: tensor.empty(%[[M]], %[[N]]) : tensor + + func.func @test_gemm_dynamic_01(%a: tensor, %b: tensor) -> tensor { + %none = "onnx.NoValue"() {value} : () -> none + %output = "onnx.Gemm"(%a, %b, %none) {transB = 1 : si64} : (tensor, tensor, none) -> tensor + return %output : tensor + } + // CHECK-LABEL: func.func @test_gemm_dynamic_01 + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C0]] + // CHECK: tensor.empty(%[[M]], %[[N]]) : tensor + + func.func @test_gemm_dynamic_11(%a: tensor<4x?xf16>, %b: tensor, %c: tensor) -> tensor { + %output = "onnx.Gemm"(%a, %b, %c) {transA = 1 : si64, transB = 1 : si64} : (tensor<4x?xf16>, tensor, tensor) -> tensor + return %output : tensor + } + // CHECK-LABEL: func.func @test_gemm_dynamic_11 + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor<4x?xf16>, %[[B:[A-Za-z0-9_]+]]: tensor + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C1]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C0]] + // CHECK: tensor.empty(%[[M]], %[[N]]) : tensor + // Dummy entry point required by generateModuleMetadata. func.func @main_graph(%arg0: tensor<1x5120xf16>, %arg1: tensor<5120x5120xf16>, %arg2: tensor<5120xf16>) -> tensor<1x5120xf16> { return %arg0 : tensor<1x5120xf16> diff --git a/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_c.mlir b/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_c.mlir new file mode 100644 index 000000000..173e3ed9e --- /dev/null +++ b/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_c.mlir @@ -0,0 +1,15 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// RUN: hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hip %s 2>&1 | FileCheck %s + +module { + func.func @main_graph(%a: tensor<2x4xf16>, %b: tensor<4x5xf16>, + %c: tensor<3xf16>) -> tensor<2x5xf16> { + // CHECK: error: gemm C dimension 3 is not broadcastable to output dimension 5 + %result = "onnx.Gemm"(%a, %b, %c) + : (tensor<2x4xf16>, tensor<4x5xf16>, tensor<3xf16>) + -> tensor<2x5xf16> + return %result : tensor<2x5xf16> + } +} diff --git a/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_transpose.mlir b/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_transpose.mlir new file mode 100644 index 000000000..cfa08156d --- /dev/null +++ b/test/lit/Conversion/onnx-to-hip/test_gemm_invalid_transpose.mlir @@ -0,0 +1,15 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// RUN: hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hip %s 2>&1 | FileCheck %s + +module { + func.func @main_graph(%a: tensor<2x4xf16>, + %b: tensor<4x5xf16>) -> tensor<2x5xf16> { + // CHECK: error: gemm transA and transB must be 0 or 1 + %none = "onnx.NoValue"() {value} : () -> none + %result = "onnx.Gemm"(%a, %b, %none) {transA = 2 : si64} + : (tensor<2x4xf16>, tensor<4x5xf16>, none) -> tensor<2x5xf16> + return %result : tensor<2x5xf16> + } +} diff --git a/test/lit/Conversion/onnx-to-hip/test_matmul.mlir b/test/lit/Conversion/onnx-to-hip/test_matmul.mlir index 1b8e00e23..6ae0d5ab6 100644 --- a/test/lit/Conversion/onnx-to-hip/test_matmul.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_matmul.mlir @@ -65,4 +65,38 @@ module { // CHECK: hip.matmul(%[[CTX]]) ins(%[[A]], %[[B]] : tensor, tensor) outs(%[[INIT]] : tensor) // CHECK-NOT: hip.alloc // CHECK-NOT: hip.copy + + // ===== Test 4: 2D x 3D -- batch comes from B, M comes from A ===== + func.func @matmul_2d_3d(%A: tensor, %B: tensor) -> tensor { + %result = "onnx.MatMul"(%A, %B) : (tensor, tensor) -> tensor + return %result : tensor + } + + // CHECK-LABEL: func.func @matmul_2d_3d + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index + // CHECK-DAG: %[[BATCH:.*]] = tensor.dim %[[B]], %[[C0]] + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C2]] + // CHECK: %[[INIT:.*]] = tensor.empty(%[[BATCH]], %[[M]], %[[N]]) : tensor + // CHECK: hip.matmul + + // ===== Test 5: 2D x 4D -- both batch axes come from B ===== + func.func @matmul_2d_4d(%A: tensor, %B: tensor) -> tensor { + %result = "onnx.MatMul"(%A, %B) : (tensor, tensor) -> tensor + return %result : tensor + } + + // CHECK-LABEL: func.func @matmul_2d_4d + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index + // CHECK-DAG: %[[B0:.*]] = tensor.dim %[[B]], %[[C0]] + // CHECK-DAG: %[[B1:.*]] = tensor.dim %[[B]], %[[C1]] + // CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] + // CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C3]] + // CHECK: %[[INIT:.*]] = tensor.empty(%[[B0]], %[[B1]], %[[M]], %[[N]]) : tensor + // CHECK: hip.matmul } diff --git a/test/lit/Conversion/onnx-to-hip/test_max.mlir b/test/lit/Conversion/onnx-to-hip/test_max.mlir index 46c4c37ca..391e1d75b 100644 --- a/test/lit/Conversion/onnx-to-hip/test_max.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_max.mlir @@ -61,8 +61,41 @@ module { // CHECK-SAME: (%[[CTX:.*]]: !hip.context, %[[A:.*]]: tensor, %[[B:.*]]: tensor) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index - // CHECK: %[[DIM0:.*]] = tensor.dim %[[A]], %[[C0]] : tensor - // CHECK: %[[DIM1:.*]] = tensor.dim %[[A]], %[[C1]] : tensor + // CHECK-DAG: %[[A0:.*]] = tensor.dim %[[A]], %[[C0]] : tensor + // CHECK-DAG: %[[A1:.*]] = tensor.dim %[[A]], %[[C1]] : tensor + // CHECK-DAG: %[[B0:.*]] = tensor.dim %[[B]], %[[C0]] : tensor + // CHECK-DAG: %[[B1:.*]] = tensor.dim %[[B]], %[[C1]] : tensor + // CHECK: %[[IS1_0:.*]] = arith.cmpi eq, %[[A0]], %[[C1]] : index + // CHECK: %[[DIM0:.*]] = arith.select %[[IS1_0]], %[[B0]], %[[A0]] : index + // CHECK: %[[IS1_1:.*]] = arith.cmpi eq, %[[A1]], %[[C1]] : index + // CHECK: %[[DIM1:.*]] = arith.select %[[IS1_1]], %[[B1]], %[[A1]] : index // CHECK: %[[INIT:.*]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor // CHECK: hip.max(%[[CTX]]) ins(%[[A]], %[[B]] : tensor, tensor) outs(%[[INIT]] : tensor) + + // --- Case 5: different operands contribute different dynamic axes --- + func.func @max_asymmetric_broadcast(%a: tensor, %b: tensor<1x?xf32>) -> tensor { + %result = "onnx.Max"(%a, %b) : (tensor, tensor<1x?xf32>) -> tensor + return %result : tensor + } + + // CHECK-LABEL: func.func @max_asymmetric_broadcast + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor<1x?xf32>) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] : tensor + // CHECK: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] : tensor<1x?xf32> + // CHECK: %[[INIT:.*]] = tensor.empty(%[[M]], %[[N]]) : tensor + // CHECK: hip.max + + // --- Case 6: variadic pairwise lowering grows intermediate rank --- + func.func @max_variadic_rank_growth(%a: tensor<4xf32>, %b: tensor<3x4xf32>, %c: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { + %result = "onnx.Max"(%a, %b, %c) : (tensor<4xf32>, tensor<3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32> + return %result : tensor<2x3x4xf32> + } + + // CHECK-LABEL: func.func @max_variadic_rank_growth + // CHECK: %[[E0:.*]] = tensor.empty() : tensor<3x4xf32> + // CHECK: %[[M0:.*]] = hip.max{{.*}}outs(%[[E0]] : tensor<3x4xf32>) + // CHECK: %[[E1:.*]] = tensor.empty() : tensor<2x3x4xf32> + // CHECK: hip.max{{.*}}ins(%[[M0]], {{.*}} : tensor<3x4xf32>, tensor<2x3x4xf32>) outs(%[[E1]] : tensor<2x3x4xf32>) } diff --git a/test/lit/Conversion/onnx-to-hip/test_min.mlir b/test/lit/Conversion/onnx-to-hip/test_min.mlir index 19da173a4..102fe8499 100644 --- a/test/lit/Conversion/onnx-to-hip/test_min.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_min.mlir @@ -65,8 +65,41 @@ module { // CHECK-SAME: (%[[CTX:.*]]: !hip.context, %[[A:.*]]: tensor, %[[B:.*]]: tensor) // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index - // CHECK: %[[DIM0:.*]] = tensor.dim %[[A]], %[[C0]] : tensor - // CHECK: %[[DIM1:.*]] = tensor.dim %[[A]], %[[C1]] : tensor + // CHECK-DAG: %[[A0:.*]] = tensor.dim %[[A]], %[[C0]] : tensor + // CHECK-DAG: %[[A1:.*]] = tensor.dim %[[A]], %[[C1]] : tensor + // CHECK-DAG: %[[B0:.*]] = tensor.dim %[[B]], %[[C0]] : tensor + // CHECK-DAG: %[[B1:.*]] = tensor.dim %[[B]], %[[C1]] : tensor + // CHECK: %[[IS1_0:.*]] = arith.cmpi eq, %[[A0]], %[[C1]] : index + // CHECK: %[[DIM0:.*]] = arith.select %[[IS1_0]], %[[B0]], %[[A0]] : index + // CHECK: %[[IS1_1:.*]] = arith.cmpi eq, %[[A1]], %[[C1]] : index + // CHECK: %[[DIM1:.*]] = arith.select %[[IS1_1]], %[[B1]], %[[A1]] : index // CHECK: %[[INIT:.*]] = tensor.empty(%[[DIM0]], %[[DIM1]]) : tensor // CHECK: hip.min(%[[CTX]]) ins(%[[A]], %[[B]] : tensor, tensor) outs(%[[INIT]] : tensor) + + // --- Case 5: different operands contribute different dynamic axes --- + func.func @min_asymmetric_broadcast(%a: tensor, %b: tensor<1x?xf32>) -> tensor { + %result = "onnx.Min"(%a, %b) : (tensor, tensor<1x?xf32>) -> tensor + return %result : tensor + } + + // CHECK-LABEL: func.func @min_asymmetric_broadcast + // CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor<1x?xf32>) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + // CHECK: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] : tensor + // CHECK: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] : tensor<1x?xf32> + // CHECK: %[[INIT:.*]] = tensor.empty(%[[M]], %[[N]]) : tensor + // CHECK: hip.min + + // --- Case 6: variadic pairwise lowering grows intermediate rank --- + func.func @min_variadic_rank_growth(%a: tensor<4xf32>, %b: tensor<3x4xf32>, %c: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { + %result = "onnx.Min"(%a, %b, %c) : (tensor<4xf32>, tensor<3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32> + return %result : tensor<2x3x4xf32> + } + + // CHECK-LABEL: func.func @min_variadic_rank_growth + // CHECK: %[[E0:.*]] = tensor.empty() : tensor<3x4xf32> + // CHECK: %[[M0:.*]] = hip.min{{.*}}outs(%[[E0]] : tensor<3x4xf32>) + // CHECK: %[[E1:.*]] = tensor.empty() : tensor<2x3x4xf32> + // CHECK: hip.min{{.*}}ins(%[[M0]], {{.*}} : tensor<3x4xf32>, tensor<2x3x4xf32>) outs(%[[E1]] : tensor<2x3x4xf32>) } diff --git a/test/lit/Dialect/hip-broadcast-reify-shapes.mlir b/test/lit/Dialect/hip-broadcast-reify-shapes.mlir new file mode 100644 index 000000000..31d7fc886 --- /dev/null +++ b/test/lit/Dialect/hip-broadcast-reify-shapes.mlir @@ -0,0 +1,74 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// RUN: hip-mlir-opt --resolve-shaped-type-result-dims %s | FileCheck %s --check-prefix=REIFY +// RUN: hip-mlir-opt --hip-infer-shapes %s | FileCheck %s --check-prefix=INFER + +// Dynamic/dynamic broadcast must select the non-unit runtime extent. +// REIFY-LABEL: func.func @dynamic_dynamic +// REIFY-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor +// REIFY-DAG: %[[C0:.*]] = arith.constant 0 : index +// REIFY-DAG: %[[C1:.*]] = arith.constant 1 : index +// REIFY-DAG: %[[AD:.*]] = tensor.dim %[[A]], %[[C0]] +// REIFY-DAG: %[[BD:.*]] = tensor.dim %[[B]], %[[C0]] +// REIFY: %[[IS1:.*]] = arith.cmpi eq, %[[AD]], %[[C1]] : index +// REIFY: %[[D:.*]] = arith.select %[[IS1]], %[[BD]], %[[AD]] : index +// REIFY: return %[[D]] +func.func @dynamic_dynamic(%ctx: !hip.context, %a: tensor, + %b: tensor, %out: tensor) -> index { + %r = hip.add(%ctx) + ins(%a, %b : tensor, tensor) + outs(%out : tensor) -> tensor + %c0 = arith.constant 0 : index + %d = tensor.dim %r, %c0 : tensor + return %d : index +} + +// Each asymmetric axis comes directly from the operand that spans it. +// REIFY-LABEL: func.func @asymmetric_axes +// REIFY-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor<1x?xf32> +// REIFY-DAG: %[[C0:.*]] = arith.constant 0 : index +// REIFY-DAG: %[[C1:.*]] = arith.constant 1 : index +// REIFY: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] +// REIFY: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] +// REIFY: return %[[M]], %[[N]] +func.func @asymmetric_axes(%ctx: !hip.context, %a: tensor, + %b: tensor<1x?xf32>, + %out: tensor) -> (index, index) { + %r = hip.max(%ctx) + ins(%a, %b : tensor, tensor<1x?xf32>) + outs(%out : tensor) : tensor + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %d0 = tensor.dim %r, %c0 : tensor + %d1 = tensor.dim %r, %c1 : tensor + return %d0, %d1 : index, index +} + +// Broadcasting zero with one yields zero; max(0, 1) would be incorrect. +// REIFY-LABEL: func.func @zero_one +// REIFY-DAG: %[[C0:.*]] = arith.constant 0 : index +// REIFY: return %[[C0]] +func.func @zero_one(%ctx: !hip.context, %a: tensor<0xf32>, + %b: tensor<1xf32>, %out: tensor) -> index { + %r = hip.min(%ctx) + ins(%a, %b : tensor<0xf32>, tensor<1xf32>) + outs(%out : tensor) : tensor + %c0 = arith.constant 0 : index + %d = tensor.dim %r, %c0 : tensor + return %d : index +} + +// Rank-zero broadcast is a successful empty reified shape. hip-infer-shapes +// must not fail merely because the result has no dimensions. +// INFER-LABEL: func.func @rank_zero +// INFER: %[[R:.*]] = hip.where +// INFER: return %[[R]] : tensor +func.func @rank_zero(%ctx: !hip.context, %cond: tensor, + %x: tensor, %y: tensor, + %out: tensor) -> tensor { + %r = hip.where(%ctx) + ins(%cond, %x, %y : tensor, tensor, tensor) + outs(%out : tensor) : tensor + return %r : tensor +} diff --git a/test/lit/Dialect/hip-gemm-reify-shapes.mlir b/test/lit/Dialect/hip-gemm-reify-shapes.mlir new file mode 100644 index 000000000..fab270f95 --- /dev/null +++ b/test/lit/Dialect/hip-gemm-reify-shapes.mlir @@ -0,0 +1,47 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// RUN: hip-mlir-opt --resolve-shaped-type-result-dims %s | FileCheck %s + +// CHECK-LABEL: func.func @gemm_dynamic_default +// CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor<4x?xf16> +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] +// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C1]] +// CHECK: return %[[M]], %[[N]] +func.func @gemm_dynamic_default( + %ctx: !hip.context, %a: tensor, %b: tensor<4x?xf16>, + %out: tensor) -> (index, index) { + %r = hip.gemm(%ctx) + ins(%a, %b : tensor, tensor<4x?xf16>) + outs(%out : tensor) + {transA = 0 : i64, transB = 0 : i64} : tensor + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %d0 = tensor.dim %r, %c0 : tensor + %d1 = tensor.dim %r, %c1 : tensor + return %d0, %d1 : index, index +} + +// Both transposed: M=A.dim1, N=B.dim0. Optional C does not supply extents. +// CHECK-LABEL: func.func @gemm_dynamic_transposed +// CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor<4x?xf16>, %[[B:[A-Za-z0-9_]+]]: tensor +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index +// CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C1]] +// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C0]] +// CHECK: return %[[M]], %[[N]] +func.func @gemm_dynamic_transposed( + %ctx: !hip.context, %a: tensor<4x?xf16>, %b: tensor, + %c: tensor, %out: tensor) -> (index, index) { + %r = hip.gemm(%ctx) + ins(%a, %b, %c : tensor<4x?xf16>, tensor, tensor) + outs(%out : tensor) + {transA = 1 : i64, transB = 1 : i64} : tensor + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %d0 = tensor.dim %r, %c0 : tensor + %d1 = tensor.dim %r, %c1 : tensor + return %d0, %d1 : index, index +} diff --git a/test/lit/Dialect/hip-matmul-reify-shapes.mlir b/test/lit/Dialect/hip-matmul-reify-shapes.mlir index d8931030f..1d7df8432 100644 --- a/test/lit/Dialect/hip-matmul-reify-shapes.mlir +++ b/test/lit/Dialect/hip-matmul-reify-shapes.mlir @@ -148,3 +148,30 @@ func.func @reify_dyn_static_batch_broadcast(%ctx: !hip.context, %d2 = tensor.dim %r, %d2_idx : tensor return %d0, %d1, %d2 : index, index, index } + +// ----- + +// Unequal ranks: B supplies the leading batch dim, A supplies M, B supplies N. +// CHECK-LABEL: func.func @reify_2d_3d +// CHECK-SAME: (%{{.*}}: !hip.context, %[[A:[A-Za-z0-9_]+]]: tensor, %[[B:[A-Za-z0-9_]+]]: tensor +// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index +// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index +// CHECK-DAG: %[[BATCH:.*]] = tensor.dim %[[B]], %[[C0]] +// CHECK-DAG: %[[M:.*]] = tensor.dim %[[A]], %[[C0]] +// CHECK-DAG: %[[N:.*]] = tensor.dim %[[B]], %[[C2]] +// CHECK: return %[[BATCH]], %[[M]], %[[N]] +func.func @reify_2d_3d(%ctx: !hip.context, + %a: tensor, + %b: tensor, + %c: tensor) -> (index, index, index) { + %r = hip.matmul(%ctx) + ins(%a, %b : tensor, tensor) + outs(%c : tensor) : tensor + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %d0 = tensor.dim %r, %c0 : tensor + %d1 = tensor.dim %r, %c1 : tensor + %d2 = tensor.dim %r, %c2 : tensor + return %d0, %d1, %d2 : index, index, index +} diff --git a/test/lit/Dialect/hip-matmul-shape-verifier.mlir b/test/lit/Dialect/hip-matmul-shape-verifier.mlir index d62717e2c..8147d28bd 100644 --- a/test/lit/Dialect/hip-matmul-shape-verifier.mlir +++ b/test/lit/Dialect/hip-matmul-shape-verifier.mlir @@ -136,3 +136,16 @@ func.func @matmul_tensor_mode_static(%ctx: !hip.context, outs(%c : tensor<2x8xf16>) : tensor<2x8xf16> return %r : tensor<2x8xf16> } + +// ----- + +func.func @matmul_partial_batch_broadcast(%ctx: !hip.context, + %a: memref<2x1x4x8xf16, 1>, + %b: memref<1x3x8x16xf16, 1>, + %c: memref<2x3x4x16xf16, 1>) { + // expected-error @+1 {{matmul partial per-axis batch broadcast is not supported by the strided-batch runtime}} + hip.matmul(%ctx) + ins(%a, %b : memref<2x1x4x8xf16, 1>, memref<1x3x8x16xf16, 1>) + outs(%c : memref<2x3x4x16xf16, 1>) + return +} diff --git a/test/lit/Pipeline/pipeline-pool-lower.mlir b/test/lit/Pipeline/pipeline-pool-lower.mlir index 98d2fb1c6..df21b6fb5 100644 --- a/test/lit/Pipeline/pipeline-pool-lower.mlir +++ b/test/lit/Pipeline/pipeline-pool-lower.mlir @@ -22,7 +22,7 @@ // CHECK: %[[DOM:.*]] = llvm.mlir.constant(0 : i32) : i32 // CHECK: llvm.call @hipdnn_ep_get_pool_base(%[[CTX]], %[[DOM]], %[[POOL_SIZE]]) : (!llvm.ptr, i32, i64) -> !llvm.ptr // CHECK: llvm.mlir.constant(256 : index) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul(%[[CTX]], +// CHECK: llvm.call @wrap_hipblasLtMatmul_v2(%[[CTX]], // CHECK: llvm.call @hip_miopen_softmax(%[[CTX]], // CHECK: llvm.return func.func @static_pool_to_llvm( @@ -47,7 +47,7 @@ func.func @static_pool_to_llvm( // CHECK: llvm.mul %arg15, %[[C32]] : i64 // CHECK: %[[DOM2:.*]] = llvm.mlir.constant(0 : i32) : i32 // CHECK: llvm.call @hipdnn_ep_get_pool_base(%[[CTX2]], %[[DOM2]], %{{[0-9]+}}) : (!llvm.ptr, i32, i64) -> !llvm.ptr -// CHECK: llvm.call @wrap_hipblasLtMatmul(%[[CTX2]], +// CHECK: llvm.call @wrap_hipblasLtMatmul_v2(%[[CTX2]], // CHECK: llvm.call @hip_miopen_softmax(%[[CTX2]], // CHECK: llvm.return func.func @dynamic_pool_to_llvm( diff --git a/test/numeric/tests/test_gemm.py b/test/numeric/tests/test_gemm.py new file mode 100644 index 000000000..529b712a8 --- /dev/null +++ b/test/numeric/tests/test_gemm.py @@ -0,0 +1,54 @@ +# +# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the MIT License. +# + +"""Numeric coverage for transpose-aware dynamic Gemm result dimensions.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper + +from framework.comparator import compare_outputs +from framework.onnx_utils import make_model_from_nodes + +M = 3 +K = 4 +N = 5 + + +def _make_gemm_model(trans_a: int, trans_b: int): + a_shape = [K, "M"] if trans_a else ["M", K] + b_shape = ["N", K] if trans_b else [K, "N"] + a = helper.make_tensor_value_info("A", TensorProto.FLOAT16, a_shape) + b = helper.make_tensor_value_info("B", TensorProto.FLOAT16, b_shape) + c = helper.make_tensor_value_info("C", TensorProto.FLOAT16, ["N"]) + output = helper.make_tensor_value_info("Y", TensorProto.FLOAT16, ["M", "N"]) + node = helper.make_node( + "Gemm", + ["A", "B", "C"], + ["Y"], + transA=trans_a, + transB=trans_b, + ) + return make_model_from_nodes([node], [a, b, c], [output]) + + +class TestGemm: + @pytest.mark.parametrize( + "trans_a,trans_b", + [(0, 0), (1, 0), (0, 1), (1, 1)], + ) + def test_dynamic_m_n_all_transposes(self, model_runner, trans_a, trans_b): + model = _make_gemm_model(trans_a, trans_b) + rng = np.random.default_rng(44) + a_shape = (K, M) if trans_a else (M, K) + b_shape = (N, K) if trans_b else (K, N) + inputs = [ + rng.uniform(-0.5, 0.5, a_shape).astype(np.float16), + rng.uniform(-0.5, 0.5, b_shape).astype(np.float16), + rng.uniform(-0.5, 0.5, (N,)).astype(np.float16), + ] + + actual, expected = model_runner.run_sample(model, inputs, reference="cpu") + compare_outputs(actual, expected, atol=1e-3) diff --git a/test/numeric/tests/test_matmul.py b/test/numeric/tests/test_matmul.py index bf1448817..29b6e590b 100644 --- a/test/numeric/tests/test_matmul.py +++ b/test/numeric/tests/test_matmul.py @@ -40,6 +40,14 @@ def _make_matmul_model(input_shape: list[int], weight_shape: list[int]): return make_model_from_nodes([node], [X], [Y], initializers=[w_init]) +def _make_unequal_rank_model(a_shape, b_shape, output_shape): + A = helper.make_tensor_value_info("A", TensorProto.FLOAT16, a_shape) + B = helper.make_tensor_value_info("B", TensorProto.FLOAT16, b_shape) + Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT16, output_shape) + node = helper.make_node("MatMul", ["A", "B"], ["Y"]) + return make_model_from_nodes([node], [A, B], [Y]) + + class TestMatMul: @pytest.mark.parametrize( "input_shape,weight_shape", @@ -89,3 +97,39 @@ def test_matmul_gate_up_proj_llama_shape(self, model_runner, seq_len): actual, expected = model_runner.run_sample(model, [x]) compare_outputs(actual, expected, atol=1e-3) + + @pytest.mark.parametrize( + "a_type_shape,b_type_shape,out_type_shape,a_shape,b_shape", + [ + ( + ["M", 4], + ["B", 4, "N"], + ["B", "M", "N"], + (3, 4), + (2, 4, 5), + ), + ( + ["M", 4], + ["B0", "B1", 4, "N"], + ["B0", "B1", "M", "N"], + (6, 4), + (2, 3, 4, 5), + ), + ], + ) + def test_matmul_unequal_rank_batch_from_b( + self, + model_runner, + a_type_shape, + b_type_shape, + out_type_shape, + a_shape, + b_shape, + ): + model = _make_unequal_rank_model(a_type_shape, b_type_shape, out_type_shape) + rng = np.random.default_rng(45) + a = rng.uniform(-1, 1, a_shape).astype(np.float16) + b = rng.uniform(-1, 1, b_shape).astype(np.float16) + + actual, expected = model_runner.run_sample(model, [a, b], reference="cpu") + compare_outputs(actual, expected, atol=1e-3) diff --git a/test/numeric/tests/test_max_min_broadcast.py b/test/numeric/tests/test_max_min_broadcast.py new file mode 100644 index 000000000..f7f9af674 --- /dev/null +++ b/test/numeric/tests/test_max_min_broadcast.py @@ -0,0 +1,55 @@ +# +# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the MIT License. +# + +"""Numeric coverage for variadic Max/Min and multidirectional broadcasting.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper + +from framework.comparator import compare_outputs +from framework.onnx_utils import make_model_from_nodes + + +def _make_binary_model(op_type: str): + lhs = helper.make_tensor_value_info("lhs", TensorProto.FLOAT, ["M", 1]) + rhs = helper.make_tensor_value_info("rhs", TensorProto.FLOAT, [1, "N"]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, ["M", "N"]) + node = helper.make_node(op_type, ["lhs", "rhs"], ["output"]) + return make_model_from_nodes([node], [lhs, rhs], [output]) + + +def _make_variadic_model(op_type: str): + a = helper.make_tensor_value_info("a", TensorProto.FLOAT, [4]) + b = helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 4]) + c = helper.make_tensor_value_info("c", TensorProto.FLOAT, [2, 3, 4]) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 3, 4]) + node = helper.make_node(op_type, ["a", "b", "c"], ["output"]) + return make_model_from_nodes([node], [a, b, c], [output]) + + +class TestMaxMinBroadcast: + @pytest.mark.parametrize("op_type", ["Max", "Min"]) + def test_asymmetric_dynamic_broadcast(self, model_runner, op_type): + model = _make_binary_model(op_type) + rng = np.random.default_rng(42) + lhs = rng.uniform(-1, 1, (3, 1)).astype(np.float32) + rhs = rng.uniform(-1, 1, (1, 5)).astype(np.float32) + + actual, expected = model_runner.run_sample(model, [lhs, rhs], reference="cpu") + compare_outputs(actual, expected, atol=1e-6) + + @pytest.mark.parametrize("op_type", ["Max", "Min"]) + def test_variadic_rank_growth(self, model_runner, op_type): + model = _make_variadic_model(op_type) + rng = np.random.default_rng(43) + inputs = [ + rng.uniform(-1, 1, (4,)).astype(np.float32), + rng.uniform(-1, 1, (3, 4)).astype(np.float32), + rng.uniform(-1, 1, (2, 3, 4)).astype(np.float32), + ] + + actual, expected = model_runner.run_sample(model, inputs, reference="cpu") + compare_outputs(actual, expected, atol=1e-6) diff --git a/test/numeric/tests/test_where.py b/test/numeric/tests/test_where.py index 7aa6a8c3b..9c385c4f2 100644 --- a/test/numeric/tests/test_where.py +++ b/test/numeric/tests/test_where.py @@ -105,6 +105,34 @@ def test_where_scalar_condition(self, model_runner, dtype): actual, expected = model_runner.run_sample(model, [cond, x, y]) compare_outputs(actual, expected, atol=1e-4) + def test_where_asymmetric_dynamic_broadcast(self, model_runner): + """Every dynamic output axis may be selected from a different input.""" + m, n = 3, 5 + model = _make_where_model( + [1, "N"], + ["M", 1], + ["M", "N"], + ["M", "N"], + np.float32, + ) + rng = np.random.default_rng(46) + cond = rng.integers(0, 2, (1, n), dtype=np.bool_) + x = rng.uniform(-1, 1, (m, 1)).astype(np.float32) + y = rng.uniform(-1, 1, (m, n)).astype(np.float32) + + actual, expected = model_runner.run_sample(model, [cond, x, y], reference="cpu") + compare_outputs(actual, expected, atol=1e-6) + + def test_where_rank_zero_output(self, model_runner): + """All-scalar Where proves rank-zero reification is a success.""" + model = _make_where_model([], [], [], [], np.float32) + cond = np.array(False, dtype=np.bool_) + x = np.array(1.25, dtype=np.float32) + y = np.array(-2.5, dtype=np.float32) + + actual, expected = model_runner.run_sample(model, [cond, x, y], reference="cpu") + compare_outputs(actual, expected, atol=0) + @pytest.mark.parametrize("seq_len", SEQ_LENS) def test_where_qwen9b_rope_mask_shape(self, model_runner, seq_len): """RoPE mask Where as it appears in Qwen3.5-9B text.onnx. From 6db8a03b641550b9bc4b184ea44f4ae80e1a00c4 Mon Sep 17 00:00:00 2001 From: fhanuman Date: Tue, 28 Jul 2026 06:22:49 -0600 Subject: [PATCH 2/9] fix(hip): diagnose unsupported MatMul batch layouts Emit explicit errors for invalid strided-batch shapes and document why repeated broadcast operands reuse their first mixed shape. Co-Authored-By: GPT-5.6 Sol Made-with: Cursor --- lib/Dialect/IR/HipShapeUtils.cpp | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index 8719dbbc8..ec5e6e08d 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -134,8 +134,10 @@ mlir::hip::inferMatmulShape(ArrayRef aShape, ArrayRef bShape, LogicalResult mlir::hip::verifyStridedBatchMatmul( ArrayRef aShape, ArrayRef bShape, function_ref emitError) { - if (aShape.size() < 2 || bShape.size() < 2) + if (aShape.size() < 2 || bShape.size() < 2) { + emitError() << "strided-batch matmul requires rank >= 2 operands"; return failure(); + } ArrayRef aBatch = aShape.drop_back(2); ArrayRef bBatch = bShape.drop_back(2); @@ -157,8 +159,11 @@ LogicalResult mlir::hip::verifyStridedBatchMatmul( } SmallVector outputBatch; - if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, outputBatch)) + if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, outputBatch)) { + emitError() << "matmul batch broadcast failure: A.batch=" + << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); return failure(); + } auto product = [](ArrayRef shape) { return std::accumulate(shape.begin(), shape.end(), int64_t{1}, std::multiplies()); @@ -289,6 +294,10 @@ FailureOr> mlir::hip::reifyBroadcastResultShape( return failure(); } bool reused = false; + // Reuse the first mixed shape for repeated SSA operands (e.g. x*x) so + // broadcastDim sees identical OpFoldResults and emits no redundant + // tensor.dim/cmpi/select chain. Broadcast arity is normally 2-3, so a + // linear scan is simpler than maintaining a side map. for (size_t j : llvm::seq(0, i)) { if (operand == operands[j]) { shapes.push_back(shapes[j]); From 6512f42f549ea162fb4601d716fa72c34b2e72e9 Mon Sep 17 00:00:00 2001 From: fhanuman Date: Wed, 29 Jul 2026 16:23:28 -0600 Subject: [PATCH 3/9] test(hip): keep result-shape coverage focused on LIT Remove Python numeric additions from the demonstration branch so the PR stays scoped to compiler changes and targeted IR coverage. Co-Authored-By: GPT-5.6 Sol Made-with: Cursor --- test/numeric/tests/test_gemm.py | 54 ------------------- test/numeric/tests/test_matmul.py | 44 ---------------- test/numeric/tests/test_max_min_broadcast.py | 55 -------------------- test/numeric/tests/test_where.py | 28 ---------- 4 files changed, 181 deletions(-) delete mode 100644 test/numeric/tests/test_gemm.py delete mode 100644 test/numeric/tests/test_max_min_broadcast.py diff --git a/test/numeric/tests/test_gemm.py b/test/numeric/tests/test_gemm.py deleted file mode 100644 index 529b712a8..000000000 --- a/test/numeric/tests/test_gemm.py +++ /dev/null @@ -1,54 +0,0 @@ -# -# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# Licensed under the MIT License. -# - -"""Numeric coverage for transpose-aware dynamic Gemm result dimensions.""" - -import numpy as np -import pytest -from onnx import TensorProto, helper - -from framework.comparator import compare_outputs -from framework.onnx_utils import make_model_from_nodes - -M = 3 -K = 4 -N = 5 - - -def _make_gemm_model(trans_a: int, trans_b: int): - a_shape = [K, "M"] if trans_a else ["M", K] - b_shape = ["N", K] if trans_b else [K, "N"] - a = helper.make_tensor_value_info("A", TensorProto.FLOAT16, a_shape) - b = helper.make_tensor_value_info("B", TensorProto.FLOAT16, b_shape) - c = helper.make_tensor_value_info("C", TensorProto.FLOAT16, ["N"]) - output = helper.make_tensor_value_info("Y", TensorProto.FLOAT16, ["M", "N"]) - node = helper.make_node( - "Gemm", - ["A", "B", "C"], - ["Y"], - transA=trans_a, - transB=trans_b, - ) - return make_model_from_nodes([node], [a, b, c], [output]) - - -class TestGemm: - @pytest.mark.parametrize( - "trans_a,trans_b", - [(0, 0), (1, 0), (0, 1), (1, 1)], - ) - def test_dynamic_m_n_all_transposes(self, model_runner, trans_a, trans_b): - model = _make_gemm_model(trans_a, trans_b) - rng = np.random.default_rng(44) - a_shape = (K, M) if trans_a else (M, K) - b_shape = (N, K) if trans_b else (K, N) - inputs = [ - rng.uniform(-0.5, 0.5, a_shape).astype(np.float16), - rng.uniform(-0.5, 0.5, b_shape).astype(np.float16), - rng.uniform(-0.5, 0.5, (N,)).astype(np.float16), - ] - - actual, expected = model_runner.run_sample(model, inputs, reference="cpu") - compare_outputs(actual, expected, atol=1e-3) diff --git a/test/numeric/tests/test_matmul.py b/test/numeric/tests/test_matmul.py index 29b6e590b..bf1448817 100644 --- a/test/numeric/tests/test_matmul.py +++ b/test/numeric/tests/test_matmul.py @@ -40,14 +40,6 @@ def _make_matmul_model(input_shape: list[int], weight_shape: list[int]): return make_model_from_nodes([node], [X], [Y], initializers=[w_init]) -def _make_unequal_rank_model(a_shape, b_shape, output_shape): - A = helper.make_tensor_value_info("A", TensorProto.FLOAT16, a_shape) - B = helper.make_tensor_value_info("B", TensorProto.FLOAT16, b_shape) - Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT16, output_shape) - node = helper.make_node("MatMul", ["A", "B"], ["Y"]) - return make_model_from_nodes([node], [A, B], [Y]) - - class TestMatMul: @pytest.mark.parametrize( "input_shape,weight_shape", @@ -97,39 +89,3 @@ def test_matmul_gate_up_proj_llama_shape(self, model_runner, seq_len): actual, expected = model_runner.run_sample(model, [x]) compare_outputs(actual, expected, atol=1e-3) - - @pytest.mark.parametrize( - "a_type_shape,b_type_shape,out_type_shape,a_shape,b_shape", - [ - ( - ["M", 4], - ["B", 4, "N"], - ["B", "M", "N"], - (3, 4), - (2, 4, 5), - ), - ( - ["M", 4], - ["B0", "B1", 4, "N"], - ["B0", "B1", "M", "N"], - (6, 4), - (2, 3, 4, 5), - ), - ], - ) - def test_matmul_unequal_rank_batch_from_b( - self, - model_runner, - a_type_shape, - b_type_shape, - out_type_shape, - a_shape, - b_shape, - ): - model = _make_unequal_rank_model(a_type_shape, b_type_shape, out_type_shape) - rng = np.random.default_rng(45) - a = rng.uniform(-1, 1, a_shape).astype(np.float16) - b = rng.uniform(-1, 1, b_shape).astype(np.float16) - - actual, expected = model_runner.run_sample(model, [a, b], reference="cpu") - compare_outputs(actual, expected, atol=1e-3) diff --git a/test/numeric/tests/test_max_min_broadcast.py b/test/numeric/tests/test_max_min_broadcast.py deleted file mode 100644 index f7f9af674..000000000 --- a/test/numeric/tests/test_max_min_broadcast.py +++ /dev/null @@ -1,55 +0,0 @@ -# -# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# Licensed under the MIT License. -# - -"""Numeric coverage for variadic Max/Min and multidirectional broadcasting.""" - -import numpy as np -import pytest -from onnx import TensorProto, helper - -from framework.comparator import compare_outputs -from framework.onnx_utils import make_model_from_nodes - - -def _make_binary_model(op_type: str): - lhs = helper.make_tensor_value_info("lhs", TensorProto.FLOAT, ["M", 1]) - rhs = helper.make_tensor_value_info("rhs", TensorProto.FLOAT, [1, "N"]) - output = helper.make_tensor_value_info("output", TensorProto.FLOAT, ["M", "N"]) - node = helper.make_node(op_type, ["lhs", "rhs"], ["output"]) - return make_model_from_nodes([node], [lhs, rhs], [output]) - - -def _make_variadic_model(op_type: str): - a = helper.make_tensor_value_info("a", TensorProto.FLOAT, [4]) - b = helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 4]) - c = helper.make_tensor_value_info("c", TensorProto.FLOAT, [2, 3, 4]) - output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [2, 3, 4]) - node = helper.make_node(op_type, ["a", "b", "c"], ["output"]) - return make_model_from_nodes([node], [a, b, c], [output]) - - -class TestMaxMinBroadcast: - @pytest.mark.parametrize("op_type", ["Max", "Min"]) - def test_asymmetric_dynamic_broadcast(self, model_runner, op_type): - model = _make_binary_model(op_type) - rng = np.random.default_rng(42) - lhs = rng.uniform(-1, 1, (3, 1)).astype(np.float32) - rhs = rng.uniform(-1, 1, (1, 5)).astype(np.float32) - - actual, expected = model_runner.run_sample(model, [lhs, rhs], reference="cpu") - compare_outputs(actual, expected, atol=1e-6) - - @pytest.mark.parametrize("op_type", ["Max", "Min"]) - def test_variadic_rank_growth(self, model_runner, op_type): - model = _make_variadic_model(op_type) - rng = np.random.default_rng(43) - inputs = [ - rng.uniform(-1, 1, (4,)).astype(np.float32), - rng.uniform(-1, 1, (3, 4)).astype(np.float32), - rng.uniform(-1, 1, (2, 3, 4)).astype(np.float32), - ] - - actual, expected = model_runner.run_sample(model, inputs, reference="cpu") - compare_outputs(actual, expected, atol=1e-6) diff --git a/test/numeric/tests/test_where.py b/test/numeric/tests/test_where.py index 9c385c4f2..7aa6a8c3b 100644 --- a/test/numeric/tests/test_where.py +++ b/test/numeric/tests/test_where.py @@ -105,34 +105,6 @@ def test_where_scalar_condition(self, model_runner, dtype): actual, expected = model_runner.run_sample(model, [cond, x, y]) compare_outputs(actual, expected, atol=1e-4) - def test_where_asymmetric_dynamic_broadcast(self, model_runner): - """Every dynamic output axis may be selected from a different input.""" - m, n = 3, 5 - model = _make_where_model( - [1, "N"], - ["M", 1], - ["M", "N"], - ["M", "N"], - np.float32, - ) - rng = np.random.default_rng(46) - cond = rng.integers(0, 2, (1, n), dtype=np.bool_) - x = rng.uniform(-1, 1, (m, 1)).astype(np.float32) - y = rng.uniform(-1, 1, (m, n)).astype(np.float32) - - actual, expected = model_runner.run_sample(model, [cond, x, y], reference="cpu") - compare_outputs(actual, expected, atol=1e-6) - - def test_where_rank_zero_output(self, model_runner): - """All-scalar Where proves rank-zero reification is a success.""" - model = _make_where_model([], [], [], [], np.float32) - cond = np.array(False, dtype=np.bool_) - x = np.array(1.25, dtype=np.float32) - y = np.array(-2.5, dtype=np.float32) - - actual, expected = model_runner.run_sample(model, [cond, x, y], reference="cpu") - compare_outputs(actual, expected, atol=0) - @pytest.mark.parametrize("seq_len", SEQ_LENS) def test_where_qwen9b_rope_mask_shape(self, model_runner, seq_len): """RoPE mask Where as it appears in Qwen3.5-9B text.onnx. From 7518cd09f376cad097b30146700ab6bf06198314 Mon Sep 17 00:00:00 2001 From: fhanuman Date: Wed, 29 Jul 2026 16:53:47 -0600 Subject: [PATCH 4/9] refactor(hip): streamline result-shape implementation Consolidate MatMul on one runtime ABI, align shape helper failure handling, and tighten compiler comments and LIT coverage. Co-Authored-By: GPT-5.6 Sol Made-with: Cursor --- CLAUDE.md | 2 +- docs/design/hip-shape-inference.md | 4 +- docs/hip_dialect_intro.md | 9 +++-- include/hip/Dialect/IR/HipShapeUtils.h | 25 +++--------- lib/Conversion/HipToLLVM/HipToLLVMUtils.h | 1 - lib/Conversion/HipToLLVM/MatmulLowering.cpp | 38 ++++++------------- lib/Conversion/OnnxToHip/GemmConversion.cpp | 16 ++++---- lib/Conversion/OnnxToHip/MatMulConversion.cpp | 20 +++++----- lib/Conversion/OnnxToHip/MaxConversion.cpp | 17 ++++----- lib/Conversion/OnnxToHip/MinConversion.cpp | 18 ++++----- lib/Conversion/OnnxToHip/OnnxToHipUtils.h | 10 +++-- lib/Conversion/OnnxToHip/WhereConversion.cpp | 24 ++++++------ lib/Dialect/IR/HipDialect.cpp | 12 +++--- lib/Dialect/IR/HipReifyResultShapesImpl.cpp | 10 +++-- lib/Dialect/IR/HipShapeUtils.cpp | 31 ++++++++------- lib/Runtime/hipdnn_ep_runtime.h | 25 ++---------- lib/Runtime/mock/mock_gpu.cpp | 19 +++------- lib/Runtime/real/matmul.cpp | 11 +----- .../Conversion/hip-to-llvm/test_matmul.mlir | 19 ++++++++-- .../test_matmul_invalid_batch.mlir | 16 ++++++++ test/lit/Pipeline/pipeline-pool-lower.mlir | 4 +- test/lit/e2e/test_matmul_model.mlir | 2 +- test/lit/e2e/test_mlp_model.mlir | 2 +- 23 files changed, 147 insertions(+), 188 deletions(-) create mode 100644 test/lit/Conversion/onnx-to-hip/test_matmul_invalid_batch.mlir diff --git a/CLAUDE.md b/CLAUDE.md index fd731e270..dc5a62576 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -144,7 +144,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co - Fully dynamic broadcast uses `select(lhs == 1, rhs, lhs)`, not integer maximum: broadcasting extents 0 and 1 produces 0. - Rank-zero success is an empty shape carried by `FailureOr`; never use an empty vector as both success and failure. - Variadic Max/Min derive every pairwise intermediate rank from the shared broadcast shape. -- MatMul uses the reified output batch product plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Per-axis partial batch broadcasting is rejected because the single-stride runtime cannot represent it. New artifacts call `wrap_hipblasLtMatmul_v2`; retain the legacy wrapper for cached artifacts compiled with the old signature. +- MatMul uses the reified output batch product plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Per-axis partial batch broadcasting is rejected because the single-stride runtime cannot represent it. `wrap_hipblasLtMatmul` carries both strides; invalidate LLVM-IR artifacts compiled against the previous wrapper ABI. - See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md). ### Allocation and memory planning diff --git a/docs/design/hip-shape-inference.md b/docs/design/hip-shape-inference.md index cbde1cb46..3906dbddb 100644 --- a/docs/design/hip-shape-inference.md +++ b/docs/design/hip-shape-inference.md @@ -155,8 +155,8 @@ operand uses stride 0, allowing the entire matrix to broadcast against a batched operand on the other side. Per-axis partial batch broadcasting where both operands contain fewer matrices than the output batch product is not yet representable by one constant stride and is rejected by the MatMul verifier. -New lowering calls the versioned `wrap_hipblasLtMatmul_v2`; the legacy wrapper -remains available so cached artifacts with the previous signature still load. +`wrap_hipblasLtMatmul` carries both strides. LLVM-IR artifacts compiled against +the previous wrapper ABI must be invalidated. ## `--hip-infer-shapes` diff --git a/docs/hip_dialect_intro.md b/docs/hip_dialect_intro.md index af6695916..d799f7b06 100644 --- a/docs/hip_dialect_intro.md +++ b/docs/hip_dialect_intro.md @@ -44,11 +44,12 @@ 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 | +| `hip.matmul` | `(%ctx) ins(%A, %B : ...) outs(%C : ...)` | `wrap_hipblasLtMatmul(state, slot, A, B, C, M, N, K, batch, elem, a_stride, b_stride)` | 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. +The output batch count is the product of the broadcasted leading dimensions. +Either operand may provide one matrix for all batches (`stride = 0`) or one +matrix per output batch. Per-axis partial batch broadcasting is rejected +because it cannot be represented by one constant stride per operand. --- diff --git a/include/hip/Dialect/IR/HipShapeUtils.h b/include/hip/Dialect/IR/HipShapeUtils.h index 1e678e9d8..263843d1b 100644 --- a/include/hip/Dialect/IR/HipShapeUtils.h +++ b/include/hip/Dialect/IR/HipShapeUtils.h @@ -18,26 +18,13 @@ namespace mlir { namespace hip { /// Compute the shape of `A @ B` for matmul with NumPy-style batch broadcast -/// over the leading dims. Last two dims of `aShape` are `[M, K]`; last two -/// dims of `bShape` are `[K, N]`. Leading dims are broadcast (right-aligned, -/// missing dims treated as 1). +/// over the leading dimensions. The matrix dimensions are `A[..., M, K]` and +/// `B[..., K, N]`. /// -/// Returns the inferred shape on success. Returns an empty `SmallVector` and -/// emits a diagnostic via `emitError` on rank-, K-, or batch-broadcast -/// mismatch. -/// -/// `ShapedType::kDynamic` is treated as a wildcard: -/// - K_a or K_b dynamic -> K match passes (result K is dropped anyway). -/// - Batch dim broadcast follows NumPy / TF / ONNX MatMul semantics -/// (delegated to `mlir::OpTrait::util::getBroadcastedShape`): -/// * 1 broadcasts against any dim. -/// * dynamic + static>1 -> static (the dynamic side must be 1 or -/// match the static side at runtime per the broadcast contract; -/// taking the static side is the strictly-correct tightening). -/// * dynamic + dynamic -> dynamic. -/// * static + static, equal -> static; unequal and neither is 1 -/// -> error. -SmallVector +/// Dynamic contraction dimensions are treated as compatible. Batch dimensions +/// use `OpTrait::util::getBroadcastedShape`. On failure, emits a diagnostic +/// through `emitError`. +FailureOr> inferMatmulShape(ArrayRef aShape, ArrayRef bShape, function_ref emitError); diff --git a/lib/Conversion/HipToLLVM/HipToLLVMUtils.h b/lib/Conversion/HipToLLVM/HipToLLVMUtils.h index 259fcc834..a32c638cd 100644 --- a/lib/Conversion/HipToLLVM/HipToLLVMUtils.h +++ b/lib/Conversion/HipToLLVM/HipToLLVMUtils.h @@ -52,7 +52,6 @@ inline constexpr const char *kMiopenConvolutionForward = inline constexpr const char *kMiopenConvolutionTranspose = "wrap_miopenConvolutionTranspose"; inline constexpr const char *kWrapHipblasltMatmul = "wrap_hipblasLtMatmul"; -inline constexpr const char *kWrapHipblasltMatmulV2 = "wrap_hipblasLtMatmul_v2"; inline constexpr const char *kWrapMiopenT5LayerNormForward = "wrap_miopenT5LayerNormForward"; inline constexpr const char *kWrapSkipSimplifiedLayerNorm = diff --git a/lib/Conversion/HipToLLVM/MatmulLowering.cpp b/lib/Conversion/HipToLLVM/MatmulLowering.cpp index bdac586f2..3947c48ef 100644 --- a/lib/Conversion/HipToLLVM/MatmulLowering.cpp +++ b/lib/Conversion/HipToLLVM/MatmulLowering.cpp @@ -11,13 +11,16 @@ namespace mlir { namespace hip { namespace { -// ===== hipBLASLt ops ========================================================= - -// hip.hipblaslt.matmul(handle) ins(A, B) outs(C) -// -> wrap_hipblasLtMatmul(state, A, B, C, M, N, K, batch, A-stride, -// B-stride) -// Rank-generic for whole-tensor batch broadcast: a rank-2 (or leading-one) -// operand uses stride 0 while the other operand supplies the output batches. +// Lower `hip.matmul` to the hipBLASLt runtime ABI. Each operand gets an +// independent batch stride so either whole matrix may be broadcast. +// +// Before: +// hip.matmul ins(%a, %b : memref<128x4096xf16>, +// memref<2x4096x1024xf16>) outs(%out : ...) +// After: +// %a_stride = llvm.mlir.constant(0 : i64) : i64 +// %b_stride = llvm.mul %k, %n : i64 +// llvm.call @wrap_hipblasLtMatmul(..., %a_stride, %b_stride) struct MatmulOpLowering : public ConvertOpToLLVMPattern { using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern; @@ -30,20 +33,17 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { Type i32Type = rewriter.getI32Type(); Type i64Type = rewriter.getI64Type(); - // Helper: create i64 constant auto createI64Const = [&](int64_t value) -> Value { return LLVM::ConstantOp::create(rewriter, loc, i64Type, rewriter.getI64IntegerAttr(value)); }; - // Extract pointers Value statePtr = adaptor.getCtx(); Value APtr = extractContiguousMemRefPtr(adaptor.getA(), rewriter, loc); Value BPtr = extractContiguousMemRefPtr(adaptor.getB(), rewriter, loc); Value outputPtr = extractContiguousMemRefPtr(adaptor.getOutput(), rewriter, loc); - // Get memref types and shapes auto AType = cast(op.getA().getType()); auto BType = cast(op.getB().getType()); auto outputType = cast(op.getOutput().getType()); @@ -54,16 +54,11 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { [&]() { return op.emitOpError(); }))) return failure(); - // === DYNAMIC SHAPE SUPPORT === - // For dynamic shapes, we compute dimensions at runtime MemRefDescriptor ADesc(adaptor.getA()); MemRefDescriptor BDesc(adaptor.getB()); MemRefDescriptor outputDesc(adaptor.getOutput()); - // Compute M, K, N from runtime dimensions - // A: [..., M, K], B: [..., K, N] or B: [K, N] - Value M = - (ARank >= 2) ? ADesc.size(rewriter, loc, ARank - 2) : createI64Const(1); + Value M = ADesc.size(rewriter, loc, ARank - 2); Value K = ADesc.size(rewriter, loc, ARank - 1); Value N = BDesc.size(rewriter, loc, BRank - 1); @@ -74,20 +69,11 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { batchCount = LLVM::MulOp::create(rewriter, loc, batchCount, dim); } - // Compute element size in bytes unsigned elemBits = AType.getElementType().getIntOrFloatBitWidth(); Value elemSize = createI64Const(elemBits / 8); // Each operand independently uses stride 0 when it contains one matrix, // or its matrix size when it contains one matrix per output batch. - // - // Before: - // hip.matmul ins(%A, %B : memref<128x4096xf16>, - // memref<2x4096x1024xf16>) - // After: - // %a_stride = llvm.mlir.constant(0 : i64) : i64 - // %b_stride = llvm.mul %K, %N : i64 - // llvm.call @wrap_hipblasLtMatmul(..., %a_stride, %b_stride) auto computeBatchStride = [&](MemRefType type, MemRefDescriptor desc, Value matrixElements) -> Value { int64_t rank = type.getRank(); @@ -147,7 +133,7 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { }; FailureOr funcOp = LLVM::lookupOrCreateFn( - rewriter, module, kWrapHipblasltMatmulV2, paramTypes, i32Type); + rewriter, module, kWrapHipblasltMatmul, paramTypes, i32Type); if (failed(funcOp)) return failure(); diff --git a/lib/Conversion/OnnxToHip/GemmConversion.cpp b/lib/Conversion/OnnxToHip/GemmConversion.cpp index a61351acb..85ac47ed3 100644 --- a/lib/Conversion/OnnxToHip/GemmConversion.cpp +++ b/lib/Conversion/OnnxToHip/GemmConversion.cpp @@ -13,10 +13,13 @@ namespace { // ONNX Gemm -> HIP Gemm // // Before: -// %init = tensor.empty(tensor.dim %A, 0, tensor.dim %A, 1) +// %r = "onnx.Gemm"(%a, %b) {transA = 1 : i64, transB = 0 : i64} +// : (tensor, tensor) -> tensor // After: -// %shape = gemm_shape(%A, %B, transA, transB) -// %init = tensor.empty(%shape.M, %shape.N) +// %m = tensor.dim %a, %c1 +// %n = tensor.dim %b, %c1 +// %init = tensor.empty(%m, %n) : tensor +// %r = hip.gemm ... outs(%init : tensor) //===----------------------------------------------------------------------===// struct GemmToHip : public mlir::RewritePattern { GemmToHip(mlir::MLIRContext *ctx) @@ -62,7 +65,6 @@ struct GemmToHip : public mlir::RewritePattern { return rewriter.notifyMatchFailure( op, "Gemm result type is incompatible with inferred shape"); - // hip.gemm llvm::SmallVector attrs; attrs.push_back( rewriter.getNamedAttr("alpha", rewriter.getF32FloatAttr(alpha))); @@ -74,12 +76,10 @@ struct GemmToHip : public mlir::RewritePattern { rewriter.getNamedAttr("transB", rewriter.getI64IntegerAttr(transB))); llvm::SmallVector operands = {context, inputA, inputB}; - if (hasInputC) { + if (hasInputC) operands.push_back(inputC); - } operands.push_back(*init); - // Result type inferred from `init` via InferTypeOpInterface — DPS contract: - // result type == outs operand type. + // Let ODS infer the result type from the DPS init. auto hipOp = mlir::hip::GemmOp::create(rewriter, loc, operands, attrs); rewriter.replaceOp(op, hipOp.getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/MatMulConversion.cpp b/lib/Conversion/OnnxToHip/MatMulConversion.cpp index f4952f956..4d0c2fe72 100644 --- a/lib/Conversion/OnnxToHip/MatMulConversion.cpp +++ b/lib/Conversion/OnnxToHip/MatMulConversion.cpp @@ -9,14 +9,17 @@ namespace mlir { namespace hip { namespace { -/// onnx.MatMul -> hip.hipblaslt.matmul +/// onnx.MatMul -> hip.matmul /// /// Before: -/// %batch = tensor.dim %A, %c0 -/// %init = tensor.empty(%batch, ...) : tensor +/// %r = "onnx.MatMul"(%a, %b) +/// : (tensor, tensor) -> tensor /// After: -/// %shape = matmul_shape(%A, %B) -/// %init = tensor.empty(%shape...) : tensor +/// %batch = tensor.dim %b, %c0 +/// %m = tensor.dim %a, %c0 +/// %n = tensor.dim %b, %c2 +/// %init = tensor.empty(%batch, %m, %n) : tensor +/// %r = hip.matmul ... outs(%init : tensor) struct MatMulToHip : public mlir::RewritePattern { MatMulToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.MatMul", /*benefit=*/1, ctx) {} @@ -50,12 +53,7 @@ MatMulToHip::matchAndRewrite(mlir::Operation *op, if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "MatMul result type is incompatible with inferred shape"); - // Inferred-type Op::create overload: result type is read from the typed - // outs operand via the auto-emitted MatmulOp::inferReturnTypes (HipOps.td - // base, autoInfer=1). Equivalent to passing `resultType` explicitly -- - // outs.getType() == resultType by construction here -- but keeps the DPS - // contract `result_type == outs_operand_type` closed by ODS rather than - // restated at the callsite. + // Let ODS infer the result type from the DPS init. auto hipOp = mlir::hip::MatmulOp::create(rewriter, loc, context, a, b, *init); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/MaxConversion.cpp b/lib/Conversion/OnnxToHip/MaxConversion.cpp index 03edf0a5b..1071d154e 100644 --- a/lib/Conversion/OnnxToHip/MaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/MaxConversion.cpp @@ -9,17 +9,15 @@ namespace mlir { namespace hip { namespace { -/// onnx.Max -> hip.max (via MIOpen miopenOpTensor with miopenTensorOpMax) -/// -/// Handles variadic inputs by pairwise chaining: -/// max(a, b, c) = max(max(a, b), c) -/// Single input is identity (pass through). +/// Lower variadic `onnx.Max` to pairwise `hip.max` operations. /// /// Before: -/// %tmp = hip.max ... outs(%same_rank_as_a) +/// %r = "onnx.Max"(%a, %b, %c) : (...) -> tensor<2x3x4xf32> /// After: -/// %tmp_shape = broadcast_shape(%a, %b) -/// %tmp = hip.max ... outs(%empty_for_tmp_shape) +/// %ab_init = tensor.empty() : tensor<3x4xf32> +/// %ab = hip.max ... outs(%ab_init : tensor<3x4xf32>) +/// %abc_init = tensor.empty() : tensor<2x3x4xf32> +/// %r = hip.max ... outs(%abc_init : tensor<2x3x4xf32>) struct MaxToHip : public mlir::RewritePattern { MaxToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.Max", /*benefit=*/1, ctx) {} @@ -50,9 +48,8 @@ MaxToHip::matchAndRewrite(mlir::Operation *op, auto resultType = mlir::cast(op->getResult(0).getType()); - // Pairwise chaining: accumulate = max(accumulate, next_input) mlir::Value accumulate = op->getOperand(0); - for (unsigned i = 1; i < numInputs; ++i) { + for (unsigned i : llvm::seq(1, numInputs)) { mlir::Value rhs = op->getOperand(i); mlir::FailureOr> stepShape = mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, diff --git a/lib/Conversion/OnnxToHip/MinConversion.cpp b/lib/Conversion/OnnxToHip/MinConversion.cpp index 9b8a0fc81..83f4b1a45 100644 --- a/lib/Conversion/OnnxToHip/MinConversion.cpp +++ b/lib/Conversion/OnnxToHip/MinConversion.cpp @@ -9,17 +9,15 @@ namespace mlir { namespace hip { namespace { -/// onnx.Min -> hip.min (via MIOpen miopenOpTensor with miopenTensorOpMin) -/// -/// Handles variadic inputs by pairwise chaining: -/// min(a, b, c) = min(min(a, b), c) -/// Single input is identity (pass through). +/// Lower variadic `onnx.Min` to pairwise `hip.min` operations. /// /// Before: -/// %tmp = hip.min ... outs(%same_rank_as_a) +/// %r = "onnx.Min"(%a, %b, %c) : (...) -> tensor<2x3x4xf32> /// After: -/// %tmp_shape = broadcast_shape(%a, %b) -/// %tmp = hip.min ... outs(%empty_for_tmp_shape) +/// %ab_init = tensor.empty() : tensor<3x4xf32> +/// %ab = hip.min ... outs(%ab_init : tensor<3x4xf32>) +/// %abc_init = tensor.empty() : tensor<2x3x4xf32> +/// %r = hip.min ... outs(%abc_init : tensor<2x3x4xf32>) struct MinToHip : public mlir::RewritePattern { MinToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.Min", /*benefit=*/1, ctx) {} @@ -36,7 +34,6 @@ MinToHip::matchAndRewrite(mlir::Operation *op, if (numInputs == 0) return rewriter.notifyMatchFailure(op, "Min requires at least 1 input"); - // Single input: identity if (numInputs == 1) { rewriter.replaceOp(op, op->getOperand(0)); return mlir::success(); @@ -51,9 +48,8 @@ MinToHip::matchAndRewrite(mlir::Operation *op, auto resultType = mlir::cast(op->getResult(0).getType()); - // Pairwise chaining: accumulate = min(accumulate, next_input) mlir::Value accumulate = op->getOperand(0); - for (unsigned i = 1; i < numInputs; ++i) { + for (unsigned i : llvm::seq(1, numInputs)) { mlir::Value rhs = op->getOperand(i); mlir::FailureOr> stepShape = mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h index d12cea994..1342bf514 100644 --- a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h @@ -192,10 +192,12 @@ getTensorTypeFromReifiedShape(llvm::ArrayRef reifiedShape, /// broadcast of \p operands. Converter destination construction delegates to /// the same dialect helper used by ReifyRankedShapedTypeOpInterface. /// -/// Use this for binary/multinary broadcast elementwise ops (Add, Mul, Where, -/// ...). Do NOT use `createEmptyTensor(resultType, source)` when operands can -/// disagree on which side supplies a dynamic extent (e.g. `[?x1] + [1x?] -> -/// [?x?]` -- dim 0 from lhs, dim 1 from rhs). +/// Before: +/// %init = tensor.empty(%lhs_dim) : tensor +/// After: +/// %lhs_is_one = arith.cmpi eq, %lhs_dim, %c1 : index +/// %extent = arith.select %lhs_is_one, %rhs_dim, %lhs_dim : index +/// %init = tensor.empty(%extent) : tensor inline mlir::FailureOr createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, mlir::RankedTensorType resultType, diff --git a/lib/Conversion/OnnxToHip/WhereConversion.cpp b/lib/Conversion/OnnxToHip/WhereConversion.cpp index 1b6cd9f5a..fa2c1d1ba 100644 --- a/lib/Conversion/OnnxToHip/WhereConversion.cpp +++ b/lib/Conversion/OnnxToHip/WhereConversion.cpp @@ -9,11 +9,18 @@ namespace mlir { namespace hip { namespace { -/// onnx.Where -> hip.where -/// ONNX Where: output[i] = condition[i] ? X[i] : Y[i]. -/// Supports multidirectional (NumPy-style) broadcasting between condition, -/// X and Y. The condition tensor is bool (i1); X and Y share the result -/// element type. +/// Lower `onnx.Where` with multidirectional broadcasting. +/// +/// Before: +/// %r = "onnx.Where"(%cond, %x, %y) +/// : (tensor, tensor<1xf32>, tensor) -> tensor +/// After: +/// %cond_dim = tensor.dim %cond, %c0 +/// %y_dim = tensor.dim %y, %c0 +/// %cond_is_one = arith.cmpi eq, %cond_dim, %c1 : index +/// %extent = arith.select %cond_is_one, %y_dim, %cond_dim : index +/// %init = tensor.empty(%extent) : tensor +/// %r = hip.where ... outs(%init : tensor) struct WhereToHip : public mlir::RewritePattern { WhereToHip(mlir::MLIRContext *ctx) : RewritePattern("onnx.Where", /*benefit=*/1, ctx) {} @@ -46,15 +53,10 @@ WhereToHip::matchAndRewrite(mlir::Operation *op, return rewriter.notifyMatchFailure( op, "onnx.Where lowering expects a ranked tensor result"); - // ONNX Where supports multidirectional (NumPy-style) broadcasting, so any - // given output dim may be contributed by a different operand. Resolve each - // dynamic result dim by scanning all three operands rather than relying on - // a single "source" tensor. mlir::FailureOr initOrFailure = createBroadcastEmptyTensor(rewriter, loc, resultType, {condition, x, y}); if (mlir::failed(initOrFailure)) - return rewriter.notifyMatchFailure( - op, "onnx.Where: no ranked operand spans dynamic result dim"); + return mlir::failure(); auto hipOp = mlir::hip::WhereOp::create(rewriter, loc, context, condition, x, y, *initOrFailure); rewriter.replaceOp(op, hipOp->getResult(0)); diff --git a/lib/Dialect/IR/HipDialect.cpp b/lib/Dialect/IR/HipDialect.cpp index 442113818..73f04d5e3 100644 --- a/lib/Dialect/IR/HipDialect.cpp +++ b/lib/Dialect/IR/HipDialect.cpp @@ -569,16 +569,14 @@ LogicalResult MatmulOp::verify() { ArrayRef aShape = getShapeOf(getA()); ArrayRef bShape = getShapeOf(getB()); - // Static shape check via the shared matmul helper. The lambda is - // invoked once; it returns `{outputShape}` on success or `{}` on shape - // mismatch (in which case it has already issued a diagnostic on `*this`). if (failed(mlir::hip::verifyHipOpShape( *this, [&]() -> SmallVector> { - SmallVector outShape = mlir::hip::inferMatmulShape( - aShape, bShape, [&]() { return this->emitOpError(); }); - if (outShape.empty()) + FailureOr> outShape = + mlir::hip::inferMatmulShape( + aShape, bShape, [&]() { return this->emitOpError(); }); + if (failed(outShape)) return {}; - return {std::move(outShape)}; + return {std::move(*outShape)}; }))) return failure(); return mlir::hip::verifyStridedBatchMatmul( diff --git a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp index 2e86ba4b0..e69ee0103 100644 --- a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp +++ b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp @@ -49,11 +49,13 @@ ArrayRef getShapeOf(Value v) { // NumPy broadcast semantics. // // Before: -// %m = hip.matmul ins(%a, %b : tensor, tensor<4096x4096xf16>) -// outs(%out : tensor) -> tensor +// %r = hip.matmul ins(%a, %b : tensor, +// tensor) +// outs(%out : tensor) -> tensor // After (reified result shape): -// dim 0 (dynamic M) -> %d0 = tensor.dim %a, %c0 -// dim 1 (static N) -> 4096 : index +// dim 0 (batch) -> tensor.dim %b, %c0 +// dim 1 (M) -> tensor.dim %a, %c0 +// dim 2 (N) -> tensor.dim %b, %c2 //===----------------------------------------------------------------------===// LogicalResult diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index ec5e6e08d..917970bfa 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -91,16 +91,16 @@ broadcastDim(OpBuilder &b, Location loc, OpFoldResult lhs, OpFoldResult rhs, } // namespace -SmallVector +FailureOr> mlir::hip::inferMatmulShape(ArrayRef aShape, ArrayRef bShape, function_ref emitError) { if (aShape.size() < 2) { emitError() << "matmul A must have rank >= 2, got rank " << aShape.size(); - return {}; + return failure(); } if (bShape.size() < 2) { emitError() << "matmul B must have rank >= 2, got rank " << bShape.size(); - return {}; + return failure(); } int64_t M = aShape[aShape.size() - 2]; @@ -112,18 +112,16 @@ mlir::hip::inferMatmulShape(ArrayRef aShape, ArrayRef bShape, if (!ShapedType::isDynamic(Ka) && !ShapedType::isDynamic(Kb) && Ka != Kb) { emitError() << "matmul contraction dim mismatch: A.shape[-1]=" << Ka << " vs B.shape[-2]=" << Kb; - return {}; + return failure(); } - // Batch broadcast (NumPy / ONNX MatMul) on the leading dims; see header - // for the full case table. ArrayRef aBatch = aShape.drop_back(2); ArrayRef bBatch = bShape.drop_back(2); SmallVector result; if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, result)) { emitError() << "matmul batch broadcast failure: A.batch=" << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); - return {}; + return failure(); } result.reserve(result.size() + 2); result.push_back(M); @@ -141,6 +139,13 @@ LogicalResult mlir::hip::verifyStridedBatchMatmul( ArrayRef aBatch = aShape.drop_back(2); ArrayRef bBatch = bShape.drop_back(2); + SmallVector outputBatch; + if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, outputBatch)) { + emitError() << "matmul batch broadcast failure: A.batch=" + << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); + return failure(); + } + auto isSingleMatrix = [](ArrayRef batch) { return llvm::all_of(batch, [](int64_t dim) { return dim == 1; }); }; @@ -153,14 +158,8 @@ LogicalResult mlir::hip::verifyStridedBatchMatmul( auto isDynamic = [](int64_t dim) { return ShapedType::isDynamic(dim); }; if (llvm::any_of(aBatch, isDynamic) || llvm::any_of(bBatch, isDynamic)) { - emitError() << "matmul with two nontrivial dynamic batch shapes is not " - "supported by the strided-batch runtime"; - return failure(); - } - - SmallVector outputBatch; - if (!OpTrait::util::getBroadcastedShape(aBatch, bBatch, outputBatch)) { - emitError() << "matmul batch broadcast failure: A.batch=" + emitError() << "matmul dynamic batch layout is not representable by one " + "constant stride per operand: A.batch=" << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); return failure(); } @@ -321,7 +320,7 @@ FailureOr> mlir::hip::reifyMatmulResultShape( emitError() << "matmul operands must be ranked tensors"; return failure(); } - if (inferMatmulShape(aType.getShape(), bType.getShape(), emitError).empty()) + if (failed(inferMatmulShape(aType.getShape(), bType.getShape(), emitError))) return failure(); if (failed(verifyStridedBatchMatmul(aType.getShape(), bType.getShape(), emitError))) diff --git a/lib/Runtime/hipdnn_ep_runtime.h b/lib/Runtime/hipdnn_ep_runtime.h index 8a6a65674..a26ce6e46 100644 --- a/lib/Runtime/hipdnn_ep_runtime.h +++ b/lib/Runtime/hipdnn_ep_runtime.h @@ -801,35 +801,16 @@ int wrap_hipblasLtGemm(void *handle, // hipBLASLt handle const void *beta, // Scalar beta void *C); // Matrix C GPU pointer (in/out) -// MatMul operation wrapper (batched matrix multiplication) -// Called by generated IR for onnx.MatMul lowering -// Computes output = A @ B for each batch -// A/B may each contain one broadcast matrix or one matrix per output batch. -// output: [batch_count x M x N] +// Compute output = A @ B for each batch. Either operand may contain one matrix +// broadcast across all output batches or one matrix per output batch. // // `a_batch_stride` / `b_batch_stride` are hipBLASLt's per-batch advances in // elements. A stride is 0 when one matrix is broadcast across all batches; // otherwise it is M*K for A or K*N for B. -// Legacy wrapper retained for cached model artifacts compiled before the A -// batch-stride parameter was added. It assumes A is per-batch. int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, - int64_t b_batch_stride); - -int wrap_hipblasLtMatmul_v2( - RuntimeState *state, - int op_state_slot, // per-instance op-state slot (shared algo table) - const void *A, // Matrix A GPU pointer - const void *B, // Matrix B GPU pointer - void *output, // Output GPU pointer - int64_t M, // Rows of A (per batch) - int64_t N, // Columns of B - int64_t K, // Columns of A / Rows of B - int64_t batch_count, // Number of batches - int64_t elem_size, // Element size in bytes (2=f16, 4=f32) - int64_t a_batch_stride, // 0 = broadcast; M*K = per-batch - int64_t b_batch_stride); // 0 = broadcast (any rank); K*N = per-batch + int64_t a_batch_stride, int64_t b_batch_stride); // GroupQueryAttention operation wrapper (Full MS spec) // Called by generated IR for onnx.Custom(GroupQueryAttention) lowering diff --git a/lib/Runtime/mock/mock_gpu.cpp b/lib/Runtime/mock/mock_gpu.cpp index 11bd83e3e..6023f7c93 100644 --- a/lib/Runtime/mock/mock_gpu.cpp +++ b/lib/Runtime/mock/mock_gpu.cpp @@ -595,18 +595,7 @@ int wrap_hipblasLtGemm(void *handle, void *stream, int64_t m, int64_t n, int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, - int64_t b_batch_stride) { - return wrap_hipblasLtMatmul_v2(state, op_state_slot, A, B, output, M, N, K, - batch_count, elem_size, M * K, b_batch_stride); -} - -int wrap_hipblasLtMatmul_v2(RuntimeState *state, int op_state_slot, - const void *A, const void *B, void *output, - int64_t M, int64_t N, int64_t K, - int64_t batch_count, int64_t elem_size, - int64_t a_batch_stride, int64_t b_batch_stride) { - (void)a_batch_stride; - (void)b_batch_stride; + int64_t a_batch_stride, int64_t b_batch_stride) { (void)op_state_slot; if (!state) { fprintf(stderr, "Invalid state in wrap_hipblasLtMatmul\n"); @@ -614,9 +603,11 @@ int wrap_hipblasLtMatmul_v2(RuntimeState *state, int op_state_slot, } MOCK_PRINT("[MOCK] wrap_hipblasLtMatmul(M=%lld, N=%lld, K=%lld, " - "batch=%lld, elem_size=%lld)\n", + "batch=%lld, elem_size=%lld, a_batch_stride=%lld, " + "b_batch_stride=%lld)\n", (long long)M, (long long)N, (long long)K, (long long)batch_count, - (long long)elem_size); + (long long)elem_size, (long long)a_batch_stride, + (long long)b_batch_stride); return 0; } diff --git a/lib/Runtime/real/matmul.cpp b/lib/Runtime/real/matmul.cpp index 85d30bbe1..8410d4537 100644 --- a/lib/Runtime/real/matmul.cpp +++ b/lib/Runtime/real/matmul.cpp @@ -419,16 +419,7 @@ static void autotuneMatmul(hipblasLtHandle_t handle, hipStream_t stream, int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, - int64_t b_batch_stride) { - return wrap_hipblasLtMatmul_v2(state, op_state_slot, A, B, output, M, N, K, - batch_count, elem_size, M * K, b_batch_stride); -} - -int wrap_hipblasLtMatmul_v2(RuntimeState *state, int op_state_slot, - const void *A, const void *B, void *output, - int64_t M, int64_t N, int64_t K, - int64_t batch_count, int64_t elem_size, - int64_t a_batch_stride, int64_t b_batch_stride) { + int64_t a_batch_stride, int64_t b_batch_stride) { OP_PROFILE( "matmul", [&] { diff --git a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir index 830f29e66..6179f35bf 100644 --- a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir +++ b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir @@ -24,7 +24,10 @@ module { // CHECK-LABEL: llvm.func @test_matmul_rank2_b // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select -// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK-NOT: @wrap_hipblasLtMatmul_v2 +// CHECK: %[[A0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 +// CHECK: %[[B0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A0_STRIDE]], %[[B0_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 // Verify 12 parameters: // - 4 pointers: state, A, B, output // - 7 i64: M, N, K, batch_count, elem_size, A stride, B stride @@ -54,7 +57,12 @@ module { // CHECK-LABEL: llvm.func @test_matmul_rank3_leading_one_b // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select -// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK-NOT: @wrap_hipblasLtMatmul_v2 +// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[A1_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[B1_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A1_STRIDE]], %[[B1_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 // ----- Rank-2 broadcast A against rank-3 B ------------------------------- // A contains one matrix and therefore uses A stride 0. B and the output carry @@ -75,4 +83,9 @@ module { // CHECK-LABEL: llvm.func @test_matmul_rank2_a // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select -// CHECK: llvm.call @wrap_hipblasLtMatmul_v2({{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK-NOT: @wrap_hipblasLtMatmul_v2 +// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[B2_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[A2_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A2_STRIDE]], %[[B2_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 diff --git a/test/lit/Conversion/onnx-to-hip/test_matmul_invalid_batch.mlir b/test/lit/Conversion/onnx-to-hip/test_matmul_invalid_batch.mlir new file mode 100644 index 000000000..2574674b4 --- /dev/null +++ b/test/lit/Conversion/onnx-to-hip/test_matmul_invalid_batch.mlir @@ -0,0 +1,16 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// RUN: hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hip %s 2>&1 | FileCheck %s + +module { + func.func @main_graph(%a: tensor<2x1x4x8xf16>, + %b: tensor<1x3x8x16xf16>) + -> tensor<2x3x4x16xf16> { + // CHECK: error: matmul partial per-axis batch broadcast is not supported by the strided-batch runtime + %result = "onnx.MatMul"(%a, %b) + : (tensor<2x1x4x8xf16>, tensor<1x3x8x16xf16>) + -> tensor<2x3x4x16xf16> + return %result : tensor<2x3x4x16xf16> + } +} diff --git a/test/lit/Pipeline/pipeline-pool-lower.mlir b/test/lit/Pipeline/pipeline-pool-lower.mlir index df21b6fb5..98d2fb1c6 100644 --- a/test/lit/Pipeline/pipeline-pool-lower.mlir +++ b/test/lit/Pipeline/pipeline-pool-lower.mlir @@ -22,7 +22,7 @@ // CHECK: %[[DOM:.*]] = llvm.mlir.constant(0 : i32) : i32 // CHECK: llvm.call @hipdnn_ep_get_pool_base(%[[CTX]], %[[DOM]], %[[POOL_SIZE]]) : (!llvm.ptr, i32, i64) -> !llvm.ptr // CHECK: llvm.mlir.constant(256 : index) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul_v2(%[[CTX]], +// CHECK: llvm.call @wrap_hipblasLtMatmul(%[[CTX]], // CHECK: llvm.call @hip_miopen_softmax(%[[CTX]], // CHECK: llvm.return func.func @static_pool_to_llvm( @@ -47,7 +47,7 @@ func.func @static_pool_to_llvm( // CHECK: llvm.mul %arg15, %[[C32]] : i64 // CHECK: %[[DOM2:.*]] = llvm.mlir.constant(0 : i32) : i32 // CHECK: llvm.call @hipdnn_ep_get_pool_base(%[[CTX2]], %[[DOM2]], %{{[0-9]+}}) : (!llvm.ptr, i32, i64) -> !llvm.ptr -// CHECK: llvm.call @wrap_hipblasLtMatmul_v2(%[[CTX2]], +// CHECK: llvm.call @wrap_hipblasLtMatmul(%[[CTX2]], // CHECK: llvm.call @hip_miopen_softmax(%[[CTX2]], // CHECK: llvm.return func.func @dynamic_pool_to_llvm( diff --git a/test/lit/e2e/test_matmul_model.mlir b/test/lit/e2e/test_matmul_model.mlir index c70dd7fde..2bc0184fb 100644 --- a/test/lit/e2e/test_matmul_model.mlir +++ b/test/lit/e2e/test_matmul_model.mlir @@ -13,7 +13,7 @@ // CHECK: module attributes { // CHECK-SAME: hipdnn.input_count = 1 // CHECK-SAME: hipdnn.output_count = 1 -// CHECK: llvm.func @wrap_hipblasLtMatmul +// CHECK: llvm.func @wrap_hipblasLtMatmul( // CHECK: llvm.func @inference_init // CHECK: llvm.func @inference_compute // CHECK: llvm.func @inference_cleanup diff --git a/test/lit/e2e/test_mlp_model.mlir b/test/lit/e2e/test_mlp_model.mlir index 8906e5b6a..6dc0e55ec 100644 --- a/test/lit/e2e/test_mlp_model.mlir +++ b/test/lit/e2e/test_mlp_model.mlir @@ -15,7 +15,7 @@ // CHECK: module attributes { // CHECK-SAME: hipdnn.input_count = 1 // CHECK-SAME: hipdnn.output_count = 1 -// CHECK: llvm.func @wrap_hipblasLtMatmul +// CHECK: llvm.func @wrap_hipblasLtMatmul( // CHECK: llvm.func @inference_init // CHECK: llvm.func @inference_compute // CHECK: llvm.func @inference_cleanup From fe4d878d770a0c615f77aff119164e8a74c8719f Mon Sep 17 00:00:00 2001 From: fhanuman Date: Fri, 31 Jul 2026 08:33:24 -0600 Subject: [PATCH 5/9] fix(hip): accept dynamic batched MatMul and unify reduction result shapes Restore support for `[?, H, M, K] @ [?, H, K, N]`, which the strided-batch check rejected for any dynamic batch extent spanning more than one batch axis even though nothing broadcasts and one stride per operand is exact. Representability is now decided per axis, so only a provable partial broadcast is rejected. Extend the shared result-shape rule to reductions, which still computed their destination independently of `reifyResultShapes`: a positional dim copy is wrong under `keepdims = 0` once a reduced axis precedes a kept one, and gating the converter on the axes operand *count* while reification gated on it being *constant* made the two disagree for opset-13+ constant axes. One `computeReductionDimMap` now backs the static shape, the mixed shape, and result type inference, replacing ReduceProd's private copy of the rule. Split each shape category into a pure `infer*` function of static shapes and a `reify*` function that validates through it before touching the builder, so a failure cannot leave stray dim ops behind; `reifyGemmResultShape` previously emitted `tensor.dim` for A and B before validating the optional C. The split also gives `GemmOp` the static shape verification it lacked. Route the remaining `hip.max` / `hip.min` destinations through the broadcast helper, share the variadic pairwise chain between Max and Min, and derive each hipBLASLt batch stride by comparing the operand's matrix count against the output's. Co-Authored-By: Claude Opus 5 Made-with: Cursor Co-authored-by: Cursor --- CLAUDE.md | 8 +- docs/design/hip-shape-inference.md | 85 ++++- docs/hip_dialect_intro.md | 6 +- include/hip/Dialect/IR/HipOps.td | 5 + include/hip/Dialect/IR/HipShapeUtils.h | 102 +++++- lib/Conversion/HipToLLVM/MatmulLowering.cpp | 69 ++-- lib/Conversion/OnnxToHip/ClipConversion.cpp | 20 +- lib/Conversion/OnnxToHip/MatMulConversion.cpp | 13 + lib/Conversion/OnnxToHip/MaxConversion.cpp | 54 +-- lib/Conversion/OnnxToHip/MinConversion.cpp | 54 +-- lib/Conversion/OnnxToHip/OnnxToHipUtils.h | 214 ++++++++++- .../OnnxToHip/ReduceMaxConversion.cpp | 26 +- .../OnnxToHip/ReduceMeanConversion.cpp | 27 +- .../OnnxToHip/ReduceMinConversion.cpp | 26 +- .../OnnxToHip/ReduceProdConversion.cpp | 158 +-------- .../OnnxToHip/ReduceSumConversion.cpp | 33 +- lib/Conversion/OnnxToHip/ReluConversion.cpp | 15 +- lib/Dialect/IR/HipDialect.cpp | 25 ++ lib/Dialect/IR/HipShapeUtils.cpp | 335 ++++++++++++------ .../Conversion/hip-to-llvm/test_matmul.mlir | 83 +++-- .../onnx-to-hip/test_reduce_sum.mlir | 45 +++ .../Dialect/hip-matmul-shape-verifier.mlir | 38 ++ 22 files changed, 883 insertions(+), 558 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index dc5a62576..b9e572940 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -140,11 +140,13 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co ### Result-shape agreement -- Converter destination construction and `reifyResultShapes` must use the same `OpFoldResult` shape helper for broadcast, Gemm, and MatMul. +- Converter destination construction, op verification, and `reifyResultShapes` must use the same `HipShapeUtils` shape rule for broadcast, Gemm, MatMul, and reductions. +- 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, so emitting IR is always the last step. - Fully dynamic broadcast uses `select(lhs == 1, rhs, lhs)`, not integer maximum: broadcasting extents 0 and 1 produces 0. - Rank-zero success is an empty shape carried by `FailureOr`; never use an empty vector as both success and failure. -- Variadic Max/Min derive every pairwise intermediate rank from the shared broadcast shape. -- MatMul uses the reified output batch product plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Per-axis partial batch broadcasting is rejected because the single-stride runtime cannot represent it. `wrap_hipblasLtMatmul` carries both strides; invalidate LLVM-IR artifacts compiled against the previous wrapper ABI. +- Variadic Max/Min share one pairwise-chain helper so every intermediate rank comes from the shared broadcast shape. +- Reductions resolve to one out-to-in dimension map. `keepdims = 0` makes the output dimension order non-positional in the input, so never copy input extents positionally when a reduced axis can precede a kept one. +- MatMul uses the reified output batch count plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Only *partial per-axis* batch broadcast is rejected; dynamic batch extents are supported and must not be rejected wholesale. `wrap_hipblasLtMatmul` carries both strides; invalidate LLVM-IR artifacts compiled against the previous wrapper ABI. - See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md). ### Allocation and memory planning diff --git a/docs/design/hip-shape-inference.md b/docs/design/hip-shape-inference.md index 3906dbddb..85d1b10bc 100644 --- a/docs/design/hip-shape-inference.md +++ b/docs/design/hip-shape-inference.md @@ -120,10 +120,26 @@ Reification is per result: `reifyResultShapes` returns one shape vector for ever ### Shared converter/reification shape helpers Converter destination construction and operation reification must not -independently implement the same shape category. Broadcast, Gemm, and MatMul -use helpers in `HipShapeUtils` that return `FailureOr>`. -The `FailureOr` is required because a valid rank-zero result has a successful -empty shape. +independently implement the same shape category. Broadcast, Gemm, MatMul, and +reductions use helpers in `HipShapeUtils` that return +`FailureOr>`. The `FailureOr` is required because a +valid rank-zero result has a successful empty shape. + +`HipShapeUtils` splits each category into a pure `infer*` function of static +shapes and a `reify*` function that may materialize index SSA. Every `reify*` +helper validates its preconditions through the matching `infer*` function +**before** it touches the builder, so a failure never leaves stray dimension ops +behind. Both the pattern-rewrite contract and +`ReifyRankedShapedTypeOpInterface` require the IR to be unchanged when a +rewrite or reification reports failure; upstream's +`ResolveShapedTypeResultDims` erases such stray ops explicitly, and this +codebase avoids creating them in the first place. Emitting IR is therefore the +last thing a helper does. + +The same split gives operation verifiers a shape rule to check against: +`MatmulOp::verify` and `GemmOp::verify` call `verifyHipOpShape` with +`inferMatmulShape` / `inferGemmShape`, so the `outs` shape, the converter +destination, and `reifyResultShapes` are all held to one shape function. Broadcast dimensions are right-aligned. Static 1 yields to the other side; equal non-unit static dimensions agree; dynamic/static-non-1 tightens to the @@ -144,17 +160,49 @@ converter materializes a constant index size. `--hip-infer-shapes` remains the single owner of later type narrowing, destination rebuilding, and cast barriers. -Variadic Max/Min derive every pairwise intermediate type from the shared -broadcast shape. Gemm derives M/N from A/B with transpose-aware indices and -checks optional C without using C as an extent source. MatMul broadcasts only -the leading batch slices, then appends M from A[-2] and N from B[-1]. - -The hipBLASLt MatMul lowering takes the batch product from the reified output -shape and carries independent A/B batch strides. A rank-2 or all-leading-one -operand uses stride 0, allowing the entire matrix to broadcast against a -batched operand on the other side. Per-axis partial batch broadcasting where -both operands contain fewer matrices than the output batch product is not yet -representable by one constant stride and is rejected by the MatMul verifier. +Variadic Max/Min share one `lowerVariadicBroadcastChain` helper that derives +every pairwise intermediate type from the shared broadcast shape. Gemm derives +M/N from A/B with transpose-aware indices and validates optional C without using +C as an extent source. MatMul broadcasts only the leading batch slices, then +appends M from A[-2] and N from B[-1]. + +Reductions resolve to one out-to-in dimension map, `computeReductionDimMap`, +which both `inferReductionShape` (static extents, used for destination types) +and `reifyReductionResultShape` (mixed extents, used for destination +construction and `reifyResultShapes`) consume. The mapping matters for +`keepdims = 0`, where dropping reduced axes makes the output dimension order +non-positional in the input: reducing axes `[1, 2]` of a rank-4 input maps +output dimension 1 to input dimension 3. A positional copy from the input is +correct only when no reduced axis precedes a kept one. When the reduced axes are +only known at runtime the mapping is data-dependent, so the converter falls back +to a positional copy and reification lifts the `outs` shape; both sides bail on +the same condition and therefore still agree. + +### 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 — +and is rejected by `verifyStridedBatchMatmul`, the MatMul verifier, and the +converter. The check is per axis: an operand is rejected only when some axis is +statically 1 while the output extent is not, *and* another axis carries batches. +Extents that are not statically 1 count as carrying batches, so an unknown +extent never hides a partial broadcast. Ordinary batched matmul with dynamic +leading extents (`[?, H, M, K] @ [?, H, K, N]`) is representable and accepted; a +blanket rejection of dynamic batch extents would fail that common layout. + +The lowering distinguishes the two strides by comparing the operand's matrix +count against the output batch count, folding at compile time when both are +static. Comparing against the output count rather than testing for "more than +one matrix" also keeps the read in bounds if a runtime-only shape ever violated +the invariant. + `wrap_hipblasLtMatmul` carries both strides. LLVM-IR artifacts compiled against the previous wrapper ABI must be invalidated. @@ -269,7 +317,9 @@ Primary regression coverage: | `test/lit/Dialect/hip-broadcast-reify-shapes.mlir` | Broadcast dynamic SSA, zero extents, and rank-zero success | | `test/lit/Dialect/hip-gemm-reify-shapes.mlir` | Transpose-aware Gemm M/N reification | | `test/lit/Dialect/hip-matmul-reify-shapes.mlir` | Per-op reification through `--resolve-shaped-type-result-dims` | -| `test/lit/Dialect/hip-matmul-shape-verifier.mlir` | Static MatMul shape validation | +| `test/lit/Dialect/hip-matmul-shape-verifier.mlir` | Static MatMul shape validation, including accepted dynamic batch layouts and rejected partial per-axis broadcast | +| `test/lit/Conversion/hip-to-llvm/test_matmul.mlir` | Per-operand batch strides: compile-time 0 / matrix size, and the runtime comparison against the output batch count | +| `test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir` | Reduction destinations, including the non-positional `keepdims = 0` dimension mapping | | `test/lit/Dialect/hip-loop-verifier.mlir` | Loop-carried type contract | | `test/lit/Dialect/hip-resolve-tensor-dims.mlir` | Production pre-bufferization dim folding | @@ -278,7 +328,8 @@ Complex operations may use dedicated files; common shape categories should exten ## Current limitations - ONNX MatMul rank-1 operands require promotion to rank 2 before constructing `hip.matmul`; the runtime and current verifier require rank at least 2. -- MatMul supports whole-matrix batch broadcast (one operand's leading product is 1) and equal flattened batch counts; per-axis partial batch broadcast remains unsupported by the strided-batch runtime. +- MatMul supports whole-matrix batch broadcast and one-matrix-per-output-batch operands, including dynamic batch extents; per-axis partial batch broadcast remains unsupported by the strided-batch runtime. +- Reduction destinations fall back to a positional copy from the input when the reduced axes are only known at runtime; reification mirrors that fallback rather than inventing a shape. - Converter migration to inferred-type builders is incremental; explicit result-type builders remain supported. - A future multi-result operation that needs custom `InferTypeOpInterface` logic may require a dedicated result-type inference implementation file. - Runtime-dependent extents without pre-execution SSA remain dynamic. diff --git a/docs/hip_dialect_intro.md b/docs/hip_dialect_intro.md index d799f7b06..424941139 100644 --- a/docs/hip_dialect_intro.md +++ b/docs/hip_dialect_intro.md @@ -48,8 +48,10 @@ Matrix multiplication backed by the hipBLASLt library (`hipblasLtMatmul`). The output batch count is the product of the broadcasted leading dimensions. Either operand may provide one matrix for all batches (`stride = 0`) or one -matrix per output batch. Per-axis partial batch broadcasting is rejected -because it cannot be represented by one constant stride per operand. +matrix per output batch, including when the leading extents are dynamic. Only +partial per-axis batch broadcasting -- some axes broadcast up while others carry +batches -- is rejected, because the resulting matrix count cannot be expressed +by one constant stride per operand. --- diff --git a/include/hip/Dialect/IR/HipOps.td b/include/hip/Dialect/IR/HipOps.td index be6288342..700174c8e 100644 --- a/include/hip/Dialect/IR/HipOps.td +++ b/include/hip/Dialect/IR/HipOps.td @@ -3474,6 +3474,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=*/[], diff --git a/include/hip/Dialect/IR/HipShapeUtils.h b/include/hip/Dialect/IR/HipShapeUtils.h index 263843d1b..7c2aa0095 100644 --- a/include/hip/Dialect/IR/HipShapeUtils.h +++ b/include/hip/Dialect/IR/HipShapeUtils.h @@ -14,9 +14,28 @@ #include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/SmallVector.h" +#include + namespace mlir { namespace hip { +//===----------------------------------------------------------------------===// +// Contract shared by every helper in this header +// +// The `infer*` helpers are pure functions of static shapes: they take no +// builder and emit no IR. The `reify*` helpers may materialize index SSA, and +// each one validates every precondition through its `infer*` counterpart +// BEFORE touching the builder. A `reify*` failure therefore never leaves +// stray ops behind, which both the pattern-rewrite contract and +// `ReifyRankedShapedTypeOpInterface` require (see +// `GreedyPatternRewriteDriver`'s expensive checks and +// `ResolveShapedTypeResultDims`). +// +// Converter destination construction and `reifyResultShapes` call the same +// helper for a given op, so the DPS `outs` shape and the shape observed by +// consumers cannot disagree. See `docs/design/hip-shape-inference.md`. +//===----------------------------------------------------------------------===// + /// Compute the shape of `A @ B` for matmul with NumPy-style batch broadcast /// over the leading dimensions. The matrix dimensions are `A[..., M, K]` and /// `B[..., K, N]`. @@ -29,13 +48,37 @@ inferMatmulShape(ArrayRef aShape, ArrayRef bShape, function_ref emitError); /// Verify that MatMul's broadcasted batches are representable by one constant -/// strided-batch offset per operand. Whole-matrix broadcast (rank 2 or all -/// leading extents 1) and equal static batch products are supported. Partial -/// per-axis broadcast and two nontrivial dynamic batch shapes are rejected. +/// strided-batch offset per operand. A stride can only express "one matrix +/// broadcast across every output batch" (stride 0) or "one matrix per output +/// batch" (stride == matrix size), so an operand is rejected only when it +/// provably needs something in between: a partial broadcast that pads some +/// batch axes up to the output extent while carrying batches on others. +/// +/// Extents that are not statically 1 count as carrying batches, so an unknown +/// extent never hides a partial broadcast. Ordinary batched matmul with +/// dynamic leading extents (`[?, H, M, K] @ [?, H, K, N]`) is representable +/// and accepted. LogicalResult verifyStridedBatchMatmul(ArrayRef aShape, ArrayRef bShape, function_ref emitError); +/// Compute the NumPy-broadcast result shape of `shapes` (right-aligned) from +/// static extents only. Folds `OpTrait::util::getBroadcastedShape` pairwise so +/// static broadcast validation is identical to the matmul batch path. +FailureOr> +inferBroadcastShape(ArrayRef> shapes, + function_ref emitError); + +/// Compute ONNX Gemm's rank-2 `{M, N}` result shape from static extents. +/// Validates that A and B are rank 2, that `transA`/`transB` are 0 or 1, that +/// the transpose-aware contraction extents agree, and that the optional C is +/// unidirectionally broadcastable onto `{M, N}`. `cShape` is `std::nullopt` +/// when C is absent; C never contributes M or N. +FailureOr> +inferGemmShape(ArrayRef aShape, ArrayRef bShape, + std::optional> cShape, int64_t transA, + int64_t transB, function_ref emitError); + /// Verify that the actual `outs` operand shapes of a DPS HIP op match the /// shapes returned by `computeExpected`. `op` must implement /// `DestinationStyleOpInterface`. @@ -168,21 +211,52 @@ SmallVector reifyGatherWithAxis(OpBuilder &b, Location loc, SmallVector reifyGatherND(OpBuilder &b, Location loc, Value data, Value indices, int64_t batchDims); +/// Sentinel in a reduction dim map: this output dimension is a reduced axis +/// retained by `keepdims=1`, so its extent is 1 rather than an input extent. +constexpr int64_t kReducedDim = -1; + +/// Map each output dimension of an ONNX reduction to the input dimension it +/// takes its extent from, or `kReducedDim` for a kept reduced axis. +/// +/// `axes` holds the already-resolved reduced axis indices (ONNX negative-axis +/// convention); an empty list means no reduction. `keepdims=0` drops reduced +/// axes from the output rank, so the output dimension order is *not* +/// positional in the input: +/// +/// data = tensor, axes = [1, 2], keepdims = 0 +/// -> map = [0, 3], i.e. output dim 1 comes from input dim 3, not dim 1. +/// +/// This mapping is the single source of truth behind `inferReductionShape` +/// (static extents, used for destination types) and +/// `reifyReductionResultShape` (mixed extents, used for destination +/// construction and `reifyResultShapes`), so the three can never disagree. +/// Returns failure when an axis is out of range for `dataRank`. +FailureOr> computeReductionDimMap(int64_t dataRank, + ArrayRef axes, + int64_t keepdims); + +/// Static ONNX reduction result shape: `computeReductionDimMap` applied to +/// `dataShape`, with kept reduced axes becoming 1. +FailureOr> inferReductionShape(ArrayRef dataShape, + ArrayRef axes, + int64_t keepdims); + +/// Mixed ONNX reduction result shape: `computeReductionDimMap` applied to +/// `data`, emitting `tensor.dim` only for dimensions that are dynamic in +/// `data`. `data` must be a `RankedTensorType`-typed Value. +FailureOr> +reifyReductionResultShape(OpBuilder &b, Location loc, Value data, + ArrayRef axes, int64_t keepdims); + /// Reify the result shape of a reduction op (reduce_sum / reduce_max / /// reduce_prod) given `data`, the `axes` operand (rank-1 i64 tensor), /// and the `keepdims` / `noop_with_empty_axes` attributes. /// -/// Tries to introspect `axes` as an `arith.constant` (the typical case -/// after the OnnxToHip converter materializes it from the ONNX -/// attribute). When successful: -/// - keepdims=1: axes-listed dims become `IndexAttr(1)`; non-axes -/// dims pass through from `data`. -/// - keepdims=0: axes-listed dims are dropped from the output rank; -/// non-axes dims pass through. -/// - Empty axes + noop_with_empty_axes=0: ALL dims become 1 -/// (keepdims=1) or output is rank-0 (keepdims=0). -/// - Empty axes + noop_with_empty_axes=1: output equals input -/// (no reduction). +/// Introspects `axes` as an `arith.constant` (the typical case after the +/// OnnxToHip converter materializes it from the ONNX attribute), resolves +/// ONNX's empty-axes semantics against `noop_with_empty_axes` — reduce every +/// axis when 0, reduce nothing when 1 — and delegates the shape rule to +/// `reifyReductionResultShape`. /// /// Returns `success()` and writes the reified dim list into `out` when /// `axes` can be introspected. Returns `failure()` when `axes` is not a diff --git a/lib/Conversion/HipToLLVM/MatmulLowering.cpp b/lib/Conversion/HipToLLVM/MatmulLowering.cpp index 3947c48ef..184202e30 100644 --- a/lib/Conversion/HipToLLVM/MatmulLowering.cpp +++ b/lib/Conversion/HipToLLVM/MatmulLowering.cpp @@ -7,6 +7,8 @@ #include "hip/Dialect/IR/HipShapeUtils.h" +#include "llvm/ADT/STLExtras.h" + namespace mlir { namespace hip { namespace { @@ -72,43 +74,40 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { unsigned elemBits = AType.getElementType().getIntOrFloatBitWidth(); Value elemSize = createI64Const(elemBits / 8); - // Each operand independently uses stride 0 when it contains one matrix, - // or its matrix size when it contains one matrix per output batch. - auto computeBatchStride = [&](MemRefType type, MemRefDescriptor desc, - Value matrixElements) -> Value { - int64_t rank = type.getRank(); - if (rank == 2) - return createI64Const(0); - - bool allLeadingStatic = true; - int64_t staticLeadingProduct = 1; - for (int64_t i : llvm::seq(0, rank - 2)) { - if (type.isDynamicDim(i)) { - allLeadingStatic = false; - break; - } - staticLeadingProduct *= type.getDimSize(i); - } - if (allLeadingStatic) - return staticLeadingProduct <= 1 ? createI64Const(0) : matrixElements; - - Value one = createI64Const(1); - Value leadingProduct = one; - for (int64_t i : llvm::seq(0, rank - 2)) { - Value dim = desc.size(rewriter, loc, i); - leadingProduct = - LLVM::MulOp::create(rewriter, loc, leadingProduct, dim).getRes(); - } - Value isBroadcast = LLVM::ICmpOp::create( - rewriter, loc, LLVM::ICmpPredicate::sle, leadingProduct, one); - return LLVM::SelectOp::create(rewriter, loc, isBroadcast, - createI64Const(0), matrixElements) + // `verifyStridedBatchMatmul` guarantees each operand holds either one + // matrix or one matrix per output batch, so one constant stride per operand + // is exact: 0 reuses a single matrix across every batch and the matrix size + // walks one matrix per batch. The two cases are distinguished by comparing + // the operand's matrix count against the output's, which also keeps the + // read in bounds if a runtime-only shape ever violated the invariant. + auto batchStride = + [&](MemRefType type, MemRefDescriptor desc, + llvm::function_ref matrixElements) -> Value { + ArrayRef batch = type.getShape().drop_back(2); + if (!ShapedType::isDynamicShape(batch)) + return llvm::product_of(batch) == 1 ? createI64Const(0) + : matrixElements(); + + Value count = createI64Const(1); + for (unsigned i : llvm::seq(0, batch.size())) + count = LLVM::MulOp::create(rewriter, loc, count, + desc.size(rewriter, loc, i)); + Value perBatch = LLVM::ICmpOp::create( + rewriter, loc, LLVM::ICmpPredicate::eq, count, batchCount); + // Materialize both arms into locals: passing calls that emit IR directly + // as arguments would leave the emission order unspecified. + Value perBatchStride = matrixElements(); + Value broadcastStride = createI64Const(0); + return LLVM::SelectOp::create(rewriter, loc, perBatch, perBatchStride, + broadcastStride) .getRes(); }; - Value aMatrixElements = LLVM::MulOp::create(rewriter, loc, M, K); - Value bMatrixElements = LLVM::MulOp::create(rewriter, loc, K, N); - Value aBatchStride = computeBatchStride(AType, ADesc, aMatrixElements); - Value bBatchStride = computeBatchStride(BType, BDesc, bMatrixElements); + Value aBatchStride = batchStride(AType, ADesc, [&] { + return LLVM::MulOp::create(rewriter, loc, M, K).getRes(); + }); + Value bBatchStride = batchStride(BType, BDesc, [&] { + return LLVM::MulOp::create(rewriter, loc, K, N).getRes(); + }); // Runtime signature: // int wrap_hipblasLtMatmul(RuntimeState* state, diff --git a/lib/Conversion/OnnxToHip/ClipConversion.cpp b/lib/Conversion/OnnxToHip/ClipConversion.cpp index fb402739d..d91cee4b9 100644 --- a/lib/Conversion/OnnxToHip/ClipConversion.cpp +++ b/lib/Conversion/OnnxToHip/ClipConversion.cpp @@ -82,17 +82,25 @@ struct ClipToHipMinMax : public mlir::RewritePattern { return mlir::failure(); mlir::Value context = *ctxOrFailure; + // `hip.max` / `hip.min` reify their result as the broadcast of both + // operands, so build each destination with the same helper. ONNX Clip + // bounds are scalars in practice, for which the broadcast collapses to + // `cur`'s own extents. mlir::Value cur = x; if (loV) { - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, cur); - cur = mlir::hip::MaxOp::create(rewriter, loc, resultType, context, cur, - loV, init) + mlir::FailureOr init = + createBroadcastEmptyTensor(rewriter, loc, resultType, {cur, loV}); + if (mlir::failed(init)) + return mlir::failure(); + cur = mlir::hip::MaxOp::create(rewriter, loc, context, cur, loV, *init) ->getResult(0); } if (hiV) { - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, cur); - cur = mlir::hip::MinOp::create(rewriter, loc, resultType, context, cur, - hiV, init) + mlir::FailureOr init = + createBroadcastEmptyTensor(rewriter, loc, resultType, {cur, hiV}); + if (mlir::failed(init)) + return mlir::failure(); + cur = mlir::hip::MinOp::create(rewriter, loc, context, cur, hiV, *init) ->getResult(0); } diff --git a/lib/Conversion/OnnxToHip/MatMulConversion.cpp b/lib/Conversion/OnnxToHip/MatMulConversion.cpp index 4d0c2fe72..1bb4410bb 100644 --- a/lib/Conversion/OnnxToHip/MatMulConversion.cpp +++ b/lib/Conversion/OnnxToHip/MatMulConversion.cpp @@ -43,6 +43,19 @@ MatMulToHip::matchAndRewrite(mlir::Operation *op, auto resultType = mlir::cast(op->getResult(0).getType()); + // The strided-batch runtime carries one constant stride per operand, so + // reject layouts it cannot express before any destination IR is emitted. + // This is a backend capability check rather than a shape rule, which is why + // it lives here and in `MatmulOp::verify` instead of in the shape helper. + auto aType = mlir::dyn_cast(a.getType()); + auto bType = mlir::dyn_cast(b.getType()); + if (!aType || !bType) + return rewriter.notifyMatchFailure(op, "MatMul operands must be ranked"); + if (mlir::failed(mlir::hip::verifyStridedBatchMatmul( + aType.getShape(), bType.getShape(), + [&]() { return op->emitError(); }))) + return mlir::failure(); + mlir::FailureOr> resultShape = mlir::hip::reifyMatmulResultShape(rewriter, loc, a, b, [&]() { return op->emitError(); }); diff --git a/lib/Conversion/OnnxToHip/MaxConversion.cpp b/lib/Conversion/OnnxToHip/MaxConversion.cpp index 1071d154e..374f8cb9c 100644 --- a/lib/Conversion/OnnxToHip/MaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/MaxConversion.cpp @@ -24,58 +24,10 @@ struct MaxToHip : public mlir::RewritePattern { mlir::LogicalResult matchAndRewrite(mlir::Operation *op, - mlir::PatternRewriter &rewriter) const override; -}; - -mlir::LogicalResult -MaxToHip::matchAndRewrite(mlir::Operation *op, - mlir::PatternRewriter &rewriter) const { - unsigned numInputs = op->getNumOperands(); - if (numInputs == 0) - return rewriter.notifyMatchFailure(op, "Max requires at least 1 input"); - - if (numInputs == 1) { - rewriter.replaceOp(op, op->getOperand(0)); - return mlir::success(); - } - - auto ctxOrFailure = getContextArg(op, rewriter); - if (mlir::failed(ctxOrFailure)) - return mlir::failure(); - mlir::Value context = *ctxOrFailure; - mlir::Location loc = op->getLoc(); - - auto resultType = - mlir::cast(op->getResult(0).getType()); - - mlir::Value accumulate = op->getOperand(0); - for (unsigned i : llvm::seq(1, numInputs)) { - mlir::Value rhs = op->getOperand(i); - mlir::FailureOr> stepShape = - mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, - [&]() { return op->emitError(); }); - if (mlir::failed(stepShape)) - return mlir::failure(); - - bool isFinal = i == numInputs - 1; - mlir::RankedTensorType stepResultType = - isFinal ? resultType - : getTensorTypeFromReifiedShape(*stepShape, - resultType.getElementType()); - mlir::FailureOr init = createEmptyTensorFromReifiedShape( - rewriter, loc, stepResultType, *stepShape); - if (mlir::failed(init)) - return rewriter.notifyMatchFailure( - op, "Max result type is incompatible with broadcast shape"); - - auto maxOp = mlir::hip::MaxOp::create(rewriter, loc, context, accumulate, - rhs, *init); - accumulate = maxOp->getResult(0); + mlir::PatternRewriter &rewriter) const override { + return lowerVariadicBroadcastChain(op, rewriter); } - - rewriter.replaceOp(op, accumulate); - return mlir::success(); -} +}; } // namespace diff --git a/lib/Conversion/OnnxToHip/MinConversion.cpp b/lib/Conversion/OnnxToHip/MinConversion.cpp index 83f4b1a45..4209eaf13 100644 --- a/lib/Conversion/OnnxToHip/MinConversion.cpp +++ b/lib/Conversion/OnnxToHip/MinConversion.cpp @@ -24,58 +24,10 @@ struct MinToHip : public mlir::RewritePattern { mlir::LogicalResult matchAndRewrite(mlir::Operation *op, - mlir::PatternRewriter &rewriter) const override; -}; - -mlir::LogicalResult -MinToHip::matchAndRewrite(mlir::Operation *op, - mlir::PatternRewriter &rewriter) const { - unsigned numInputs = op->getNumOperands(); - if (numInputs == 0) - return rewriter.notifyMatchFailure(op, "Min requires at least 1 input"); - - if (numInputs == 1) { - rewriter.replaceOp(op, op->getOperand(0)); - return mlir::success(); + mlir::PatternRewriter &rewriter) const override { + return lowerVariadicBroadcastChain(op, rewriter); } - - auto ctxOrFailure = getContextArg(op, rewriter); - if (mlir::failed(ctxOrFailure)) - return mlir::failure(); - mlir::Value context = *ctxOrFailure; - mlir::Location loc = op->getLoc(); - - auto resultType = - mlir::cast(op->getResult(0).getType()); - - mlir::Value accumulate = op->getOperand(0); - for (unsigned i : llvm::seq(1, numInputs)) { - mlir::Value rhs = op->getOperand(i); - mlir::FailureOr> stepShape = - mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, - [&]() { return op->emitError(); }); - if (mlir::failed(stepShape)) - return mlir::failure(); - - bool isFinal = i == numInputs - 1; - mlir::RankedTensorType stepResultType = - isFinal ? resultType - : getTensorTypeFromReifiedShape(*stepShape, - resultType.getElementType()); - mlir::FailureOr initOrFailure = - createEmptyTensorFromReifiedShape(rewriter, loc, stepResultType, - *stepShape); - if (mlir::failed(initOrFailure)) - return rewriter.notifyMatchFailure( - op, "Min result type is incompatible with broadcast shape"); - auto minOp = mlir::hip::MinOp::create(rewriter, loc, context, accumulate, - rhs, *initOrFailure); - accumulate = minOp->getResult(0); - } - - rewriter.replaceOp(op, accumulate); - return mlir::success(); -} +}; } // namespace diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h index 1342bf514..4af0f534b 100644 --- a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h @@ -99,6 +99,92 @@ inline mlir::Value createEmptyTensor(mlir::OpBuilder &builder, resultType.getElementType(), dynSizes); } +/// Recognize \p value as a compile-time constant integer vector, covering the +/// forms an ONNX attribute can still take once conversion runs: an inlined +/// `arith.constant`, a defining op carrying a `value` attribute, or a constant +/// that `lowerOnnxConstants` externalized into a `memref.global`. +/// +/// This must recognize at least everything the reification helpers can +/// introspect (an inlined `arith.constant`). A converter that saw *fewer* +/// constants than reification would build its destination from a weaker rule +/// than the shape consumers observe, and the two would disagree. +inline bool extractConstantIntVector(mlir::Value value, + llvm::SmallVectorImpl &out) { + out.clear(); + mlir::Operation *defOp = value.getDefiningOp(); + if (!defOp) + return false; + + mlir::DenseElementsAttr dense; + if (auto cst = mlir::dyn_cast(defOp)) + dense = mlir::dyn_cast(cst.getValue()); + if (!dense) + if (auto attr = defOp->getAttr("value")) + dense = mlir::dyn_cast(attr); + if (!dense) { + // Externalized constant: to_tensor(get_global) whose global still carries + // an initial value. + if (auto toTensor = mlir::dyn_cast(defOp)) + if (auto getGlobal = + toTensor.getBuffer().getDefiningOp()) + if (auto module = getGlobal->getParentOfType()) + if (auto global = module.lookupSymbol( + getGlobal.getNameAttr())) + dense = mlir::dyn_cast_or_null( + global.getInitialValueAttr()); + } + if (!dense) + return false; + + auto denseType = mlir::dyn_cast(dense.getType()); + if (!denseType || denseType.getRank() > 1) + return false; + mlir::Type elemType = denseType.getElementType(); + if (!elemType.isInteger(64) && !elemType.isInteger(32)) + return false; + for (mlir::APInt entry : dense.getValues()) + out.push_back(entry.getSExtValue()); + return true; +} + +/// Resolve the reduced axis list of an ONNX reduction op into \p axes. +/// +/// The axes arrive either as an `axes` attribute (opset < 13) or as an operand +/// (opset 13+) that may still be a compile-time constant. ONNX's empty-axes +/// semantics are applied here: with `noop_with_empty_axes = 0` an absent or +/// empty list reduces every axis, and with 1 it reduces nothing. +/// +/// Returns false when the axes are only known at runtime. The axis mapping is +/// then data-dependent and no shape rule applies, so both destination +/// construction and reification fall back to the `outs` shape. Deciding this +/// from one predicate is what keeps those two paths in agreement -- gating the +/// converter on the operand *count* while reification gates on the operand +/// being *constant* is exactly how they drift apart. +inline bool resolveReductionAxes(mlir::Operation *op, mlir::Value data, + int64_t noopWithEmptyAxes, + llvm::SmallVectorImpl &axes) { + axes.clear(); + bool hasAxesOperand = op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType()); + if (hasAxesOperand) { + if (!extractConstantIntVector(op->getOperand(1), axes)) + return false; + } else if (auto axesAttr = op->getAttrOfType("axes")) { + for (mlir::Attribute entry : axesAttr) + axes.push_back( + mlir::cast(entry).getValue().getSExtValue()); + } + + if (!axes.empty() || noopWithEmptyAxes != 0) + return true; + // Empty axes with noop_with_empty_axes = 0 reduces every axis. + auto dataType = mlir::dyn_cast(data.getType()); + if (!dataType) + return false; + axes = llvm::to_vector(llvm::seq(0, dataType.getRank())); + return true; +} + /// Resolve the ranked result type of an ONNX reduction op (ReduceMax / Sum / /// Mean / Prod / ...). /// @@ -113,9 +199,13 @@ inline mlir::Value createEmptyTensor(mlir::OpBuilder &builder, /// keepdims=1: reduced axes become size 1, other dims preserved. /// keepdims=0: reduced axes are dropped. /// +/// The shape rule itself is shared with destination construction and +/// `reifyResultShapes` through `mlir::hip::inferReductionShape`. +/// /// \p reducedAxes reduced axis indices (may be negative; normalized -/// here). For the all-axes default the caller passes -/// every axis; for a noop (empty axes) it passes none. +/// by the shared helper). For the all-axes default the +/// caller passes every axis; for a noop (empty axes) +/// it passes none. /// \p axesStaticallyKnown false when axes are only known at runtime, in which /// case an unranked result cannot be inferred. /// Returns failure only when the result is unranked AND cannot be inferred @@ -130,20 +220,12 @@ inferReduceResultType(mlir::Operation *op, mlir::Value data, auto inputType = mlir::dyn_cast(data.getType()); if (!inputType || !axesStaticallyKnown) return mlir::failure(); - int64_t rank = inputType.getRank(); - llvm::SmallVector reduced(rank, false); - for (int64_t a : reducedAxes) - reduced[a < 0 ? a + rank : a] = true; - llvm::SmallVector outShape; - for (int64_t i = 0; i < rank; ++i) { - if (reduced[i]) { - if (keepdims) - outShape.push_back(1); - } else { - outShape.push_back(inputType.getDimSize(i)); - } - } - return mlir::RankedTensorType::get(outShape, inputType.getElementType()); + mlir::FailureOr> outShape = + mlir::hip::inferReductionShape(inputType.getShape(), reducedAxes, + keepdims); + if (mlir::failed(outShape)) + return mlir::failure(); + return mlir::RankedTensorType::get(*outShape, inputType.getElementType()); } /// Build a tensor.empty with the imported result type and the dynamic sizes @@ -210,6 +292,43 @@ createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); } +/// Create a tensor.empty for the DPS init of an ONNX reduction op. +/// +/// When the reduced axes are known at compile time the extents come from +/// `reifyReductionResultShape` — the same helper that backs +/// `reifyResultShapes` — so the destination and the shape observed by +/// consumers implement one ONNX reduction shape function. This matters for +/// `keepdims=0`, where the output dimension order is not positional in the +/// input: +/// +/// Before (positional, wrong once a reduced axis precedes a kept one): +/// %d0 = tensor.dim %data, %c0 +/// %d1 = tensor.dim %data, %c1 +/// %init = tensor.empty(%d0, %d1) : tensor +/// After (axes = [1, 2], keepdims = 0 on a rank-4 input): +/// %d0 = tensor.dim %data, %c0 +/// %d3 = tensor.dim %data, %c3 +/// %init = tensor.empty(%d0, %d3) : tensor +/// +/// When the axes are only known at runtime the mapping is data-dependent and +/// no shape function applies, so the destination falls back to a positional +/// copy from `data`. `reifyReductionShape` bails on the same condition and +/// lifts the `outs` shape instead, so the two still agree. +inline mlir::FailureOr +createReductionEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, + mlir::RankedTensorType resultType, mlir::Value data, + llvm::ArrayRef reducedAxes, bool axesKnown, + int64_t keepdims) { + if (!axesKnown) + return createEmptyTensor(builder, loc, resultType, data); + mlir::FailureOr> shape = + mlir::hip::reifyReductionResultShape(builder, loc, data, reducedAxes, + keepdims); + if (mlir::failed(shape)) + return mlir::failure(); + return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); +} + /// Get !hip.context from function argument 0. Returns failure if the /// function has no arguments or the first argument is not !hip.context. inline mlir::FailureOr @@ -227,6 +346,69 @@ getContextArg(mlir::Operation *op, mlir::PatternRewriter &rewriter) { return ctx; } +/// Lower a variadic ONNX elementwise op to a left-associated chain of pairwise +/// broadcasting HIP ops: +/// +/// Max(a, b, c) -> hip.max(hip.max(a, b), c) +/// +/// Each step's destination comes from `reifyBroadcastResultShape`, the same +/// helper backing `Hip_DpsOp_Broadcast`'s `reifyResultShapes`, so no step's +/// `outs` shape can disagree with the shape its consumers observe. Intermediate +/// steps take the type of their own broadcast shape; only the final step has to +/// match the imported ONNX result type. A single operand is the identity. +template +mlir::LogicalResult +lowerVariadicBroadcastChain(mlir::Operation *op, + mlir::PatternRewriter &rewriter) { + llvm::StringRef opName = op->getName().getStringRef(); + unsigned numInputs = op->getNumOperands(); + if (numInputs == 0) + return rewriter.notifyMatchFailure(op, llvm::Twine(opName) + + " requires at least 1 input"); + + if (numInputs == 1) { + rewriter.replaceOp(op, op->getOperand(0)); + return mlir::success(); + } + + auto ctxOrFailure = getContextArg(op, rewriter); + if (mlir::failed(ctxOrFailure)) + return mlir::failure(); + mlir::Value context = *ctxOrFailure; + mlir::Location loc = op->getLoc(); + + auto resultType = + mlir::cast(op->getResult(0).getType()); + + mlir::Value accumulate = op->getOperand(0); + for (unsigned i : llvm::seq(1, numInputs)) { + mlir::Value rhs = op->getOperand(i); + mlir::FailureOr> stepShape = + mlir::hip::reifyBroadcastResultShape(rewriter, loc, {accumulate, rhs}, + [&]() { return op->emitError(); }); + if (mlir::failed(stepShape)) + return mlir::failure(); + + bool isFinal = i == numInputs - 1; + mlir::RankedTensorType stepResultType = + isFinal ? resultType + : getTensorTypeFromReifiedShape(*stepShape, + resultType.getElementType()); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, stepResultType, *stepShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, llvm::Twine(opName) + + " result type is incompatible with the broadcast shape"); + + accumulate = HipOpTy::create(rewriter, loc, context, accumulate, rhs, *init) + ->getResult(0); + } + + rewriter.replaceOp(op, accumulate); + return mlir::success(); +} + /// Build a hip.gqa op for the Whisper-MHA / Whisper-encoder-Attention paths. /// /// Emits one `hip.gqa` op with the 19-slot AttrSizedOperandSegments layout diff --git a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp index e62e3cc34..9e612bea1 100644 --- a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp @@ -45,19 +45,9 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, // Statically-known reduced axes (only when axes is NOT a runtime operand). // Used both to materialize the axes constant below and to infer the result // type when the ONNX importer left the result unranked (see below). - bool axesStaticallyKnown = op->getNumOperands() <= 1; llvm::SmallVector axesVec; - auto inputType = mlir::dyn_cast(data.getType()); - if (axesStaticallyKnown) { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - } else if (noopWithEmptyAxes == 0 && inputType) { - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - } - } + bool axesStaticallyKnown = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); // The ONNX importer can leave the ReduceMax result unranked (e.g. Phi's // pos_ids_reformat ReduceMax(position_ids) feeding GreaterOrEqual); infer a @@ -70,7 +60,11 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, data); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceMax result type is incompatible with the reduction shape"); mlir::Value axesOperand; if (op->getNumOperands() > 1) { @@ -86,9 +80,9 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopWithEmptyAxesAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); - auto hipOp = - mlir::hip::ReduceMaxOp::create(rewriter, loc, context, data, axesOperand, - init, keepdimsAttr, noopWithEmptyAxesAttr); + auto hipOp = mlir::hip::ReduceMaxOp::create(rewriter, loc, context, data, + axesOperand, *init, keepdimsAttr, + noopWithEmptyAxesAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp index c7bfc1f52..bf38e802c 100644 --- a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp @@ -62,22 +62,11 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, keepdims = keepdimsAttr.getSInt(); } - // Statically-known reduced axes (only when axes is NOT a runtime operand). - bool axesStaticallyKnown = op->getNumOperands() <= 1; - auto inputType = mlir::dyn_cast(data.getType()); + // Reduced axes, resolved from the attribute or a compile-time-constant + // operand. llvm::SmallVector axesVec; - if (axesStaticallyKnown) { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - } else if (noopWithEmptyAxes == 0 && inputType) { - // Default: reduce all axes (when noop_with_empty_axes is 0) - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - } - // noop_with_empty_axes == 1 with no axes -> axesVec stays empty (identity). - } + bool axesStaticallyKnown = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); // Resolve the result type (infer if the importer left it unranked). auto resultTypeOr = @@ -87,7 +76,11 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, op, "ReduceMean: cannot infer unranked result (need ranked input and " "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, data); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceMean result type is incompatible with the reduction shape"); // axes is always required in HIP dialect; create empty tensor<0xi64> when not // provided @@ -109,7 +102,7 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopWithEmptyAxesAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); auto hipOp = mlir::hip::ReduceMeanOp::create(rewriter, loc, context, data, - axesOperand, init, keepdimsAttr, + axesOperand, *init, keepdimsAttr, noopWithEmptyAxesAttr); rewriter.replaceOp(op, hipOp->getResult(0)); diff --git a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp index f0ff3dba6..438557954 100644 --- a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp @@ -42,19 +42,9 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, keepdims = keepdimsAttr.getSInt(); } - bool axesStaticallyKnown = op->getNumOperands() <= 1; llvm::SmallVector axesVec; - auto inputType = mlir::dyn_cast(data.getType()); - if (axesStaticallyKnown) { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - } else if (noopWithEmptyAxes == 0 && inputType) { - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - } - } + bool axesStaticallyKnown = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); auto resultTypeOr = inferReduceResultType(op, data, axesVec, axesStaticallyKnown, keepdims); @@ -64,7 +54,11 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, data); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceMin result type is incompatible with the reduction shape"); mlir::Value axesOperand; if (op->getNumOperands() > 1) { @@ -80,9 +74,9 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopWithEmptyAxesAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); - auto hipOp = - mlir::hip::ReduceMinOp::create(rewriter, loc, context, data, axesOperand, - init, keepdimsAttr, noopWithEmptyAxesAttr); + auto hipOp = mlir::hip::ReduceMinOp::create(rewriter, loc, context, data, + axesOperand, *init, keepdimsAttr, + noopWithEmptyAxesAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp b/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp index 293566c9a..e29b68918 100644 --- a/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp @@ -5,135 +5,10 @@ #include "OnnxToHipUtils.h" -#include "llvm/ADT/APInt.h" -#include "llvm/ADT/SmallSet.h" - namespace mlir { namespace hip { namespace { -/// Try to recognise \p v as a compile-time 1-D integer constant tensor. -/// Mirrors the helper in SliceConversion.cpp / PadConversion.cpp. -static mlir::DenseElementsAttr getCompileTimeConstantTensor(mlir::Value value) { - mlir::Operation *defOp = value.getDefiningOp(); - if (!defOp) - return nullptr; - if (auto cst = mlir::dyn_cast(defOp)) - return mlir::dyn_cast(cst.getValue()); - if (auto attr = defOp->getAttr("value")) - if (auto dense = mlir::dyn_cast(attr)) - return dense; - if (auto toTensor = mlir::dyn_cast(defOp)) { - auto bufDef = - toTensor.getBuffer().getDefiningOp(); - if (!bufDef) - return nullptr; - auto module = bufDef->getParentOfType(); - if (!module) - return nullptr; - auto global = - module.lookupSymbol(bufDef.getNameAttr()); - if (!global) - return nullptr; - return mlir::dyn_cast_or_null( - global.getInitialValueAttr()); - } - return nullptr; -} - -static mlir::LogicalResult -extractIntVector(mlir::Value v, llvm::SmallVectorImpl &out) { - if (!v) - return mlir::failure(); - auto dense = getCompileTimeConstantTensor(v); - if (!dense) - return mlir::failure(); - auto tensorType = mlir::dyn_cast(dense.getType()); - if (!tensorType || tensorType.getRank() != 1) - return mlir::failure(); - auto elemTy = tensorType.getElementType(); - if (!elemTy.isInteger(64) && !elemTy.isInteger(32)) - return mlir::failure(); - for (mlir::APInt entry : dense.getValues()) - out.push_back(entry.getSExtValue()); - return mlir::success(); -} - -/// Build the destination `tensor.empty` for ReduceProd. -/// -/// Output shape semantics: -/// * keepdims=1: out_rank == in_rank; out[i] = in[i] when i is NOT -/// reduced, otherwise 1. -/// * keepdims=0: out_rank == in_rank - #axes; out skips the reduced -/// axes; out[j] = in[non_reduced_axis_at_position_j]. -/// -/// For dynamic output dims we map back to the source `data` dim using -/// the known axes set (which must be a compile-time constant -- a dynamic -/// `axes` would make the mapping data-dependent). When `axes` is not -/// known, we fall back to positional alignment (correct only for -/// keepdims=1 and the all-reduce case). -static mlir::Value buildReduceProdInit(mlir::PatternRewriter &rewriter, - mlir::Location loc, - mlir::RankedTensorType resultType, - mlir::Value data, - llvm::ArrayRef axesVec, - bool axesKnown, int64_t keepdims) { - auto dataType = mlir::cast(data.getType()); - int64_t inRank = dataType.getRank(); - - // Build the lookup: for each output dim, which input dim is its source? - // Static -> reduced axis (resulting size 1) gets no entry. Identity- - // mapped output dim points back to the corresponding input dim. - llvm::SmallSet reducedAxes; - for (int64_t a : axesVec) { - if (a < 0) - a += inRank; - reducedAxes.insert(a); - } - - // Map outIdx -> inIdx (or -1 for "reduced -> size 1"). - llvm::SmallVector outToIn(resultType.getRank(), -1); - if (axesKnown) { - if (keepdims) { - for (int64_t i = 0; i < resultType.getRank(); ++i) - outToIn[i] = reducedAxes.contains(i) ? -1 : i; - } else { - int64_t outIdx = 0; - for (int64_t i = 0; i < inRank; ++i) { - if (reducedAxes.contains(i)) - continue; - if (outIdx < resultType.getRank()) - outToIn[outIdx] = i; - ++outIdx; - } - } - } else { - // Fallback: positional alignment. - for (int64_t i = 0; i < resultType.getRank(); ++i) - outToIn[i] = i < inRank ? i : -1; - } - - llvm::SmallVector dynSizes; - for (int64_t i = 0; i < resultType.getRank(); ++i) { - if (!resultType.isDynamicDim(i)) - continue; - int64_t inIdx = outToIn[i]; - if (inIdx < 0) { - // Reduced axis -- size 1. - dynSizes.push_back( - mlir::arith::ConstantIndexOp::create(rewriter, loc, 1)); - } else if (dataType.isDynamicDim(inIdx)) { - dynSizes.push_back( - mlir::tensor::DimOp::create(rewriter, loc, data, inIdx)); - } else { - dynSizes.push_back(mlir::arith::ConstantIndexOp::create( - rewriter, loc, dataType.getDimSize(inIdx))); - } - } - return mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); -} - /// onnx.ReduceProd -> hip.reduce_prod /// /// Mirrors the existing ReduceSum/ReduceMax conversion: lifts an optional @@ -170,31 +45,17 @@ ReduceProdToHip::matchAndRewrite(mlir::Operation *op, if (auto keepdimsAttr = op->getAttrOfType("keepdims")) keepdims = keepdimsAttr.getSInt(); - // Materialise the axes vector (and remember whether it is known at - // compile time) so we can build the destination tensor with the right - // dynamic dim sources even when the result is partially dynamic. + // Reduced axes, resolved from the attribute or a compile-time-constant + // operand, so the destination shape can map each output dimension back to the + // input dimension it comes from. llvm::SmallVector axesVec; - bool axesKnown = false; + bool axesKnown = resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + mlir::Value axesOperand; if (op->getNumOperands() > 1 && !mlir::isa(op->getOperand(1).getType())) { axesOperand = op->getOperand(1); - if (mlir::succeeded(extractIntVector(axesOperand, axesVec))) - axesKnown = true; } else { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - axesKnown = true; - } else if (noopWithEmptyAxes == 0) { - auto inputType = mlir::cast(data.getType()); - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - axesKnown = true; - } else { - axesKnown = true; // empty axes, noop - } auto axesType = mlir::RankedTensorType::get( {static_cast(axesVec.size())}, rewriter.getI64Type()); auto axesAttr = @@ -213,13 +74,16 @@ ReduceProdToHip::matchAndRewrite(mlir::Operation *op, "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = buildReduceProdInit(rewriter, loc, resultType, data, - axesVec, axesKnown, keepdims); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, axesVec, axesKnown, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceProd result type is incompatible with the reduction shape"); auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); auto hipOp = mlir::hip::ReduceProdOp::create( - rewriter, loc, context, data, axesOperand, init, keepdimsAttr, noopAttr); + rewriter, loc, context, data, axesOperand, *init, keepdimsAttr, noopAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp index a4ca17df9..47166da96 100644 --- a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp @@ -43,23 +43,12 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, keepdims = keepdimsAttr.getSInt(); } - // Statically-known reduced axes (only when axes is NOT a runtime operand). - // Used to materialize the axes constant and to infer an unranked result. - bool axesStaticallyKnown = op->getNumOperands() <= 1; - auto inputType = mlir::dyn_cast(data.getType()); + // Reduced axes, resolved from the attribute or a compile-time-constant + // operand. Drives the axes constant we materialize, the unranked-result + // inference, and the destination shape. llvm::SmallVector axesVec; - if (axesStaticallyKnown) { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - } else if (noopWithEmptyAxes == 0 && inputType) { - // Default: reduce all axes (when noop_with_empty_axes is 0) - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - } - // noop_with_empty_axes == 1 with no axes -> axesVec stays empty (identity). - } + bool axesStaticallyKnown = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); // Resolve the result type (infer if the importer left it unranked). auto resultTypeOr = @@ -69,7 +58,11 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, op, "ReduceSum: cannot infer unranked result (need ranked input and " "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, data); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceSum result type is incompatible with the reduction shape"); // axes is always required in HIP dialect; create empty tensor<0xi64> when not // provided @@ -90,9 +83,9 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, // Create hip.reduce_sum operation (axes always provided, may be empty) auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopWithEmptyAxesAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); - auto hipOp = - mlir::hip::ReduceSumOp::create(rewriter, loc, context, data, axesOperand, - init, keepdimsAttr, noopWithEmptyAxesAttr); + auto hipOp = mlir::hip::ReduceSumOp::create(rewriter, loc, context, data, + axesOperand, *init, keepdimsAttr, + noopWithEmptyAxesAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/ReluConversion.cpp b/lib/Conversion/OnnxToHip/ReluConversion.cpp index bf9e5709e..443ebd722 100644 --- a/lib/Conversion/OnnxToHip/ReluConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReluConversion.cpp @@ -83,13 +83,16 @@ struct ReluToHipMax : public mlir::RewritePattern { auto resultType = mlir::cast(op->getResult(0).getType()); - // Init tensor for DPS output. Use `x` as the shape source so dynamic - // dims of the result are tied to the same SSA values as `x` (the scalar - // zero has rank 0 and carries no dim info). - mlir::Value init = createEmptyTensor(rewriter, loc, resultType, x); + // `hip.max` reifies its result as the broadcast of both operands, so build + // the destination with the same helper. The rank-0 zero contributes no + // extent, leaving `x`'s own dims. + mlir::FailureOr init = + createBroadcastEmptyTensor(rewriter, loc, resultType, {x, zero}); + if (mlir::failed(init)) + return mlir::failure(); - auto maxOp = mlir::hip::MaxOp::create(rewriter, loc, resultType, context, x, - zero, init); + auto maxOp = + mlir::hip::MaxOp::create(rewriter, loc, context, x, zero, *init); rewriter.replaceOp(op, maxOp->getResult(0)); return mlir::success(); } diff --git a/lib/Dialect/IR/HipDialect.cpp b/lib/Dialect/IR/HipDialect.cpp index 73f04d5e3..3a1984906 100644 --- a/lib/Dialect/IR/HipDialect.cpp +++ b/lib/Dialect/IR/HipDialect.cpp @@ -1296,6 +1296,31 @@ void GemmOp::getEffects( emitDpsMemoryEffects(getDpsInputOperands(), getDpsInitsMutable(), effects); } +LogicalResult GemmOp::verify() { + SmallVector dataOperands = {getInputA(), getInputB()}; + if (getInputC()) + dataOperands.push_back(getInputC()); + dataOperands.push_back(getOutput()); + // The cross-cutting DPS contract first (all-tensor-or-all-memref + + // result-count parity); it also rules out non-shaped data operands, so the + // shape check below can rely on getShapeOf(). + if (failed(verifyDpsComputeOp(*this, dataOperands, /*numInits=*/1))) + return failure(); + + std::optional> cShape; + if (getInputC()) + cShape = getShapeOf(getInputC()); + return mlir::hip::verifyHipOpShape( + *this, [&]() -> SmallVector> { + FailureOr> outShape = mlir::hip::inferGemmShape( + getShapeOf(getInputA()), getShapeOf(getInputB()), cShape, + getTransA(), getTransB(), [&]() { return this->emitOpError(); }); + if (failed(outShape)) + return {}; + return {std::move(*outShape)}; + }); +} + //===----------------------------------------------------------------------===// // GqaOp: Full MS spec implementation // ins(query, [key, value, past_key, past_value], seqlens_k, diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index 917970bfa..15a36ce53 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -31,8 +31,7 @@ #include "llvm/Support/raw_ostream.h" #include -#include -#include +#include using namespace mlir; using namespace mlir::hip; @@ -51,7 +50,7 @@ std::string formatShape(ArrayRef shape) { os << d; }); os << "]"; - return os.str(); + return out; } FailureOr @@ -129,6 +128,38 @@ mlir::hip::inferMatmulShape(ArrayRef aShape, ArrayRef bShape, return result; } +namespace { + +/// Whether `batch` -- one operand's leading extents, right-aligned against +/// `outputBatch` -- holds a matrix count that one constant stride can express. +/// +/// A stride of 0 reuses a single matrix for every output batch and a stride of +/// the matrix size walks one matrix per output batch, so the operand's matrix +/// count must be either 1 or the output's. A partial broadcast lands strictly +/// between the two: `[2, 1]` against an output batch of `[2, 3]` holds 2 +/// matrices where the output needs 6, so neither stride is correct. +/// +/// Extents that are not statically 1 count as carrying batches and output +/// extents that are not statically 1 count as broadcast targets, so an unknown +/// extent is never assumed away. +bool isSingleStrideBatchLayout(ArrayRef batch, + ArrayRef outputBatch) { + size_t pad = outputBatch.size() - batch.size(); + bool broadcastsUp = false; + bool carriesBatches = false; + for (size_t i : llvm::seq(0, outputBatch.size())) { + // Axes below `pad` are the implicit leading ones of right-alignment. + int64_t extent = i < pad ? 1 : batch[i - pad]; + if (extent == 1) + broadcastsUp |= outputBatch[i] != 1; + else + carriesBatches = true; + } + return !(broadcastsUp && carriesBatches); +} + +} // namespace + LogicalResult mlir::hip::verifyStridedBatchMatmul( ArrayRef aShape, ArrayRef bShape, function_ref emitError) { @@ -146,34 +177,82 @@ LogicalResult mlir::hip::verifyStridedBatchMatmul( return failure(); } - auto isSingleMatrix = [](ArrayRef batch) { - return llvm::all_of(batch, [](int64_t dim) { return dim == 1; }); - }; - if (isSingleMatrix(aBatch) || isSingleMatrix(bBatch)) - return success(); - // With at most one batch axis, every valid broadcast has operand matrix - // counts in {1, output_count}, which one constant stride can represent. - if (std::max(aBatch.size(), bBatch.size()) <= 1) + if (isSingleStrideBatchLayout(aBatch, outputBatch) && + isSingleStrideBatchLayout(bBatch, outputBatch)) return success(); - auto isDynamic = [](int64_t dim) { return ShapedType::isDynamic(dim); }; - if (llvm::any_of(aBatch, isDynamic) || llvm::any_of(bBatch, isDynamic)) { - emitError() << "matmul dynamic batch layout is not representable by one " - "constant stride per operand: A.batch=" - << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); + emitError() << "matmul partial per-axis batch broadcast is not supported by " + "the strided-batch runtime: A.batch=" + << formatShape(aBatch) << " B.batch=" << formatShape(bBatch); + return failure(); +} + +FailureOr> +mlir::hip::inferBroadcastShape(ArrayRef> shapes, + function_ref emitError) { + if (shapes.empty()) { + emitError() << "broadcast requires at least one input shape"; return failure(); } - auto product = [](ArrayRef shape) { - return std::accumulate(shape.begin(), shape.end(), int64_t{1}, - std::multiplies()); - }; - int64_t outputCount = product(outputBatch); - if (product(aBatch) == outputCount && product(bBatch) == outputCount) - return success(); - emitError() << "matmul partial per-axis batch broadcast is not supported by " - "the strided-batch runtime"; - return failure(); + SmallVector result(shapes.front()); + for (ArrayRef shape : shapes.drop_front()) { + SmallVector merged; + if (!OpTrait::util::getBroadcastedShape(result, shape, merged)) { + emitError() << "incompatible broadcast shapes " << formatShape(result) + << " and " << formatShape(shape); + return failure(); + } + result = std::move(merged); + } + return result; +} + +FailureOr> +mlir::hip::inferGemmShape(ArrayRef aShape, ArrayRef bShape, + std::optional> cShape, + int64_t transA, int64_t transB, + function_ref emitError) { + if (aShape.size() != 2 || bShape.size() != 2) { + emitError() << "gemm A and B must be rank-2 tensors"; + return failure(); + } + if ((transA != 0 && transA != 1) || (transB != 0 && transB != 1)) { + emitError() << "gemm transA and transB must be 0 or 1"; + return failure(); + } + + // Contraction K must agree (kDynamic on either side is a wildcard). + int64_t aK = aShape[transA ? 0 : 1]; + int64_t bK = bShape[transB ? 1 : 0]; + if (!ShapedType::isDynamic(aK) && !ShapedType::isDynamic(bK) && aK != bK) { + emitError() << "gemm contraction dim mismatch: A has " << aK + << " but B has " << bK; + return failure(); + } + + SmallVector result = {aShape[transA ? 1 : 0], + bShape[transB ? 0 : 1]}; + if (!cShape) + return result; + + if (cShape->size() > 2) { + emitError() << "gemm C must be a ranked tensor of rank at most 2"; + return failure(); + } + // ONNX Gemm broadcasts C onto `{M, N}` unidirectionally: C never widens the + // result, so every C extent must be 1 or equal to the output extent. + size_t pad = result.size() - cShape->size(); + for (size_t i : llvm::seq(pad, result.size())) { + int64_t cDim = (*cShape)[i - pad]; + if (ShapedType::isDynamic(cDim) || ShapedType::isDynamic(result[i]) || + cDim == 1 || cDim == result[i]) + continue; + emitError() << "gemm C dimension " << cDim + << " is not broadcastable to output dimension " << result[i]; + return failure(); + } + return result; } LogicalResult mlir::hip::verifyHipOpShape( @@ -189,12 +268,16 @@ LogicalResult mlir::hip::verifyHipOpShape( return failure(); // Each helper returns one shape per DPS init by construction; assert in - // debug, fail-safe in release to avoid OOB on `expected[i]` below. + // debug, and in release emit a diagnostic rather than failing silently (a + // verifier that returns failure must say why) before `expected[i]` below + // could read out of bounds. auto inits = dpsOp.getDpsInits(); assert(expected.size() == inits.size() && "shape helper must produce one expected shape per DPS init operand"); if (expected.size() != inits.size()) - return failure(); + return op->emitOpError("shape helper produced ") + << expected.size() << " expected shapes for " << inits.size() + << " DPS init operands"; for (auto [i, init] : llvm::enumerate(inits)) { auto initType = dyn_cast(init.getType()); @@ -284,14 +367,25 @@ FailureOr> mlir::hip::reifyBroadcastResultShape( return failure(); } + SmallVector> staticShapes; + staticShapes.reserve(operands.size()); + for (Value operand : operands) { + auto operandType = dyn_cast(operand.getType()); + if (!operandType) { + emitError() << "broadcast operand must be a ranked tensor"; + return failure(); + } + staticShapes.push_back(operandType.getShape()); + } + // Validate broadcastability before emitting any `tensor.dim`, so a failure + // leaves the IR unchanged (see the contract in HipShapeUtils.h). + if (failed(inferBroadcastShape(staticShapes, emitError))) + return failure(); + SmallVector> shapes; shapes.reserve(operands.size()); for (size_t i : llvm::seq(0, operands.size())) { Value operand = operands[i]; - if (!isa(operand.getType())) { - emitError() << "broadcast operand must be a ranked tensor"; - return failure(); - } bool reused = false; // Reuse the first mixed shape for repeated SSA operands (e.g. x*x) so // broadcastDim sees identical OpFoldResults and emits no redundant @@ -320,11 +414,12 @@ FailureOr> mlir::hip::reifyMatmulResultShape( emitError() << "matmul operands must be ranked tensors"; return failure(); } + // Validate before emitting any `tensor.dim`, so a failure leaves the IR + // unchanged (see the contract in HipShapeUtils.h). Strided-batch + // representability is a backend capability rather than a shape rule, so it + // stays with `MatmulOp::verify` and the converter. if (failed(inferMatmulShape(aType.getShape(), bType.getShape(), emitError))) return failure(); - if (failed(verifyStridedBatchMatmul(aType.getShape(), bType.getShape(), - emitError))) - return failure(); SmallVector aSizes = tensor::getMixedSizes(b, loc, A); SmallVector bSizes = tensor::getMixedSizes(b, loc, B); @@ -348,53 +443,31 @@ mlir::hip::reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, function_ref emitError) { auto aType = dyn_cast(A.getType()); auto bType = dyn_cast(B.getType()); - if (!aType || !bType || aType.getRank() != 2 || bType.getRank() != 2) { - emitError() << "gemm A and B must be rank-2 tensors"; - return failure(); - } - if ((transA != 0 && transA != 1) || (transB != 0 && transB != 1)) { - emitError() << "gemm transA and transB must be 0 or 1"; - return failure(); - } - - int64_t aKDim = transA ? 0 : 1; - int64_t bKDim = transB ? 1 : 0; - int64_t aK = aType.getDimSize(aKDim); - int64_t bK = bType.getDimSize(bKDim); - if (!ShapedType::isDynamic(aK) && !ShapedType::isDynamic(bK) && aK != bK) { - emitError() << "gemm contraction dim mismatch: A has " << aK - << " but B has " << bK; + if (!aType || !bType) { + emitError() << "gemm A and B must be ranked tensors"; return failure(); } - - SmallVector aSizes = tensor::getMixedSizes(b, loc, A); - SmallVector bSizes = tensor::getMixedSizes(b, loc, B); - SmallVector result = {aSizes[transA ? 1 : 0], - bSizes[transB ? 0 : 1]}; - + std::optional> cShape; if (optionalC) { auto cType = dyn_cast(optionalC.getType()); - if (!cType || cType.getRank() > 2) { - emitError() << "gemm C must be a ranked tensor of rank at most 2"; + if (!cType) { + emitError() << "gemm C must be a ranked tensor"; return failure(); } - SmallVector cSizes = tensor::getMixedSizes(b, loc, optionalC); - size_t pad = result.size() - cSizes.size(); - for (size_t i : llvm::seq(0, result.size())) { - if (i < pad) - continue; - std::optional cStatic = getConstantIntValue(cSizes[i - pad]); - std::optional resultStatic = getConstantIntValue(result[i]); - if (cStatic && resultStatic && *cStatic != 1 && - *cStatic != *resultStatic) { - emitError() << "gemm C dimension " << *cStatic - << " is not broadcastable to output dimension " - << *resultStatic; - return failure(); - } - } + cShape = cType.getShape(); } - return result; + // Validate before emitting any `tensor.dim`, so a failure leaves the IR + // unchanged (see the contract in HipShapeUtils.h). + if (failed(inferGemmShape(aType.getShape(), bType.getShape(), cShape, transA, + transB, emitError))) + return failure(); + + // C is validated above but never contributes an extent: Gemm's result is + // exactly `{M, N}` from the transpose-aware A and B dimensions. + SmallVector aSizes = tensor::getMixedSizes(b, loc, A); + SmallVector bSizes = tensor::getMixedSizes(b, loc, B); + return SmallVector{aSizes[transA ? 1 : 0], + bSizes[transB ? 0 : 1]}; } SmallVector @@ -519,6 +592,70 @@ bool extractConstantInts(Value v, SmallVectorImpl &out) { } // namespace +FailureOr> +mlir::hip::computeReductionDimMap(int64_t dataRank, ArrayRef axes, + int64_t keepdims) { + // Normalize negative axes (ONNX convention) into a set. + llvm::SmallSet reduced; + for (int64_t axis : axes) { + if (axis < 0) + axis += dataRank; + if (axis < 0 || axis >= dataRank) + return failure(); + reduced.insert(axis); + } + + SmallVector dimMap; + dimMap.reserve(dataRank); + for (int64_t i : llvm::seq(0, dataRank)) { + if (!reduced.contains(i)) + dimMap.push_back(i); + else if (keepdims) + dimMap.push_back(kReducedDim); + // keepdims=0: the reduced axis leaves the output rank entirely. + } + return dimMap; +} + +FailureOr> +mlir::hip::inferReductionShape(ArrayRef dataShape, + ArrayRef axes, int64_t keepdims) { + FailureOr> dimMap = + computeReductionDimMap(dataShape.size(), axes, keepdims); + if (failed(dimMap)) + return failure(); + + SmallVector shape; + shape.reserve(dimMap->size()); + for (int64_t sourceDim : *dimMap) + shape.push_back(sourceDim == kReducedDim ? 1 : dataShape[sourceDim]); + return shape; +} + +FailureOr> +mlir::hip::reifyReductionResultShape(OpBuilder &b, Location loc, Value data, + ArrayRef axes, int64_t keepdims) { + auto dataType = dyn_cast(data.getType()); + if (!dataType) + return failure(); + // Validate before emitting any `tensor.dim`, so a failure leaves the IR + // unchanged (see the contract in HipShapeUtils.h). + FailureOr> dimMap = + computeReductionDimMap(dataType.getRank(), axes, keepdims); + if (failed(dimMap)) + return failure(); + + ArrayRef dataShape = dataType.getShape(); + SmallVector dims; + dims.reserve(dimMap->size()); + for (int64_t sourceDim : *dimMap) + dims.push_back(sourceDim == kReducedDim + ? OpFoldResult(b.getIndexAttr(1)) + : reifyDimOrConstant(b, loc, dataShape[sourceDim], data, + sourceDim)); + return dims; +} + LogicalResult mlir::hip::reifyReductionWithKeepdims( OpBuilder &b, Location loc, Value data, Value axes, int64_t keepdims, int64_t noopWithEmptyAxes, SmallVectorImpl &out) { @@ -526,8 +663,6 @@ LogicalResult mlir::hip::reifyReductionWithKeepdims( auto dataType = dyn_cast(data.getType()); if (!dataType) return failure(); - ArrayRef dataShape = dataType.getShape(); - int64_t dataRank = dataType.getRank(); // Axes operand: try to fold to a constant int vector. ONNX semantics // allow a size-0 vector to mean "no axes specified" — combined with the @@ -536,45 +671,17 @@ LogicalResult mlir::hip::reifyReductionWithKeepdims( SmallVector axesList; if (!extractConstantInts(axes, axesList)) return failure(); - - // Empty axes branch: "no axes specified" semantics. - if (axesList.empty()) { - if (noopWithEmptyAxes != 0) { - // No-op: output == data shape. - out.reserve(dataRank); - for (int64_t i : llvm::seq(0, dataRank)) - out.push_back(reifyDimOrConstant(b, loc, dataShape[i], data, i)); - return success(); - } - // Reduce all axes: every dim is reduced. - if (keepdims) { - out.append(dataRank, b.getIndexAttr(1)); - return success(); - } - // keepdims=0: output is rank-0 — `out` stays empty (a valid result). - return success(); + if (axesList.empty() && noopWithEmptyAxes == 0) { + // Reduce every axis. `noop_with_empty_axes = 1` instead means "reduce + // nothing", which the empty list already expresses. + axesList = llvm::to_vector(llvm::seq(0, dataType.getRank())); } - // Normalize negative axes (ONNX convention). - llvm::SmallSet reducedSet; - for (int64_t a : axesList) { - if (a < 0) - a += dataRank; - if (a < 0 || a >= dataRank) - return failure(); - reducedSet.insert(a); - } - - out.reserve(dataRank); - for (int64_t i : llvm::seq(0, dataRank)) { - if (reducedSet.contains(i)) { - if (keepdims) - out.push_back(b.getIndexAttr(1)); - // else: drop the dim from the output. - } else { - out.push_back(reifyDimOrConstant(b, loc, dataShape[i], data, i)); - } - } + FailureOr> dims = + reifyReductionResultShape(b, loc, data, axesList, keepdims); + if (failed(dims)) + return failure(); + out.assign(dims->begin(), dims->end()); return success(); } diff --git a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir index 6179f35bf..44fa9095f 100644 --- a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir +++ b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir @@ -3,11 +3,20 @@ // RUN: hip-mlir-opt --convert-hip-to-llvm %s | FileCheck %s -// ----- Rank-2 broadcast B [K, N] --------------------------------------------- -// b_batch_stride MUST be a constant 0 (no leading dims to multiply). -// `CHECK-NOT` between the label and the call asserts the compile-time -// constant path was taken (no runtime `llvm.icmp` / `llvm.select` over a -// leading-dim product). +// `wrap_hipblasLtMatmul` takes 12 arguments: +// - 4 pointers: state, A, B, output +// - 7 i64: M, N, K, batch_count, elem_size, A stride, B stride +// - 1 i32: op_state_slot (-1 here; --assign-op-state-slots is not in this RUN) +// +// Each operand's stride is 0 when it holds a single matrix broadcast across the +// output batches, and the matrix size when it holds one matrix per output batch. +// The two strides are the last arguments, so each case anchors on the +// `elem_size` constant (2 for f16) that immediately precedes them. + +// ----- Both operands hold a single matrix ------------------------------------ +// A's leading dim is statically 1 and B is rank 2, so both strides fold to a +// compile-time 0. `CHECK-NOT` asserts no runtime `llvm.icmp` / `llvm.select` +// over a matrix count was emitted. module { func.func @test_matmul_rank2_b(%ctx: !hip.context, @@ -25,22 +34,16 @@ module { // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 +// CHECK: %[[ELEM0:.*]] = llvm.mlir.constant(2 : i64) : i64 // CHECK: %[[A0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 // CHECK: %[[B0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A0_STRIDE]], %[[B0_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 -// Verify 12 parameters: -// - 4 pointers: state, A, B, output -// - 7 i64: M, N, K, batch_count, elem_size, A stride, B stride -// (B is rank-2 [K, N] = broadcast weight → stride = 0, compile-time const) -// - 1 i32: op_state_slot (-1 here — --assign-op-state-slots not run in this RUN) +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM0]], %[[A0_STRIDE]], %[[B0_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 // ----- Rank-3 leading-one B [1, K, N] ---------------------------------------- -// Leading dim is statically 1: still ONE [K, N] matrix in the buffer, so -// b_batch_stride must be 0 (NOT K*N). The compiler folds this at compile -// time when all leading dims are static — verified via CHECK-NOT for icmp / -// select. Encoding this case as "rank > 2 ⇒ per-batch" (the prior `b_batched` -// bool rule) caused hipBLASLt to step K*N elements past the end of the -// weight buffer on batch > 0 and feed garbage into the GEMM. +// B's leading dim is statically 1, so the buffer still holds ONE [K, N] matrix +// and its stride must be 0, not K*N. Treating "rank > 2" as per-batch made +// hipBLASLt step K*N elements past the end of the weight on batch > 0 and feed +// garbage into the GEMM. A carries two batches and gets stride M*K. module { func.func @test_matmul_rank3_leading_one_b(%ctx: !hip.context, @@ -58,15 +61,15 @@ module { // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 -// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[ELEM1:.*]] = llvm.mlir.constant(2 : i64) : i64 // CHECK: %[[A1_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 -// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 // CHECK: %[[B1_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A1_STRIDE]], %[[B1_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM1]], %[[A1_STRIDE]], %[[B1_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 -// ----- Rank-2 broadcast A against rank-3 B ------------------------------- -// A contains one matrix and therefore uses A stride 0. B and the output carry -// two batches. +// ----- Rank-2 broadcast A against rank-3 B ----------------------------------- +// The mirror image of the case above: A holds one matrix (stride 0) while B and +// the output carry two batches (stride K*N). Independent per-operand strides are +// what let either whole matrix broadcast across the other's batches. module { func.func @test_matmul_rank2_a(%ctx: !hip.context, @@ -84,8 +87,34 @@ module { // CHECK-NOT: llvm.icmp // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 -// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 -// CHECK: %{{.*}} = llvm.mul %{{.*}}, %{{.*}} : i64 -// CHECK: %[[B2_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: %[[ELEM2:.*]] = llvm.mlir.constant(2 : i64) : i64 // CHECK: %[[A2_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[A2_STRIDE]], %[[B2_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: %[[B2_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM2]], %[[A2_STRIDE]], %[[B2_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 + +// ----- Dynamic batch on both operands ---------------------------------------- +// `[?, 8, M, K] @ [?, 8, K, N]` is an ordinary batched matmul: neither operand +// broadcasts, so both are representable and must NOT be rejected. Because the +// leading extents are dynamic, each stride is chosen at runtime by comparing the +// operand's matrix count against the output's batch count -- which also keeps +// the read in bounds if the counts ever disagree. + +module { + func.func @test_matmul_dynamic_batch(%ctx: !hip.context, + %A: memref, + %B: memref, + %output: memref) { + hip.matmul(%ctx) + ins(%A, %B : memref, memref) + outs(%output : memref) + return + } +} + +// CHECK-LABEL: llvm.func @test_matmul_dynamic_batch +// CHECK: %[[ELEM3:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[A3_CMP:.*]] = llvm.icmp "eq" +// CHECK: %[[A3_STRIDE:.*]] = llvm.select %[[A3_CMP]] +// CHECK: %[[B3_CMP:.*]] = llvm.icmp "eq" +// CHECK: %[[B3_STRIDE:.*]] = llvm.select %[[B3_CMP]] +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM3]], %[[A3_STRIDE]], %[[B3_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 diff --git a/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir b/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir index e65fbd82e..4ed90572d 100644 --- a/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir @@ -68,4 +68,49 @@ module { // CHECK: %[[INIT:.*]] = tensor.empty(%{{.*}}, %{{.*}}) : tensor // CHECK: hip.reduce_sum(%[[CTX]]) ins(%[[DATA]], %[[AXES]] : tensor, tensor) outs(%[[INIT]] : tensor) {keepdims = 0 : i64} // CHECK-NOT: hip.alloc + + // keepdims = 0 with compile-time axes and dynamic extents. Dropping the + // reduced axes makes the output dimension order non-positional in the input: + // output dim 1 must come from `data` dim 3, not dim 1. The destination is + // built from the same shape helper that backs `reifyResultShapes`, so both + // agree on that mapping. + func.func @reduce_sum_no_keepdims_static_axes(%data: tensor) + -> tensor { + %output = "onnx.ReduceSum"(%data) + {axes = [1 : si64, 2 : si64], keepdims = 0 : si64, + noop_with_empty_axes = 0 : si64} + : (tensor) -> tensor + return %output : tensor + } + + // CHECK-LABEL: func.func @reduce_sum_no_keepdims_static_axes + // CHECK-SAME: (%[[CTX:.*]]: !hip.context, %[[DATA:.*]]: tensor) + // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index + // CHECK: %[[D0:.*]] = tensor.dim %[[DATA]], %[[C0]] : tensor + // CHECK: %[[D3:.*]] = tensor.dim %[[DATA]], %[[C3]] : tensor + // CHECK: %[[INIT:.*]] = tensor.empty(%[[D0]], %[[D3]]) : tensor + // CHECK: hip.reduce_sum(%[[CTX]]) ins(%[[DATA]], %{{.*}} : tensor, tensor<2xi64>) outs(%[[INIT]] : tensor) {keepdims = 0 : i64} + + // Same mapping when the axes arrive as a compile-time-constant OPERAND + // (opset 13+) rather than an attribute. Gating destination construction on the + // operand count instead of on the operand being constant would build a + // positional destination here while reification computed the real mapping. + func.func @reduce_sum_no_keepdims_constant_axes_operand( + %data: tensor) -> tensor { + %axes = arith.constant dense<[1, 2]> : tensor<2xi64> + %output = "onnx.ReduceSum"(%data, %axes) + {keepdims = 0 : si64, noop_with_empty_axes = 0 : si64} + : (tensor, tensor<2xi64>) -> tensor + return %output : tensor + } + + // CHECK-LABEL: func.func @reduce_sum_no_keepdims_constant_axes_operand + // CHECK-SAME: (%[[CTX:.*]]: !hip.context, %[[DATA:.*]]: tensor) + // CHECK-DAG: %[[OC0:.*]] = arith.constant 0 : index + // CHECK-DAG: %[[OC3:.*]] = arith.constant 3 : index + // CHECK: %[[OD0:.*]] = tensor.dim %[[DATA]], %[[OC0]] : tensor + // CHECK: %[[OD3:.*]] = tensor.dim %[[DATA]], %[[OC3]] : tensor + // CHECK: %[[OINIT:.*]] = tensor.empty(%[[OD0]], %[[OD3]]) : tensor + // CHECK: hip.reduce_sum(%[[CTX]]) ins(%[[DATA]], %{{.*}} : tensor, tensor<2xi64>) outs(%[[OINIT]] : tensor) {keepdims = 0 : i64} } diff --git a/test/lit/Dialect/hip-matmul-shape-verifier.mlir b/test/lit/Dialect/hip-matmul-shape-verifier.mlir index 8147d28bd..029de7382 100644 --- a/test/lit/Dialect/hip-matmul-shape-verifier.mlir +++ b/test/lit/Dialect/hip-matmul-shape-verifier.mlir @@ -59,6 +59,44 @@ func.func @matmul_dynamic_k(%ctx: !hip.context, // ----- +// An ordinary batched matmul with a dynamic leading extent and a second static +// batch axis on both operands: nothing broadcasts, so each operand holds one +// matrix per output batch and a single stride per operand is exact. Rejecting +// every dynamic batch extent with more than one batch axis would fail this +// legal `[?, H, M, K] @ [?, H, K, N]` layout. + +// CHECK-LABEL: func.func @matmul_dynamic_batch_two_axes +// CHECK: hip.matmul +func.func @matmul_dynamic_batch_two_axes(%ctx: !hip.context, + %a: memref, + %b: memref, + %c: memref) { + hip.matmul(%ctx) + ins(%a, %b : memref, memref) + outs(%c : memref) + return +} + +// ----- + +// Whole-matrix broadcast of A across B's batches, with a dynamic batch extent. +// A's batch extents are all statically 1, so A uses stride 0 regardless of what +// the dynamic output batch turns out to be. + +// CHECK-LABEL: func.func @matmul_dynamic_batch_broadcast_a +// CHECK: hip.matmul +func.func @matmul_dynamic_batch_broadcast_a(%ctx: !hip.context, + %a: memref<1x1x4x16xf16, 1>, + %b: memref, + %c: memref) { + hip.matmul(%ctx) + ins(%a, %b : memref<1x1x4x16xf16, 1>, memref) + outs(%c : memref) + return +} + +// ----- + func.func @matmul_k_mismatch(%ctx: !hip.context, %a: memref<2x4xf16, 1>, %b: memref<8x16xf16, 1>, From 454059a7ef0f184a8a43ff45c64159a3d9f4de17 Mon Sep 17 00:00:00 2001 From: fhanuman Date: Fri, 31 Jul 2026 09:10:25 -0600 Subject: [PATCH 6/9] refactor(hip): tighten shape-helper APIs to upstream MLIR conventions Replace the sentinel error channel in `verifyHipOpShape`, where an empty outer vector meant "the callee already diagnosed", with a `FailureOr` returning helper. Both callers are single-destination ops, so the speculative one-shape-per-init generality goes away too and a verifier now pairs directly with its `infer*` rule instead of restating an adaptor lambda. Express "not known at compile time" as `std::optional` rather than a parallel `bool`: `resolveReductionAxes` returns an optional view over caller-owned storage, and `inferReduceResultType` / `createReductionEmptyTensor` take it directly, so no caller can pass axes that contradict a flag. Narrow the public surface to shape rules. The reduction dimension map and the static broadcast fold have no callers outside `HipShapeUtils.cpp` and become file-static; the map now uses `llvm::SmallBitVector` over the closed axis domain and `std::optional` in place of a `-1` sentinel, which was ambiguous next to ONNX negative-axis indices. Move the non-template destination helpers out of the header that all 87 conversion units include, and assert the right-alignment precondition whose unsigned subtraction would otherwise wrap silently. Co-Authored-By: Claude Opus 5 Made-with: Cursor Co-authored-by: Cursor --- CLAUDE.md | 2 + docs/design/hip-shape-inference.md | 49 +++-- include/hip/Dialect/IR/HipShapeUtils.h | 71 +++---- lib/Conversion/OnnxToHip/CMakeLists.txt | 1 + lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp | 168 ++++++++++++++++ lib/Conversion/OnnxToHip/OnnxToHipUtils.h | 184 ++++-------------- .../OnnxToHip/ReduceMaxConversion.cpp | 17 +- .../OnnxToHip/ReduceMeanConversion.cpp | 17 +- .../OnnxToHip/ReduceMinConversion.cpp | 17 +- .../OnnxToHip/ReduceProdConversion.cpp | 19 +- .../OnnxToHip/ReduceSumConversion.cpp | 17 +- lib/Dialect/IR/HipDialect.cpp | 27 +-- lib/Dialect/IR/HipShapeUtils.cpp | 146 +++++++------- .../Dialect/hip-matmul-shape-verifier.mlir | 6 +- 14 files changed, 401 insertions(+), 340 deletions(-) create mode 100644 lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp diff --git a/CLAUDE.md b/CLAUDE.md index b9e572940..5928ca7c2 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -142,6 +142,8 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co - Converter destination construction, op verification, and `reifyResultShapes` must use the same `HipShapeUtils` shape rule for broadcast, Gemm, MatMul, and reductions. - 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, so emitting IR is always the last step. +- 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. - Fully dynamic broadcast uses `select(lhs == 1, rhs, lhs)`, not integer maximum: broadcasting extents 0 and 1 produces 0. - Rank-zero success is an empty shape carried by `FailureOr`; never use an empty vector as both success and failure. - Variadic Max/Min share one pairwise-chain helper so every intermediate rank comes from the shared broadcast shape. diff --git a/docs/design/hip-shape-inference.md b/docs/design/hip-shape-inference.md index 85d1b10bc..7cfdac36f 100644 --- a/docs/design/hip-shape-inference.md +++ b/docs/design/hip-shape-inference.md @@ -136,10 +136,24 @@ rewrite or reification reports failure; upstream's codebase avoids creating them in the first place. Emitting IR is therefore the last thing a helper does. -The same split gives operation verifiers a shape rule to check against: -`MatmulOp::verify` and `GemmOp::verify` call `verifyHipOpShape` with -`inferMatmulShape` / `inferGemmShape`, so the `outs` shape, the converter -destination, and `reifyResultShapes` are all held to one shape function. +The same split gives operation verifiers a shape rule to check against. +`verifyHipOpShape` takes an `infer*` helper directly, so a single-destination op +pairs the two in one expression and the `outs` shape, the converter destination, +and `reifyResultShapes` are all held to one shape function: + +```c++ +LogicalResult MatmulOp::verify() { + // ... DPS contract first ... + return verifyHipOpShape(*this, [&] { + return inferMatmulShape(aShape, bShape, + [&] { return this->emitOpError(); }); + }); +} +``` + +Only the `infer*` helpers are public. The dimension mappings and static +broadcast folds they are built from stay internal to `HipShapeUtils.cpp`, so the +header exposes shape *rules* rather than the machinery behind them. Broadcast dimensions are right-aligned. Static 1 yields to the other side; equal non-unit static dimensions agree; dynamic/static-non-1 tightens to the @@ -166,17 +180,22 @@ M/N from A/B with transpose-aware indices and validates optional C without using C as an extent source. MatMul broadcasts only the leading batch slices, then appends M from A[-2] and N from B[-1]. -Reductions resolve to one out-to-in dimension map, `computeReductionDimMap`, -which both `inferReductionShape` (static extents, used for destination types) -and `reifyReductionResultShape` (mixed extents, used for destination -construction and `reifyResultShapes`) consume. The mapping matters for -`keepdims = 0`, where dropping reduced axes makes the output dimension order -non-positional in the input: reducing axes `[1, 2]` of a rank-4 input maps -output dimension 1 to input dimension 3. A positional copy from the input is -correct only when no reduced axis precedes a kept one. When the reduced axes are -only known at runtime the mapping is data-dependent, so the converter falls back -to a positional copy and reification lifts the `outs` shape; both sides bail on -the same condition and therefore still agree. +Reductions resolve to one internal out-to-in dimension map, consumed by both +`inferReductionShape` (static extents, used for destination types) and +`reifyReductionResultShape` (mixed extents, used for destination construction +and `reifyResultShapes`). The mapping matters for `keepdims = 0`, where dropping +reduced axes makes the output dimension order non-positional in the input: +reducing axes `[1, 2]` of a rank-4 input maps output dimension 1 to input +dimension 3. A positional copy from the input is correct only when no reduced +axis precedes a kept one. + +Whether the axes are usable at all is decided once, by +`resolveReductionAxes`, which returns `std::nullopt` when they are only known at +runtime. Both the destination and reification key off that single answer: the +converter falls back to a positional copy and reification lifts the `outs` +shape. Deciding it twice is how the two drift apart — gating the converter on +the axes operand *count* while reification gates on the operand being +*constant* leaves opset-13+ constant axes handled inconsistently. ### MatMul strided-batch representability diff --git a/include/hip/Dialect/IR/HipShapeUtils.h b/include/hip/Dialect/IR/HipShapeUtils.h index 7c2aa0095..aec94fa5f 100644 --- a/include/hip/Dialect/IR/HipShapeUtils.h +++ b/include/hip/Dialect/IR/HipShapeUtils.h @@ -62,13 +62,6 @@ LogicalResult verifyStridedBatchMatmul(ArrayRef aShape, ArrayRef bShape, function_ref emitError); -/// Compute the NumPy-broadcast result shape of `shapes` (right-aligned) from -/// static extents only. Folds `OpTrait::util::getBroadcastedShape` pairwise so -/// static broadcast validation is identical to the matmul batch path. -FailureOr> -inferBroadcastShape(ArrayRef> shapes, - function_ref emitError); - /// Compute ONNX Gemm's rank-2 `{M, N}` result shape from static extents. /// Validates that A and B are rank 2, that `transA`/`transB` are 0 or 1, that /// the transpose-aware contraction extents agree, and that the optional C is @@ -79,22 +72,25 @@ inferGemmShape(ArrayRef aShape, ArrayRef bShape, std::optional> cShape, int64_t transA, int64_t transB, function_ref emitError); -/// Verify that the actual `outs` operand shapes of a DPS HIP op match the -/// shapes returned by `computeExpected`. `op` must implement -/// `DestinationStyleOpInterface`. +/// Verify that the `outs` shape of a single-destination DPS HIP op matches the +/// shape returned by `inferShape`. `op` must implement +/// `DestinationStyleOpInterface` and have exactly one DPS init. +/// +/// `inferShape` is invoked once and returns `failure()` when the underlying +/// `infer*` helper already emitted a diagnostic, which this function +/// propagates without re-emitting. Pair it directly with an `infer*` helper: /// -/// `computeExpected` is invoked once and must return one shape per init -/// operand (== one per `OpResult` for tensor mode; same count for memref -/// mode, just no SSA result). An empty outer vector signals that the -/// shape-arithmetic helper already emitted a diagnostic — this function -/// returns `failure()` without re-emitting. +/// return verifyHipOpShape(*this, [&] { +/// return inferMatmulShape(aShape, bShape, +/// [&] { return this->emitOpError(); }); +/// }); /// /// Element-type checks are intentionally not handled here: dtype-changing /// ops (cast, equal, less, not, and) keep their own element-type checks in /// their op-local verifiers. -LogicalResult verifyHipOpShape( - Operation *op, - function_ref>()> computeExpected); +LogicalResult +verifyHipOpShape(Operation *op, + function_ref>()> inferShape); /// Build an `OpFoldResult` for one dimension of a reify-callable op's /// result: @@ -211,39 +207,26 @@ SmallVector reifyGatherWithAxis(OpBuilder &b, Location loc, SmallVector reifyGatherND(OpBuilder &b, Location loc, Value data, Value indices, int64_t batchDims); -/// Sentinel in a reduction dim map: this output dimension is a reduced axis -/// retained by `keepdims=1`, so its extent is 1 rather than an input extent. -constexpr int64_t kReducedDim = -1; - -/// Map each output dimension of an ONNX reduction to the input dimension it -/// takes its extent from, or `kReducedDim` for a kept reduced axis. +/// ONNX reduction result shape over `axes`, from static extents only. /// /// `axes` holds the already-resolved reduced axis indices (ONNX negative-axis -/// convention); an empty list means no reduction. `keepdims=0` drops reduced -/// axes from the output rank, so the output dimension order is *not* -/// positional in the input: +/// convention); an empty list means no reduction. `keepdims = 0` drops reduced +/// axes from the output rank, so the output dimension order is *not* positional +/// in the input: reducing axes `[1, 2]` of a rank-4 input maps output dimension +/// 1 to input dimension 3. /// -/// data = tensor, axes = [1, 2], keepdims = 0 -/// -> map = [0, 3], i.e. output dim 1 comes from input dim 3, not dim 1. -/// -/// This mapping is the single source of truth behind `inferReductionShape` -/// (static extents, used for destination types) and -/// `reifyReductionResultShape` (mixed extents, used for destination -/// construction and `reifyResultShapes`), so the three can never disagree. -/// Returns failure when an axis is out of range for `dataRank`. -FailureOr> computeReductionDimMap(int64_t dataRank, - ArrayRef axes, - int64_t keepdims); - -/// Static ONNX reduction result shape: `computeReductionDimMap` applied to -/// `dataShape`, with kept reduced axes becoming 1. +/// This and `reifyReductionResultShape` share one internal output-to-input +/// dimension mapping, so a destination built from either cannot disagree with +/// the shape `reifyResultShapes` reports. Returns failure when an axis is out +/// of range for `dataShape`. FailureOr> inferReductionShape(ArrayRef dataShape, ArrayRef axes, int64_t keepdims); -/// Mixed ONNX reduction result shape: `computeReductionDimMap` applied to -/// `data`, emitting `tensor.dim` only for dimensions that are dynamic in -/// `data`. `data` must be a `RankedTensorType`-typed Value. +/// ONNX reduction result shape over `axes` as mixed extents, emitting +/// `tensor.dim` only for dimensions that are dynamic in `data`. Same mapping +/// as `inferReductionShape`; see it for the `keepdims` semantics. `data` must +/// be a `RankedTensorType`-typed Value. FailureOr> reifyReductionResultShape(OpBuilder &b, Location loc, Value data, ArrayRef axes, int64_t keepdims); diff --git a/lib/Conversion/OnnxToHip/CMakeLists.txt b/lib/Conversion/OnnxToHip/CMakeLists.txt index 318d42bcf..a4a0ccdc8 100644 --- a/lib/Conversion/OnnxToHip/CMakeLists.txt +++ b/lib/Conversion/OnnxToHip/CMakeLists.txt @@ -5,6 +5,7 @@ add_library(OnnxToHip STATIC OnnxToHip.cpp + OnnxToHipUtils.cpp SimplifyOnnx.cpp MatMulConversion.cpp TransposeConversion.cpp diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp b/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp new file mode 100644 index 000000000..9c4fae2fe --- /dev/null +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp @@ -0,0 +1,168 @@ +/* + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. + * Licensed under the MIT License. + */ +//===- OnnxToHipUtils.cpp - Shared helpers for ONNX-to-HIP patterns ------===// +// +// Out-of-line definitions for the destination-construction helpers declared in +// `OnnxToHipUtils.h`. That header is included by every per-operator conversion +// file, so only the pattern-facing templates stay inline there. +// +// See `docs/design/hip-shape-inference.md` for how these helpers relate to the +// shared `HipShapeUtils` shape rules and to `reifyResultShapes`. +// +//===----------------------------------------------------------------------===// + +#include "OnnxToHipUtils.h" + +namespace mlir { +namespace hip { + +bool extractConstantIntVector(mlir::Value value, + llvm::SmallVectorImpl &out) { + out.clear(); + mlir::Operation *defOp = value.getDefiningOp(); + if (!defOp) + return false; + + mlir::DenseElementsAttr dense; + if (auto cst = mlir::dyn_cast(defOp)) + dense = mlir::dyn_cast(cst.getValue()); + if (!dense) + if (auto attr = defOp->getAttr("value")) + dense = mlir::dyn_cast(attr); + if (!dense) { + // Externalized constant: to_tensor(get_global) whose global still carries + // an initial value. + if (auto toTensor = mlir::dyn_cast(defOp)) + if (auto getGlobal = + toTensor.getBuffer().getDefiningOp()) + if (auto module = getGlobal->getParentOfType()) + if (auto global = module.lookupSymbol( + getGlobal.getNameAttr())) + dense = mlir::dyn_cast_or_null( + global.getInitialValueAttr()); + } + if (!dense) + return false; + + auto denseType = mlir::dyn_cast(dense.getType()); + if (!denseType || denseType.getRank() > 1) + return false; + mlir::Type elemType = denseType.getElementType(); + if (!elemType.isInteger(64) && !elemType.isInteger(32)) + return false; + for (mlir::APInt entry : dense.getValues()) + out.push_back(entry.getSExtValue()); + return true; +} + +std::optional> +resolveReductionAxes(mlir::Operation *op, mlir::Value data, + int64_t noopWithEmptyAxes, + llvm::SmallVectorImpl &storage) { + storage.clear(); + bool hasAxesOperand = op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType()); + if (hasAxesOperand) { + if (!extractConstantIntVector(op->getOperand(1), storage)) + return std::nullopt; + } else if (auto axesAttr = op->getAttrOfType("axes")) { + for (mlir::Attribute entry : axesAttr) + storage.push_back( + mlir::cast(entry).getValue().getSExtValue()); + } + + if (storage.empty() && noopWithEmptyAxes == 0) { + // Empty axes with noop_with_empty_axes = 0 reduces every axis. + auto dataType = mlir::dyn_cast(data.getType()); + if (!dataType) + return std::nullopt; + llvm::append_range(storage, llvm::seq(0, dataType.getRank())); + } + return llvm::ArrayRef(storage); +} + +mlir::FailureOr +inferReduceResultType(mlir::Operation *op, mlir::Value data, + std::optional> reducedAxes, + int64_t keepdims) { + if (auto ranked = + mlir::dyn_cast(op->getResult(0).getType())) + return ranked; + auto inputType = mlir::dyn_cast(data.getType()); + if (!inputType || !reducedAxes) + return mlir::failure(); + mlir::FailureOr> outShape = + mlir::hip::inferReductionShape(inputType.getShape(), *reducedAxes, + keepdims); + if (mlir::failed(outShape)) + return mlir::failure(); + return mlir::RankedTensorType::get(*outShape, inputType.getElementType()); +} + +mlir::FailureOr createEmptyTensorFromReifiedShape( + mlir::OpBuilder &builder, mlir::Location loc, + mlir::RankedTensorType resultType, + llvm::ArrayRef reifiedShape) { + if (static_cast(reifiedShape.size()) != resultType.getRank()) + return mlir::failure(); + + llvm::SmallVector dynSizes; + for (int64_t dimIdx : llvm::seq(resultType.getRank())) { + std::optional reifiedStatic = + mlir::getConstantIntValue(reifiedShape[dimIdx]); + if (!resultType.isDynamicDim(dimIdx)) { + if (reifiedStatic && *reifiedStatic != resultType.getDimSize(dimIdx)) + return mlir::failure(); + continue; + } + dynSizes.push_back(mlir::getValueOrCreateConstantIndexOp( + builder, loc, reifiedShape[dimIdx])); + } + return mlir::Value( + mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), + resultType.getElementType(), dynSizes)); +} + +mlir::RankedTensorType +getTensorTypeFromReifiedShape(llvm::ArrayRef reifiedShape, + mlir::Type elementType, + mlir::Attribute encoding) { + llvm::SmallVector shape; + shape.reserve(reifiedShape.size()); + for (mlir::OpFoldResult dim : reifiedShape) + shape.push_back( + mlir::getConstantIntValue(dim).value_or(mlir::ShapedType::kDynamic)); + return mlir::RankedTensorType::get(shape, elementType, encoding); +} + +mlir::FailureOr +createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, + mlir::RankedTensorType resultType, + mlir::ValueRange operands) { + mlir::FailureOr> shape = + mlir::hip::reifyBroadcastResultShape( + builder, loc, operands, [&] { return mlir::emitError(loc); }); + if (mlir::failed(shape)) + return mlir::failure(); + return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); +} + +mlir::FailureOr +createReductionEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, + mlir::RankedTensorType resultType, mlir::Value data, + std::optional> reducedAxes, + int64_t keepdims) { + if (!reducedAxes) + return createEmptyTensor(builder, loc, resultType, data); + mlir::FailureOr> shape = + mlir::hip::reifyReductionResultShape(builder, loc, data, *reducedAxes, + keepdims); + if (mlir::failed(shape)) + return mlir::failure(); + return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); +} + +} // namespace hip +} // namespace mlir diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h index 4af0f534b..676661e71 100644 --- a/lib/Conversion/OnnxToHip/OnnxToHipUtils.h +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.h @@ -108,82 +108,29 @@ inline mlir::Value createEmptyTensor(mlir::OpBuilder &builder, /// introspect (an inlined `arith.constant`). A converter that saw *fewer* /// constants than reification would build its destination from a weaker rule /// than the shape consumers observe, and the two would disagree. -inline bool extractConstantIntVector(mlir::Value value, - llvm::SmallVectorImpl &out) { - out.clear(); - mlir::Operation *defOp = value.getDefiningOp(); - if (!defOp) - return false; - - mlir::DenseElementsAttr dense; - if (auto cst = mlir::dyn_cast(defOp)) - dense = mlir::dyn_cast(cst.getValue()); - if (!dense) - if (auto attr = defOp->getAttr("value")) - dense = mlir::dyn_cast(attr); - if (!dense) { - // Externalized constant: to_tensor(get_global) whose global still carries - // an initial value. - if (auto toTensor = mlir::dyn_cast(defOp)) - if (auto getGlobal = - toTensor.getBuffer().getDefiningOp()) - if (auto module = getGlobal->getParentOfType()) - if (auto global = module.lookupSymbol( - getGlobal.getNameAttr())) - dense = mlir::dyn_cast_or_null( - global.getInitialValueAttr()); - } - if (!dense) - return false; - - auto denseType = mlir::dyn_cast(dense.getType()); - if (!denseType || denseType.getRank() > 1) - return false; - mlir::Type elemType = denseType.getElementType(); - if (!elemType.isInteger(64) && !elemType.isInteger(32)) - return false; - for (mlir::APInt entry : dense.getValues()) - out.push_back(entry.getSExtValue()); - return true; -} +bool extractConstantIntVector(mlir::Value value, + llvm::SmallVectorImpl &out); -/// Resolve the reduced axis list of an ONNX reduction op into \p axes. +/// Resolve the reduced axis list of an ONNX reduction op. /// /// The axes arrive either as an `axes` attribute (opset < 13) or as an operand /// (opset 13+) that may still be a compile-time constant. ONNX's empty-axes /// semantics are applied here: with `noop_with_empty_axes = 0` an absent or /// empty list reduces every axis, and with 1 it reduces nothing. /// -/// Returns false when the axes are only known at runtime. The axis mapping is -/// then data-dependent and no shape rule applies, so both destination -/// construction and reification fall back to the `outs` shape. Deciding this -/// from one predicate is what keeps those two paths in agreement -- gating the -/// converter on the operand *count* while reification gates on the operand -/// being *constant* is exactly how they drift apart. -inline bool resolveReductionAxes(mlir::Operation *op, mlir::Value data, - int64_t noopWithEmptyAxes, - llvm::SmallVectorImpl &axes) { - axes.clear(); - bool hasAxesOperand = op->getNumOperands() > 1 && - !mlir::isa(op->getOperand(1).getType()); - if (hasAxesOperand) { - if (!extractConstantIntVector(op->getOperand(1), axes)) - return false; - } else if (auto axesAttr = op->getAttrOfType("axes")) { - for (mlir::Attribute entry : axesAttr) - axes.push_back( - mlir::cast(entry).getValue().getSExtValue()); - } - - if (!axes.empty() || noopWithEmptyAxes != 0) - return true; - // Empty axes with noop_with_empty_axes = 0 reduces every axis. - auto dataType = mlir::dyn_cast(data.getType()); - if (!dataType) - return false; - axes = llvm::to_vector(llvm::seq(0, dataType.getRank())); - return true; -} +/// Returns `std::nullopt` when the axes are only known at runtime. The axis +/// mapping is then data-dependent and no shape rule applies, so both +/// destination construction and reification fall back to the `outs` shape. +/// Deciding this from one predicate is what keeps those two paths in agreement +/// -- gating the converter on the operand *count* while reification gates on +/// the operand being *constant* is exactly how they drift apart. +/// +/// \p storage is caller-owned scratch that backs the returned view; it must +/// outlive the result and must not be modified while the view is in use. +std::optional> +resolveReductionAxes(mlir::Operation *op, mlir::Value data, + int64_t noopWithEmptyAxes, + llvm::SmallVectorImpl &storage); /// Resolve the ranked result type of an ONNX reduction op (ReduceMax / Sum / /// Mean / Prod / ...). @@ -202,73 +149,32 @@ inline bool resolveReductionAxes(mlir::Operation *op, mlir::Value data, /// The shape rule itself is shared with destination construction and /// `reifyResultShapes` through `mlir::hip::inferReductionShape`. /// -/// \p reducedAxes reduced axis indices (may be negative; normalized -/// by the shared helper). For the all-axes default the -/// caller passes every axis; for a noop (empty axes) -/// it passes none. -/// \p axesStaticallyKnown false when axes are only known at runtime, in which -/// case an unranked result cannot be inferred. +/// \p reducedAxes reduced axis indices (may be negative; normalized by the +/// shared helper), or `std::nullopt` when the axes are only +/// known at runtime, in which case an unranked result cannot be +/// inferred. /// Returns failure only when the result is unranked AND cannot be inferred /// (unranked/absent input type, or runtime-only axes). -inline mlir::FailureOr +mlir::FailureOr inferReduceResultType(mlir::Operation *op, mlir::Value data, - llvm::ArrayRef reducedAxes, - bool axesStaticallyKnown, int64_t keepdims) { - if (auto ranked = - mlir::dyn_cast(op->getResult(0).getType())) - return ranked; - auto inputType = mlir::dyn_cast(data.getType()); - if (!inputType || !axesStaticallyKnown) - return mlir::failure(); - mlir::FailureOr> outShape = - mlir::hip::inferReductionShape(inputType.getShape(), reducedAxes, - keepdims); - if (mlir::failed(outShape)) - return mlir::failure(); - return mlir::RankedTensorType::get(*outShape, inputType.getElementType()); -} + std::optional> reducedAxes, + int64_t keepdims); /// Build a tensor.empty with the imported result type and the dynamic sizes /// described by `reifiedShape`. Static reified dimensions are materialized as /// constant index operands when the imported type keeps that dimension /// dynamic; the existing `hip-infer-shapes` pass owns later type refinement. -inline mlir::FailureOr createEmptyTensorFromReifiedShape( +mlir::FailureOr createEmptyTensorFromReifiedShape( mlir::OpBuilder &builder, mlir::Location loc, mlir::RankedTensorType resultType, - llvm::ArrayRef reifiedShape) { - if (static_cast(reifiedShape.size()) != resultType.getRank()) - return mlir::failure(); - - llvm::SmallVector dynSizes; - for (int64_t dimIdx : llvm::seq(resultType.getRank())) { - std::optional reifiedStatic = - mlir::getConstantIntValue(reifiedShape[dimIdx]); - if (!resultType.isDynamicDim(dimIdx)) { - if (reifiedStatic && *reifiedStatic != resultType.getDimSize(dimIdx)) - return mlir::failure(); - continue; - } - dynSizes.push_back(mlir::getValueOrCreateConstantIndexOp( - builder, loc, reifiedShape[dimIdx])); - } - return mlir::Value( - mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), - resultType.getElementType(), dynSizes)); -} + llvm::ArrayRef reifiedShape); /// Derive a tensor type for a synthesized intermediate from a reified shape. /// Constant dimensions become static; all other dimensions stay dynamic. -inline mlir::RankedTensorType +mlir::RankedTensorType getTensorTypeFromReifiedShape(llvm::ArrayRef reifiedShape, mlir::Type elementType, - mlir::Attribute encoding = {}) { - llvm::SmallVector shape; - shape.reserve(reifiedShape.size()); - for (mlir::OpFoldResult dim : reifiedShape) - shape.push_back( - mlir::getConstantIntValue(dim).value_or(mlir::ShapedType::kDynamic)); - return mlir::RankedTensorType::get(shape, elementType, encoding); -} + mlir::Attribute encoding = {}); /// Create a tensor.empty for a DPS init whose shape is the NumPy-style /// broadcast of \p operands. Converter destination construction delegates to @@ -280,21 +186,14 @@ getTensorTypeFromReifiedShape(llvm::ArrayRef reifiedShape, /// %lhs_is_one = arith.cmpi eq, %lhs_dim, %c1 : index /// %extent = arith.select %lhs_is_one, %rhs_dim, %lhs_dim : index /// %init = tensor.empty(%extent) : tensor -inline mlir::FailureOr +mlir::FailureOr createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, mlir::RankedTensorType resultType, - mlir::ValueRange operands) { - mlir::FailureOr> shape = - mlir::hip::reifyBroadcastResultShape( - builder, loc, operands, [&]() { return mlir::emitError(loc); }); - if (mlir::failed(shape)) - return mlir::failure(); - return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); -} + mlir::ValueRange operands); /// Create a tensor.empty for the DPS init of an ONNX reduction op. /// -/// When the reduced axes are known at compile time the extents come from +/// When \p reducedAxes carries a value the extents come from /// `reifyReductionResultShape` — the same helper that backs /// `reifyResultShapes` — so the destination and the shape observed by /// consumers implement one ONNX reduction shape function. This matters for @@ -310,24 +209,15 @@ createBroadcastEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, /// %d3 = tensor.dim %data, %c3 /// %init = tensor.empty(%d0, %d3) : tensor /// -/// When the axes are only known at runtime the mapping is data-dependent and -/// no shape function applies, so the destination falls back to a positional -/// copy from `data`. `reifyReductionShape` bails on the same condition and -/// lifts the `outs` shape instead, so the two still agree. -inline mlir::FailureOr +/// `std::nullopt` means the axes are only known at runtime: the mapping is then +/// data-dependent and no shape function applies, so the destination falls back +/// to a positional copy from `data`. `reifyReductionShape` bails on the same +/// condition and lifts the `outs` shape instead, so the two still agree. +mlir::FailureOr createReductionEmptyTensor(mlir::OpBuilder &builder, mlir::Location loc, mlir::RankedTensorType resultType, mlir::Value data, - llvm::ArrayRef reducedAxes, bool axesKnown, - int64_t keepdims) { - if (!axesKnown) - return createEmptyTensor(builder, loc, resultType, data); - mlir::FailureOr> shape = - mlir::hip::reifyReductionResultShape(builder, loc, data, reducedAxes, - keepdims); - if (mlir::failed(shape)) - return mlir::failure(); - return createEmptyTensorFromReifiedShape(builder, loc, resultType, *shape); -} + std::optional> reducedAxes, + int64_t keepdims); /// Get !hip.context from function argument 0. Returns failure if the /// function has no arguments or the first argument is not !hip.context. diff --git a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp index 9e612bea1..febb3ff95 100644 --- a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp @@ -45,15 +45,14 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, // Statically-known reduced axes (only when axes is NOT a runtime operand). // Used both to materialize the axes constant below and to infer the result // type when the ONNX importer left the result unranked (see below). - llvm::SmallVector axesVec; - bool axesStaticallyKnown = - resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); // The ONNX importer can leave the ReduceMax result unranked (e.g. Phi's // pos_ids_reformat ReduceMax(position_ids) feeding GreaterOrEqual); infer a // ranked result type in that case (see inferReduceResultType). - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesStaticallyKnown, keepdims); + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceMax: cannot infer unranked result (need ranked input and " @@ -61,7 +60,7 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, mlir::RankedTensorType resultType = *resultTypeOr; mlir::FailureOr init = createReductionEmptyTensor( - rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + rewriter, loc, resultType, data, reducedAxes, keepdims); if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "ReduceMax result type is incompatible with the reduction shape"); @@ -71,9 +70,9 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, axesOperand = op->getOperand(1); } else { auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } diff --git a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp index bf38e802c..f7c347755 100644 --- a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp @@ -64,20 +64,19 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, // Reduced axes, resolved from the attribute or a compile-time-constant // operand. - llvm::SmallVector axesVec; - bool axesStaticallyKnown = - resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); // Resolve the result type (infer if the importer left it unranked). - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesStaticallyKnown, keepdims); + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceMean: cannot infer unranked result (need ranked input and " "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; mlir::FailureOr init = createReductionEmptyTensor( - rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + rewriter, loc, resultType, data, reducedAxes, keepdims); if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "ReduceMean result type is incompatible with the reduction shape"); @@ -91,9 +90,9 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, } else { // Create constant tensor for axes auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } diff --git a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp index 438557954..61aabb3fc 100644 --- a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp @@ -42,12 +42,11 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, keepdims = keepdimsAttr.getSInt(); } - llvm::SmallVector axesVec; - bool axesStaticallyKnown = - resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesStaticallyKnown, keepdims); + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceMin: cannot infer unranked result (need ranked input and " @@ -55,7 +54,7 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, mlir::RankedTensorType resultType = *resultTypeOr; mlir::FailureOr init = createReductionEmptyTensor( - rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + rewriter, loc, resultType, data, reducedAxes, keepdims); if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "ReduceMin result type is incompatible with the reduction shape"); @@ -65,9 +64,9 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, axesOperand = op->getOperand(1); } else { auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } diff --git a/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp b/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp index e29b68918..e39be24f2 100644 --- a/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceProdConversion.cpp @@ -48,8 +48,9 @@ ReduceProdToHip::matchAndRewrite(mlir::Operation *op, // Reduced axes, resolved from the attribute or a compile-time-constant // operand, so the destination shape can map each output dimension back to the // input dimension it comes from. - llvm::SmallVector axesVec; - bool axesKnown = resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); mlir::Value axesOperand; if (op->getNumOperands() > 1 && @@ -57,17 +58,15 @@ ReduceProdToHip::matchAndRewrite(mlir::Operation *op, axesOperand = op->getOperand(1); } else { auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } - // Resolve the result type (infer if the importer left it unranked). axesKnown - // also covers the compile-time-constant axes-operand case extracted above. - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesKnown, keepdims); + // Resolve the result type (infer if the importer left it unranked). + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceProd: cannot infer unranked result (need ranked input and " @@ -75,7 +74,7 @@ ReduceProdToHip::matchAndRewrite(mlir::Operation *op, mlir::RankedTensorType resultType = *resultTypeOr; mlir::FailureOr init = createReductionEmptyTensor( - rewriter, loc, resultType, data, axesVec, axesKnown, keepdims); + rewriter, loc, resultType, data, reducedAxes, keepdims); if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "ReduceProd result type is incompatible with the reduction shape"); diff --git a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp index 47166da96..f9d1e3317 100644 --- a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp @@ -46,20 +46,19 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, // Reduced axes, resolved from the attribute or a compile-time-constant // operand. Drives the axes constant we materialize, the unranked-result // inference, and the destination shape. - llvm::SmallVector axesVec; - bool axesStaticallyKnown = - resolveReductionAxes(op, data, noopWithEmptyAxes, axesVec); + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); // Resolve the result type (infer if the importer left it unranked). - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesStaticallyKnown, keepdims); + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceSum: cannot infer unranked result (need ranked input and " "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; mlir::FailureOr init = createReductionEmptyTensor( - rewriter, loc, resultType, data, axesVec, axesStaticallyKnown, keepdims); + rewriter, loc, resultType, data, reducedAxes, keepdims); if (mlir::failed(init)) return rewriter.notifyMatchFailure( op, "ReduceSum result type is incompatible with the reduction shape"); @@ -73,9 +72,9 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, } else { // Create constant tensor for axes auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } diff --git a/lib/Dialect/IR/HipDialect.cpp b/lib/Dialect/IR/HipDialect.cpp index 3a1984906..2efa89489 100644 --- a/lib/Dialect/IR/HipDialect.cpp +++ b/lib/Dialect/IR/HipDialect.cpp @@ -569,15 +569,10 @@ LogicalResult MatmulOp::verify() { ArrayRef aShape = getShapeOf(getA()); ArrayRef bShape = getShapeOf(getB()); - if (failed(mlir::hip::verifyHipOpShape( - *this, [&]() -> SmallVector> { - FailureOr> outShape = - mlir::hip::inferMatmulShape( - aShape, bShape, [&]() { return this->emitOpError(); }); - if (failed(outShape)) - return {}; - return {std::move(*outShape)}; - }))) + if (failed(mlir::hip::verifyHipOpShape(*this, [&] { + return mlir::hip::inferMatmulShape(aShape, bShape, + [&] { return this->emitOpError(); }); + }))) return failure(); return mlir::hip::verifyStridedBatchMatmul( aShape, bShape, [&]() { return this->emitOpError(); }); @@ -1310,15 +1305,11 @@ LogicalResult GemmOp::verify() { std::optional> cShape; if (getInputC()) cShape = getShapeOf(getInputC()); - return mlir::hip::verifyHipOpShape( - *this, [&]() -> SmallVector> { - FailureOr> outShape = mlir::hip::inferGemmShape( - getShapeOf(getInputA()), getShapeOf(getInputB()), cShape, - getTransA(), getTransB(), [&]() { return this->emitOpError(); }); - if (failed(outShape)) - return {}; - return {std::move(*outShape)}; - }); + return mlir::hip::verifyHipOpShape(*this, [&] { + return mlir::hip::inferGemmShape( + getShapeOf(getInputA()), getShapeOf(getInputB()), cShape, getTransA(), + getTransB(), [&] { return this->emitOpError(); }); + }); } //===----------------------------------------------------------------------===// diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index 15a36ce53..a2c0e31db 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -27,7 +27,7 @@ #include "mlir/Interfaces/DestinationStyleOpInterface.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/Sequence.h" -#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/SmallBitVector.h" #include "llvm/Support/raw_ostream.h" #include @@ -144,6 +144,10 @@ namespace { /// extent is never assumed away. bool isSingleStrideBatchLayout(ArrayRef batch, ArrayRef outputBatch) { + // Guaranteed by getBroadcastedShape, and load-bearing: the unsigned `pad` + // below would wrap and silently report "representable" if it were violated. + assert(batch.size() <= outputBatch.size() && + "operand batch cannot outrank the broadcasted output batch"); size_t pad = outputBatch.size() - batch.size(); bool broadcastsUp = false; bool carriesBatches = false; @@ -187,9 +191,14 @@ LogicalResult mlir::hip::verifyStridedBatchMatmul( return failure(); } +namespace { + +/// NumPy-broadcast result shape of `shapes` (right-aligned) from static extents +/// only. Folds `OpTrait::util::getBroadcastedShape` pairwise so static +/// broadcast validation is identical to the matmul batch path. FailureOr> -mlir::hip::inferBroadcastShape(ArrayRef> shapes, - function_ref emitError) { +inferBroadcastShape(ArrayRef> shapes, + function_ref emitError) { if (shapes.empty()) { emitError() << "broadcast requires at least one input shape"; return failure(); @@ -208,6 +217,8 @@ mlir::hip::inferBroadcastShape(ArrayRef> shapes, return result; } +} // namespace + FailureOr> mlir::hip::inferGemmShape(ArrayRef aShape, ArrayRef bShape, std::optional> cShape, @@ -256,54 +267,44 @@ mlir::hip::inferGemmShape(ArrayRef aShape, ArrayRef bShape, } LogicalResult mlir::hip::verifyHipOpShape( - Operation *op, - function_ref>()> computeExpected) { + Operation *op, function_ref>()> inferShape) { // Asserting cast: every op wired to verifyHipOpShape also implements DPS // via TableGen; a missing interface is a programmer error in the op def. auto dpsOp = cast(op); - - // Empty outer vector means the shape helper already emitted a diagnostic. - SmallVector> expected = computeExpected(); - if (expected.empty()) - return failure(); - - // Each helper returns one shape per DPS init by construction; assert in - // debug, and in release emit a diagnostic rather than failing silently (a - // verifier that returns failure must say why) before `expected[i]` below - // could read out of bounds. auto inits = dpsOp.getDpsInits(); - assert(expected.size() == inits.size() && - "shape helper must produce one expected shape per DPS init operand"); - if (expected.size() != inits.size()) - return op->emitOpError("shape helper produced ") - << expected.size() << " expected shapes for " << inits.size() - << " DPS init operands"; - - for (auto [i, init] : llvm::enumerate(inits)) { - auto initType = dyn_cast(init.getType()); - if (!initType) - return op->emitOpError("init #") << i << " is not a shaped type"; - ArrayRef actualShape = initType.getShape(); - ArrayRef expShape = expected[i]; - if (actualShape.size() != expShape.size()) - return op->emitOpError("rank mismatch on result #") - << i << ": expected rank " << expShape.size() << " " - << formatShape(expShape) << " but outs has rank " - << actualShape.size() << " " << formatShape(actualShape); - for (size_t d : llvm::seq(0, actualShape.size())) { - // kDynamic on either side is a wildcard. - if (ShapedType::isDynamic(actualShape[d]) || - ShapedType::isDynamic(expShape[d])) - continue; - if (actualShape[d] != expShape[d]) - return op->emitOpError("dim ") - << d << " of result #" << i << " mismatch: expected " - << expShape[d] << " " << formatShape(expShape) - << " but outs has " << actualShape[d] << " " - << formatShape(actualShape); - } + assert(inits.size() == 1 && + "verifyHipOpShape covers single-destination ops; a multi-destination " + "op needs one expected shape per init"); + if (inits.size() != 1) + return op->emitOpError("expected a single DPS init operand, got ") + << inits.size(); + + // The shape helper has already emitted a diagnostic on failure. + FailureOr> expected = inferShape(); + if (failed(expected)) + return failure(); + + auto initType = dyn_cast(inits.front().getType()); + if (!initType) + return op->emitOpError("init operand is not a shaped type"); + ArrayRef actual = initType.getShape(); + if (actual.size() != expected->size()) + return op->emitOpError("rank mismatch on result: expected rank ") + << expected->size() << " " << formatShape(*expected) + << " but outs has rank " << actual.size() << " " + << formatShape(actual); + + for (size_t d : llvm::seq(0, actual.size())) { + // kDynamic on either side is a wildcard. + if (ShapedType::isDynamic(actual[d]) || + ShapedType::isDynamic((*expected)[d])) + continue; + if (actual[d] != (*expected)[d]) + return op->emitOpError("dim ") + << d << " of result mismatch: expected " << (*expected)[d] << " " + << formatShape(*expected) << " but outs has " << actual[d] << " " + << formatShape(actual); } - return success(); } @@ -592,43 +593,55 @@ bool extractConstantInts(Value v, SmallVectorImpl &out) { } // namespace -FailureOr> -mlir::hip::computeReductionDimMap(int64_t dataRank, ArrayRef axes, - int64_t keepdims) { - // Normalize negative axes (ONNX convention) into a set. - llvm::SmallSet reduced; +namespace { + +/// Map each output dimension of an ONNX reduction to the input dimension it +/// takes its extent from; `std::nullopt` marks a reduced axis retained by +/// `keepdims`, whose extent is the literal 1 rather than an input extent. +/// +/// `axes` holds reduced axis indices in the ONNX negative-axis convention; an +/// empty list means no reduction. Returns failure when an axis is out of range. +/// +/// This is the single source of truth behind `inferReductionShape` and +/// `reifyReductionResultShape`, so the static and mixed forms cannot disagree. +FailureOr>> +computeReductionDimMap(int64_t dataRank, ArrayRef axes, + int64_t keepdims) { + // Axis membership over the closed domain [0, dataRank). + llvm::SmallBitVector reduced(dataRank); for (int64_t axis : axes) { - if (axis < 0) - axis += dataRank; - if (axis < 0 || axis >= dataRank) + int64_t normalized = axis < 0 ? axis + dataRank : axis; + if (normalized < 0 || normalized >= dataRank) return failure(); - reduced.insert(axis); + reduced.set(normalized); } - SmallVector dimMap; + SmallVector> dimMap; dimMap.reserve(dataRank); for (int64_t i : llvm::seq(0, dataRank)) { - if (!reduced.contains(i)) + if (!reduced.test(i)) dimMap.push_back(i); else if (keepdims) - dimMap.push_back(kReducedDim); + dimMap.push_back(std::nullopt); // keepdims=0: the reduced axis leaves the output rank entirely. } return dimMap; } +} // namespace + FailureOr> mlir::hip::inferReductionShape(ArrayRef dataShape, ArrayRef axes, int64_t keepdims) { - FailureOr> dimMap = + FailureOr>> dimMap = computeReductionDimMap(dataShape.size(), axes, keepdims); if (failed(dimMap)) return failure(); SmallVector shape; shape.reserve(dimMap->size()); - for (int64_t sourceDim : *dimMap) - shape.push_back(sourceDim == kReducedDim ? 1 : dataShape[sourceDim]); + for (std::optional sourceDim : *dimMap) + shape.push_back(sourceDim ? dataShape[*sourceDim] : 1); return shape; } @@ -640,7 +653,7 @@ mlir::hip::reifyReductionResultShape(OpBuilder &b, Location loc, Value data, return failure(); // Validate before emitting any `tensor.dim`, so a failure leaves the IR // unchanged (see the contract in HipShapeUtils.h). - FailureOr> dimMap = + FailureOr>> dimMap = computeReductionDimMap(dataType.getRank(), axes, keepdims); if (failed(dimMap)) return failure(); @@ -648,11 +661,10 @@ mlir::hip::reifyReductionResultShape(OpBuilder &b, Location loc, Value data, ArrayRef dataShape = dataType.getShape(); SmallVector dims; dims.reserve(dimMap->size()); - for (int64_t sourceDim : *dimMap) - dims.push_back(sourceDim == kReducedDim - ? OpFoldResult(b.getIndexAttr(1)) - : reifyDimOrConstant(b, loc, dataShape[sourceDim], data, - sourceDim)); + for (std::optional sourceDim : *dimMap) + dims.push_back(sourceDim ? reifyDimOrConstant(b, loc, dataShape[*sourceDim], + data, *sourceDim) + : OpFoldResult(b.getIndexAttr(1))); return dims; } diff --git a/test/lit/Dialect/hip-matmul-shape-verifier.mlir b/test/lit/Dialect/hip-matmul-shape-verifier.mlir index 029de7382..32069d53a 100644 --- a/test/lit/Dialect/hip-matmul-shape-verifier.mlir +++ b/test/lit/Dialect/hip-matmul-shape-verifier.mlir @@ -114,7 +114,7 @@ func.func @matmul_m_mismatch(%ctx: !hip.context, %a: memref<2x4xf16, 1>, %b: memref<4x8xf16, 1>, %c: memref<3x8xf16, 1>) { - // expected-error @below {{dim 0 of result #0 mismatch: expected 2}} + // expected-error @below {{dim 0 of result mismatch: expected 2}} hip.matmul(%ctx) ins(%a, %b : memref<2x4xf16, 1>, memref<4x8xf16, 1>) outs(%c : memref<3x8xf16, 1>) @@ -127,7 +127,7 @@ func.func @matmul_n_mismatch(%ctx: !hip.context, %a: memref<2x4xf16, 1>, %b: memref<4x8xf16, 1>, %c: memref<2x9xf16, 1>) { - // expected-error @below {{dim 1 of result #0 mismatch: expected 8}} + // expected-error @below {{dim 1 of result mismatch: expected 8}} hip.matmul(%ctx) ins(%a, %b : memref<2x4xf16, 1>, memref<4x8xf16, 1>) outs(%c : memref<2x9xf16, 1>) @@ -153,7 +153,7 @@ func.func @matmul_rank_mismatch(%ctx: !hip.context, %a: memref<2x4x8xf16, 1>, %b: memref<8x16xf16, 1>, %c: memref<4x16xf16, 1>) { - // expected-error @below {{rank mismatch on result #0: expected rank 3}} + // expected-error @below {{rank mismatch on result: expected rank 3}} hip.matmul(%ctx) ins(%a, %b : memref<2x4x8xf16, 1>, memref<8x16xf16, 1>) outs(%c : memref<4x16xf16, 1>) From f226186a2e1ee798455b550aceb9fb9a298b272e Mon Sep 17 00:00:00 2001 From: fhanuman Date: Fri, 31 Jul 2026 10:11:56 -0600 Subject: [PATCH 7/9] refactor(hip): finish narrowing the shape-helper surface `reifyBroadcastShape` and `reifyReductionWithKeepdims` have no callers outside `HipShapeUtils.cpp`, so they join the dimension map and the static broadcast fold as file-static. Every remaining declaration in the header now has a real external consumer, and the header exposes only shape rules. Check every dimension in `createEmptyTensorFromReifiedShape` before materializing any index value: bailing out partway through left stray constants behind, which is what the surrounding no-stray-ops contract forbids. Guard `extractConstantIntVector` against the null Value that optional ONNX operands arrive as, since `getDefiningOp` would dereference it. Pass the operand's row and column extents to the batch-stride helper instead of a `function_ref` that materializes the product, which drops the lazily-invoked callback and the two lambdas at its call sites while keeping the matrix size off the paths that do not need it. Co-Authored-By: Claude Opus 5 Made-with: Cursor Co-authored-by: Cursor --- include/hip/Dialect/IR/HipOps.td | 9 ++--- include/hip/Dialect/IR/HipShapeUtils.h | 39 ++----------------- lib/Conversion/HipToLLVM/MatmulLowering.cpp | 35 +++++++++-------- lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp | 25 ++++++++---- lib/Dialect/IR/HipShapeUtils.cpp | 42 +++++++++++++++------ 5 files changed, 72 insertions(+), 78 deletions(-) diff --git a/include/hip/Dialect/IR/HipOps.td b/include/hip/Dialect/IR/HipOps.td index 700174c8e..17ba5d30a 100644 --- a/include/hip/Dialect/IR/HipOps.td +++ b/include/hip/Dialect/IR/HipOps.td @@ -164,11 +164,10 @@ class Hip_DpsOp reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source); -/// Compute a NumPy-broadcast result shape from already-reified operand shapes. -/// Static extents remain `IndexAttr`; dynamic/dynamic pairs materialize the -/// runtime broadcast rule as ordinary index SSA. +/// Compute the NumPy-broadcast result shape over ranked tensor `operands`, +/// using `tensor::getMixedSizes` for each. Static extents remain `IndexAttr`; +/// dynamic/dynamic pairs materialize the runtime broadcast rule as index SSA. /// /// Before (choosing either dynamic operand is incorrect when it is 1): /// %lhs_dim = tensor.dim %lhs, %c0 @@ -133,14 +133,6 @@ SmallVector reifyElementwiseSameShape(OpBuilder &b, Location loc, /// %init = tensor.empty(%extent) : tensor /// /// The `FailureOr` distinguishes failure from a successful rank-zero shape. -FailureOr> -reifyBroadcastShape(OpBuilder &b, Location loc, - ArrayRef> inputShapes, - function_ref emitError); - -/// Compute the NumPy-broadcast result shape over ranked tensor `operands`. -/// This is the ValueRange convenience wrapper around the mixed-shape helper -/// above and uses `tensor::getMixedSizes` for each operand. /// /// Used by elementwise ops that take broadcast-shape operands and write /// the broadcast result into their `outs` (add, mul, sub, div, min, mod, @@ -231,31 +223,6 @@ FailureOr> reifyReductionResultShape(OpBuilder &b, Location loc, Value data, ArrayRef axes, int64_t keepdims); -/// Reify the result shape of a reduction op (reduce_sum / reduce_max / -/// reduce_prod) given `data`, the `axes` operand (rank-1 i64 tensor), -/// and the `keepdims` / `noop_with_empty_axes` attributes. -/// -/// Introspects `axes` as an `arith.constant` (the typical case after the -/// OnnxToHip converter materializes it from the ONNX attribute), resolves -/// ONNX's empty-axes semantics against `noop_with_empty_axes` — reduce every -/// axis when 0, reduce nothing when 1 — and delegates the shape rule to -/// `reifyReductionResultShape`. -/// -/// Returns `success()` and writes the reified dim list into `out` when -/// `axes` can be introspected. Returns `failure()` when `axes` is not a -/// recognised constant — the caller should then fall back to -/// `reifyElementwiseSameShape(output)` to keep the reify interface -/// non-failing. -/// -/// Uses `LogicalResult` (rather than the empty-vector sentinel used by -/// the other helpers in this header) because a valid rank-0 reduction -/// result has an empty dim list, which would otherwise be -/// indistinguishable from the bail path. -LogicalResult reifyReductionWithKeepdims(OpBuilder &b, Location loc, Value data, - Value axes, int64_t keepdims, - int64_t noopWithEmptyAxes, - SmallVectorImpl &out); - /// One-shot reify body for ONNX-style reduction ops (reduce_sum, /// reduce_max, reduce_prod). Tries `reifyReductionWithKeepdims` first /// to recover per-input-dim mappings from a constant `axes` operand. diff --git a/lib/Conversion/HipToLLVM/MatmulLowering.cpp b/lib/Conversion/HipToLLVM/MatmulLowering.cpp index 184202e30..31be2915e 100644 --- a/lib/Conversion/HipToLLVM/MatmulLowering.cpp +++ b/lib/Conversion/HipToLLVM/MatmulLowering.cpp @@ -80,13 +80,16 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { // walks one matrix per batch. The two cases are distinguished by comparing // the operand's matrix count against the output's, which also keeps the // read in bounds if a runtime-only shape ever violated the invariant. - auto batchStride = - [&](MemRefType type, MemRefDescriptor desc, - llvm::function_ref matrixElements) -> Value { + // `rows` x `cols` is the operand's matrix size in elements, materialized + // only on the paths that need it. + auto batchStride = [&](MemRefType type, MemRefDescriptor desc, Value rows, + Value cols) -> Value { ArrayRef batch = type.getShape().drop_back(2); - if (!ShapedType::isDynamicShape(batch)) - return llvm::product_of(batch) == 1 ? createI64Const(0) - : matrixElements(); + if (!ShapedType::isDynamicShape(batch)) { + if (llvm::product_of(batch) == 1) + return createI64Const(0); + return LLVM::MulOp::create(rewriter, loc, rows, cols).getRes(); + } Value count = createI64Const(1); for (unsigned i : llvm::seq(0, batch.size())) @@ -94,20 +97,16 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { desc.size(rewriter, loc, i)); Value perBatch = LLVM::ICmpOp::create( rewriter, loc, LLVM::ICmpPredicate::eq, count, batchCount); - // Materialize both arms into locals: passing calls that emit IR directly - // as arguments would leave the emission order unspecified. - Value perBatchStride = matrixElements(); - Value broadcastStride = createI64Const(0); - return LLVM::SelectOp::create(rewriter, loc, perBatch, perBatchStride, - broadcastStride) + // Both arms in locals, so the emission order does not depend on the + // unspecified evaluation order of call arguments. + Value matrixElements = LLVM::MulOp::create(rewriter, loc, rows, cols); + Value zero = createI64Const(0); + return LLVM::SelectOp::create(rewriter, loc, perBatch, matrixElements, + zero) .getRes(); }; - Value aBatchStride = batchStride(AType, ADesc, [&] { - return LLVM::MulOp::create(rewriter, loc, M, K).getRes(); - }); - Value bBatchStride = batchStride(BType, BDesc, [&] { - return LLVM::MulOp::create(rewriter, loc, K, N).getRes(); - }); + Value aBatchStride = batchStride(AType, ADesc, M, K); + Value bBatchStride = batchStride(BType, BDesc, K, N); // Runtime signature: // int wrap_hipblasLtMatmul(RuntimeState* state, diff --git a/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp b/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp index 9c4fae2fe..fe56bc764 100644 --- a/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp +++ b/lib/Conversion/OnnxToHip/OnnxToHipUtils.cpp @@ -21,6 +21,10 @@ namespace hip { bool extractConstantIntVector(mlir::Value value, llvm::SmallVectorImpl &out) { out.clear(); + // Optional ONNX operands arrive as a null Value; `getDefiningOp` would + // dereference it. + if (!value) + return false; mlir::Operation *defOp = value.getDefiningOp(); if (!defOp) return false; @@ -108,18 +112,23 @@ mlir::FailureOr createEmptyTensorFromReifiedShape( if (static_cast(reifiedShape.size()) != resultType.getRank()) return mlir::failure(); - llvm::SmallVector dynSizes; + // Check every dimension before materializing any index value. Bailing out + // partway through the loop below would leave stray constants behind, which + // the contract in HipShapeUtils.h forbids. for (int64_t dimIdx : llvm::seq(resultType.getRank())) { + if (resultType.isDynamicDim(dimIdx)) + continue; std::optional reifiedStatic = mlir::getConstantIntValue(reifiedShape[dimIdx]); - if (!resultType.isDynamicDim(dimIdx)) { - if (reifiedStatic && *reifiedStatic != resultType.getDimSize(dimIdx)) - return mlir::failure(); - continue; - } - dynSizes.push_back(mlir::getValueOrCreateConstantIndexOp( - builder, loc, reifiedShape[dimIdx])); + if (reifiedStatic && *reifiedStatic != resultType.getDimSize(dimIdx)) + return mlir::failure(); } + + llvm::SmallVector dynSizes; + for (int64_t dimIdx : llvm::seq(resultType.getRank())) + if (resultType.isDynamicDim(dimIdx)) + dynSizes.push_back(mlir::getValueOrCreateConstantIndexOp( + builder, loc, reifiedShape[dimIdx])); return mlir::Value( mlir::tensor::EmptyOp::create(builder, loc, resultType.getShape(), resultType.getElementType(), dynSizes)); diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index a2c0e31db..54363567e 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -331,10 +331,15 @@ mlir::hip::reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source) { return dims; } +namespace { + +/// NumPy-broadcast result shape from already-reified operand shapes. Callers +/// must have validated broadcastability against the static shapes first, since +/// this materializes index SSA as it folds. FailureOr> -mlir::hip::reifyBroadcastShape(OpBuilder &b, Location loc, - ArrayRef> inputShapes, - function_ref emitError) { +reifyBroadcastShape(OpBuilder &b, Location loc, + ArrayRef> inputShapes, + function_ref emitError) { if (inputShapes.empty()) { emitError() << "broadcast requires at least one input shape"; return failure(); @@ -360,6 +365,8 @@ mlir::hip::reifyBroadcastShape(OpBuilder &b, Location loc, return result; } +} // namespace + FailureOr> mlir::hip::reifyBroadcastResultShape( OpBuilder &b, Location loc, ValueRange operands, function_ref emitError) { @@ -668,18 +675,29 @@ mlir::hip::reifyReductionResultShape(OpBuilder &b, Location loc, Value data, return dims; } -LogicalResult mlir::hip::reifyReductionWithKeepdims( - OpBuilder &b, Location loc, Value data, Value axes, int64_t keepdims, - int64_t noopWithEmptyAxes, SmallVectorImpl &out) { +namespace { + +/// Reduction result shape recovered from a constant `axes` operand. +/// +/// Introspects `axes` as an `arith.constant` (the typical case after the +/// OnnxToHip converter materializes it from the ONNX attribute), resolves +/// ONNX's empty-axes semantics against `noop_with_empty_axes` — reduce every +/// axis when 0, reduce nothing when 1 — and delegates the shape rule to +/// `reifyReductionResultShape`. +/// +/// Returns failure when `axes` is not a recognised constant, which is why this +/// returns `LogicalResult` and writes through `out`: a valid rank-0 reduction +/// result is a successful *empty* dim list and would otherwise be +/// indistinguishable from the bail path. +LogicalResult reifyReductionWithKeepdims(OpBuilder &b, Location loc, Value data, + Value axes, int64_t keepdims, + int64_t noopWithEmptyAxes, + SmallVectorImpl &out) { out.clear(); auto dataType = dyn_cast(data.getType()); if (!dataType) return failure(); - // Axes operand: try to fold to a constant int vector. ONNX semantics - // allow a size-0 vector to mean "no axes specified" — combined with the - // `noop_with_empty_axes` attribute that selects between "no-op" and - // "reduce all axes". SmallVector axesList; if (!extractConstantInts(axes, axesList)) return failure(); @@ -690,13 +708,15 @@ LogicalResult mlir::hip::reifyReductionWithKeepdims( } FailureOr> dims = - reifyReductionResultShape(b, loc, data, axesList, keepdims); + mlir::hip::reifyReductionResultShape(b, loc, data, axesList, keepdims); if (failed(dims)) return failure(); out.assign(dims->begin(), dims->end()); return success(); } +} // namespace + LogicalResult mlir::hip::reifyBroadcastShapeFor(OpBuilder &b, Location loc, ValueRange operands, Operation *op, From 96acf9d5cb6da77d68d0600d4fb8d8acbce2525f Mon Sep 17 00:00:00 2001 From: fhanuman Date: Fri, 31 Jul 2026 10:42:50 -0600 Subject: [PATCH 8/9] fix(hip): port ReduceL2 onto the shared reduction shape rule MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `reduce_l2` landed on main after this branch diverged, carrying its own copy of the reduction dimension map — the same `outToIn` / `SmallSet` / `axesKnown` structure this branch removed from ReduceProd — and calling `inferReduceResultType` through the boolean-flag signature that no longer exists. The merge therefore did not compile, which is why CI failed on all three build jobs while each side built on its own. Route it through `resolveReductionAxes` and `createReductionEmptyTensor` like the other five reduction converters, so its destination and `reifyResultShapes` share one rule and `keepdims = 0` maps output dimensions back to the input dimensions they actually come from. No private copy of the dimension map remains in the tree. Co-Authored-By: Claude Opus 5 Made-with: Cursor Co-authored-by: Cursor --- .../OnnxToHip/ReduceL2Conversion.cpp | 157 ++---------------- 1 file changed, 17 insertions(+), 140 deletions(-) diff --git a/lib/Conversion/OnnxToHip/ReduceL2Conversion.cpp b/lib/Conversion/OnnxToHip/ReduceL2Conversion.cpp index fb8d18ba0..31df437fd 100644 --- a/lib/Conversion/OnnxToHip/ReduceL2Conversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceL2Conversion.cpp @@ -5,125 +5,10 @@ #include "OnnxToHipUtils.h" -#include "llvm/ADT/APInt.h" -#include "llvm/ADT/SmallSet.h" - namespace mlir { namespace hip { namespace { -/// Try to recognise \p v as a compile-time 1-D integer constant tensor. -/// Mirrors the helper in SliceConversion.cpp / ReduceProdConversion.cpp. -static mlir::DenseElementsAttr getCompileTimeConstantTensor(mlir::Value value) { - mlir::Operation *defOp = value.getDefiningOp(); - if (!defOp) - return nullptr; - if (auto cst = mlir::dyn_cast(defOp)) - return mlir::dyn_cast(cst.getValue()); - if (auto attr = defOp->getAttr("value")) - if (auto dense = mlir::dyn_cast(attr)) - return dense; - if (auto toTensor = mlir::dyn_cast(defOp)) { - auto bufDef = - toTensor.getBuffer().getDefiningOp(); - if (!bufDef) - return nullptr; - auto module = bufDef->getParentOfType(); - if (!module) - return nullptr; - auto global = - module.lookupSymbol(bufDef.getNameAttr()); - if (!global) - return nullptr; - return mlir::dyn_cast_or_null( - global.getInitialValueAttr()); - } - return nullptr; -} - -static mlir::LogicalResult -extractIntVector(mlir::Value v, llvm::SmallVectorImpl &out) { - if (!v) - return mlir::failure(); - auto dense = getCompileTimeConstantTensor(v); - if (!dense) - return mlir::failure(); - auto tensorType = mlir::dyn_cast(dense.getType()); - if (!tensorType || tensorType.getRank() != 1) - return mlir::failure(); - auto elemTy = tensorType.getElementType(); - if (!elemTy.isInteger(64) && !elemTy.isInteger(32)) - return mlir::failure(); - for (mlir::APInt entry : dense.getValues()) - out.push_back(entry.getSExtValue()); - return mlir::success(); -} - -/// Build the destination `tensor.empty` for ReduceL2. -/// -/// Output shape semantics match other ONNX reduce ops: -/// * keepdims=1: out_rank == in_rank; reduced axes become size 1. -/// * keepdims=0: out_rank == in_rank - #axes; reduced axes are dropped. -/// -/// For dynamic output dims we map back to the source `data` dim using the -/// known axes set (compile-time constant). When axes are not known, fall -/// back to positional alignment (correct for keepdims=1 and all-reduce). -static mlir::Value buildReduceL2Init(mlir::PatternRewriter &rewriter, - mlir::Location loc, - mlir::RankedTensorType resultType, - mlir::Value data, - llvm::ArrayRef axesVec, - bool axesKnown, int64_t keepdims) { - auto dataType = mlir::cast(data.getType()); - int64_t inRank = dataType.getRank(); - - llvm::SmallSet reducedAxes; - for (int64_t a : axesVec) { - if (a < 0) - a += inRank; - reducedAxes.insert(a); - } - - llvm::SmallVector outToIn(resultType.getRank(), -1); - if (axesKnown) { - if (keepdims) { - for (int64_t i = 0; i < resultType.getRank(); ++i) - outToIn[i] = reducedAxes.contains(i) ? -1 : i; - } else { - int64_t outIdx = 0; - for (int64_t i = 0; i < inRank; ++i) { - if (reducedAxes.contains(i)) - continue; - if (outIdx < resultType.getRank()) - outToIn[outIdx] = i; - ++outIdx; - } - } - } else { - for (int64_t i = 0; i < resultType.getRank(); ++i) - outToIn[i] = i < inRank ? i : -1; - } - - llvm::SmallVector dynSizes; - for (int64_t i = 0; i < resultType.getRank(); ++i) { - if (!resultType.isDynamicDim(i)) - continue; - int64_t inIdx = outToIn[i]; - if (inIdx < 0) { - dynSizes.push_back( - mlir::arith::ConstantIndexOp::create(rewriter, loc, 1)); - } else if (dataType.isDynamicDim(inIdx)) { - dynSizes.push_back( - mlir::tensor::DimOp::create(rewriter, loc, data, inIdx)); - } else { - dynSizes.push_back(mlir::arith::ConstantIndexOp::create( - rewriter, loc, dataType.getDimSize(inIdx))); - } - } - return mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); -} - /// onnx.ReduceL2 -> hip.reduce_l2 /// /// Direct, dim-tolerant conversion: sqrt(sum(x^2)) along the reduced axes @@ -161,52 +46,44 @@ ReduceL2ToHip::matchAndRewrite(mlir::Operation *op, keepdims = keepdimsAttr.getSInt(); } - llvm::SmallVector axesVec; - bool axesKnown = false; + // Reduced axes, resolved from the attribute or a compile-time-constant + // operand, so the destination shape can map each output dimension back to the + // input dimension it comes from. + llvm::SmallVector axesStorage; + std::optional> reducedAxes = + resolveReductionAxes(op, data, noopWithEmptyAxes, axesStorage); + mlir::Value axesOperand; if (op->getNumOperands() > 1 && !mlir::isa(op->getOperand(1).getType())) { axesOperand = op->getOperand(1); - if (mlir::succeeded(extractIntVector(axesOperand, axesVec))) - axesKnown = true; } else { - if (auto axesAttr = op->getAttrOfType("axes")) { - for (auto a : axesAttr) - axesVec.push_back( - mlir::cast(a).getValue().getSExtValue()); - axesKnown = true; - } else if (noopWithEmptyAxes == 0) { - auto inputType = mlir::cast(data.getType()); - for (int64_t i : llvm::seq(inputType.getRank())) - axesVec.push_back(i); - axesKnown = true; - } else { - axesKnown = true; // empty axes, noop - } auto axesType = mlir::RankedTensorType::get( - {static_cast(axesVec.size())}, rewriter.getI64Type()); - auto axesAttr = - mlir::DenseIntElementsAttr::get(axesType, llvm::ArrayRef(axesVec)); + {static_cast(axesStorage.size())}, rewriter.getI64Type()); + auto axesAttr = mlir::DenseIntElementsAttr::get( + axesType, llvm::ArrayRef(axesStorage)); axesOperand = mlir::arith::ConstantOp::create(rewriter, loc, axesType, axesAttr); } - auto resultTypeOr = - inferReduceResultType(op, data, axesVec, axesKnown, keepdims); + auto resultTypeOr = inferReduceResultType(op, data, reducedAxes, keepdims); if (mlir::failed(resultTypeOr)) return rewriter.notifyMatchFailure( op, "ReduceL2: cannot infer unranked result (need ranked input and " "static axes)"); mlir::RankedTensorType resultType = *resultTypeOr; - mlir::Value init = buildReduceL2Init(rewriter, loc, resultType, data, axesVec, - axesKnown, keepdims); + mlir::FailureOr init = createReductionEmptyTensor( + rewriter, loc, resultType, data, reducedAxes, keepdims); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "ReduceL2 result type is incompatible with the reduction shape"); auto keepdimsAttr = rewriter.getI64IntegerAttr(keepdims); auto noopWithEmptyAxesAttr = rewriter.getI64IntegerAttr(noopWithEmptyAxes); auto hipOp = mlir::hip::ReduceL2Op::create(rewriter, loc, context, data, axesOperand, - init, keepdimsAttr, noopWithEmptyAxesAttr); + *init, keepdimsAttr, noopWithEmptyAxesAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); From 237e3a31ae202dfb3b2a3899a267d47e4a08d4aa Mon Sep 17 00:00:00 2001 From: fhanuman Date: Fri, 31 Jul 2026 15:45:24 -0600 Subject: [PATCH 9/9] fix(hip): harden shared shape rules and dynamic MatMul batches Make rank-zero shape success unambiguous by returning FailureOr from the transpose, gather, GatherND, and shape-preserving reifiers, and validate all preconditions before emitting dimension IR. Route the corresponding converters through those shared rules, including MatMulNBits' N-derived output width. Close the dynamic partial-broadcast hole without aborting the host process: pass both operand batch counts to the MatMul runtime wrapper, reject counts that are neither one nor the output batch count, and report the failure through the existing device error flag so the generated interface returns a recoverable non-zero status to ORT. Also make verifier diagnostics reachable, normalize NoneType reduction axes, and align the documentation and LIT ABI coverage. Co-Authored-By: GPT-5.6 Sol Made-with: Cursor --- CLAUDE.md | 4 +- docs/design/hip-shape-inference.md | 46 +++++------ docs/hip_dialect_intro.md | 2 +- include/hip/Dialect/IR/HipShapeUtils.h | 58 ++++++++------ lib/Conversion/HipToLLVM/MatmulLowering.cpp | 58 +++++++++----- lib/Conversion/OnnxToHip/GatherConversion.cpp | 50 ++++-------- .../OnnxToHip/GatherElementsConversion.cpp | 25 +++--- .../OnnxToHip/GatherNDConversion.cpp | 55 +++++--------- .../OnnxToHip/MatMulNBitsConversion.cpp | 18 ++++- .../OnnxToHip/ReduceMaxConversion.cpp | 3 +- .../OnnxToHip/ReduceMeanConversion.cpp | 3 +- .../OnnxToHip/ReduceMinConversion.cpp | 3 +- .../OnnxToHip/ReduceSumConversion.cpp | 3 +- .../OnnxToHip/TransposeConversion.cpp | 25 +++--- lib/Dialect/IR/HipReifyResultShapesImpl.cpp | 67 +++++++--------- lib/Dialect/IR/HipShapeUtils.cpp | 76 ++++++++++++------- lib/Runtime/hipdnn_ep_runtime.h | 17 ++++- lib/Runtime/hipdnn_ep_runtime_state.cpp | 13 ++++ lib/Runtime/mock/mock_gpu.cpp | 13 +++- lib/Runtime/real/matmul.cpp | 23 +++++- .../Conversion/hip-to-llvm/test_matmul.mlir | 58 +++++++++++--- .../onnx-to-hip/test_reduce_sum.mlir | 17 +++++ 22 files changed, 374 insertions(+), 263 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 5928ca7c2..9a23b26fa 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -140,7 +140,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co ### Result-shape agreement -- Converter destination construction, op verification, and `reifyResultShapes` must use the same `HipShapeUtils` shape rule for broadcast, Gemm, MatMul, and reductions. +- Converter destination construction and `reifyResultShapes` must use the same `HipShapeUtils` shape rule for broadcast, Gemm, MatMul, and reductions. MatMul and Gemm also use their shared rule for static verification; broadcast and reduction verifiers are future work. - 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, so emitting IR is always the last step. - 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. @@ -148,7 +148,7 @@ See [CONTRIBUTING.md](CONTRIBUTING.md) for PR, formatting, AI-disclosure, and co - Rank-zero success is an empty shape carried by `FailureOr`; never use an empty vector as both success and failure. - Variadic Max/Min share one pairwise-chain helper so every intermediate rank comes from the shared broadcast shape. - Reductions resolve to one out-to-in dimension map. `keepdims = 0` makes the output dimension order non-positional in the input, so never copy input extents positionally when a reduced axis can precede a kept one. -- MatMul uses the reified output batch count plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Only *partial per-axis* batch broadcast is rejected; dynamic batch extents are supported and must not be rejected wholesale. `wrap_hipblasLtMatmul` carries both strides; invalidate LLVM-IR artifacts compiled against the previous wrapper ABI. +- MatMul uses the reified output batch count plus independent A/B strides, so either whole matrix may broadcast across the other's batches. Static partial per-axis broadcast is rejected by the verifier; dynamic layouts are validated at runtime and return a recoverable inference error when an operand contains neither one matrix nor one matrix per output batch. `wrap_hipblasLtMatmul` carries both operand batch counts and strides; invalidate LLVM-IR artifacts compiled against the previous wrapper ABI. - See [docs/design/hip-shape-inference.md](docs/design/hip-shape-inference.md). ### Allocation and memory planning diff --git a/docs/design/hip-shape-inference.md b/docs/design/hip-shape-inference.md index 7cfdac36f..cf83c9e62 100644 --- a/docs/design/hip-shape-inference.md +++ b/docs/design/hip-shape-inference.md @@ -136,10 +136,11 @@ rewrite or reification reports failure; upstream's codebase avoids creating them in the first place. Emitting IR is therefore the last thing a helper does. -The same split gives operation verifiers a shape rule to check against. -`verifyHipOpShape` takes an `infer*` helper directly, so a single-destination op -pairs the two in one expression and the `outs` shape, the converter destination, -and `reifyResultShapes` are all held to one shape function: +The same split gives targeted operation verifiers a shape rule to check against. +MatMul and Gemm use `verifyHipOpShape` with their `infer*` helper, so their +`outs` shape, converter destination, and `reifyResultShapes` are all held to one +shape function. Broadcast and reduction ops share their converter/reify rule +but do not yet have shared static shape verifiers: ```c++ LogicalResult MatmulOp::verify() { @@ -207,23 +208,22 @@ 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 — -and is rejected by `verifyStridedBatchMatmul`, the MatMul verifier, and the -converter. The check is per axis: an operand is rejected only when some axis is -statically 1 while the output extent is not, *and* another axis carries batches. -Extents that are not statically 1 count as carrying batches, so an unknown -extent never hides a partial broadcast. Ordinary batched matmul with dynamic -leading extents (`[?, H, M, K] @ [?, H, K, N]`) is representable and accepted; a -blanket rejection of dynamic batch extents would fail that common layout. - -The lowering distinguishes the two strides by comparing the operand's matrix -count against the output batch count, folding at compile time when both are -static. Comparing against the output count rather than testing for "more than -one matrix" also keeps the read in bounds if a runtime-only shape ever violated -the invariant. - -`wrap_hipblasLtMatmul` carries both strides. LLVM-IR artifacts compiled against -the previous wrapper ABI must be invalidated. +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, so the lowering also computes each +operand's runtime matrix count. `wrap_hipblasLtMatmul` dispatches only when each +count is either 1 or the output batch count; otherwise it records an error in +the runtime state's device error flag and skips hipBLASLt. The generated +interface observes that flag after its existing stream synchronization and +returns a recoverable non-zero inference status to ORT. + +This preserves ordinary dynamic batched matmul +(`[?, H, M, K] @ [?, H, K, N]`) without treating every non-output matrix count +as whole-matrix broadcast. The stride is 0 only for one matrix and the matrix +size only for one matrix per output batch. + +`wrap_hipblasLtMatmul` carries both operand batch counts and both strides. +LLVM-IR artifacts compiled against the previous wrapper ABI must be invalidated. ## `--hip-infer-shapes` @@ -337,7 +337,7 @@ Primary regression coverage: | `test/lit/Dialect/hip-gemm-reify-shapes.mlir` | Transpose-aware Gemm M/N reification | | `test/lit/Dialect/hip-matmul-reify-shapes.mlir` | Per-op reification through `--resolve-shaped-type-result-dims` | | `test/lit/Dialect/hip-matmul-shape-verifier.mlir` | Static MatMul shape validation, including accepted dynamic batch layouts and rejected partial per-axis broadcast | -| `test/lit/Conversion/hip-to-llvm/test_matmul.mlir` | Per-operand batch strides: compile-time 0 / matrix size, and the runtime comparison against the output batch count | +| `test/lit/Conversion/hip-to-llvm/test_matmul.mlir` | Per-operand batch counts and strides: compile-time 0 / matrix size, dynamic count comparison, and the runtime-validation ABI | | `test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir` | Reduction destinations, including the non-positional `keepdims = 0` dimension mapping | | `test/lit/Dialect/hip-loop-verifier.mlir` | Loop-carried type contract | | `test/lit/Dialect/hip-resolve-tensor-dims.mlir` | Production pre-bufferization dim folding | @@ -347,7 +347,7 @@ Complex operations may use dedicated files; common shape categories should exten ## Current limitations - ONNX MatMul rank-1 operands require promotion to rank 2 before constructing `hip.matmul`; the runtime and current verifier require rank at least 2. -- MatMul supports whole-matrix batch broadcast and one-matrix-per-output-batch operands, including dynamic batch extents; per-axis partial batch broadcast remains unsupported by the strided-batch runtime. +- MatMul supports whole-matrix batch broadcast and one-matrix-per-output-batch operands, including dynamic batch extents. Static partial per-axis broadcast is rejected during compilation; a dynamically concealed partial layout returns a recoverable runtime error. - Reduction destinations fall back to a positional copy from the input when the reduced axes are only known at runtime; reification mirrors that fallback rather than inventing a shape. - Converter migration to inferred-type builders is incremental; explicit result-type builders remain supported. - A future multi-result operation that needs custom `InferTypeOpInterface` logic may require a dedicated result-type inference implementation file. diff --git a/docs/hip_dialect_intro.md b/docs/hip_dialect_intro.md index 424941139..d49be977c 100644 --- a/docs/hip_dialect_intro.md +++ b/docs/hip_dialect_intro.md @@ -44,7 +44,7 @@ Matrix multiplication backed by the hipBLASLt library (`hipblasLtMatmul`). | Op | DPS Syntax | Runtime | Status | |---|---|---|---| -| `hip.matmul` | `(%ctx) ins(%A, %B : ...) outs(%C : ...)` | `wrap_hipblasLtMatmul(state, slot, A, B, C, M, N, K, batch, elem, a_stride, b_stride)` | Full impl | +| `hip.matmul` | `(%ctx) ins(%A, %B : ...) outs(%C : ...)` | `wrap_hipblasLtMatmul(state, slot, A, B, C, M, N, K, batch, elem, a_batches, b_batches, a_stride, b_stride)` | Full impl | The output batch count is the product of the broadcasted leading dimensions. Either operand may provide one matrix for all batches (`stride = 0`) or one diff --git a/include/hip/Dialect/IR/HipShapeUtils.h b/include/hip/Dialect/IR/HipShapeUtils.h index 5f6233e16..1ddf14cce 100644 --- a/include/hip/Dialect/IR/HipShapeUtils.h +++ b/include/hip/Dialect/IR/HipShapeUtils.h @@ -54,10 +54,10 @@ inferMatmulShape(ArrayRef aShape, ArrayRef bShape, /// provably needs something in between: a partial broadcast that pads some /// batch axes up to the output extent while carrying batches on others. /// -/// Extents that are not statically 1 count as carrying batches, so an unknown -/// extent never hides a partial broadcast. Ordinary batched matmul with -/// dynamic leading extents (`[?, H, M, K] @ [?, H, K, N]`) is representable -/// and accepted. +/// This helper rejects layouts that are provably partial from static types. +/// Dynamic extents can conceal a partial broadcast, so the HIP-to-LLVM lowering +/// also passes each operand's runtime batch count to the wrapper. The wrapper +/// dispatches only when that count is 1 or equals the output batch count. LogicalResult verifyStridedBatchMatmul(ArrayRef aShape, ArrayRef bShape, function_ref emitError); @@ -112,11 +112,10 @@ OpFoldResult reifyDimOrConstant(OpBuilder &b, Location loc, int64_t staticDim, /// dim becomes `tensor.dim %source, %i`. Used by ops whose result has /// the same shape as one designated input (e.g. rope, rms_norm, qmoe). /// -/// `source` must be a `RankedTensorType`-typed Value -- this helper is -/// called from `reifyResultShapes` impls, which are invoked only in -/// tensor mode per the interface contract. -SmallVector reifyElementwiseSameShape(OpBuilder &b, Location loc, - Value source); +/// Returns failure when `source` is not a ranked tensor. The `FailureOr` +/// distinguishes that failure from a successful rank-zero shape. +FailureOr> +reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source); /// Compute the NumPy-broadcast result shape over ranked tensor `operands`, /// using `tensor::getMixedSizes` for each. Static extents remain `IndexAttr`; @@ -154,6 +153,15 @@ FailureOr> reifyMatmulResultShape(OpBuilder &b, Location loc, Value A, Value B, function_ref emitError); +/// Reify MatMulNBits' result shape: `A`'s leading dimensions followed by the +/// static `N` attribute. The quantized weight is stored transposed, so the +/// contraction dimension never appears in the result -- which is why the last +/// dimension comes from `N` and not from `A`. +/// +/// `A` must be a `RankedTensorType`-typed Value of rank at least 1. +FailureOr> +reifyMatMulNBitsResultShape(OpBuilder &b, Location loc, Value A, int64_t N); + /// Reify ONNX Gemm's rank-2 `{M, N}` result using transpose-aware dimensions. /// Optional C is checked for static unidirectional broadcast compatibility but /// never supplies M or N. @@ -163,41 +171,41 @@ reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, function_ref emitError); /// Reify the result shape of a transpose op as `output[i] = input[perm[i]]`. -/// `perm` must be a permutation of `[0, rank-1)` and have the same length -/// as `input`'s rank — the verifier should already guarantee this; the -/// helper bails (returns empty) on mismatch. +/// `perm` must be a permutation of `[0, rank)` and have the same length as +/// `input`'s rank. Returns failure on malformed input, before emitting IR. /// /// Each output dim `i`: /// - emits `IndexAttr(input.shape[perm[i]])` when that dim is static, /// - emits `tensor.dim %input, perm[i]` otherwise. /// -/// `input` must be a `RankedTensorType`-typed Value. -SmallVector reifyTransposeByPerm(OpBuilder &b, Location loc, - Value input, - ArrayRef perm); +/// The `FailureOr` distinguishes failure from a successful rank-zero shape. +FailureOr> +reifyTransposeByPerm(OpBuilder &b, Location loc, Value input, + ArrayRef perm); /// Reify the result shape of a gather op as /// `output = data.shape[:axis] ++ indices.shape ++ data.shape[axis+1:]`. /// `axis` is normalized into `[0, data.rank)` (negative axis follows ONNX -/// convention). The helper bails (returns empty) on a malformed axis. +/// convention). Returns failure on a malformed axis. /// -/// `data` and `indices` must be `RankedTensorType`-typed Values. -SmallVector reifyGatherWithAxis(OpBuilder &b, Location loc, - Value data, Value indices, - int64_t axis); +/// The `FailureOr` distinguishes failure from a successful rank-zero shape. +FailureOr> +reifyGatherWithAxis(OpBuilder &b, Location loc, Value data, Value indices, + int64_t axis); /// Reify the result shape of a `gather_nd` op as /// `batch_dims_from_data ++ indices.shape[batch_dims:-1] ++ /// data.shape[batch_dims + indices.shape[-1]:]`. /// Per ONNX GatherND semantics, output rank = /// `q + r - indices.shape[-1] - 1 - batch_dims`, where `q = rank(indices)` -/// and `r = rank(data)`. The helper bails (returns empty) when the +/// and `r = rank(data)`. Returns failure when the /// trailing index-tuple width (`indices.shape[-1]`) is dynamic — the /// output rank itself is then unknown and reify cannot run. /// -/// `data` and `indices` must be `RankedTensorType`-typed Values. -SmallVector reifyGatherND(OpBuilder &b, Location loc, Value data, - Value indices, int64_t batchDims); +/// The `FailureOr` distinguishes failure from a successful rank-zero shape. +FailureOr> reifyGatherND(OpBuilder &b, Location loc, + Value data, Value indices, + int64_t batchDims); /// ONNX reduction result shape over `axes`, from static extents only. /// diff --git a/lib/Conversion/HipToLLVM/MatmulLowering.cpp b/lib/Conversion/HipToLLVM/MatmulLowering.cpp index 31be2915e..23345852e 100644 --- a/lib/Conversion/HipToLLVM/MatmulLowering.cpp +++ b/lib/Conversion/HipToLLVM/MatmulLowering.cpp @@ -40,12 +40,6 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { rewriter.getI64IntegerAttr(value)); }; - Value statePtr = adaptor.getCtx(); - Value APtr = extractContiguousMemRefPtr(adaptor.getA(), rewriter, loc); - Value BPtr = extractContiguousMemRefPtr(adaptor.getB(), rewriter, loc); - Value outputPtr = - extractContiguousMemRefPtr(adaptor.getOutput(), rewriter, loc); - auto AType = cast(op.getA().getType()); auto BType = cast(op.getB().getType()); auto outputType = cast(op.getOutput().getType()); @@ -56,6 +50,13 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { [&]() { return op.emitOpError(); }))) return failure(); + // No IR is emitted before all static representability checks pass. + Value statePtr = adaptor.getCtx(); + Value APtr = extractContiguousMemRefPtr(adaptor.getA(), rewriter, loc); + Value BPtr = extractContiguousMemRefPtr(adaptor.getB(), rewriter, loc); + Value outputPtr = + extractContiguousMemRefPtr(adaptor.getOutput(), rewriter, loc); + MemRefDescriptor ADesc(adaptor.getA()); MemRefDescriptor BDesc(adaptor.getB()); MemRefDescriptor outputDesc(adaptor.getOutput()); @@ -74,15 +75,28 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { unsigned elemBits = AType.getElementType().getIntOrFloatBitWidth(); Value elemSize = createI64Const(elemBits / 8); - // `verifyStridedBatchMatmul` guarantees each operand holds either one - // matrix or one matrix per output batch, so one constant stride per operand - // is exact: 0 reuses a single matrix across every batch and the matrix size - // walks one matrix per batch. The two cases are distinguished by comparing - // the operand's matrix count against the output's, which also keeps the - // read in bounds if a runtime-only shape ever violated the invariant. + auto operandBatchCount = [&](MemRefType type, + MemRefDescriptor desc) -> Value { + ArrayRef batch = type.getShape().drop_back(2); + if (!ShapedType::isDynamicShape(batch)) + return createI64Const(llvm::product_of(batch)); + + Value count = createI64Const(1); + for (unsigned i : llvm::seq(0, batch.size())) + count = LLVM::MulOp::create(rewriter, loc, count, + desc.size(rewriter, loc, i)); + return count; + }; + Value aBatchCount = operandBatchCount(AType, ADesc); + Value bBatchCount = operandBatchCount(BType, BDesc); + + // A stride can represent exactly one matrix (stride 0) or one matrix per + // output batch (matrix-size stride). Dynamic extents are checked again by + // the runtime wrapper because static types cannot rule out a partial + // per-axis broadcast in every invocation. // `rows` x `cols` is the operand's matrix size in elements, materialized // only on the paths that need it. - auto batchStride = [&](MemRefType type, MemRefDescriptor desc, Value rows, + auto batchStride = [&](MemRefType type, Value count, Value rows, Value cols) -> Value { ArrayRef batch = type.getShape().drop_back(2); if (!ShapedType::isDynamicShape(batch)) { @@ -91,10 +105,6 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { return LLVM::MulOp::create(rewriter, loc, rows, cols).getRes(); } - Value count = createI64Const(1); - for (unsigned i : llvm::seq(0, batch.size())) - count = LLVM::MulOp::create(rewriter, loc, count, - desc.size(rewriter, loc, i)); Value perBatch = LLVM::ICmpOp::create( rewriter, loc, LLVM::ICmpPredicate::eq, count, batchCount); // Both arms in locals, so the emission order does not depend on the @@ -105,17 +115,19 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { zero) .getRes(); }; - Value aBatchStride = batchStride(AType, ADesc, M, K); - Value bBatchStride = batchStride(BType, BDesc, K, N); + Value aBatchStride = batchStride(AType, aBatchCount, M, K); + Value bBatchStride = batchStride(BType, bBatchCount, K, N); // Runtime signature: // int wrap_hipblasLtMatmul(RuntimeState* state, // const void* A, const void* B, void* output, // int64_t M, int64_t N, int64_t K, // int64_t batch_count, int64_t elem_size, + // int64_t a_batch_count, + // int64_t b_batch_count, // int64_t a_batch_stride, // int64_t b_batch_stride) - SmallVector paramTypes = { + SmallVector paramTypes = { ptrType, // state i32Type, // op_state_slot ptrType, // A @@ -126,6 +138,8 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { i64Type, // K i64Type, // batch_count i64Type, // elem_size + i64Type, // a_batch_count + i64Type, // b_batch_count i64Type, // a_batch_stride i64Type // b_batch_stride }; @@ -135,7 +149,7 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { if (failed(funcOp)) return failure(); - SmallVector args = {statePtr, + SmallVector args = {statePtr, getOpStateSlotValue(op, rewriter, loc), APtr, BPtr, @@ -145,6 +159,8 @@ struct MatmulOpLowering : public ConvertOpToLLVMPattern { K, batchCount, elemSize, + aBatchCount, + bBatchCount, aBatchStride, bBatchStride}; diff --git a/lib/Conversion/OnnxToHip/GatherConversion.cpp b/lib/Conversion/OnnxToHip/GatherConversion.cpp index 20034cbe0..a4783e155 100644 --- a/lib/Conversion/OnnxToHip/GatherConversion.cpp +++ b/lib/Conversion/OnnxToHip/GatherConversion.cpp @@ -36,47 +36,23 @@ struct GatherToHip : public mlir::RewritePattern { // Get result type auto resultType = mlir::cast(op->getResult(0).getType()); - auto dataType = mlir::cast(data.getType()); - auto indicesType = mlir::cast(indices.getType()); - // Normalize negative axis for dimension calculations only - int64_t normalizedAxis = axis < 0 ? axis + dataType.getRank() : axis; - - // Create output tensor with dynamic shape support - // Output shape: [data[0:axis], indices.shape, data[axis+1:]] - llvm::SmallVector dynSizes; - int64_t outDimIdx = 0; - - // Copy dimensions before axis from data - for (auto i : llvm::seq(0, normalizedAxis)) { - if (outDimIdx < resultType.getRank() && - resultType.isDynamicDim(outDimIdx)) - dynSizes.push_back(mlir::tensor::DimOp::create(rewriter, loc, data, i)); - outDimIdx++; - } - // Copy all dimensions from indices - for (auto i : llvm::seq(0, indicesType.getRank())) { - if (outDimIdx < resultType.getRank() && - resultType.isDynamicDim(outDimIdx)) - dynSizes.push_back( - mlir::tensor::DimOp::create(rewriter, loc, indices, i)); - outDimIdx++; - } - // Copy dimensions after axis from data - for (auto i : llvm::seq(normalizedAxis + 1, dataType.getRank())) { - if (outDimIdx < resultType.getRank() && - resultType.isDynamicDim(outDimIdx)) - dynSizes.push_back(mlir::tensor::DimOp::create(rewriter, loc, data, i)); - outDimIdx++; - } - - mlir::Value init = - mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); + // Output shape is `data[:axis] ++ indices.shape ++ data[axis+1:]`. Use the + // same helper that backs `GatherOp::reifyResultShapes` so the destination + // and the shape consumers observe cannot disagree. + mlir::FailureOr> resultShape = + mlir::hip::reifyGatherWithAxis(rewriter, loc, data, indices, axis); + if (mlir::failed(resultShape)) + return rewriter.notifyMatchFailure(op, "Gather axis is not reifiable"); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "Gather result type is incompatible with the gathered shape"); // Create hip.gather operation auto gatherOp = mlir::hip::GatherOp::create(rewriter, loc, context, data, - indices, init, axisAttr); + indices, *init, axisAttr); rewriter.replaceOp(op, gatherOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/GatherElementsConversion.cpp b/lib/Conversion/OnnxToHip/GatherElementsConversion.cpp index 82f4c03fa..01b706841 100644 --- a/lib/Conversion/OnnxToHip/GatherElementsConversion.cpp +++ b/lib/Conversion/OnnxToHip/GatherElementsConversion.cpp @@ -37,21 +37,22 @@ GatherElementsToHip::matchAndRewrite(mlir::Operation *op, auto resultType = mlir::cast(op->getResult(0).getType()); - auto indicesType = mlir::cast(indices.getType()); - llvm::SmallVector dynSizes; - for (auto i : llvm::seq(0, resultType.getRank())) { - if (resultType.isDynamicDim(i)) - dynSizes.push_back( - mlir::tensor::DimOp::create(rewriter, loc, indices, i)); - } - - mlir::Value init = - mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); + // ONNX GatherElements produces `indices`'s shape. Use the same helper that + // backs `GatherElementsOp::reifyResultShapes` so the two cannot disagree. + mlir::FailureOr> resultShape = + mlir::hip::reifyElementwiseSameShape(rewriter, loc, indices); + if (mlir::failed(resultShape)) + return rewriter.notifyMatchFailure( + op, "GatherElements indices must be a ranked tensor"); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "GatherElements result type is incompatible with indices shape"); auto gatherOp = mlir::hip::GatherElementsOp::create( - rewriter, loc, context, data, indices, init, + rewriter, loc, context, data, indices, *init, rewriter.getI64IntegerAttr(axis)); rewriter.replaceOp(op, gatherOp->getResult(0)); diff --git a/lib/Conversion/OnnxToHip/GatherNDConversion.cpp b/lib/Conversion/OnnxToHip/GatherNDConversion.cpp index 842660f25..7481125ee 100644 --- a/lib/Conversion/OnnxToHip/GatherNDConversion.cpp +++ b/lib/Conversion/OnnxToHip/GatherNDConversion.cpp @@ -46,55 +46,34 @@ struct GatherNDToHip : public mlir::RewritePattern { auto resultType = mlir::cast(op->getResult(0).getType()); - auto dataType = mlir::cast(data.getType()); auto indicesType = mlir::cast(indices.getType()); int64_t batchDims = 0; if (auto attr = op->getAttrOfType("batch_dims")) batchDims = attr.getValue().getSExtValue(); - int64_t q = indicesType.getRank(); - int64_t r = dataType.getRank(); - int64_t outRank = resultType.getRank(); - - // The last dim of `indices` (call it `k`) controls how many leading data - // dims each index tuple consumes. It must be statically known so we can - // compute the data-tail mapping; dynamic-k GatherND is not expressible - // with a single `tensor.empty` dynsize list. - int64_t k = indicesType.getDimSize(q - 1); - if (k == mlir::ShapedType::kDynamic) + // The trailing index-tuple width must be statically known: it determines + // the output rank, so a dynamic one leaves nothing to build a destination + // from. The shared helper bails on the same condition. + if (indicesType.isDynamicDim(indicesType.getRank() - 1)) return rewriter.notifyMatchFailure( op, "GatherND requires static indices.shape[-1]"); - llvm::SmallVector dynSizes; - for (int64_t i = 0; i < outRank; ++i) { - if (!resultType.isDynamicDim(i)) - continue; - - mlir::Value dim; - if (i < batchDims) { - // Batch dim -- ONNX requires data and indices to agree; prefer data - // because we already have it bound for the tail case below. - dim = mlir::tensor::DimOp::create(rewriter, loc, data, i); - } else if (i < q - 1) { - // Indices-outer dim: result[i] == indices[i] for batch_dims <= i < q-1. - dim = mlir::tensor::DimOp::create(rewriter, loc, indices, i); - } else { - // Data-tail dim: result[i] == data[i - (q-1) + batch_dims + k]. - int64_t dataIdx = i - (q - 1) + batchDims + k; - if (dataIdx < 0 || dataIdx >= r) - return rewriter.notifyMatchFailure( - op, "cannot resolve dynamic result dim from data tail"); - dim = mlir::tensor::DimOp::create(rewriter, loc, data, dataIdx); - } - dynSizes.push_back(dim); - } - mlir::Value init = - mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); + // Same helper that backs `GatherNDOp::reifyResultShapes`, so the + // destination and the shape consumers observe cannot disagree. + mlir::FailureOr> resultShape = + mlir::hip::reifyGatherND(rewriter, loc, data, indices, batchDims); + if (mlir::failed(resultShape)) + return rewriter.notifyMatchFailure( + op, "GatherND batch_dims/indices layout is not reifiable"); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "GatherND result type is incompatible with the gathered shape"); auto hipOp = mlir::hip::GatherNDOp::create( - rewriter, loc, context, data, indices, init, + rewriter, loc, context, data, indices, *init, rewriter.getI64IntegerAttr(batchDims)); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); diff --git a/lib/Conversion/OnnxToHip/MatMulNBitsConversion.cpp b/lib/Conversion/OnnxToHip/MatMulNBitsConversion.cpp index fc774bf82..39fda3ea3 100644 --- a/lib/Conversion/OnnxToHip/MatMulNBitsConversion.cpp +++ b/lib/Conversion/OnnxToHip/MatMulNBitsConversion.cpp @@ -109,13 +109,25 @@ MatMulNBitsToHip::matchAndRewrite(mlir::Operation *op, auto zpElemSizeAttr = rewriter.getI64IntegerAttr(zpElemSize); auto rt = mlir::cast(op->getResult(0).getType()); - mlir::Value init = createEmptyTensor(rewriter, loc, rt, A); + // Same helper that backs `MatMulNBitsOp::reifyResultShapes`. A positional + // copy from A would size the last dimension from A's contraction extent (K) + // instead of the output width (N) whenever the imported type leaves it + // dynamic. + mlir::FailureOr> resultShape = + mlir::hip::reifyMatMulNBitsResultShape(rewriter, loc, A, NAttr.getInt()); + if (mlir::failed(resultShape)) + return rewriter.notifyMatchFailure(op, "MatMulNBits A must be ranked"); + mlir::FailureOr init = + createEmptyTensorFromReifiedShape(rewriter, loc, rt, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "MatMulNBits result type is incompatible with the inferred shape"); // Result type inferred from `init` via InferTypeOpInterface — DPS contract: // result type == outs operand type. auto hipOp = mlir::hip::MatMulNBitsOp::create( - rewriter, loc, context, A, B, scales, zeroPoints, gIdx, bias, init, KAttr, - NAttr, bitsAttr, blockSizeAttr, accuracyLevelAttr, zpElemSizeAttr); + rewriter, loc, context, A, B, scales, zeroPoints, gIdx, bias, *init, + KAttr, NAttr, bitsAttr, blockSizeAttr, accuracyLevelAttr, zpElemSizeAttr); rewriter.replaceOp(op, hipOp->getResults()); return mlir::success(); } diff --git a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp index febb3ff95..6517f93a3 100644 --- a/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMaxConversion.cpp @@ -66,7 +66,8 @@ ReduceMaxToHip::matchAndRewrite(mlir::Operation *op, op, "ReduceMax result type is incompatible with the reduction shape"); mlir::Value axesOperand; - if (op->getNumOperands() > 1) { + if (op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType())) { axesOperand = op->getOperand(1); } else { auto axesType = mlir::RankedTensorType::get( diff --git a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp index f7c347755..2091faa84 100644 --- a/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMeanConversion.cpp @@ -84,7 +84,8 @@ ReduceMeanToHip::matchAndRewrite(mlir::Operation *op, // axes is always required in HIP dialect; create empty tensor<0xi64> when not // provided mlir::Value axesOperand; - if (op->getNumOperands() > 1) { + if (op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType())) { // Axes provided as operand (opset 18+) axesOperand = op->getOperand(1); } else { diff --git a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp index 61aabb3fc..a84a05077 100644 --- a/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceMinConversion.cpp @@ -60,7 +60,8 @@ ReduceMinToHip::matchAndRewrite(mlir::Operation *op, op, "ReduceMin result type is incompatible with the reduction shape"); mlir::Value axesOperand; - if (op->getNumOperands() > 1) { + if (op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType())) { axesOperand = op->getOperand(1); } else { auto axesType = mlir::RankedTensorType::get( diff --git a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp index f9d1e3317..043d78ee3 100644 --- a/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp +++ b/lib/Conversion/OnnxToHip/ReduceSumConversion.cpp @@ -66,7 +66,8 @@ ReduceSumToHip::matchAndRewrite(mlir::Operation *op, // axes is always required in HIP dialect; create empty tensor<0xi64> when not // provided mlir::Value axesOperand; - if (op->getNumOperands() > 1) { + if (op->getNumOperands() > 1 && + !mlir::isa(op->getOperand(1).getType())) { // Axes provided as operand (opset 13+) axesOperand = op->getOperand(1); } else { diff --git a/lib/Conversion/OnnxToHip/TransposeConversion.cpp b/lib/Conversion/OnnxToHip/TransposeConversion.cpp index 446f6781e..302b1d6d9 100644 --- a/lib/Conversion/OnnxToHip/TransposeConversion.cpp +++ b/lib/Conversion/OnnxToHip/TransposeConversion.cpp @@ -78,22 +78,21 @@ TransposeToHip::matchAndRewrite(mlir::Operation *op, seen[p] = true; } - // Build dynamic sizes for the DPS init tensor: output dim i = input dim - // perm[i]. - llvm::SmallVector dynSizes; - for (auto [outDimIdx, srcDim] : llvm::enumerate(perm)) { - if (resultType.isDynamicDim(outDimIdx)) - dynSizes.push_back(mlir::tensor::DimOp::create( - rewriter, loc, data, static_cast(srcDim))); - } - - mlir::Value init = - mlir::tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), - resultType.getElementType(), dynSizes); + // Same helper that backs `TransposeOp::reifyResultShapes`, so the + // destination and the shape consumers observe cannot disagree. + mlir::FailureOr> resultShape = + mlir::hip::reifyTransposeByPerm(rewriter, loc, data, perm); + if (mlir::failed(resultShape)) + return rewriter.notifyMatchFailure(op, "Transpose perm is not reifiable"); + mlir::FailureOr init = createEmptyTensorFromReifiedShape( + rewriter, loc, resultType, *resultShape); + if (mlir::failed(init)) + return rewriter.notifyMatchFailure( + op, "Transpose result type is incompatible with the permuted shape"); mlir::ArrayAttr permArrayAttr = rewriter.getI64ArrayAttr(perm); auto hipOp = mlir::hip::TransposeOp::create(rewriter, loc, context, data, - init, permArrayAttr); + *init, permArrayAttr); rewriter.replaceOp(op, hipOp->getResult(0)); return mlir::success(); } diff --git a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp index e69ee0103..2a5e38c19 100644 --- a/lib/Dialect/IR/HipReifyResultShapesImpl.cpp +++ b/lib/Dialect/IR/HipReifyResultShapesImpl.cpp @@ -95,10 +95,11 @@ RopeOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { if (getNumResults() == 0) return failure(); - if (!isa(getInput().getType())) + FailureOr> dims = + mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput()); + if (failed(dims)) return failure(); - reifiedReturnShapes.assign( - {mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput())}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -124,10 +125,11 @@ RmsNormOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { if (getNumResults() == 0) return failure(); - if (!isa(getInput().getType())) + FailureOr> dims = + mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput()); + if (failed(dims)) return failure(); - reifiedReturnShapes.assign( - {mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput())}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -156,10 +158,11 @@ QMoEOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { if (getNumResults() == 0) return failure(); - if (!isa(getInput().getType())) + FailureOr> dims = + mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput()); + if (failed(dims)) return failure(); - reifiedReturnShapes.assign( - {mlir::hip::reifyElementwiseSameShape(b, getLoc(), getInput())}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -189,20 +192,11 @@ LogicalResult MatMulNBitsOp::reifyResultShapes( OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) { if (getNumResults() == 0) return failure(); - ArrayRef aShape = getShapeOf(getA()); - if (aShape.empty()) + FailureOr> dims = + mlir::hip::reifyMatMulNBitsResultShape(b, getLoc(), getA(), getN()); + if (failed(dims)) return failure(); - - Location loc = getLoc(); - Value A = getA(); - size_t aRank = aShape.size(); - SmallVector dims; - dims.reserve(aRank); - // Leading dims (rank-1 of them) from A; final dim is the static N attr. - for (size_t i : llvm::seq(0, aRank - 1)) - dims.push_back(mlir::hip::reifyDimOrConstant(b, loc, aShape[i], A, i)); - dims.push_back(b.getIndexAttr(getN())); - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -333,11 +327,11 @@ LogicalResult TransposeOp::reifyResultShapes( perm.push_back(ia.getInt()); } - SmallVector dims = + FailureOr> dims = mlir::hip::reifyTransposeByPerm(b, getLoc(), getInput(), perm); - if (dims.empty()) + if (failed(dims)) return failure(); - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -350,11 +344,11 @@ GatherOp::reifyResultShapes(OpBuilder &b, !isa(getIndices().getType())) return failure(); - SmallVector dims = mlir::hip::reifyGatherWithAxis( + FailureOr> dims = mlir::hip::reifyGatherWithAxis( b, getLoc(), getData(), getIndices(), getAxis()); - if (dims.empty()) + if (failed(dims)) return failure(); - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } @@ -363,14 +357,11 @@ GatherElementsOp::reifyResultShapes(OpBuilder &b, ReifiedRankedShapedTypeDims &reified) { if (getNumResults() == 0) return failure(); - auto indicesType = dyn_cast(getIndices().getType()); - if (!indicesType) + FailureOr> dims = + mlir::hip::reifyElementwiseSameShape(b, getLoc(), getIndices()); + if (failed(dims)) return failure(); - - SmallVector dims; - for (auto i : llvm::seq(0, indicesType.getRank())) - dims.push_back(tensor::getMixedSize(b, getLoc(), getIndices(), i)); - reified.assign({std::move(dims)}); + reified.assign({std::move(*dims)}); return success(); } @@ -478,11 +469,11 @@ LogicalResult GatherNDOp::reifyResultShapes( !isa(getIndices().getType())) return failure(); - SmallVector dims = mlir::hip::reifyGatherND( + FailureOr> dims = mlir::hip::reifyGatherND( b, getLoc(), getData(), getIndices(), getBatchDims()); - if (dims.empty()) + if (failed(dims)) return failure(); - reifiedReturnShapes.assign({std::move(dims)}); + reifiedReturnShapes.assign({std::move(*dims)}); return success(); } diff --git a/lib/Dialect/IR/HipShapeUtils.cpp b/lib/Dialect/IR/HipShapeUtils.cpp index 54363567e..f55a08001 100644 --- a/lib/Dialect/IR/HipShapeUtils.cpp +++ b/lib/Dialect/IR/HipShapeUtils.cpp @@ -268,13 +268,11 @@ mlir::hip::inferGemmShape(ArrayRef aShape, ArrayRef bShape, LogicalResult mlir::hip::verifyHipOpShape( Operation *op, function_ref>()> inferShape) { - // Asserting cast: every op wired to verifyHipOpShape also implements DPS - // via TableGen; a missing interface is a programmer error in the op def. - auto dpsOp = cast(op); + auto dpsOp = dyn_cast(op); + if (!dpsOp) + return op->emitOpError( + "shape verification requires DestinationStyleOpInterface"); auto inits = dpsOp.getDpsInits(); - assert(inits.size() == 1 && - "verifyHipOpShape covers single-destination ops; a multi-destination " - "op needs one expected shape per init"); if (inits.size() != 1) return op->emitOpError("expected a single DPS init operand, got ") << inits.size(); @@ -318,11 +316,11 @@ OpFoldResult mlir::hip::reifyDimOrConstant(OpBuilder &b, Location loc, return tensor::getMixedSize(b, loc, source, sourceDim); } -SmallVector +FailureOr> mlir::hip::reifyElementwiseSameShape(OpBuilder &b, Location loc, Value source) { - // Caller must hand a ranked tensor; reify is only invoked in tensor mode - // per the ReifyRankedShapedTypeOpInterface contract. - auto sourceType = cast(source.getType()); + auto sourceType = dyn_cast(source.getType()); + if (!sourceType) + return failure(); ArrayRef shape = sourceType.getShape(); SmallVector dims; dims.reserve(shape.size()); @@ -445,6 +443,22 @@ FailureOr> mlir::hip::reifyMatmulResultShape( return result; } +FailureOr> +mlir::hip::reifyMatMulNBitsResultShape(OpBuilder &b, Location loc, Value A, + int64_t N) { + auto aType = dyn_cast(A.getType()); + if (!aType || aType.getRank() < 1) + return failure(); + + ArrayRef aShape = aType.getShape(); + SmallVector dims; + dims.reserve(aShape.size()); + for (size_t i : llvm::seq(0, aShape.size() - 1)) + dims.push_back(reifyDimOrConstant(b, loc, aShape[i], A, i)); + dims.push_back(b.getIndexAttr(N)); + return dims; +} + FailureOr> mlir::hip::reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, Value optionalC, int64_t transA, int64_t transB, @@ -478,41 +492,47 @@ mlir::hip::reifyGemmResultShape(OpBuilder &b, Location loc, Value A, Value B, bSizes[transB ? 0 : 1]}; } -SmallVector +FailureOr> mlir::hip::reifyTransposeByPerm(OpBuilder &b, Location loc, Value input, ArrayRef perm) { auto inputType = dyn_cast(input.getType()); if (!inputType) - return {}; + return failure(); ArrayRef inputShape = inputType.getShape(); int64_t rank = inputType.getRank(); if (static_cast(perm.size()) != rank) - return {}; + return failure(); + + // Validate the entire permutation before materializing any tensor.dim. A + // failure must leave the IR unchanged. + llvm::SmallBitVector seen(rank); + for (int64_t dim : perm) { + if (dim < 0 || dim >= rank || seen.test(dim)) + return failure(); + seen.set(dim); + } SmallVector dims; dims.reserve(perm.size()); - for (int64_t pi : perm) { - if (pi < 0 || pi >= rank) - return {}; - dims.push_back(reifyDimOrConstant(b, loc, inputShape[pi], input, pi)); - } + for (int64_t dim : perm) + dims.push_back(reifyDimOrConstant(b, loc, inputShape[dim], input, dim)); return dims; } -SmallVector +FailureOr> mlir::hip::reifyGatherWithAxis(OpBuilder &b, Location loc, Value data, Value indices, int64_t axis) { auto dataType = dyn_cast(data.getType()); auto indicesType = dyn_cast(indices.getType()); if (!dataType || !indicesType) - return {}; + return failure(); int64_t dataRank = dataType.getRank(); int64_t indicesRank = indicesType.getRank(); // Negative-axis normalization (ONNX convention). if (axis < 0) axis += dataRank; if (axis < 0 || axis >= dataRank) - return {}; + return failure(); ArrayRef dataShape = dataType.getShape(); ArrayRef indicesShape = indicesType.getShape(); @@ -529,17 +549,17 @@ mlir::hip::reifyGatherWithAxis(OpBuilder &b, Location loc, Value data, return dims; } -SmallVector mlir::hip::reifyGatherND(OpBuilder &b, Location loc, - Value data, Value indices, - int64_t batchDims) { +FailureOr> +mlir::hip::reifyGatherND(OpBuilder &b, Location loc, Value data, Value indices, + int64_t batchDims) { auto dataType = dyn_cast(data.getType()); auto indicesType = dyn_cast(indices.getType()); if (!dataType || !indicesType) - return {}; + return failure(); int64_t dataRank = dataType.getRank(); int64_t indicesRank = indicesType.getRank(); if (indicesRank < 1) - return {}; + return failure(); ArrayRef dataShape = dataType.getShape(); ArrayRef indicesShape = indicesType.getShape(); @@ -548,10 +568,10 @@ SmallVector mlir::hip::reifyGatherND(OpBuilder &b, Location loc, // synthesise a rank with a dynamic count of dim entries. int64_t tupleWidth = indicesShape[indicesRank - 1]; if (ShapedType::isDynamic(tupleWidth)) - return {}; + return failure(); if (batchDims < 0 || batchDims > indicesRank - 1 || batchDims + tupleWidth > dataRank) - return {}; + return failure(); // Output = data.shape[:batch_dims] (the shared batch prefix) ++ // indices.shape[batch_dims:-1] (the gathered tuple count) ++ diff --git a/lib/Runtime/hipdnn_ep_runtime.h b/lib/Runtime/hipdnn_ep_runtime.h index 370f6d022..0c44f4cd8 100644 --- a/lib/Runtime/hipdnn_ep_runtime.h +++ b/lib/Runtime/hipdnn_ep_runtime.h @@ -418,11 +418,16 @@ int hipdnn_ep_state_ensure_matmul_dp4a_scratch(RuntimeState *state, // walks the array via each object's deletor. bool hipdnn_ep_op_states_alloc(RuntimeState *state, int64_t n); -// Device-side runtime error flag (set by kernels, observed by wrappers). -// Intended for operators that detect runtime-invalid inputs on GPU (e.g. Range -// delta==0) and need to propagate an error code back through main_graph. +// Device-side runtime error flag (set by kernels or queued by wrappers). +// Intended for operators that detect runtime-invalid inputs (e.g. Range +// delta==0 or a dynamically concealed partial MatMul batch broadcast) and need +// to propagate an error code back through main_graph. void *hipdnn_ep_state_get_error_flag_device_ptr(RuntimeState *state); int hipdnn_ep_state_reset_error_flag(RuntimeState *state); +// Record a host-detected runtime error on the stream. The generated interface +// observes it at the same boundary as device-detected errors and returns a +// recoverable non-zero status to ORT. +int hipdnn_ep_state_set_error_flag(RuntimeState *state); int hipdnn_ep_state_read_and_clear_error_flag(RuntimeState *state); // Mark the start of a new Compute() call. Invalidates per-forward-pass // runtime caches -- today: the GQA seqlens_k cache (see @@ -807,9 +812,15 @@ int wrap_hipblasLtGemm(void *handle, // hipBLASLt handle // `a_batch_stride` / `b_batch_stride` are hipBLASLt's per-batch advances in // elements. A stride is 0 when one matrix is broadcast across all batches; // otherwise it is M*K for A or K*N for B. +// +// `a_batch_count` / `b_batch_count` are validated at runtime because dynamic +// batch extents can conceal a partial per-axis broadcast from the static +// verifier. Each count must be either 1 or `batch_count`; otherwise the wrapper +// records a recoverable runtime error and does not dispatch hipBLASLt. int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, + int64_t a_batch_count, int64_t b_batch_count, int64_t a_batch_stride, int64_t b_batch_stride); // GroupQueryAttention operation wrapper (Full MS spec) diff --git a/lib/Runtime/hipdnn_ep_runtime_state.cpp b/lib/Runtime/hipdnn_ep_runtime_state.cpp index e29a06791..02f8e7fe4 100644 --- a/lib/Runtime/hipdnn_ep_runtime_state.cpp +++ b/lib/Runtime/hipdnn_ep_runtime_state.cpp @@ -1429,6 +1429,19 @@ int hipdnn_ep_state_reset_error_flag(RuntimeState *state) { return (err == hipSuccess) ? 0 : -1; } +int hipdnn_ep_state_set_error_flag(RuntimeState *state) { + if (!state || !state->device_error_flag || !state->stream) { + fprintf(stderr, "hipdnn_ep_state_set_error_flag: invalid state\n"); + return -1; + } + // Byte-wise 0xff is the two's-complement representation of the runtime's + // generic error code (-1). Queue it on the inference stream so the interface + // boundary observes it after the existing stream synchronization. + hipError_t err = hipMemsetAsync(state->device_error_flag, 0xff, sizeof(int), + state->stream); + return (err == hipSuccess) ? 0 : -1; +} + int hipdnn_ep_state_read_and_clear_error_flag(RuntimeState *state) { if (!state || !state->device_error_flag || !state->stream) { fprintf(stderr, diff --git a/lib/Runtime/mock/mock_gpu.cpp b/lib/Runtime/mock/mock_gpu.cpp index 2f02ceed8..7a7f40c45 100644 --- a/lib/Runtime/mock/mock_gpu.cpp +++ b/lib/Runtime/mock/mock_gpu.cpp @@ -595,18 +595,25 @@ int wrap_hipblasLtGemm(void *handle, void *stream, int64_t m, int64_t n, int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, + int64_t a_batch_count, int64_t b_batch_count, int64_t a_batch_stride, int64_t b_batch_stride) { (void)op_state_slot; if (!state) { fprintf(stderr, "Invalid state in wrap_hipblasLtMatmul\n"); return -1; } + if ((a_batch_count != 1 && a_batch_count != batch_count) || + (b_batch_count != 1 && b_batch_count != batch_count)) { + (void)hipdnn_ep_state_set_error_flag(state); + return -1; + } MOCK_PRINT("[MOCK] wrap_hipblasLtMatmul(M=%lld, N=%lld, K=%lld, " - "batch=%lld, elem_size=%lld, a_batch_stride=%lld, " - "b_batch_stride=%lld)\n", + "batch=%lld, elem_size=%lld, a_batches=%lld, b_batches=%lld, " + "a_batch_stride=%lld, b_batch_stride=%lld)\n", (long long)M, (long long)N, (long long)K, (long long)batch_count, - (long long)elem_size, (long long)a_batch_stride, + (long long)elem_size, (long long)a_batch_count, + (long long)b_batch_count, (long long)a_batch_stride, (long long)b_batch_stride); return 0; diff --git a/lib/Runtime/real/matmul.cpp b/lib/Runtime/real/matmul.cpp index 8410d4537..802b2e2ab 100644 --- a/lib/Runtime/real/matmul.cpp +++ b/lib/Runtime/real/matmul.cpp @@ -419,6 +419,7 @@ static void autotuneMatmul(hipblasLtHandle_t handle, hipStream_t stream, int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const void *B, void *output, int64_t M, int64_t N, int64_t K, int64_t batch_count, int64_t elem_size, + int64_t a_batch_count, int64_t b_batch_count, int64_t a_batch_stride, int64_t b_batch_stride) { OP_PROFILE( "matmul", @@ -434,6 +435,22 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, return -1; } + auto isRepresentable = [batch_count](int64_t operandBatchCount) { + return operandBatchCount == 1 || operandBatchCount == batch_count; + }; + if (!isRepresentable(a_batch_count) || !isRepresentable(b_batch_count)) { + fprintf(stderr, + "wrap_hipblasLtMatmul: runtime batch layout is not representable " + "by one stride per operand (A batches=%lld, B batches=%lld, " + "output batches=%lld)\n", + (long long)a_batch_count, (long long)b_batch_count, + (long long)batch_count); + // The generated interface reads this flag after stream synchronization and + // returns a non-zero inference status to ORT. + (void)hipdnn_ep_state_set_error_flag(state); + return -1; + } + hipblasLtHandle_t handle = static_cast(hipdnn_ep_state_get_hipblas_handle(state)); hipStream_t stream = @@ -452,10 +469,12 @@ int wrap_hipblasLtMatmul(RuntimeState *state, int op_state_slot, const void *A, const char *type_name = (elem_size == 2) ? "f16" : "f32"; RUNTIME_DEBUG_LOG("[REAL] wrap_hipblasLtMatmul: M=%lld, N=%lld, K=%lld, " - "batch=%lld, a_batch_stride=%lld, b_batch_stride=%lld, " + "batch=%lld, a_batches=%lld, b_batches=%lld, " + "a_batch_stride=%lld, b_batch_stride=%lld, " "elem_size=%lld (%s), total_bytes=%lld\n", (long long)M, (long long)N, (long long)K, - (long long)batch_count, (long long)a_batch_stride, + (long long)batch_count, (long long)a_batch_count, + (long long)b_batch_count, (long long)a_batch_stride, (long long)b_batch_stride, (long long)elem_size, type_name, (long long)(batch_count * M * N * elem_size)); diff --git a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir index 44fa9095f..b529e3227 100644 --- a/test/lit/Conversion/hip-to-llvm/test_matmul.mlir +++ b/test/lit/Conversion/hip-to-llvm/test_matmul.mlir @@ -3,9 +3,10 @@ // RUN: hip-mlir-opt --convert-hip-to-llvm %s | FileCheck %s -// `wrap_hipblasLtMatmul` takes 12 arguments: +// `wrap_hipblasLtMatmul` takes 14 arguments: // - 4 pointers: state, A, B, output -// - 7 i64: M, N, K, batch_count, elem_size, A stride, B stride +// - 9 i64: M, N, K, output batch count, element size, +// A/B operand batch counts, A/B strides // - 1 i32: op_state_slot (-1 here; --assign-op-state-slots is not in this RUN) // // Each operand's stride is 0 when it holds a single matrix broadcast across the @@ -35,9 +36,11 @@ module { // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 // CHECK: %[[ELEM0:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[A0_COUNT:.*]] = llvm.mlir.constant(1 : i64) : i64 +// CHECK: %[[B0_COUNT:.*]] = llvm.mlir.constant(1 : i64) : i64 // CHECK: %[[A0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 // CHECK: %[[B0_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM0]], %[[A0_STRIDE]], %[[B0_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM0]], %[[A0_COUNT]], %[[B0_COUNT]], %[[A0_STRIDE]], %[[B0_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64, i64, i64) -> i32 // ----- Rank-3 leading-one B [1, K, N] ---------------------------------------- // B's leading dim is statically 1, so the buffer still holds ONE [K, N] matrix @@ -62,9 +65,11 @@ module { // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 // CHECK: %[[ELEM1:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[A1_COUNT:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[B1_COUNT:.*]] = llvm.mlir.constant(1 : i64) : i64 // CHECK: %[[A1_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 // CHECK: %[[B1_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM1]], %[[A1_STRIDE]], %[[B1_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM1]], %[[A1_COUNT]], %[[B1_COUNT]], %[[A1_STRIDE]], %[[B1_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64, i64, i64) -> i32 // ----- Rank-2 broadcast A against rank-3 B ----------------------------------- // The mirror image of the case above: A holds one matrix (stride 0) while B and @@ -88,16 +93,19 @@ module { // CHECK-NOT: llvm.select // CHECK-NOT: @wrap_hipblasLtMatmul_v2 // CHECK: %[[ELEM2:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[A2_COUNT:.*]] = llvm.mlir.constant(1 : i64) : i64 +// CHECK: %[[B2_COUNT:.*]] = llvm.mlir.constant(2 : i64) : i64 // CHECK: %[[A2_STRIDE:.*]] = llvm.mlir.constant(0 : i64) : i64 // CHECK: %[[B2_STRIDE:.*]] = llvm.mul %{{.*}}, %{{.*}} : i64 -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM2]], %[[A2_STRIDE]], %[[B2_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM2]], %[[A2_COUNT]], %[[B2_COUNT]], %[[A2_STRIDE]], %[[B2_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64, i64, i64) -> i32 // ----- Dynamic batch on both operands ---------------------------------------- // `[?, 8, M, K] @ [?, 8, K, N]` is an ordinary batched matmul: neither operand // broadcasts, so both are representable and must NOT be rejected. Because the // leading extents are dynamic, each stride is chosen at runtime by comparing the -// operand's matrix count against the output's batch count -- which also keeps -// the read in bounds if the counts ever disagree. +// operand's matrix count against the output's batch count. Both counts are also +// passed to the wrapper, which reports a recoverable error if either is neither +// 1 nor the output count (a dynamically concealed partial broadcast). module { func.func @test_matmul_dynamic_batch(%ctx: !hip.context, @@ -113,8 +121,38 @@ module { // CHECK-LABEL: llvm.func @test_matmul_dynamic_batch // CHECK: %[[ELEM3:.*]] = llvm.mlir.constant(2 : i64) : i64 -// CHECK: %[[A3_CMP:.*]] = llvm.icmp "eq" +// CHECK: %[[A3_PART:.*]] = llvm.mul +// CHECK: %[[A3_COUNT:.*]] = llvm.mul %[[A3_PART]], +// CHECK: %[[B3_PART:.*]] = llvm.mul +// CHECK: %[[B3_COUNT:.*]] = llvm.mul %[[B3_PART]], +// CHECK: %[[A3_CMP:.*]] = llvm.icmp "eq" %[[A3_COUNT]], // CHECK: %[[A3_STRIDE:.*]] = llvm.select %[[A3_CMP]] -// CHECK: %[[B3_CMP:.*]] = llvm.icmp "eq" +// CHECK: %[[B3_CMP:.*]] = llvm.icmp "eq" %[[B3_COUNT]], // CHECK: %[[B3_STRIDE:.*]] = llvm.select %[[B3_CMP]] -// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM3]], %[[A3_STRIDE]], %[[B3_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64) -> i32 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM3]], %[[A3_COUNT]], %[[B3_COUNT]], %[[A3_STRIDE]], %[[B3_STRIDE]]) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64, i64, i64) -> i32 + +// ----- Dynamic extent can conceal a partial per-axis broadcast --------------- +// A's runtime batch count is either 2 (when ? == 1, not representable) or 6 +// (when ? == 3, one matrix per output batch). The static verifier cannot +// distinguish those invocations, so the wrapper receives both operand counts +// and reports a recoverable error for the count-2 case. + +module { + func.func @test_matmul_dynamic_partial_guard( + %ctx: !hip.context, + %A: memref<2x?x4x16xf16, 1>, + %B: memref<2x3x16x32xf16, 1>, + %output: memref<2x3x4x32xf16, 1>) { + hip.matmul(%ctx) + ins(%A, %B : memref<2x?x4x16xf16, 1>, memref<2x3x16x32xf16, 1>) + outs(%output : memref<2x3x4x32xf16, 1>) + return + } +} + +// CHECK-LABEL: llvm.func @test_matmul_dynamic_partial_guard +// CHECK: %[[ELEM4:.*]] = llvm.mlir.constant(2 : i64) : i64 +// CHECK: %[[A4_PART:.*]] = llvm.mul +// CHECK: %[[A4_COUNT:.*]] = llvm.mul %[[A4_PART]], +// CHECK: %[[B4_COUNT:.*]] = llvm.mlir.constant(6 : i64) : i64 +// CHECK: llvm.call @wrap_hipblasLtMatmul({{.*}}, %[[ELEM4]], %[[A4_COUNT]], %[[B4_COUNT]], %{{.*}}, %{{.*}}) : (!llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, i64, i64, i64) -> i32 diff --git a/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir b/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir index 4ed90572d..df1a8e837 100644 --- a/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir +++ b/test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir @@ -113,4 +113,21 @@ module { // CHECK: %[[OD3:.*]] = tensor.dim %[[DATA]], %[[OC3]] : tensor // CHECK: %[[OINIT:.*]] = tensor.empty(%[[OD0]], %[[OD3]]) : tensor // CHECK: hip.reduce_sum(%[[CTX]]) ins(%[[DATA]], %{{.*}} : tensor, tensor<2xi64>) outs(%[[OINIT]] : tensor) {keepdims = 0 : i64} + + // An explicit ONNX NoValue is semantically the absent axes operand. The HIP + // op still requires an axes tensor, so conversion must materialize all axes + // rather than forwarding a `none` value. + func.func @reduce_sum_none_axes(%data: tensor) -> tensor { + %none = "onnx.NoValue"() {value} : () -> none + %output = "onnx.ReduceSum"(%data, %none) + {keepdims = 0 : si64, noop_with_empty_axes = 0 : si64} + : (tensor, none) -> tensor + return %output : tensor + } + + // CHECK-LABEL: func.func @reduce_sum_none_axes + // CHECK: %[[AXES:.*]] = arith.constant dense<[0, 1]> : tensor<2xi64> + // CHECK: %[[INIT:.*]] = tensor.empty() : tensor + // CHECK: hip.reduce_sum({{.*}}) ins({{.*}}, %[[AXES]] : tensor, tensor<2xi64>) outs(%[[INIT]] : tensor) {keepdims = 0 : i64} + // CHECK-NOT: onnx.NoValue }