[Mlir-commits] [mlir] e18d3f7 - [AMDGPU][MLIR] Fix legacy gfx942 bf16 sparse mfma lowering in gfx950 (#195977)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 6 07:55:34 PDT 2026
Author: Eric
Date: 2026-05-06T07:55:28-07:00
New Revision: e18d3f7ccc0aa6a791ade3d0521f3645c247ccc6
URL: https://github.com/llvm/llvm-project/commit/e18d3f7ccc0aa6a791ade3d0521f3645c247ccc6
DIFF: https://github.com/llvm/llvm-project/commit/e18d3f7ccc0aa6a791ade3d0521f3645c247ccc6.diff
LOG: [AMDGPU][MLIR] Fix legacy gfx942 bf16 sparse mfma lowering in gfx950 (#195977)
Previously, we treated all gfx950 BF16 sparse MFMA operands as native
bf16. But we could still select the legacy/shorter gfx942 variants of
the intrinsics (`16x16x32.bf16` / `32x32x16.bf16`) in gfx950. At the
ROCDL/LLVM level, these still expect the operands to be in the old `i16`
form, so bitcast them too.
---------
Signed-off-by: Eric Feng <Eric.Feng at amd.com>
Added:
Modified:
mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 78a87e3972d92..8b9b3e194adf9 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -1697,7 +1697,17 @@ struct SparseMFMAOpLowering : public ConvertOpToLLVMPattern<SparseMFMAOp> {
// smfmac is supported on gfx942 and gfx950.
if (chipset.majorVersion != 9 || chipset < kGfx942)
return op->emitOpError("sparse MFMA (smfmac) only supported on gfx942+");
- bool isGfx950 = chipset >= kGfx950;
+
+ std::optional<StringRef> maybeIntrinsic = smfmacOpToIntrinsic(op, chipset);
+ if (!maybeIntrinsic.has_value())
+ return op.emitOpError(
+ "no intrinsic matching sparse MFMA on the given chipset");
+ bool isGfx942BF16 =
+ (*maybeIntrinsic ==
+ ROCDL::smfmac_f32_16x16x32_bf16::getOperationName() ||
+ *maybeIntrinsic ==
+ ROCDL::smfmac_f32_32x32x16_bf16::getOperationName());
+ bool isGfx950 = (chipset >= kGfx950) && !isGfx942BF16;
Value a = convertPackedVectorOperand(rewriter, loc, adaptor.getSourceA(),
isGfx950);
@@ -1705,11 +1715,6 @@ struct SparseMFMAOpLowering : public ConvertOpToLLVMPattern<SparseMFMAOp> {
isGfx950);
Value c = adaptor.getDestC();
- std::optional<StringRef> maybeIntrinsic = smfmacOpToIntrinsic(op, chipset);
- if (!maybeIntrinsic.has_value())
- return op.emitOpError(
- "no intrinsic matching sparse MFMA on the given chipset");
-
// Bitcast sparse indices from vector<4xi8> or vector<2xi16> to i32.
// gfx950 8-bit variants already carry the index as i32; skip the bitcast.
Value sparseIdx = adaptor.getSparseIdx();
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir b/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
index 0f2a74917c757..acf9df948f26a 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/sparse-mfma-gfx950.mlir
@@ -6,7 +6,8 @@ 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>,
- %arg15 : vector<2xi16>, %arg16 : i32) {
+ %arg15 : vector<2xi16>, %arg16 : i32,
+ %arg17 : vector<4xbf16>, %arg18 : vector<4xi8>) {
// 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(%arg15 : vector<2xi16>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<8xf16>, vector<16xf16>, vector<4xf32>
@@ -14,6 +15,11 @@ func.func @sparse_mfma_to_rocdl(%arg0 : vector<8xf16>, %arg1 : vector<16xf16>,
// CHECK: rocdl.smfmac.f32.16x16x64.bf16{{.*}}: (vector<8xbf16>, vector<16xbf16>, vector<4xf32>, i32) -> 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<4xbf16> to vector<4xi16>
+ // CHECK: llvm.bitcast {{.*}} : vector<8xbf16> to vector<8xi16>
+ // CHECK: rocdl.smfmac.f32.16x16x32.bf16{{.*}}: (vector<4xi16>, vector<8xi16>, vector<4xf32>, i32) -> vector<4xf32>
+ amdgpu.sparse_mfma 16x16x32 %arg17 * %arg4 + %arg2 sparse(%arg18 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<4xbf16>, vector<8xbf16>, vector<4xf32>
+
// CHECK: llvm.bitcast {{.*}} : vector<16xi8> to vector<4xi32>
// CHECK: llvm.bitcast {{.*}} : vector<32xi8> to vector<8xi32>
// CHECK: rocdl.smfmac.i32.16x16x128.i8{{.*}}: (vector<4xi32>, vector<8xi32>, vector<4xi32>, i32) -> vector<4xi32>
@@ -41,6 +47,9 @@ func.func @sparse_mfma_to_rocdl(%arg0 : vector<8xf16>, %arg1 : vector<16xf16>,
// CHECK: rocdl.smfmac.f32.32x32x32.bf16{{.*}}: (vector<8xbf16>, vector<16xbf16>, vector<16xf32>, i32) -> 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.f32.32x32x16.bf16{{.*}}: (vector<4xi16>, vector<8xi16>, vector<16xf32>, i32) -> vector<16xf32>
+ amdgpu.sparse_mfma 32x32x16 %arg17 * %arg4 + %arg3 sparse(%arg18 : vector<4xi8>) { abid = 0 : i32, cbsz = 0 : i32 } : vector<4xbf16>, vector<8xbf16>, 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(%arg16 : i32) { abid = 0 : i32, cbsz = 0 : i32 } : vector<16xi8>, vector<32xi8>, vector<16xi32>
More information about the Mlir-commits
mailing list