Skip to content

Add gfx950 mfma instructions to ROCDL dialect #123361

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jungpark-mlir merged 5 commits into from
Jan 21, 2025
Merged

Add gfx950 mfma instructions to ROCDL dialect #123361

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
aac972c
Add gfx950 mfma instructions to ROCDL dialect
yiqian1 Jan 16, 2025
05785e6
Remove an old comment
yiqian1 Jan 17, 2025
342573d
Update a comment.
yiqian1 Jan 20, 2025
1cf0e2d
Merge branch 'main' into gfx950-mfma-rocdl
jungpark-mlir Jan 21, 2025
96bd837
Merge branch 'main' into gfx950-mfma-rocdl
jungpark-mlir Jan 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -379,7 +379,6 @@ 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">;
// fp8, only on gfx940
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">;
Expand All @@ -388,6 +387,13 @@ def ROCDL_mfma_f32_32x32x16_bf8_bf8 : ROCDL_Mfma_IntrOp<"mfma.f32.32x32x16.bf8.b
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">;
// New in gfx950
Copy link
Member

@Groverkss Groverkss Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The above comment says the fp8 instructions are only supported on gfx940. Do we want to update that to "new in gfx940"?

Copy link
Contributor

@krzysz00 krzysz00 Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Part of what's new is the instructions supporting OCP FP8 instead of the gfx940 format

Copy link
Contributor

@krzysz00 krzysz00 Jan 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So yeah, comment needs updating

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">;

//===---------------------------------------------------------------------===//
// WMMA intrinsics
Expand Down
35 changes: 34 additions & 1 deletion mlir/test/Target/LLVMIR/rocdl.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,9 @@ llvm.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
%arg4 : vector<16 x f32>, %arg5 : vector<4xf32>,
%arg6 : vector<4xf16>, %arg7 : vector<32 x i32>,
%arg8 : vector<16 x i32>, %arg9 : vector<4xi32>,
%arg10 : vector<2xi16>, %arg11 : i64) -> vector<32 x f32> {
%arg10 : vector<2xi16>, %arg11 : i64,
%arg12 : vector<8xbf16>, %arg13 : vector<4xi32>,
%arg14 : vector<8xf16>) -> vector<32 x f32> {
%csti32 = llvm.mlir.constant(42 : i32) : i32

// CHECK-LABEL: rocdl.xdlops
Expand Down Expand Up @@ -362,6 +364,37 @@ llvm.func @rocdl.xdlops(%arg0 : f32, %arg1 : f32,
%r27 = rocdl.mfma.f32.32x32x16.bf8.bf8 %arg11, %arg11, %arg4, %csti32, %csti32, %csti32 :
(i64, i64, vector<16xf32>,
i32, i32, i32) -> vector<16xf32>

// CHECK: call <4 x float> @llvm.amdgcn.mfma.f32.16x16x32.bf16(<8 x bfloat> %{{.*}}, <8 x bfloat> %{{.*}}, <4 x float> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r28 = rocdl.mfma.f32.16x16x32.bf16 %arg12, %arg12, %arg5, %csti32, %csti32, %csti32 :
(vector<8xbf16>, vector<8xbf16>, vector<4xf32>,
i32, i32, i32) -> vector<4xf32>

// CHECK: call <4 x i32> @llvm.amdgcn.mfma.i32.16x16x64.i8(<4 x i32> %{{.*}}, <4 x i32> %{{.*}}, <4 x i32> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r29 = rocdl.mfma.i32.16x16x64.i8 %arg9, %arg9, %arg9, %csti32, %csti32, %csti32 :
(vector<4xi32>, vector<4xi32>, vector<4xi32>,
i32, i32, i32) -> vector<4xi32>

// 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>

// 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 :
(vector<8xbf16>, vector<8xbf16>, vector<16xf32>,
i32, i32, i32) -> vector<16xf32>

// CHECK: call <16 x i32> @llvm.amdgcn.mfma.i32.32x32x32.i8(<4 x i32> %{{.*}}, <4 x i32> %{{.*}}, <16 x i32> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r32 = rocdl.mfma.i32.32x32x32.i8 %arg9, %arg9, %arg8, %csti32, %csti32, %csti32 :
(vector<4xi32>, vector<4xi32>, vector<16xi32>,
i32, i32, i32) -> vector<16xi32>

// CHECK: call <16 x float> @llvm.amdgcn.mfma.f32.32x32x16.f16(<8 x half> %{{.*}}, <8 x half> %{{.*}}, <16 x float> %{{.*}}, i32 {{.*}}, i32 {{.*}}, i32 {{.*}})
%r33 = rocdl.mfma.f32.32x32x16.f16 %arg14, %arg14, %arg4, %csti32, %csti32, %csti32 :
(vector<8xf16>, vector<8xf16>, vector<16xf32>,
i32, i32, i32) -> vector<16xf32>

llvm.return %r0 : vector<32 x f32>
}

Expand Down
Loading