[Mlir-commits] [mlir] d2adf78 - [MLIR][NVVM] Support kind::ti16 for tcgen05.mma Ops (#216588)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Aug 19 06:22:41 PDT 2026
Author: Kirill Vedernikov
Date: 2026-08-19T15:22:36+02:00
New Revision: d2adf78b4492cfe166f905addc74d8c424208358
URL: https://github.com/llvm/llvm-project/commit/d2adf78b4492cfe166f905addc74d8c424208358
DIFF: https://github.com/llvm/llvm-project/commit/d2adf78b4492cfe166f905addc74d8c424208358.diff
LOG: [MLIR][NVVM] Support kind::ti16 for tcgen05.mma Ops (#216588)
This change adds `kind::ti16` support for tcgen05.mma MLIR Ops.
Added:
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-shared-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-shared-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-sp-tensor-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-tensor-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-shared-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-shared-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-tensor-ti16.mlir
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-tensor-ti16.mlir
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 22781cdae7324..b1c817b87a79e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6241,13 +6241,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;
}
@@ -6303,8 +6305,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 2a59ffc2062d7..74faabf81893f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5892,6 +5892,28 @@ 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;
+ case NVVM::Tcgen05MMAKind::MXF8F6F4:
+ case NVVM::Tcgen05MMAKind::MXF4:
+ case NVVM::Tcgen05MMAKind::MXF4NVF4:
+ // Block-scale kinds are handled by the tcgen05.mma.block_scale
+ // lowering paths and are not valid for plain tcgen05.mma.
+ llvm_unreachable("Unsupported tcgen05.mma kind");
+ }
+}
+
mlir::NVVM::IDArgPair
Tcgen05MMAOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
llvm::IRBuilderBase &builder) {
@@ -6023,7 +6045,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));
@@ -6211,7 +6234,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));
@@ -6444,7 +6468,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(
@@ -6485,7 +6510,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..0964a3686b01b
--- /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 = ti16, 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 71f6bb9a9d2b5..617cd8335210e 100644
--- a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-invalid.mlir
@@ -125,7 +125,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 = mxf8f6f4, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
llvm.return
}
@@ -135,7 +135,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 = mxf4, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
llvm.return
}
@@ -145,7 +145,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 = mxf4nvf4 {ctaGroup = #nvvm.cta_group<cta_1>}
: (!llvm.ptr<6>, i64, i64, i32, i1)
@@ -157,7 +157,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 = mxf8f6f4 {ctaGroup = #nvvm.cta_group<cta_1>}
: (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
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..37915177e5ce9
--- /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 = ti16, 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 = ti16, cta_group = <cta_1> collector_a = fill collector_b = 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..eab0ced98452f
--- /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 = ti16, 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 = ti16, cta_group = <cta_1> collector_a = fill collector_b = 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..382a9bbb55c47
--- /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 = ti16, 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 = ti16, cta_group = <cta_1> collector_a = fill collector_b = 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..c4115b3c71e40
--- /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 = ti16, 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 = ti16, cta_group = <cta_1> collector_a = fill collector_b = 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..9c16961734f11
--- /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 = 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 = 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..c219fb5fd6a75
--- /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 = 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 = 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..abade0631667b
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-sp-tensor-ti16.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_sp_tensor_ti16
+llvm.func @nvvm_tcgen05_mma_ws_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>) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.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, /* collector_b_buffer=b0 */ i32 0, /* collector_b=discard */ i32 0)
+ nvvm.tcgen05.mma.ws.sp %d_tmem, %a_tmem, %b_desc, %idesc, %enable_input_d, %spmetadata
+ kind = ti16 : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1, !llvm.ptr<6>)
+
+ llvm.return
+}
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_sp_tensor_zero_col_mask_ti16
+llvm.func @nvvm_tcgen05_mma_ws_sp_tensor_zero_col_mask_ti16(%d_tmem : !llvm.ptr<6>, %a_tmem: !llvm.ptr<6>, %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.tensor.zero_col_mask(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[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_tmem, %b_desc, %idesc, %enable_input_d, %spmetadata, %zero_col_mask
+ kind = ti16 : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1, !llvm.ptr<6>, i64)
+
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-tensor-ti16.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-tensor-ti16.mlir
new file mode 100644
index 0000000000000..db2c2ef1be17c
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-mma-ws-tensor-ti16.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-translate --mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_tensor_ti16
+llvm.func @nvvm_tcgen05_mma_ws_tensor_ti16(%d_tmem : !llvm.ptr<6>, %a_tmem: !llvm.ptr<6>, %b_desc: i64, %idesc: i32, %enable_input_d: i1) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.tensor(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[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_tmem, %b_desc, %idesc, %enable_input_d
+ kind = ti16 : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1)
+
+ llvm.return
+}
+
+// CHECK-LABEL: @nvvm_tcgen05_mma_ws_tensor_zero_col_mask_ti16
+llvm.func @nvvm_tcgen05_mma_ws_tensor_zero_col_mask_ti16(%d_tmem : !llvm.ptr<6>, %a_tmem: !llvm.ptr<6>, %b_desc: i64, %idesc: i32, %enable_input_d: i1, %zero_col_mask: i64) {
+
+ // CHECK: call void @llvm.nvvm.tcgen05.mma.ws.tensor.zero_col_mask(ptr addrspace(6) {{%[0-9]+}}, ptr addrspace(6) {{%[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_tmem, %b_desc, %idesc, %enable_input_d, %zero_col_mask
+ kind = ti16 : (!llvm.ptr<6>, !llvm.ptr<6>, i64, i32, i1, i64)
+
+ llvm.return
+}
More information about the Mlir-commits
mailing list