You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
feat(onnx-to-hip): lower MatMulBnb4, outer-product Einsum, and empty slices - #1255
Lower com.microsoft.MatMulBnb4 to hip.dequant_bnb4 plus hip.matmul so the op bufferizes. FP4 and NF4 unpack high nibble first. Static shapes are absorbed into the rocMLIR matmul kernel; a dynamic activation keeps the dequant on a HIP kernel and the GEMM on hipBLASLt.
Lower a binary Einsum with no contracted index, such as bc,d->bcd, to a permute, unit-axis reshape, and hip.mul for f16 and f32. Rank stays at most 4. f64 stays onnx.Einsum.
Keep an empty onnx.Slice and a static zero-length onnx.Concat input inside the tensor.extract_slice / tensor.insert_slice verifier. A zero-length view uses offset 0, and a trailing empty concat input is skipped.
Related issue or design
None.
Why
MatMulBnb4 stayed as an ONNX op, so one-shot bufferization rejected the graph. The existing Einsum lowering only accepts a contraction with static M, N, and K, so an equation with no contracted index never bufferized. An empty shape-vector slice (tensor<3xi64> from index 3 through INT64_MAX) then failed verification as offset 0 is out-of-bounds: 3 >= 3: the leading 0 is the dimension index, and the offset value 3 is not strictly inside a length-3 dimension. Concatenating that empty tensor<0xi64> hits the same rule when the insert lands at offset == result dim.
What
hip.dequant_bnb4 writes static [N, K] f32. quant_type 0 is FP4 and 1 is NF4. Block sizes are 16, 32, 64, 128, and 256. transB must be 1 and training_mode must be 0. The op verifier requires ui8 packed weights, f32 absmax and output, a static [N, K] result, and packed and absmax lengths of N * K / 2 and N * K / block_size. The HIP kernel clears the last-error state before launch and returns hipGetLastError() afterward. The LUT matches the ONNX Runtime blockwise BNB4 tables. This path does not use tosa.matmul_t_block_scaled.
FuseROCMlir absorbs a single-use dequant into the matmul rock kernel when the activation shape is static. HipToTosa unpacks the high nibble, gathers the LUT, tiles the block absmax, and multiplies.
An outer-product Einsum permutes each operand into output-letter order, inserts size-1 axes for missing letters, and multiplies. The element type is f16 or f32. Contractions with a K dimension still go through transpose, collapse, and hip.matmul.
A size-0 extract_slice of a positive dimension uses offset 0. A source dimension of 0 becomes tensor.empty. The hip.slice zero-extent path is unchanged: a runtime empty slice still allocates full capacity so a loop accumulator can append. A concat input whose concat-axis extent is statically 0 is not inserted.
Test plan
hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hip with FileCheck: test_matmul_bnb4.mlir, test_einsum.mlir (contraction and outer product, including f64 left as onnx.Einsum), test_slice.mlir (past-the-end and already-empty), test_concat.mlir (leading and trailing empty)
hip-mlir-opt --verify-diagnostics: hip-dequant-bnb4-verifier.mlir for valid tensor and memref forms, and for an f16 output, a short packed buffer, and an unsupported block size
hip-mlir-opt --hipdnn-pipeline: test_matmul_bnb4_model.mlir and test_einsum_outer_model.mlir
test/numeric/tests/test_matmul_bnb4.py was added and not run on a GPU
Notes for reviewers
A dynamic MatMulBnb4 activation is not fused. training_mode != 0, transB != 1, and block sizes outside 16/32/64/128/256 stay unconverted.
Einsum with ellipsis, repeated letters, a dynamic extent, rank above 4, an f64 outer product, or an empty M or N on a real contraction stays onnx.Einsum.
Cursor Grok 4.7 drafted the conversions, the verifier, the tests, and this description.
Checklist
The change is focused, or links a design/series explaining its scope.
Relevant tests were added or updated and the results are documented.
User-facing or design documentation was updated when needed.
Substantial AI assistance is disclosed, and I reviewed and understand the result.
This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.
Before requesting review, please check that:
The change is focused. Substantial work links the relevant issue
or design discussion.
The PR documents relevant test results and updates affected
documentation.
If AI tools provided substantial assistance, the description
explains what was assisted and how it was validated, and commit
trailers identify the tool. The contributor has reviewed and
understands the result.
Reviewers are assigned through CODEOWNERS where ownership
is configured.
…slices
MatMulBnb4 decomposes into a high-nibble dequant plus matmul so the op
bufferizes on the HIP and rocMLIR paths. An Einsum with no contracted
index lowers to a reshape and hip.mul. Empty Slice and Concat views no
longer place the offset at the dimension length.
Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com>
Made-with: Cursor
hip.mul's runtime multiply accepts f16 and f32. An f64 equation such as i,j->ij was rewritten anyway, and wrap_elementwise then returned failure without writing the output.
Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com>
Made-with: Cursor
The HIP-to-LLVM lowering forwards the packed, absmax, and output pointers without checking them. An f16 output with the same element count is only half as large as the f32 buffer the kernel writes.
Co-Authored-By: Cursor Grok 4.7 <cursoragent@cursor.com>
Made-with: Cursor
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The new rocMLIR fusion path lacks regression coverage, and GPU numerical validation remains unrun.
No open issues 🎉
1 previously missed issue
In code that hasn't changed since last review
Add LIT coverage for dequantization fusion capture and remapping
lib/Dialect/Transforms/FuseROCMlir.cpp:138
Please add LIT coverage for this dequantization capture/remapping, following test/lit/Dialect/hip-fuse-rocmlir.mlir. The new --hipdnn-pipeline test does not run this pass, and the HipToTosa test starts inside an already outlined function, so neither checks whether fusion produces a valid isolated function. Test static single-use fusion, verify that packed weights and absmax become kernel arguments, re-verify the outlined function, and run HipToTosa on it. Also check that dynamic activations bypass fusion and that a dequantization result with multiple uses is not absorbed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
large-prSoft size threshold (>500 LOC or >10 files). Reviewer signal, not a block.
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
github-actions
Bot
added
the
large-pr
Summary
com.microsoft.MatMulBnb4tohip.dequant_bnb4pluship.matmulso the op bufferizes. FP4 and NF4 unpack high nibble first. Static shapes are absorbed into the rocMLIR matmul kernel; a dynamic activation keeps the dequant on a HIP kernel and the GEMM on hipBLASLt.bc,d->bcd, to a permute, unit-axis reshape, andhip.mulfor f16 and f32. Rank stays at most 4. f64 staysonnx.Einsum.onnx.Sliceand a static zero-lengthonnx.Concatinput inside thetensor.extract_slice/tensor.insert_sliceverifier. A zero-length view uses offset 0, and a trailing empty concat input is skipped.Related issue or design
None.
Why
MatMulBnb4stayed as an ONNX op, so one-shot bufferization rejected the graph. The existing Einsum lowering only accepts a contraction with static M, N, and K, so an equation with no contracted index never bufferized. An empty shape-vector slice (tensor<3xi64>from index 3 throughINT64_MAX) then failed verification asoffset 0 is out-of-bounds: 3 >= 3: the leading 0 is the dimension index, and the offset value 3 is not strictly inside a length-3 dimension. Concatenating that emptytensor<0xi64>hits the same rule when the insert lands atoffset == result dim.What
hip.dequant_bnb4writes static[N, K]f32.quant_type0 is FP4 and 1 is NF4. Block sizes are 16, 32, 64, 128, and 256.transBmust be 1 andtraining_modemust be 0. The op verifier requires ui8 packed weights, f32 absmax and output, a static[N, K]result, and packed and absmax lengths ofN * K / 2andN * K / block_size. The HIP kernel clears the last-error state before launch and returnshipGetLastError()afterward. The LUT matches the ONNX Runtime blockwise BNB4 tables. This path does not usetosa.matmul_t_block_scaled.hip.matmul.extract_sliceof a positive dimension uses offset 0. A source dimension of 0 becomestensor.empty. Thehip.slicezero-extent path is unchanged: a runtime empty slice still allocates full capacity so a loop accumulator can append. A concat input whose concat-axis extent is statically 0 is not inserted.Test plan
hip-mlir-opt --hip-add-context-arg --convert-onnx-to-hipwith FileCheck:test_matmul_bnb4.mlir,test_einsum.mlir(contraction and outer product, including f64 left asonnx.Einsum),test_slice.mlir(past-the-end and already-empty),test_concat.mlir(leading and trailing empty)hip-mlir-opt --convert-hip-to-tosa:test_dequant_bnb4.mlirhip-mlir-opt --verify-diagnostics:hip-dequant-bnb4-verifier.mlirfor valid tensor and memref forms, and for an f16 output, a short packed buffer, and an unsupported block sizehip-mlir-opt --hipdnn-pipeline:test_matmul_bnb4_model.mlirandtest_einsum_outer_model.mlirtest/numeric/tests/test_matmul_bnb4.pywas added and not run on a GPUNotes for reviewers
training_mode != 0,transB != 1, and block sizes outside 16/32/64/128/256 stay unconverted.onnx.Einsum.Checklist