[Mlir-commits] [mlir] ef739b9 - [AMDGPU] Correct gfx950 smfmac sparse index verifier (#193541)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Apr 24 13:47:22 PDT 2026
Author: Eric Feng
Date: 2026-04-24T13:47:17-07:00
New Revision: ef739b97b108d81de23f9fc0f6ca6de001482164
URL: https://github.com/llvm/llvm-project/commit/ef739b97b108d81de23f9fc0f6ca6de001482164
DIFF: https://github.com/llvm/llvm-project/commit/ef739b97b108d81de23f9fc0f6ca6de001482164.diff
LOG: [AMDGPU] Correct gfx950 smfmac sparse index verifier (#193541)
Originally, the smfmac verifier expects for the sparse indices, which
describe which the positions of the non-zero elements per lane, the
following:
```
8 bit source -> require vector<2xi16>, ABID range [0, 1]
16 bit source -> require vector<4xi8>, ABID range [0, 3]
```
which is correct for CDNA3 and what was stated in the CDNA4 ISA
description as well. However, because the CDNA4 variants have double K
of the CDNA3 variants, meaning e.g, for 16 bit variants, each lane
carries 8 non-zero values rather than 4, we need 16 bit sparse indices
to express the full range of non-zero elements. This is in line with the
layout tables presented in the CDNA4 ISA.
Direct comparison for 16 bit elements:
On gfx942; we can select from one of four 8-bit sets of sparse indices
with ABID. Each set represents the location of four non-zero values per
8 following 4:2 structured sparsity. For example:
```
a0 a1 0 0 a3 a4 0 0 | a5 a6 0 0 0 0 a7 a8 | a9 0 0 a10 0 a11 0 a12 | 0 a13 0 a14 0 a15 0
```
On gfx950; because each lane carries 8 non-zero values; we can only
specify the full range of 8 non-zero values per 16 from one of two
16-bit sets. For example:
```
a0 a1 0 0 a3 a4 0 0 a5 a6 0 0 0 0 a7 a8 | a9 0 0 a10 0 a11 0 a12 0 a13 0 a14 0 a15 0
```
Similarly, for 8 bit variants on gfx950, we would need the full 32 bits
to describe the full range of the locations for the 16 non-zero 8 bit
elements. In this case, there is no option to select from different sets
of indices.
The issue arises in downstream use cases if we want to use use a set of
sparse indices targeting gfx950; because we are unable to specify the
full range of the non-zero values at the moment, we will get numerical
issues.
Assisted by: Claude
---------
Signed-off-by: Eric Feng <Eric.Feng at amd.com>
Added:
Modified:
mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
mlir/test/Dialect/AMDGPU/invalid.mlir
mlir/test/Dialect/AMDGPU/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
index 1b396e484d807..4112ea281bb96 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
@@ -938,7 +938,8 @@ def SMFMACOutTypes : AnyTypeOf<[
def SMFMACIdxTypes : AnyTypeOf<[
FixedVectorOfLengthAndType<[4], [I8]>,
- FixedVectorOfLengthAndType<[2], [I16]>
+ FixedVectorOfLengthAndType<[2], [I16]>,
+ I32
]>;
// scaled_mfma
@@ -1177,14 +1178,24 @@ def AMDGPU_SparseMFMAOp :
- M=N=16, K=64 and M=N=32, K=32 for f16 and bf16 sources
- M=N=16, K=128 and M=N=32, K=64 for i8 and fp8 sources
- The `sparseIdx` parameter contains packed indices identifying the positions
- of non-zero elements in the 2:4 sparse matrix A. For 16-bit source data,
- use `vector<4xi8>` (four 8-bit indices). For 8-bit source data, use
- `vector<2xi16>` (two 16-bit indices).
-
- The `cbsz` and `abid` parameters are repurposed to select the index set.
- If `cbsz == 0`, then `abid[1:0]` selects which index set to use.
- If `cbsz != 0`, then the very first is selected.
+ The `sparseIdx` parameter contains packed 2-bit indices identifying which
+ of every 4 dense-K positions are non-zero in the 2:4 sparse matrix A.
+ The required `sparseIdx` type depends on the variant:
+ - gfx942 16-bit (`(m,k)` in `{(16,32), (32,16)}`): 8 bits per lane,
+ carried as `vector<4xi8>` (one 8-bit set per i8 element).
+ - gfx942 8-bit (`(m,k)` in `{(16,64), (32,32)}`) and gfx950 16-bit
+ (`(m,k)` in `{(16,64), (32,32)}`): 16 bits per lane, carried as
+ `vector<2xi16>` (one 16-bit set per i16 element).
+ - gfx950 8-bit (`(m,k)` in `{(16,128), (32,64)}`): 32 bits per lane (a
+ full VGPR with no internal set structure), carried as `i32`.
+
+ The `cbsz` and `abid` parameters select which index set within the VGPR is
+ used:
+ - gfx942 16-bit: `cbsz == 0` selects one of four 8-bit sets via
+ `abid[1:0]` (range `[0, 3]`); `cbsz != 0` selects the first set.
+ - gfx942 8-bit and gfx950 16-bit: `cbsz == 0` selects one of two 16-bit
+ sets via `abid[0]` (range `[0, 1]`); `cbsz != 0` selects the first set.
+ - gfx950 8-bit: hardware ignores both `cbsz` and `abid`; both must be 0.
Example:
```mlir
@@ -1197,6 +1208,9 @@ def AMDGPU_SparseMFMAOp :
%2 = amdgpu.sparse_mfma 16x16x64 %matA * %matB + %matC sparse(%idx : vector<2xi16>)
{ cbsz = 0 : i32, abid = 1 : i32 }
: vector<8xf8E4M3FNUZ>, vector<16xf8E4M3FNUZ>, vector<4xf32>
+
+ %3 = amdgpu.sparse_mfma 16x16x128 %matA * %matB + %matC sparse(%idx : i32)
+ : vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<4xf32>
```
}];
let assemblyFormat = [{
@@ -1257,7 +1271,7 @@ def AMDGPU_SparseWMMAOp :
UnitAttr:$clamp,
UnitAttr:$wave64)>,
Results<(outs SWMMACOutTypes: $destD)> {
- let summary = "MLIR wrapper for CDNA sparse mfma (smfmac) instructions";
+ let summary = "MLIR wrapper for gfx12+ sparse wmma instructions";
let description = [{
The `amdgpu.sparse_wmma` op is an MLIR wrapper around intrinsics for various
`swmmac` instructions in the AMDGPU architecture, which perform matrix
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 66b726fe797eb..8464d1e29f0aa 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -1710,8 +1710,11 @@ struct SparseMFMAOpLowering : public ConvertOpToLLVMPattern<SparseMFMAOp> {
"no intrinsic matching sparse MFMA on the given chipset");
// Bitcast sparse indices from vector<4xi8> or vector<2xi16> to i32.
- Value sparseIdx = LLVM::BitcastOp::create(
- rewriter, loc, rewriter.getI32Type(), adaptor.getSparseIdx());
+ // gfx950 8-bit variants already carry the index as i32; skip the bitcast.
+ Value sparseIdx = adaptor.getSparseIdx();
+ Type i32Type = rewriter.getI32Type();
+ if (sparseIdx.getType() != i32Type)
+ sparseIdx = LLVM::BitcastOp::create(rewriter, loc, i32Type, sparseIdx);
OperationState loweredOp(loc, maybeIntrinsic.value());
loweredOp.addTypes(outType);
diff --git a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
index f19b0f3f9edf5..2f6f59194fba3 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
@@ -627,32 +627,51 @@ LogicalResult SparseMFMAOp::verify() {
return emitOpError(
"expected source operands to have the same element type");
- // When CBSZ == 0, ABID selects the index set within the sparse index VGPR.
- // When CBSZ != 0, the first index set is always used (ABID ignored).
+ // Classify the sparse MFMA variant. The three flavors
diff er in CBSZ/ABID
+ // handling and in the sparse-index layout:
+ // - gfx942 16-bit: max ABID = 3, sparse idx = vector<4xi8>
+ // - gfx950 16-bit / gfx942 8-bit: max ABID = 1, sparse idx = vector<2xi16>
+ // - gfx950 8-bit: CBSZ/ABID ignored by hw, sparse idx = i32
+ uint32_t m = getM(), k = getK();
bool is8BitSource = sparseElem.isFloat(8) || sparseElem.isInteger(8);
- // 8-bit source: ABID selects one of two 16-bit index sets.
- if (getCbsz() == 0 && is8BitSource && getAbid() > 1)
- return emitOpError("ABID must be 0 or 1 for 8-bit source data");
- // 16-bit source: ABID selects one of four 8-bit index sets (0-3 all valid).
- if (getCbsz() == 0 && !is8BitSource && getAbid() > 3)
- return emitOpError("ABID must be between 0 and 3 for 16-bit source data");
-
- // Validate sparseIdx type matches source element type.
- auto sparseIdxType = cast<VectorType>(getSparseIdx().getType());
- if (is8BitSource) {
- // 8-bit source data requires vector<2xi16> sparse indices.
- if (sparseIdxType.getNumElements() != 2 ||
- !sparseIdxType.getElementType().isInteger(16))
- return emitOpError("expected vector<2xi16> sparse indices for 8-bit "
- "source data, but got ")
- << getSparseIdx().getType();
+ bool is16BitGfx942 =
+ !is8BitSource && ((m == 16 && k == 32) || (m == 32 && k == 16));
+ bool is8BitGfx950 =
+ is8BitSource && ((m == 16 && k == 128) || (m == 32 && k == 64));
+
+ // CBSZ/ABID range check. On gfx950 8-bit the hardware always uses the first
+ // set and ignores these fields, so require zeros in IR. Otherwise ABID is
+ // only meaningful when CBSZ == 0 (when CBSZ != 0 the first set is always
+ // used and ABID is irrelevant, so the verifier accepts any value).
+ if (is8BitGfx950) {
+ if (getCbsz() != 0)
+ return emitOpError(
+ "CBSZ must be 0 for this variant (field is ignored by hardware)");
+ if (getAbid() != 0)
+ return emitOpError(
+ "ABID must be 0 for this variant (field is ignored by hardware)");
+ } else if (getCbsz() == 0) {
+ unsigned maxAbid = is16BitGfx942 ? 3u : 1u;
+ if (getAbid() > maxAbid)
+ return emitOpError("ABID must be in [0, ")
+ << maxAbid << "] for this variant";
+ }
+
+ Type sparseIdxType = getSparseIdx().getType();
+ if (is8BitGfx950) {
+ if (!sparseIdxType.isInteger(32))
+ return emitOpError("expected i32 sparse indices for this variant "
+ "(no internal set structure), but got ")
+ << sparseIdxType;
} else {
- // 16-bit source data requires vector<4xi8> sparse indices.
- if (sparseIdxType.getNumElements() != 4 ||
- !sparseIdxType.getElementType().isInteger(8))
- return emitOpError("expected vector<4xi8> sparse indices for 16-bit "
- "source data, but got ")
- << getSparseIdx().getType();
+ unsigned expectedIdxElems = is16BitGfx942 ? 4 : 2;
+ unsigned expectedIdxBits = is16BitGfx942 ? 8 : 16;
+ auto vecType = dyn_cast<VectorType>(sparseIdxType);
+ if (!vecType || vecType.getNumElements() != expectedIdxElems ||
+ !vecType.getElementType().isInteger(expectedIdxBits))
+ return emitOpError("expected vector<")
+ << expectedIdxElems << "xi" << expectedIdxBits
+ << "> sparse indices for this variant, but got " << sparseIdxType;
}
int64_t expectedSourceElems = (getM() * getK()) / waveSize;
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir b/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
index 557c71c32ee54..0f2a74917c757 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
@@ -6,56 +6,55 @@ func.func @sparse_mfma_to_rocdl(%arg0 : vector<8xf16>, %arg1 : vector<16xf16>,
%arg8 : vector<4xi32>, %arg9 : vector<16xi32>,
%arg10 : vector<16xf8E4M3FN>, %arg11 : vector<16xf8E5M2>,
%arg12 : vector<32xf8E4M3FN>, %arg13 : vector<32xf8E5M2>,
- %arg14 : vector<4xi8>, %arg15 : vector<2xi16>) {
- // CHECK: llvm.bitcast %{{.*}} : vector<4xi8> to i32
+ %arg15 : vector<2xi16>, %arg16 : i32) {
+ // CHECK: llvm.bitcast %{{.*}} : vector<2xi16> to i32
// CHECK: rocdl.smfmac.f32.16x16x64.f16{{.*}}: (vector<8xf16>, vector<16xf16>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x64 %arg0 * %arg1 + %arg2 sparse(%arg14 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x64 %arg0 * %arg1 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<4xf32>
// CHECK: rocdl.smfmac.f32.16x16x64.bf16{{.*}}: (vector<8xbf16>, vector<16xbf16>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x64 %arg4 * %arg5 + %arg2 sparse(%arg14 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xbf16>, vector<16xbf16>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x64 %arg4 * %arg5 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xbf16>, vector<16xbf16>, vector<4xf32>
// CHECK: llvm.bitcast {{.*}} : vector<16xi8> to vector<4xi32>
// CHECK: llvm.bitcast {{.*}} : vector<32xi8> to vector<8xi32>
- // CHECK: llvm.bitcast %{{.*}} : vector<2xi16> to i32
// CHECK: rocdl.smfmac.i32.16x16x128.i8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xi32>, i32) -> vector<4xi32>
- amdgpu.sparse_mfma 16x16x128 %arg6 * %arg7 + %arg8 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<4xi32>
+ amdgpu.sparse_mfma 16x16x128 %arg6 * %arg7 + %arg8 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<4xi32>
// CHECK: llvm.bitcast {{.*}} : vector<16xi8> to vector<4xi32>
// CHECK: llvm.bitcast {{.*}} : vector<32xi8> to vector<8xi32>
// CHECK: rocdl.smfmac.f32.16x16x128.fp8.fp8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x128 %arg10 * %arg12 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x128 %arg10 * %arg12 + %arg2 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<4xf32>
// CHECK: llvm.bitcast {{.*}} : vector<16xi8> to vector<4xi32>
// CHECK: llvm.bitcast {{.*}} : vector<32xi8> to vector<8xi32>
// CHECK: rocdl.smfmac.f32.16x16x128.bf8.bf8 {{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x128 %arg11 * %arg13 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E5M2>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x128 %arg11 * %arg13 + %arg2 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E5M2>, vector<4xf32>
// CHECK: rocdl.smfmac.f32.16x16x128.fp8.bf8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x128 %arg10 * %arg13 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E5M2>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x128 %arg10 * %arg13 + %arg2 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E5M2>, vector<4xf32>
// CHECK: rocdl.smfmac.f32.16x16x128.bf8.fp8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xf32>, i32) -> vector<4xf32>
- amdgpu.sparse_mfma 16x16x128 %arg11 * %arg12 + %arg2 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E4M3FN>, vector<4xf32>
+ amdgpu.sparse_mfma 16x16x128 %arg11 * %arg12 + %arg2 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E4M3FN>, vector<4xf32>
// CHECK: rocdl.smfmac.f32.32x32x32.f16{{.*}}: (vector<8xf16>, vector<16xf16>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x32 %arg0 * %arg1 + %arg3 sparse(%arg14 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x32 %arg0 * %arg1 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<16xf32>
// CHECK: rocdl.smfmac.f32.32x32x32.bf16{{.*}}: (vector<8xbf16>, vector<16xbf16>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x32 %arg4 * %arg5 + %arg3 sparse(%arg14 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xbf16>, vector<16xbf16>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x32 %arg4 * %arg5 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xbf16>, vector<16xbf16>, vector<16xf32>
// CHECK: rocdl.smfmac.i32.32x32x64.i8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<16xi32>, i32) -> vector<16xi32>
- amdgpu.sparse_mfma 32x32x64 %arg6 * %arg7 + %arg9 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<16xi32>
+ amdgpu.sparse_mfma 32x32x64 %arg6 * %arg7 + %arg9 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<16xi32>
// CHECK: rocdl.smfmac.f32.32x32x64.fp8.fp8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x64 %arg10 * %arg12 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x64 %arg10 * %arg12 + %arg3 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<16xf32>
// CHECK: rocdl.smfmac.f32.32x32x64.bf8.bf8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x64 %arg11 * %arg13 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E5M2>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x64 %arg11 * %arg13 + %arg3 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E5M2>, vector<16xf32>
// CHECK: rocdl.smfmac.f32.32x32x64.fp8.bf8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x64 %arg10 * %arg13 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E5M2>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x64 %arg10 * %arg13 + %arg3 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E4M3FN>, vector<32xf8E5M2>, vector<16xf32>
// CHECK: rocdl.smfmac.f32.32x32x64.bf8.fp8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<16xf32>, i32) -> vector<16xf32>
- amdgpu.sparse_mfma 32x32x64 %arg11 * %arg12 + %arg3 sparse(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E4M3FN>, vector<16xf32>
+ amdgpu.sparse_mfma 32x32x64 %arg11 * %arg12 + %arg3 sparse(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xf8E5M2>, vector<32xf8E4M3FN>, vector<16xf32>
func.return
}
diff --git a/mlir/test/Dialect/AMDGPU/invalid.mlir b/mlir/test/Dialect/AMDGPU/invalid.mlir
index 4f59ec642171f..2958b0fe2bc51 100644
--- a/mlir/test/Dialect/AMDGPU/invalid.mlir
+++ b/mlir/test/Dialect/AMDGPU/invalid.mlir
@@ -507,7 +507,7 @@ func.func @sparse_mfma_mismatched_source_types(%a: vector<4xf16>, %b: vector<8xb
// -----
func.func @sparse_mfma_abid_invalid_for_8bit(%a: vector<8xi8>, %b: vector<16xi8>, %c: vector<4xi32>, %idx: vector<2xi16>) -> vector<4xi32> {
- // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be 0 or 1 for 8-bit source data}}
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be in [0, 1] for this variant}}
%d = amdgpu.sparse_mfma 16x16x64 %a * %b + %c sparse(%idx : vector<2xi16>) { abid = 2 : i32, cbsz = 0 : i32 } : vector<8xi8>, vector<16xi8>, vector<4xi32>
func.return %d : vector<4xi32>
}
@@ -515,15 +515,47 @@ func.func @sparse_mfma_abid_invalid_for_8bit(%a: vector<8xi8>, %b: vector<16xi8>
// -----
func.func @sparse_mfma_abid_invalid_for_16bit(%a: vector<4xf16>, %b: vector<8xf16>, %c: vector<4xf32>, %idx: vector<4xi8>) -> vector<4xf32> {
- // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be between 0 and 3 for 16-bit source data}}
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be in [0, 3] for this variant}}
%d = amdgpu.sparse_mfma 16x16x32 %a * %b + %c sparse(%idx : vector<4xi8>) { abid = 4 : i32, cbsz = 0 : i32 } : vector<4xf16>, vector<8xf16>, vector<4xf32>
func.return %d : vector<4xf32>
}
// -----
+func.func @sparse_mfma_abid_invalid_for_gfx950_16bit(%a: vector<8xf16>, %b: vector<16xf16>, %c: vector<4xf32>, %idx: vector<2xi16>) -> vector<4xf32> {
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be in [0, 1] for this variant}}
+ %d = amdgpu.sparse_mfma 16x16x64 %a * %b + %c sparse(%idx : vector<2xi16>) { abid = 2 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<4xf32>
+ func.return %d : vector<4xf32>
+}
+
+// -----
+
+func.func @sparse_mfma_gfx950_8bit_nonzero_cbsz(%a: vector<16xi8>, %b: vector<32xi8>, %c: vector<4xi32>, %idx: i32) -> vector<4xi32> {
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op CBSZ must be 0 for this variant (field is ignored by hardware)}}
+ %d = amdgpu.sparse_mfma 16x16x128 %a * %b + %c sparse(%idx : i32) { abid = 0 : i32, cbsz = 1 : i32 } : vector<16xi8>, vector<32xi8>, vector<4xi32>
+ func.return %d : vector<4xi32>
+}
+
+// -----
+
+func.func @sparse_mfma_gfx950_8bit_nonzero_abid(%a: vector<16xi8>, %b: vector<32xi8>, %c: vector<4xi32>, %idx: i32) -> vector<4xi32> {
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op ABID must be 0 for this variant (field is ignored by hardware)}}
+ %d = amdgpu.sparse_mfma 16x16x128 %a * %b + %c sparse(%idx : i32) { abid = 1 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<4xi32>
+ func.return %d : vector<4xi32>
+}
+
+// -----
+
+func.func @sparse_mfma_wrong_idx_type_for_gfx950_8bit(%a: vector<16xi8>, %b: vector<32xi8>, %c: vector<4xi32>, %idx: vector<2xi16>) -> vector<4xi32> {
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op expected i32 sparse indices for this variant (no internal set structure), but got 'vector<2xi16>'}}
+ %d = amdgpu.sparse_mfma 16x16x128 %a * %b + %c sparse(%idx : vector<2xi16>) : vector<16xi8>, vector<32xi8>, vector<4xi32>
+ func.return %d : vector<4xi32>
+}
+
+// -----
+
func.func @sparse_mfma_wrong_idx_type_for_8bit(%a: vector<8xi8>, %b: vector<16xi8>, %c: vector<4xi32>, %idx: vector<4xi8>) -> vector<4xi32> {
- // expected-error at +1 {{'amdgpu.sparse_mfma' op expected vector<2xi16> sparse indices for 8-bit source data, but got 'vector<4xi8>'}}
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op expected vector<2xi16> sparse indices for this variant, but got 'vector<4xi8>'}}
%d = amdgpu.sparse_mfma 16x16x64 %a * %b + %c sparse(%idx : vector<4xi8>) : vector<8xi8>, vector<16xi8>, vector<4xi32>
func.return %d : vector<4xi32>
}
@@ -531,16 +563,24 @@ func.func @sparse_mfma_wrong_idx_type_for_8bit(%a: vector<8xi8>, %b: vector<16xi
// -----
func.func @sparse_mfma_wrong_idx_type_for_16bit(%a: vector<4xf16>, %b: vector<8xf16>, %c: vector<4xf32>, %idx: vector<2xi16>) -> vector<4xf32> {
- // expected-error at +1 {{'amdgpu.sparse_mfma' op expected vector<4xi8> sparse indices for 16-bit source data, but got 'vector<2xi16>'}}
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op expected vector<4xi8> sparse indices for this variant, but got 'vector<2xi16>'}}
%d = amdgpu.sparse_mfma 16x16x32 %a * %b + %c sparse(%idx : vector<2xi16>) : vector<4xf16>, vector<8xf16>, vector<4xf32>
func.return %d : vector<4xf32>
}
// -----
-func.func @sparse_mfma_wrong_source_count(%a: vector<4xf16>, %b: vector<8xf16>, %c: vector<16xf32>, %idx: vector<4xi8>) -> vector<16xf32> {
+func.func @sparse_mfma_wrong_idx_type_for_gfx950_16bit(%a: vector<8xf16>, %b: vector<16xf16>, %c: vector<4xf32>, %idx: vector<4xi8>) -> vector<4xf32> {
+ // expected-error at +1 {{'amdgpu.sparse_mfma' op expected vector<2xi16> sparse indices for this variant, but got 'vector<4xi8>'}}
+ %d = amdgpu.sparse_mfma 16x16x64 %a * %b + %c sparse(%idx : vector<4xi8>) : vector<8xf16>, vector<16xf16>, vector<4xf32>
+ func.return %d : vector<4xf32>
+}
+
+// -----
+
+func.func @sparse_mfma_wrong_source_count(%a: vector<4xf16>, %b: vector<8xf16>, %c: vector<16xf32>, %idx: vector<2xi16>) -> vector<16xf32> {
// expected-error at +1 {{'amdgpu.sparse_mfma' op expected 16 source values for this operation but got 8}}
- %d = amdgpu.sparse_mfma 32x32x32 %a * %b + %c sparse(%idx : vector<4xi8>) : vector<4xf16>, vector<8xf16>, vector<16xf32>
+ %d = amdgpu.sparse_mfma 32x32x32 %a * %b + %c sparse(%idx : vector<2xi16>) : vector<4xf16>, vector<8xf16>, vector<16xf32>
func.return %d : vector<16xf32>
}
diff --git a/mlir/test/Dialect/AMDGPU/ops.mlir b/mlir/test/Dialect/AMDGPU/ops.mlir
index a34550dc25420..606a7768974bf 100644
--- a/mlir/test/Dialect/AMDGPU/ops.mlir
+++ b/mlir/test/Dialect/AMDGPU/ops.mlir
@@ -649,6 +649,25 @@ func.func @scaled_mfma(%arg0 : f8E8M0FNU, %arg1 : vector<32xf6E2M3FN>, %arg2 : v
func.return %0 : vector<16xf32>
}
+// CHECK-LABEL: func @sparse_mfma
+func.func @sparse_mfma(%a16_4 : vector<4xf16>, %b16_8 : vector<8xf16>,
+ %a16_8 : vector<8xf16>, %b16_16 : vector<16xf16>,
+ %a8_8 : vector<8xi8>, %b8_16 : vector<16xi8>,
+ %a8_16 : vector<16xi8>, %b8_32 : vector<32xi8>,
+ %c4f : vector<4xf32>, %c4i : vector<4xi32>,
+ %idx4xi8 : vector<4xi8>, %idx2xi16 : vector<2xi16>,
+ %idxI32 : i32) {
+ // CHECK: amdgpu.sparse_mfma 16x16x32 {{.*}} sparse({{.*}} : vector<4xi8>)
+ %0 = amdgpu.sparse_mfma 16x16x32 %a16_4 * %b16_8 + %c4f sparse(%idx4xi8 : vector<4xi8>) { abid = 3 : i32, cbsz = 0 : i32 } : vector<4xf16>, vector<8xf16>, vector<4xf32>
+ // CHECK: amdgpu.sparse_mfma 16x16x64 {{.*}} sparse({{.*}} : vector<2xi16>)
+ %1 = amdgpu.sparse_mfma 16x16x64 %a8_8 * %b8_16 + %c4i sparse(%idx2xi16 : vector<2xi16>) { abid = 1 : i32, cbsz = 0 : i32 } : vector<8xi8>, vector<16xi8>, vector<4xi32>
+ // CHECK: amdgpu.sparse_mfma 16x16x64 {{.*}} sparse({{.*}} : vector<2xi16>)
+ %2 = amdgpu.sparse_mfma 16x16x64 %a16_8 * %b16_16 + %c4f sparse(%idx2xi16 : vector<2xi16>) { abid = 1 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<4xf32>
+ // CHECK: amdgpu.sparse_mfma 16x16x128 {{.*}} sparse({{.*}} : i32)
+ %3 = amdgpu.sparse_mfma 16x16x128 %a8_16 * %b8_32 + %c4i sparse(%idxI32 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<4xi32>
+ func.return
+}
+
// CHECK-LABEL: func @transpose_load
func.func @transpose_load(%idx1 : index, %idx2 : index, %mem : memref<128x32xf16, 3>) -> vector<4xf16> {
// CHECK: amdgpu.transpose_load
More information about the Mlir-commits
mailing list