[Mlir-commits] [mlir] [MLIR][NVVM] Support kind::ti16 for tcgen05.mma Ops (PR #216588)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Aug 16 10:53:01 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-llvm
Author: Kirill Vedernikov (kvederni)
<details>
<summary>Changes</summary>
This change adds `kind::ti16` support for tcgen05.mma MLIR Ops.
---
Patch is 23.29 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/216588.diff
12 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+6-3)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+25-4)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid-ti16.mlir (+10)
- (modified) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir (+4-4)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-ti16.mlir (+17)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-shared-ti16.mlir (+17)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-tensor-ti16.mlir (+17)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-ti16.mlir (+17)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-shared-ti16.mlir (+21)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-shared-ti16.mlir (+21)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-tensor-ti16.mlir (+21)
- (added) mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-tensor-ti16.mlir (+21)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 46c38bcb5475d..2b6ba76e3c271 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6033,13 +6033,15 @@ def Tcgen05MMAKindI8 : I32EnumAttrCase<"I8", 3, "i8">;
def Tcgen05MMAKindMXF8F6F4 : I32EnumAttrCase<"MXF8F6F4", 4, "mxf8f6f4">;
def Tcgen05MMAKindMXF4 : I32EnumAttrCase<"MXF4", 5, "mxf4">;
def Tcgen05MMAKindMXF4NVF4 : I32EnumAttrCase<"MXF4NVF4", 6, "mxf4nvf4">;
+def Tcgen05MMAKindTI16 : I32EnumAttrCase<"TI16", 7, "ti16">;
def Tcgen05MMAKind : I32EnumAttr<
"Tcgen05MMAKind",
"tcgen05 MMA Supported Types",
[Tcgen05MMAKindF16, Tcgen05MMAKindTF32, Tcgen05MMAKindF8F6F4,
Tcgen05MMAKindI8, Tcgen05MMAKindMXF8F6F4, Tcgen05MMAKindMXF4,
- Tcgen05MMAKindMXF4NVF4]> {
+ Tcgen05MMAKindMXF4NVF4, Tcgen05MMAKindTI16
+ ]> {
let cppNamespace = "::mlir::NVVM";
let genSpecializedAttr = 0;
}
@@ -6095,8 +6097,9 @@ def Tcgen05MMACollectorOpAttr : EnumAttr<NVVM_Dialect, Tcgen05MMACollectorOp, "t
}
defvar Tcgen05MMANonBlockScaleKindList =
- [Tcgen05MMAKindF16, Tcgen05MMAKindTF32,
- Tcgen05MMAKindF8F6F4, Tcgen05MMAKindI8];
+ [Tcgen05MMAKindF16, Tcgen05MMAKindTF32, Tcgen05MMAKindF8F6F4,
+ Tcgen05MMAKindI8, Tcgen05MMAKindTI16
+ ];
defvar Tcgen05MMABlockScaleKindList =
[Tcgen05MMAKindMXF8F6F4, Tcgen05MMAKindMXF4, Tcgen05MMAKindMXF4NVF4];
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 31fed0b25990d..7a499515f0165 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5569,6 +5569,23 @@ mlir::NVVM::IDArgPair TensormapReplaceOp::getIntrinsicIDAndArgs(
// NVVM tcgen05.mma functions
//===----------------------------------------------------------------------===//
+static llvm::nvvm::Tcgen05MMAKind
+getNVVMTcgen05MMAKind(NVVM::Tcgen05MMAKind kind) {
+ switch (kind) {
+ case NVVM::Tcgen05MMAKind::F16:
+ return llvm::nvvm::Tcgen05MMAKind::F16;
+ case NVVM::Tcgen05MMAKind::TF32:
+ return llvm::nvvm::Tcgen05MMAKind::TF32;
+ case NVVM::Tcgen05MMAKind::F8F6F4:
+ return llvm::nvvm::Tcgen05MMAKind::F8F6F4;
+ case NVVM::Tcgen05MMAKind::I8:
+ return llvm::nvvm::Tcgen05MMAKind::I8;
+ case NVVM::Tcgen05MMAKind::TI16:
+ return llvm::nvvm::Tcgen05MMAKind::TI16;
+ }
+ llvm_unreachable("unknown tcgen05.mma kind");
+}
+
mlir::NVVM::IDArgPair
Tcgen05MMAOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
@@ -5700,7 +5717,8 @@ Tcgen05MMAOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
if (hasDisableOutputLane)
args.push_back(DisableOutputLane);
- args.push_back(builder.getInt32(static_cast<unsigned>(thisOp.getKind())));
+ args.push_back(builder.getInt32(
+ static_cast<unsigned>(getNVVMTcgen05MMAKind(thisOp.getKind()))));
if (!hasDisableOutputLane)
args.push_back(builder.getInt32(ctaGroup));
@@ -5888,7 +5906,8 @@ mlir::NVVM::IDArgPair Tcgen05MMASparseOp::getIntrinsicIDAndArgs(
if (hasDisableOutputLane)
args.push_back(DisableOutputLane);
- args.push_back(builder.getInt32(static_cast<unsigned>(thisOp.getKind())));
+ args.push_back(builder.getInt32(
+ static_cast<unsigned>(getNVVMTcgen05MMAKind(thisOp.getKind()))));
if (!hasDisableOutputLane)
args.push_back(builder.getInt32(ctaGroup));
@@ -6121,7 +6140,8 @@ mlir::NVVM::IDArgPair Tcgen05MMAWsOp::getIntrinsicIDAndArgs(
ID = isATensor ? llvm::Intrinsic::nvvm_tcgen05_mma_ws_tensor
: llvm::Intrinsic::nvvm_tcgen05_mma_ws_shared;
- args.push_back(builder.getInt32(static_cast<unsigned>(thisOp.getKind())));
+ args.push_back(builder.getInt32(
+ static_cast<unsigned>(getNVVMTcgen05MMAKind(thisOp.getKind()))));
args.push_back(
builder.getInt32(static_cast<unsigned>(thisOp.getCollectorBBuffer())));
args.push_back(
@@ -6162,7 +6182,8 @@ mlir::NVVM::IDArgPair Tcgen05MMAWsSparseOp::getIntrinsicIDAndArgs(
ID = isATensor ? llvm::Intrinsic::nvvm_tcgen05_mma_ws_sp_tensor
: llvm::Intrinsic::nvvm_tcgen05_mma_ws_sp_shared;
- args.push_back(builder.getInt32(static_cast<unsigned>(thisOp.getKind())));
+ args.push_back(builder.getInt32(
+ static_cast<unsigned>(getNVVMTcgen05MMAKind(thisOp.getKind()))));
args.push_back(
builder.getInt32(static_cast<unsigned>(thisOp.getCollectorBBuffer())));
args.push_back(
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid-ti16.mlir
new file mode 100644
index 0000000000000..918cfbb881e58
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid-ti16.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-translate --mlir-to-llvmir -verify-diagnostics -split-input-file %s
+
+// Invalid Tcgen05MMAKind for tcgen05.mma.block_scale: ti16 is valid only for
+// mma/mma.sp/ws/ws.sp, not for block_scale ops.
+// CHECK-LABEL: @nvvm_tcgen05_mma_block_scale_invalid_kind_ti16
+llvm.func @nvvm_tcgen05_mma_block_scale_invalid_kind_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %scale_a: !llvm.ptr<6>, %scale_b: !llvm.ptr<6>) {
+ // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {mxf8f6f4, mxf4, mxf4nvf4}}}
+ nvvm.tcgen05.mma.block_scale %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %scale_a, %scale_b {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>)
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
index 3dd0d7a6c9aaa..606983d8fcf07 100644
--- a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
@@ -123,7 +123,7 @@ llvm.func @nvvm_tcgen05_mma_sp_mxf4_block_scale_default(%d_tmem : !llvm.ptr<6>,
// Invalid Tcgen05MMAKind for tcgen05.mma: mxf8f6f4 is only for block_scale ops.
// CHECK-LABEL: @nvvm_tcgen05_mma_invalid_kind_mxf8f6f4
llvm.func @nvvm_tcgen05_mma_invalid_kind_mxf8f6f4(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
- // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8}}}
+ // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8, ti16}}}
nvvm.tcgen05.mma %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d {kind = #nvvm.tcgen05_mma_kind<mxf8f6f4>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1)
llvm.return
}
@@ -133,7 +133,7 @@ llvm.func @nvvm_tcgen05_mma_invalid_kind_mxf8f6f4(%d_tmem : !llvm.ptr<6>, %a_des
// Invalid Tcgen05MMAKind for tcgen05.mma.sp: mxf4 is only for block_scale ops.
// CHECK-LABEL: @nvvm_tcgen05_mma_sp_invalid_kind_mxf4
llvm.func @nvvm_tcgen05_mma_sp_invalid_kind_mxf4(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>) {
- // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8}}}
+ // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8, ti16}}}
nvvm.tcgen05.mma.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata {kind = #nvvm.tcgen05_mma_kind<mxf4>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
llvm.return
}
@@ -143,7 +143,7 @@ llvm.func @nvvm_tcgen05_mma_sp_invalid_kind_mxf4(%d_tmem : !llvm.ptr<6>, %a_desc
// Invalid Tcgen05MMAKind for tcgen05.mma.ws: mxf4nvf4 is only for block_scale ops.
// CHECK-LABEL: @nvvm_tcgen05_mma_ws_invalid_kind_mxf4nvf4
llvm.func @nvvm_tcgen05_mma_ws_invalid_kind_mxf4nvf4(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
- // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8}}}
+ // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8, ti16}}}
nvvm.tcgen05.mma.ws %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d {kind = #nvvm.tcgen05_mma_kind<mxf4nvf4>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1)
llvm.return
}
@@ -153,7 +153,7 @@ llvm.func @nvvm_tcgen05_mma_ws_invalid_kind_mxf4nvf4(%d_tmem : !llvm.ptr<6>, %a_
// Invalid Tcgen05MMAKind for tcgen05.mma.ws.sp: mxf8f6f4 is only for block_scale ops.
// CHECK-LABEL: @nvvm_tcgen05_mma_ws_sp_invalid_kind_mxf8f6f4
llvm.func @nvvm_tcgen05_mma_ws_sp_invalid_kind_mxf8f6f4(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>) {
- // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8}}}
+ // expected-error @below {{attribute 'kind' failed to satisfy constraint: tcgen05 MMA Supported Types whose value is one of {f16, tf32, f8f6f4, i8, ti16}}}
nvvm.tcgen05.mma.ws.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata {kind = #nvvm.tcgen05_mma_kind<mxf8f6f4>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
llvm.return
}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-ti16.mlir
new file mode 100644
index 0000000000000..4b2ba39d98fe6
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-ti16.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ti16
+llvm.func @nvvm_tcgen05_mma_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
+
+ // ti16 enum value is 7, but maps to intrinsic kind 4.
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1)
+
+ // collector_a=fill(2), collector_b=use(3), kind=ti16
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>, collectorOp = #nvvm.tcgen05_mma_collectorop<fill>, collectorOpB = #nvvm.tcgen05_mma_collectorop<use>} : (!llvm.ptr<6>, i64, i64, i32, i1)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-shared-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-shared-ti16.mlir
new file mode 100644
index 0000000000000..038b6bacfaed7
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-shared-ti16.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_sp_ti16
+llvm.func @nvvm_tcgen05_mma_sp_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>) {
+
+ // ti16 enum value is 7, but maps to intrinsic kind 4.
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.sp.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
+
+ // collector_a=fill(2), collector_b=use(3), kind=ti16
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.sp.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>, collectorOp = #nvvm.tcgen05_mma_collectorop<fill>, collectorOpB = #nvvm.tcgen05_mma_collectorop<use>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-tensor-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-tensor-ti16.mlir
new file mode 100644
index 0000000000000..3445668a0b97b
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-tensor-ti16.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_sp_tensor_ti16
+llvm.func @nvvm_tcgen05_mma_sp_tensor_ti16(%d_tmem : !llvm.ptr<6>, %a_tmem: !llvm.ptr<6>, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>) {
+
+ // ti16 enum value is 7, but maps to intrinsic kind 4.
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.sp.tensor(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.sp %d_tmem, %a_tmem, %b_desc, %idesc, %enable_input_d, %spmetadata
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1, !llvm.ptr<6>)
+
+ // collector_a=fill(2), collector_b=use(3), kind=ti16
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.sp.tensor(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma.sp %d_tmem, %a_tmem, %b_desc, %idesc, %enable_input_d, %spmetadata
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>, collectorOp = #nvvm.tcgen05_mma_collectorop<fill>, collectorOpB = #nvvm.tcgen05_mma_collectorop<use>} : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1, !llvm.ptr<6>)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-ti16.mlir
new file mode 100644
index 0000000000000..df45d325e9999
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-ti16.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_tensor_ti16
+llvm.func @nvvm_tcgen05_mma_tensor_ti16(%d_tmem : !llvm.ptr<6>, %a_tmem: !llvm.ptr<6>, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
+
+ // ti16 enum value is 7, but maps to intrinsic kind 4.
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.tensor(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=discard */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma %d_tmem, %a_tmem, %b_desc, %idesc, %enable_input_d
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>} : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1)
+
+ // collector_a=fill(2), collector_b=use(3), kind=ti16
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.tensor(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* cta_group= */ i32 1, /* collector_a=fill */ i32 2, /* collector_b=use */ i32 3)
+ nvvm.tcgen05.mma %d_tmem, %a_tmem, %b_desc, %idesc, %enable_input_d
+ {kind = #nvvm.tcgen05_mma_kind<ti16>, ctaGroup = #nvvm.cta_group<cta_1>, collectorOp = #nvvm.tcgen05_mma_collectorop<fill>, collectorOpB = #nvvm.tcgen05_mma_collectorop<use>} : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-shared-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-shared-ti16.mlir
new file mode 100644
index 0000000000000..78b3b00805244
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-shared-ti16.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_ti16
+llvm.func @nvvm_tcgen05_mma_ws_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* collector_b_buffer=b0 */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.ws %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d
+ {kind = #nvvm.tcgen05_mma_kind<ti16>} : (!llvm.ptr<6>, i64, i64, i32, i1)
+
+ llvm.return
+}
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_zero_col_mask_ti16
+llvm.func @nvvm_tcgen05_mma_ws_zero_col_mask_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %zero_col_mask: i64) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.shared.zero_col_mask(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, i64 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* collector_b_buffer=b0 */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.ws %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %zero_col_mask
+ {kind = #nvvm.tcgen05_mma_kind<ti16>} : (!llvm.ptr<6>, i64, i64, i32, i1, i64)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-shared-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-shared-ti16.mlir
new file mode 100644
index 0000000000000..db5971bcd8873
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-shared-ti16.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_sp_ti16
+llvm.func @nvvm_tcgen05_mma_ws_sp_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.sp.shared(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* collector_b_buffer=b0 */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.ws.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata
+ {kind = #nvvm.tcgen05_mma_kind<ti16>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
+
+ llvm.return
+}
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_sp_zero_col_mask_ti16
+llvm.func @nvvm_tcgen05_mma_ws_sp_zero_col_mask_ti16(%d_tmem : !llvm.ptr<6>, %a_desc: i64, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %spmetadata: !llvm.ptr<6>, %zero_col_mask: i64) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.sp.shared.zero_col_mask(ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, i64 {{%[0-9]+}}, i32 {{%[0-9]+}}, i1 {{%[0-9]+}}, ptr addrspace(6) {{%[0-9]+}}, i64 {{%[0-9]+}}, /* kind=ti16 */ i32 4, /* collector_b_buffer=b0 */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.ws.sp %d_tmem, %a_desc, %b_desc, %idesc, %enable_input_d, %spmetadata, %zero_col_mask
+ {kind = #nvvm.tcgen05_mma_kind<ti16>} : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, i64)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-tensor-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-tensor-ti16.mlir
new file mode 100644
index 0000000000000..cf16aaa0f0b1a
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-w...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/216588
More information about the Mlir-commits
mailing list