[Mlir-commits] [mlir] [mlir][nvvm] Verify kind-specific target requirements of tcgen05.mma ops (PR #235628)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 00:23:28 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-nvptx
Author: Vaisman
<details>
<summary>Changes</summary>
The tcgen05.mma ops were checked against the gpu.module target only at the op level (sm_100f, sm_101f or sm_110f), but the NVPTX backend gates several of their variants more narrowly. MLIR accepted these variants on targets where llc then fails with "Cannot select":
- kind::i8 needs sm_100a or sm_110a.
- kind::ti16 and a collector_b qualifier other than discard need the Rubin family (sm_107f).
- An immediate scale_input_d needs sm_100f.
- Sparse block-scaled mxf4 and mxf4nvf4 need sm_100a, sm_103a, sm_107a or sm_110a.
Have the tcgen05.mma ops implement getRequiredMinSMVersion() instead of using the static trait, and derive the requirement from these attributes, mirroring the NVPTX predicates. When two attributes impose requirements that no target meets together, such as collector_b with kind::i8, the op is rejected on every target.
Tests cover each target restriction on an unsupported and a supported target.
---
Full diff: https://github.com/llvm/llvm-project/pull/235628.diff
3 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+6-6)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+88)
- (modified) mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir (+177)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index ccc7805dd6ae19..7b2523106a3a27 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6560,7 +6560,7 @@ defvar Tcgen05MMABlockScaleKindAttr =
[EnumAttrIsOneOf<Tcgen05MMAKindAttr, Tcgen05MMABlockScaleKindList>]>;
def NVVM_Tcgen05MMAOp : NVVM_Op<"tcgen05.mma",
- [AttrSizedOperandSegments, NVVMRequiresSMf<[100, 101, 110]>]> {
+ [AttrSizedOperandSegments, DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs MMA operation on 5th-gen tensor cores";
let description = [{
@@ -6653,7 +6653,7 @@ def NVVM_Tcgen05MMAOp : NVVM_Op<"tcgen05.mma",
}
def NVVM_Tcgen05MMASparseOp : NVVM_Op<"tcgen05.mma.sp",
- [AttrSizedOperandSegments, NVVMRequiresSMf<[100, 101, 110]>]> {
+ [AttrSizedOperandSegments, DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs MMA operation with sparse A matrix on 5th-gen tensor cores";
let description = [{
@@ -6743,7 +6743,7 @@ def Tcgen05MMABlockScaleAttr : EnumAttr<NVVM_Dialect, Tcgen05MMABlockScale,
"tcgen05_mma_block_scale">;
def NVVM_Tcgen05MMABlockScaleOp : NVVM_Op<"tcgen05.mma.block_scale",
- [NVVMRequiresSMf<[100, 101, 110]>]> {
+ [DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs block scaled MMA operation on 5th-gen tensor cores";
let description = [{
@@ -6826,7 +6826,7 @@ def NVVM_Tcgen05MMABlockScaleOp : NVVM_Op<"tcgen05.mma.block_scale",
}
def NVVM_Tcgen05MMASparseBlockScaleOp : NVVM_Op<"tcgen05.mma.sp.block_scale",
- [NVVMRequiresSMf<[100, 101, 110]>]> {
+ [DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs block scaled MMA operation with sparse A matrix on 5th-gen tensor cores";
let description = [{
@@ -6915,7 +6915,7 @@ def Tcgen05MMACollectorBBuffer : I32EnumAttr<
def Tcgen05MMACollectorBBufferAttr : EnumAttr<NVVM_Dialect, Tcgen05MMACollectorBBuffer, "tcgen05_mma_collectorb">;
def NVVM_Tcgen05MMAWsOp : NVVM_Op<"tcgen05.mma.ws",
- [NVVMRequiresSMf<[100, 101, 110]>]> {
+ [DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs weight stationary convolution MMA operation on 5th-gen tensor cores";
let description = [{
@@ -6990,7 +6990,7 @@ def NVVM_Tcgen05MMAWsOp : NVVM_Op<"tcgen05.mma.ws",
}
def NVVM_Tcgen05MMAWsSparseOp : NVVM_Op<"tcgen05.mma.ws.sp",
- [NVVMRequiresSMf<[100, 101, 110]>]> {
+ [DeclareOpInterfaceMethods<RequiresSMInterface>]> {
let summary = "Performs weight stationary convolution MMA with sparse A matrix on 5th-gen tensor cores";
let description = [{
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index d7d0e6f329f138..b1bfd33cf88cac 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -6228,6 +6228,94 @@ mlir::NVVM::IDArgPair TensormapReplaceOp::getIntrinsicIDAndArgs(
return {IDs[fieldIndex], args};
}
+//===----------------------------------------------------------------------===//
+// NVVM tcgen05.mma target requirements
+//===----------------------------------------------------------------------===//
+
+// Full SM versions as encoded by NVVMCheckSMVersion: the SM version times 10,
+// plus 2 for the family-specific ('f') and 3 for the architecture-specific
+// ('a') variant.
+static constexpr unsigned kSM100a = 1003, kSM100f = 1002, kSM101f = 1012,
+ kSM103a = 1033, kSM107a = 1073, kSM107f = 1072,
+ kSM110a = 1103, kSM110f = 1102;
+
+// The requirements below mirror the predicates of the tcgen05.mma instructions
+// in the NVPTX backend.
+static constexpr unsigned kTcgen05MMABaseSM[] = {kSM100f, kSM101f, kSM110f};
+
+static NVVMCheckSMVersion makeSMRequirement(ArrayRef<unsigned> versions) {
+ NVVMCheckSMVersion requirement;
+ requirement.fullSmVersionList.assign(versions.begin(), versions.end());
+ return requirement;
+}
+
+static ArrayRef<unsigned> getTcgen05MMAKindRequiredSM(Tcgen05MMAKind kind,
+ bool hasScaleInputD) {
+ static constexpr unsigned kRubin[] = {kSM107f};
+ static constexpr unsigned kScaleInputD[] = {kSM100f};
+ static constexpr unsigned kI8[] = {kSM100a, kSM110a};
+ if (kind == Tcgen05MMAKind::TI16)
+ return kRubin;
+ if (hasScaleInputD)
+ return kScaleInputD;
+ if (kind == Tcgen05MMAKind::I8)
+ return kI8;
+ return kTcgen05MMABaseSM;
+}
+
+// A collector_b qualifier other than discard is a Rubin (sm_107f) extension.
+// Keep the part of `kindRequirement` that a Rubin target also satisfies. An
+// empty result rejects every target.
+static NVVMCheckSMVersion
+getTcgen05MMARequiredSM(Tcgen05MMACollectorOp collectorOpB,
+ ArrayRef<unsigned> kindRequirement) {
+ NVVMCheckSMVersion requirement = makeSMRequirement(kindRequirement);
+ if (collectorOpB == Tcgen05MMACollectorOp::DISCARD)
+ return requirement;
+ if (requirement.isCompatibleWith(kSM107f))
+ return NVVMCheckSMVersion(kSM107f);
+ if (requirement.isCompatibleWith(kSM107a))
+ return NVVMCheckSMVersion(kSM107a);
+ return NVVMCheckSMVersion();
+}
+
+NVVMCheckSMVersion Tcgen05MMAOp::getRequiredMinSMVersion() {
+ return getTcgen05MMARequiredSM(
+ getCollectorOpB(), getTcgen05MMAKindRequiredSM(
+ getKind(), static_cast<bool>(getScaleInputD())));
+}
+
+NVVMCheckSMVersion Tcgen05MMASparseOp::getRequiredMinSMVersion() {
+ return getTcgen05MMARequiredSM(
+ getCollectorOpB(), getTcgen05MMAKindRequiredSM(
+ getKind(), static_cast<bool>(getScaleInputD())));
+}
+
+NVVMCheckSMVersion Tcgen05MMABlockScaleOp::getRequiredMinSMVersion() {
+ return getTcgen05MMARequiredSM(getCollectorOpB(), kTcgen05MMABaseSM);
+}
+
+NVVMCheckSMVersion Tcgen05MMASparseBlockScaleOp::getRequiredMinSMVersion() {
+ static constexpr unsigned kSparseMxf4[] = {kSM100a, kSM103a, kSM110a,
+ kSM107a};
+ Tcgen05MMAKind kind = getKind();
+ bool isMxf4 =
+ kind == Tcgen05MMAKind::MXF4 || kind == Tcgen05MMAKind::MXF4NVF4;
+ return getTcgen05MMARequiredSM(
+ getCollectorOpB(), isMxf4 ? ArrayRef<unsigned>(kSparseMxf4)
+ : ArrayRef<unsigned>(kTcgen05MMABaseSM));
+}
+
+NVVMCheckSMVersion Tcgen05MMAWsOp::getRequiredMinSMVersion() {
+ return makeSMRequirement(
+ getTcgen05MMAKindRequiredSM(getKind(), /*hasScaleInputD=*/false));
+}
+
+NVVMCheckSMVersion Tcgen05MMAWsSparseOp::getRequiredMinSMVersion() {
+ return makeSMRequirement(
+ getTcgen05MMAKindRequiredSM(getKind(), /*hasScaleInputD=*/false));
+}
+
//===----------------------------------------------------------------------===//
// NVVM tcgen05.mma functions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
index a9b1e6d0fd3b5f..906616fab84996 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
@@ -433,3 +433,180 @@ gpu.module @mbarrier_try_wait_parity_sm80 [#nvvm.target<chip = "sm_80">] {
return
}
}
+
+// -----
+
+// The tcgen05.mma ops narrow their requirement depending on the kind and on
+// the collector_b and scale qualifiers, mirroring the NVPTX backend.
+
+gpu.module @tcgen05_mma_i8_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_i8_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ // expected-error @below {{'nvvm.tcgen05.mma' op is not supported on sm_103a}}
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = i8, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_i8_sm100a [#nvvm.target<chip = "sm_100a">] {
+ func.func @tcgen05_mma_i8_sm100a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = i8, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_ti16_sm100a [#nvvm.target<chip = "sm_100a">] {
+ func.func @tcgen05_mma_ti16_sm100a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ // expected-error @below {{'nvvm.tcgen05.mma' op is not supported on sm_100a}}
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = ti16, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_ti16_sm107a [#nvvm.target<chip = "sm_107a">] {
+ func.func @tcgen05_mma_ti16_sm107a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = ti16, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_collector_b_sm100a [#nvvm.target<chip = "sm_100a">] {
+ func.func @tcgen05_mma_collector_b_sm100a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ // expected-error @below {{'nvvm.tcgen05.mma' op is not supported on sm_100a}}
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = f16, cta_group = <cta_1> collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_collector_b_sm107f [#nvvm.target<chip = "sm_107f">] {
+ func.func @tcgen05_mma_collector_b_sm107f(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = f16, cta_group = <cta_1> collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+// No target supports both collector_b and kind::i8.
+
+gpu.module @tcgen05_mma_collector_b_i8_sm107a [#nvvm.target<chip = "sm_107a">] {
+ func.func @tcgen05_mma_collector_b_i8_sm107a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ // expected-error @below {{'nvvm.tcgen05.mma' op is not supported on sm_107a}}
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = i8, cta_group = <cta_1> collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_scale_d_sm110a [#nvvm.target<chip = "sm_110a">] {
+ func.func @tcgen05_mma_scale_d_sm110a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ %scale = llvm.mlir.constant(0 : i64) : i64
+ // expected-error @below {{'nvvm.tcgen05.mma' op is not supported on sm_110a}}
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en scale = %scale
+ , kind = f16, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, i64)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_scale_d_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_scale_d_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ %scale = llvm.mlir.constant(0 : i64) : i64
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en scale = %scale
+ , kind = f16, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, i64)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_f16_sm103f [#nvvm.target<chip = "sm_103f">] {
+ func.func @tcgen05_mma_f16_sm103f(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ nvvm.tcgen05.mma %d, %a, %b, %idesc, %en
+ , kind = f16, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_sp_i8_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_sp_i8_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1, %sp: !llvm.ptr<6>) {
+ // expected-error @below {{'nvvm.tcgen05.mma.sp' op is not supported on sm_103a}}
+ nvvm.tcgen05.mma.sp %d, %a, %b, %idesc, %en, %sp
+ , kind = i8, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_block_scale_collector_b_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_block_scale_collector_b_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1, %sa: !llvm.ptr<6>, %sb: !llvm.ptr<6>) {
+ // expected-error @below {{'nvvm.tcgen05.mma.block_scale' op is not supported on sm_103a}}
+ nvvm.tcgen05.mma.block_scale %d, %a, %b, %idesc, %en, %sa, %sb
+ , kind = mxf8f6f4, cta_group = <cta_1> collector_b = lastuse : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_sp_block_scale_mxf4_sm100f [#nvvm.target<chip = "sm_100f">] {
+ func.func @tcgen05_mma_sp_block_scale_mxf4_sm100f(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1, %sp: !llvm.ptr<6>, %sa: !llvm.ptr<6>, %sb: !llvm.ptr<6>) {
+ // expected-error @below {{'nvvm.tcgen05.mma.sp.block_scale' op is not supported on sm_100f}}
+ nvvm.tcgen05.mma.sp.block_scale %d, %a, %b, %idesc, %en, %sp, %sa, %sb
+ , kind = mxf4, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_sp_block_scale_mxf4_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_sp_block_scale_mxf4_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1, %sp: !llvm.ptr<6>, %sa: !llvm.ptr<6>, %sb: !llvm.ptr<6>) {
+ nvvm.tcgen05.mma.sp.block_scale %d, %a, %b, %idesc, %en, %sp, %sa, %sb
+ , kind = mxf4, cta_group = <cta_1> : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>, !llvm.ptr<6>, !llvm.ptr<6>)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_ws_i8_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_ws_i8_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1) {
+ // expected-error @below {{'nvvm.tcgen05.mma.ws' op is not supported on sm_103a}}
+ nvvm.tcgen05.mma.ws %d, %a, %b, %idesc, %en
+ kind = i8 : (!llvm.ptr<6>, i64, i64, i32, i1)
+ return
+ }
+}
+
+// -----
+
+gpu.module @tcgen05_mma_ws_sp_i8_sm103a [#nvvm.target<chip = "sm_103a">] {
+ func.func @tcgen05_mma_ws_sp_i8_sm103a(%d: !llvm.ptr<6>, %a: i64, %b: i64, %idesc: i32, %en: i1, %sp: !llvm.ptr<6>) {
+ // expected-error @below {{'nvvm.tcgen05.mma.ws.sp' op is not supported on sm_103a}}
+ nvvm.tcgen05.mma.ws.sp %d, %a, %b, %idesc, %en, %sp
+ kind = i8 : (!llvm.ptr<6>, i64, i64, i32, i1, !llvm.ptr<6>)
+ return
+ }
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/235628
More information about the Mlir-commits
mailing list