[Mlir-commits] [mlir] [ROCDL] Refactored MFMA ops in ODS; added constraints (PR #175775)
Ravil Dorozhinskii
llvmlistbot at llvm.org
Tue Jan 13 07:11:00 PST 2026
https://github.com/ravil-mobile created https://github.com/llvm/llvm-project/pull/175775
None
>From c171c9b388722b51e92ad45dca96115d5507cb65 Mon Sep 17 00:00:00 2001
From: ravil-mobile <ravil.aviva.com at gmail.com>
Date: Tue, 13 Jan 2026 15:08:51 +0000
Subject: [PATCH] [ROCDL] Refactored MFMA ops in ODS; added constraints
---
mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td | 105 +++++++++++--------
mlir/test/Dialect/LLVMIR/rocdl.mlir | 2 +-
mlir/test/Target/LLVMIR/rocdl.mlir | 2 +-
3 files changed, 62 insertions(+), 47 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index 265c2e99f52d6..ef070d26e2451 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -617,6 +617,21 @@ class ROCDL_Mfma_IntrOp<string mnemonic, list<Trait> traits = []> :
"$args attr-dict `:` functional-type($args, $res)";
}
+class ROCDL_Mfma_IntrOpV0<string mnemonic, Type AB, Type CD> :
+ ROCDL_IntrOp<mnemonic, [], [], [], 1, 0, 0, 0, [], []>,
+ Arguments<(ins
+ LLVM_ScalarOrVectorOf<AB>:$a,
+ LLVM_ScalarOrVectorOf<AB>:$b,
+ LLVM_ScalarOrVectorOf<CD>:$c,
+ I32:$cbsz,
+ I32:$abid,
+ I32:$blgp)> {
+ let results = (outs LLVM_ScalarOrVectorOf<CD>:$res);
+ let assemblyFormat = [{
+ $a `,` $b `,` $c `,` $cbsz `,` $abid `,` $blgp attr-dict `:` functional-type(operands, $res)
+ }];
+}
+
//===---------------------------------------------------------------------===//
// MFMA intrinsics with overloaded operands
class ROCDL_Mfma_OO_IntrOp<string mnemonic, list<int> overloadedOperands,
@@ -628,56 +643,56 @@ class ROCDL_Mfma_OO_IntrOp<string mnemonic, list<int> overloadedOperands,
}
// Available on all CDNA.
-def ROCDL_mfma_f32_32x32x1f32 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x1f32">;
-def ROCDL_mfma_f32_16x16x1f32 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x1f32">;
-def ROCDL_mfma_f32_4x4x1f32 : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x1f32">;
-def ROCDL_mfma_f32_32x32x2f32 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x2f32">;
-def ROCDL_mfma_f32_16x16x4f32 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x4f32">;
-def ROCDL_mfma_f32_32x32x4f16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x4f16">;
-def ROCDL_mfma_f32_16x16x4f16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x4f16">;
-def ROCDL_mfma_f32_4x4x4f16 : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x4f16">;
-def ROCDL_mfma_f32_32x32x8f16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x8f16">;
-def ROCDL_mfma_f32_16x16x16f16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x16f16">;
-def ROCDL_mfma_i32_32x32x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x4i8">;
-def ROCDL_mfma_i32_16x16x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x4i8">;
-def ROCDL_mfma_i32_4x4x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.4x4x4i8">;
-def ROCDL_mfma_i32_32x32x8i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x8i8">;
-def ROCDL_mfma_i32_16x16x16i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x16i8">;
-def ROCDL_mfma_f32_32x32x2bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x2bf16">;
-def ROCDL_mfma_f32_16x16x2bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x2bf16">;
-def ROCDL_mfma_f32_4x4x2bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x2bf16">;
-def ROCDL_mfma_f32_32x32x4bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x4bf16">;
-def ROCDL_mfma_f32_16x16x8bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x8bf16">;
+def ROCDL_mfma_f32_32x32x1f32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x1f32", /*Type AB=*/F32, /*Type CD=*/F32>;
+def ROCDL_mfma_f32_16x16x1f32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x1f32", F32, F32>;
+def ROCDL_mfma_f32_4x4x1f32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.4x4x1f32", F32, F32>;
+def ROCDL_mfma_f32_32x32x2f32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x2f32", F32, F32>;
+def ROCDL_mfma_f32_16x16x4f32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x4f32", F32, F32>;
+def ROCDL_mfma_f32_32x32x4f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x4f16", F16, F32>;
+def ROCDL_mfma_f32_16x16x4f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x4f16", F16, F32>;
+def ROCDL_mfma_f32_4x4x4f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.4x4x4f16", F16, F32>;
+def ROCDL_mfma_f32_32x32x8f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x8f16", F16, F32>;
+def ROCDL_mfma_f32_16x16x16f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x16f16", F16, F32>;
+def ROCDL_mfma_i32_32x32x4i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.32x32x4i8", I32, I32>;
+def ROCDL_mfma_i32_16x16x4i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.16x16x4i8", I32, I32>;
+def ROCDL_mfma_i32_4x4x4i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.4x4x4i8", I32, I32>;
+def ROCDL_mfma_i32_32x32x8i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.32x32x8i8", I32, I32>;
+def ROCDL_mfma_i32_16x16x16i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.16x16x16i8", I32, I32>;
+def ROCDL_mfma_f32_32x32x2bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x2bf16", I16, F32>;
+def ROCDL_mfma_f32_16x16x2bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x2bf16", I16, F32>;
+def ROCDL_mfma_f32_4x4x2bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.4x4x2bf16", I16, F32>;
+def ROCDL_mfma_f32_32x32x4bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x4bf16", I16, F32>;
+def ROCDL_mfma_f32_16x16x8bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x8bf16", I16, F32>;
// New in gfx90a.
-def ROCDL_mfma_f32_32x32x4bf16_1k : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x4bf16.1k">;
-def ROCDL_mfma_f32_16x16x4bf16_1k : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x4bf16.1k">;
-def ROCDL_mfma_f32_4x4x4bf16_1k : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x4bf16.1k">;
-def ROCDL_mfma_f32_32x32x8bf16_1k : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x8bf16.1k">;
-def ROCDL_mfma_f32_16x16x16bf16_1k : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x16bf16.1k">;
+def ROCDL_mfma_f32_32x32x4bf16_1k : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x4bf16.1k", I16, F32>;
+def ROCDL_mfma_f32_16x16x4bf16_1k : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x4bf16.1k", I16, F32>;
+def ROCDL_mfma_f32_4x4x4bf16_1k : ROCDL_Mfma_IntrOpV0<"mfma.f32.4x4x4bf16.1k", I16, F32>;
+def ROCDL_mfma_f32_32x32x8bf16_1k : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x8bf16.1k", I16, F32>;
+def ROCDL_mfma_f32_16x16x16bf16_1k : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x16bf16.1k", I16, F32>;
// Note: in gfx94x, unlike in gfx90a, the f64 xdlops use the "blgp" argument as
// a NEG bitfield. See IntrinsicsAMDGPU.td for more info.
-def ROCDL_mfma_f64_16x16x4f64 : ROCDL_Mfma_IntrOp<"mfma.f64.16x16x4f64">;
-def ROCDL_mfma_f64_4x4x4f64 : ROCDL_Mfma_IntrOp<"mfma.f64.4x4x4f64">;
+def ROCDL_mfma_f64_16x16x4f64 : ROCDL_Mfma_IntrOpV0<"mfma.f64.16x16x4f64", F64, F64>;
+def ROCDL_mfma_f64_4x4x4f64 : ROCDL_Mfma_IntrOpV0<"mfma.f64.4x4x4f64", F64, F64>;
// New in gfx94x.
-def ROCDL_mfma_i32_16x16x32_i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x32.i8">;
-def ROCDL_mfma_i32_32x32x16_i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x16.i8">;
-def ROCDL_mfma_f32_16x16x8_xf32 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x8.xf32">;
-def ROCDL_mfma_f32_32x32x4_xf32 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x4.xf32">;
-def ROCDL_mfma_f32_16x16x32_bf8_bf8 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.bf8.bf8">;
-def ROCDL_mfma_f32_16x16x32_bf8_fp8 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.bf8.fp8">;
-def ROCDL_mfma_f32_16x16x32_fp8_bf8 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.fp8.bf8">;
-def ROCDL_mfma_f32_16x16x32_fp8_fp8 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.fp8.fp8">;
-def ROCDL_mfma_f32_32x32x16_bf8_bf8 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.bf8.bf8">;
-def ROCDL_mfma_f32_32x32x16_bf8_fp8 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.bf8.fp8">;
-def ROCDL_mfma_f32_32x32x16_fp8_bf8 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.fp8.bf8">;
-def ROCDL_mfma_f32_32x32x16_fp8_fp8 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.fp8.fp8">;
+def ROCDL_mfma_i32_16x16x32_i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.16x16x32.i8", I64, I32>;
+def ROCDL_mfma_i32_32x32x16_i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.32x32x16.i8", I64, I32>;
+def ROCDL_mfma_f32_16x16x8_xf32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x8.xf32", F32, F32>;
+def ROCDL_mfma_f32_32x32x4_xf32 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x4.xf32", F32, F32>;
+def ROCDL_mfma_f32_16x16x32_bf8_bf8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.bf8.bf8", I64, F32>;
+def ROCDL_mfma_f32_16x16x32_bf8_fp8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.bf8.fp8", I64, F32>;
+def ROCDL_mfma_f32_16x16x32_fp8_bf8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.fp8.bf8", I64, F32>;
+def ROCDL_mfma_f32_16x16x32_fp8_fp8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.fp8.fp8", I64, F32>;
+def ROCDL_mfma_f32_32x32x16_bf8_bf8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.bf8.bf8", I64, F32>;
+def ROCDL_mfma_f32_32x32x16_bf8_fp8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.bf8.fp8", I64, F32>;
+def ROCDL_mfma_f32_32x32x16_fp8_bf8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.fp8.bf8", I64, F32>;
+def ROCDL_mfma_f32_32x32x16_fp8_fp8 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.fp8.fp8", I64, F32>;
// New in gfx950.
-def ROCDL_mfma_f32_16x16x32_bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.bf16">;
-def ROCDL_mfma_i32_16x16x64_i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x64.i8">;
-def ROCDL_mfma_f32_16x16x32_f16 : ROCDL_Mfma_IntrOp<"mfma.f32.16x16x32.f16">;
-def ROCDL_mfma_f32_32x32x16_bf16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.bf16">;
-def ROCDL_mfma_i32_32x32x32_i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x32.i8">;
-def ROCDL_mfma_f32_32x32x16_f16 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.f16">;
+def ROCDL_mfma_f32_16x16x32_bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.bf16", BF16, F32>;
+def ROCDL_mfma_i32_16x16x64_i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.16x16x64.i8", I32, I32>;
+def ROCDL_mfma_f32_16x16x32_f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.16x16x32.f16", F16, F32>;
+def ROCDL_mfma_f32_32x32x16_bf16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.bf16", BF16, F32>;
+def ROCDL_mfma_i32_32x32x32_i8 : ROCDL_Mfma_IntrOpV0<"mfma.i32.32x32x32.i8", I32, I32>;
+def ROCDL_mfma_f32_32x32x16_f16 : ROCDL_Mfma_IntrOpV0<"mfma.f32.32x32x16.f16", F16, F32>;
def ROCDL_mfma_scale_f32_16x16x128_f8f6f4 : ROCDL_Mfma_OO_IntrOp<"mfma.scale.f32.16x16x128.f8f6f4", [0,1]>;
def ROCDL_mfma_scale_f32_32x32x64_f8f6f4 : ROCDL_Mfma_OO_IntrOp<"mfma.scale.f32.32x32x64.f8f6f4", [0,1]>;
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index cf2b144219f36..daa46b72057ca 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -311,7 +311,7 @@ func.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
// CHECK: rocdl.mfma.f32.16x16x32.f16 {{.*}} : (vector<8xf16>, vector<8xf16>, vector<4xf32>, i32, i32, i32) -> vector<4xi32>
%r33 = rocdl.mfma.f32.16x16x32.f16 %arg17, %arg17, %arg5, %arg3, %arg3, %arg3 :
(vector<8xf16>, vector<8xf16>, vector<4xf32>,
- i32, i32, i32) -> vector<4xi32>
+ i32, i32, i32) -> vector<4xf32>
// CHECK: rocdl.mfma.f32.32x32x16.bf16 {{.*}} : (vector<8xbf16>, vector<8xbf16>, vector<16xf32>, i32, i32, i32) -> vector<16xf32>
%r34 = rocdl.mfma.f32.32x32x16.bf16 %arg16, %arg16, %arg4, %arg3, %arg3, %arg3 :
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index dc6a00e19afc3..af77f09c091d5 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -566,7 +566,7 @@ llvm.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
// CHECK: call <4 x float> @llvm.amdgcn.mfma.f32.16x16x32.f16(<8 x half> %{{.*}}, <8 x half> %{{.*}}, <4 x float> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r30 = rocdl.mfma.f32.16x16x32.f16 %arg14, %arg14, %arg5, %csti32, %csti32, %csti32 :
(vector<8xf16>, vector<8xf16>, vector<4xf32>,
- i32, i32, i32) -> vector<4xi32>
+ i32, i32, i32) -> vector<4xf32>
// CHECK: call <16 x float> @llvm.amdgcn.mfma.f32.32x32x16.bf16(<8 x bfloat> %1{{.*}}, <8 x bfloat> %{{.*}}, <16 x float> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r31 = rocdl.mfma.f32.32x32x16.bf16 %arg12, %arg12, %arg4, %csti32, %csti32, %csti32 :
More information about the Mlir-commits
mailing list