[Mlir-commits] [mlir] [mlir][nvvm] Verify kind-specific target requirements of tcgen05.mma ops (PR #235628)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Nov 3 11:53:43 PST 2026
https://github.com/Vaisman updated https://github.com/llvm/llvm-project/pull/235628
>From ef36bb6faacb2d232fcf9821b4e51c2c4a59eaa9 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Sat, 24 Oct 2026 15:51:30 +0200
Subject: [PATCH 1/3] [mlir][nvvm] Verify kind-specific target requirements of
tcgen05.mma ops
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.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 12 +-
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 88 +++++++++
.../Dialect/LLVMIR/nvvm_check_target_sm.mlir | 177 ++++++++++++++++++
3 files changed, 271 insertions(+), 6 deletions(-)
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
+ }
+}
>From 364f4fac1d853ab3398e8e54412ccf1dd339d97e Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Tue, 3 Nov 2026 17:21:46 +0100
Subject: [PATCH 2/3] [mlir][nvvm] Accept sm_101a for tcgen05.mma i8 and sparse
mxf4
sm_101a is an alias of sm_110a in the NVPTX backend, so kind::i8 and the
sparse block-scaled mxf4 variants are also valid on it. Add it to both
requirement lists.
Also add an ArrayRef constructor to NVVMCheckSMVersion and use it instead
of the local makeSMRequirement helper.
---
.../Dialect/LLVMIR/NVVMRequiresSMTraits.h | 4 ++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 24 +++++++------------
.../Dialect/LLVMIR/nvvm_check_target_sm.mlir | 22 +++++++++++++++++
3 files changed, 35 insertions(+), 15 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
index d60b7c8bb68671..bbc66000d9cf23 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
@@ -15,6 +15,7 @@
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/StorageUniquerSupport.h"
+#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include <limits>
@@ -45,6 +46,9 @@ struct NVVMCheckSMVersion {
NVVMCheckSMVersion(Versions... fullSmVersions)
: fullSmVersionList({fullSmVersions...}) {}
+ explicit NVVMCheckSMVersion(llvm::ArrayRef<unsigned> fullSmVersions)
+ : fullSmVersionList(fullSmVersions.begin(), fullSmVersions.end()) {}
+
bool isCompatibleWith(const unsigned &targetFullSmVersion) const {
return llvm::any_of(
fullSmVersionList, [&](const unsigned &requiredFullSmVersion) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index b1bfd33cf88cac..f019b27e460941 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -6235,25 +6235,19 @@ mlir::NVVM::IDArgPair TensormapReplaceOp::getIntrinsicIDAndArgs(
// 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;
+static constexpr unsigned kSM100a = 1003, kSM100f = 1002, kSM101a = 1013,
+ 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};
+ static constexpr unsigned kI8[] = {kSM100a, kSM101a, kSM110a};
if (kind == Tcgen05MMAKind::TI16)
return kRubin;
if (hasScaleInputD)
@@ -6269,7 +6263,7 @@ static ArrayRef<unsigned> getTcgen05MMAKindRequiredSM(Tcgen05MMAKind kind,
static NVVMCheckSMVersion
getTcgen05MMARequiredSM(Tcgen05MMACollectorOp collectorOpB,
ArrayRef<unsigned> kindRequirement) {
- NVVMCheckSMVersion requirement = makeSMRequirement(kindRequirement);
+ NVVMCheckSMVersion requirement(kindRequirement);
if (collectorOpB == Tcgen05MMACollectorOp::DISCARD)
return requirement;
if (requirement.isCompatibleWith(kSM107f))
@@ -6296,8 +6290,8 @@ NVVMCheckSMVersion Tcgen05MMABlockScaleOp::getRequiredMinSMVersion() {
}
NVVMCheckSMVersion Tcgen05MMASparseBlockScaleOp::getRequiredMinSMVersion() {
- static constexpr unsigned kSparseMxf4[] = {kSM100a, kSM103a, kSM110a,
- kSM107a};
+ static constexpr unsigned kSparseMxf4[] = {kSM100a, kSM101a, kSM103a, kSM107a,
+ kSM110a};
Tcgen05MMAKind kind = getKind();
bool isMxf4 =
kind == Tcgen05MMAKind::MXF4 || kind == Tcgen05MMAKind::MXF4NVF4;
@@ -6307,12 +6301,12 @@ NVVMCheckSMVersion Tcgen05MMASparseBlockScaleOp::getRequiredMinSMVersion() {
}
NVVMCheckSMVersion Tcgen05MMAWsOp::getRequiredMinSMVersion() {
- return makeSMRequirement(
+ return NVVMCheckSMVersion(
getTcgen05MMAKindRequiredSM(getKind(), /*hasScaleInputD=*/false));
}
NVVMCheckSMVersion Tcgen05MMAWsSparseOp::getRequiredMinSMVersion() {
- return makeSMRequirement(
+ return NVVMCheckSMVersion(
getTcgen05MMAKindRequiredSM(getKind(), /*hasScaleInputD=*/false));
}
diff --git a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
index 906616fab84996..23071317f2a568 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
@@ -460,6 +460,18 @@ gpu.module @tcgen05_mma_i8_sm100a [#nvvm.target<chip = "sm_100a">] {
// -----
+// sm_101a is an alias of sm_110a in the NVPTX backend.
+
+gpu.module @tcgen05_mma_i8_sm101a [#nvvm.target<chip = "sm_101a">] {
+ func.func @tcgen05_mma_i8_sm101a(%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}}
@@ -591,6 +603,16 @@ gpu.module @tcgen05_mma_sp_block_scale_mxf4_sm103a [#nvvm.target<chip = "sm_103a
// -----
+gpu.module @tcgen05_mma_sp_block_scale_mxf4_sm101a [#nvvm.target<chip = "sm_101a">] {
+ func.func @tcgen05_mma_sp_block_scale_mxf4_sm101a(%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}}
>From 4cd3acc94d54d5305d38af02ca042255402a8c39 Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Tue, 3 Nov 2026 20:44:13 +0100
Subject: [PATCH 3/3] [mlir][nvvm] Treat sm_101 as an alias of sm_110 in target
checks
NVPTX treats sm_101* as aliases of sm_110*, but NVVMCheckSMVersion
compared them by numeric family, so sm_101a/f passed the sm_100f
requirement of tcgen05.mma scale_input_d and llc failed to select it.
Map sm_101* to sm_110* before comparing.
---
.../Dialect/LLVMIR/NVVMRequiresSMTraits.h | 16 ++++++++++--
.../Dialect/LLVMIR/nvvm_check_target_sm.mlir | 26 +++++++++++++++++++
2 files changed, 40 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
index bbc66000d9cf23..d8a872158d0ce2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
@@ -49,9 +49,12 @@ struct NVVMCheckSMVersion {
explicit NVVMCheckSMVersion(llvm::ArrayRef<unsigned> fullSmVersions)
: fullSmVersionList(fullSmVersions.begin(), fullSmVersions.end()) {}
- bool isCompatibleWith(const unsigned &targetFullSmVersion) const {
+ bool isCompatibleWith(const unsigned &fullSmVersion) const {
+ unsigned targetFullSmVersion = canonicalizeFullSmVersion(fullSmVersion);
return llvm::any_of(
- fullSmVersionList, [&](const unsigned &requiredFullSmVersion) {
+ fullSmVersionList, [&](unsigned requiredFullSmVersion) {
+ requiredFullSmVersion =
+ canonicalizeFullSmVersion(requiredFullSmVersion);
if (hasArchAcceleratedFeatures(requiredFullSmVersion))
return hasArchAcceleratedFeatures(targetFullSmVersion) &&
(getSMVersion(targetFullSmVersion) ==
@@ -114,6 +117,15 @@ struct NVVMCheckSMVersion {
static unsigned getSMFamily(unsigned fullSmVersion) {
return fullSmVersion / 100;
}
+
+ // sm_101 is the name PTX 8.6 through 8.8 used for what PTX 9.0 renamed to
+ // sm_110. The NVPTX backend treats the two as aliases, so compare sm_101
+ // and its variants as sm_110.
+ static unsigned canonicalizeFullSmVersion(unsigned fullSmVersion) {
+ if (getSMVersion(fullSmVersion) == 101)
+ return 1100 + fullSmVersion % 10;
+ return fullSmVersion;
+ }
};
} // namespace NVVM
diff --git a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
index 23071317f2a568..728dfa8c74f6dd 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm_check_target_sm.mlir
@@ -539,6 +539,32 @@ gpu.module @tcgen05_mma_scale_d_sm110a [#nvvm.target<chip = "sm_110a">] {
// -----
+// sm_101 is an alias of sm_110, which has no scale_input_d.
+
+gpu.module @tcgen05_mma_scale_d_sm101a [#nvvm.target<chip = "sm_101a">] {
+ func.func @tcgen05_mma_scale_d_sm101a(%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_101a}}
+ 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_sm101f [#nvvm.target<chip = "sm_101f">] {
+ func.func @tcgen05_mma_scale_d_sm101f(%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_101f}}
+ 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
More information about the Mlir-commits
mailing list