[Mlir-commits] [mlir] cb69027 - [mlir][nvgpu] Preserve unsigned types in MMA lowering to NVVM (#224837)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 00:51:16 PDT 2026
Author: Vaisman
Date: 2026-10-15T18:51:09+11:00
New Revision: cb6902729b96665543a40906f2f19eba312cd7eb
URL: https://github.com/llvm/llvm-project/commit/cb6902729b96665543a40906f2f19eba312cd7eb
DIFF: https://github.com/llvm/llvm-project/commit/cb6902729b96665543a40906f2f19eba312cd7eb.diff
LOG: [mlir][nvgpu] Preserve unsigned types in MMA lowering to NVVM (#224837)
Type::isInteger(width) matches signed, unsigned, and signless integers.
As a result, getNvvmMmaType lowered ui8 and ui4 multiplicands to the
signed PTX types s8 and s4. Unsigned values with the high bit set were
therefore interpreted as negative by mma.sync and mma.sp.sync, with no
diagnostic.
Select u8 and u4 for unsigned operands while retaining the existing
signed PTX types for signed and signless operands. The downstream NVVM
and PTX lowering already supports these types. Add coverage for dense
ui8/ui4 and sparse ui8/i4/ui4 MMA lowering.
Added:
Modified:
mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index c2f74baf1ff89e..1877be1c1dc425 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -319,9 +319,9 @@ struct MmaLdMatrixOpToNVVM : public ConvertOpToLLVMPattern<nvgpu::LdMatrixOp> {
static FailureOr<NVVM::MMATypes> getNvvmMmaType(Type t) {
Type elType = getElementTypeOrSelf(t);
if (elType.isInteger(8))
- return NVVM::MMATypes::s8;
+ return elType.isUnsignedInteger() ? NVVM::MMATypes::u8 : NVVM::MMATypes::s8;
if (elType.isInteger(4))
- return NVVM::MMATypes::s4;
+ return elType.isUnsignedInteger() ? NVVM::MMATypes::u4 : NVVM::MMATypes::s4;
if (elType.isF16())
return NVVM::MMATypes::f16;
if (elType.isBF16())
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
index fe11c4f038af81..074b646cdc0c71 100644
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
@@ -117,6 +117,17 @@ func.func @m16n8k32_int8(%arg0: vector<4x4xi8>, %arg1: vector<2x4xi8>, %arg2: ve
return %d : vector<2x2xi32>
}
+// CHECK-LABEL: @m16n8k32_uint8
+func.func @m16n8k32_uint8(%arg0: vector<4x4xui8>, %arg1: vector<2x4xui8>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {
+ // CHECK: [[d:%.+]] = nvvm.mma.sync
+ // CHECK-SAME: shape = <m = 16, n = 8, k = 32>
+ // CHECK-SAME: int_overflow = satfinite
+ // CHECK-SAME: multiplicand_a_ptx_type = u8
+ // CHECK-SAME: multiplicand_b_ptx_type = u8
+ %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) mmaShape = [16, 8, 32] : (vector<4x4xui8>, vector<2x4xui8>, vector<2x2xi32>) -> vector<2x2xi32>
+ return %d : vector<2x2xi32>
+}
+
// CHECK-LABEL: @m16n8k32_i4
func.func @m16n8k32_i4(%arg0: vector<2x8xi4>, %arg1: vector<1x8xi4>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {
// CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<2 x vector<8xi4>>
@@ -136,6 +147,17 @@ func.func @m16n8k32_i4(%arg0: vector<2x8xi4>, %arg1: vector<1x8xi4>, %arg2: vect
return %d : vector<2x2xi32>
}
+// CHECK-LABEL: @m16n8k32_ui4
+func.func @m16n8k32_ui4(%arg0: vector<2x8xui4>, %arg1: vector<1x8xui4>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {
+ // CHECK: [[d:%.+]] = nvvm.mma.sync
+ // CHECK-SAME: shape = <m = 16, n = 8, k = 32>
+ // CHECK-SAME: int_overflow = satfinite
+ // CHECK-SAME: multiplicand_a_ptx_type = u4
+ // CHECK-SAME: multiplicand_b_ptx_type = u4
+ %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) mmaShape = [16, 8, 32] : (vector<2x8xui4>, vector<1x8xui4>, vector<2x2xi32>) -> vector<2x2xi32>
+ return %d : vector<2x2xi32>
+}
+
// CHECK-LABEL: @m16n8k64_i4
func.func @m16n8k64_i4(%arg0: vector<4x8xi4>, %arg1: vector<2x8xi4>, %arg2: vector<2x2xi32>) -> vector<2x2xi32> {
// CHECK: [[el:%.+]] = llvm.extractvalue %{{.*}}[{{.*}}] : !llvm.array<4 x vector<8xi4>>
@@ -492,6 +514,48 @@ func.func @mma_sp_sync_i8_16864(%arg0: vector<4x4xi8>,
return %d : vector<2x2xi32>
}
+// CHECK-LABEL: func @mma_sp_sync_ui8_16864(
+func.func @mma_sp_sync_ui8_16864(%arg0: vector<4x4xui8>,
+ %arg1: vector<4x4xui8>,
+ %arg2: vector<2x2xi32>,
+ %arg3: vector<2xi16>) -> vector<2x2xi32> {
+
+ // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att
+ // CHECK-SAME: "mma.sp.sync.aligned.m16n8k64.row.col.satfinite.s32.u8.u8.s32 {$0,$1,$2,$3},{$4,$5,$6,$7},{$8,$9,$10,$11},{$12,$13,$14,$15},$16,0x0;"
+
+ %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) mmaShape = [16, 8, 64] :
+ (vector<4x4xui8>, vector<4x4xui8>, vector<2x2xi32>) -> vector<2x2xi32>
+ return %d : vector<2x2xi32>
+}
+
+// CHECK-LABEL: func @mma_sp_sync_i4_16864(
+func.func @mma_sp_sync_i4_16864(%arg0: vector<2x8xi4>,
+ %arg1: vector<2x8xi4>,
+ %arg2: vector<2x2xi32>,
+ %arg3: vector<2xi16>) -> vector<2x2xi32> {
+
+ // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att
+ // CHECK-SAME: "mma.sp.sync.aligned.m16n8k64.row.col.satfinite.s32.s4.s4.s32 {$0,$1,$2,$3},{$4,$5},{$6,$7},{$8,$9,$10,$11},$12,0x0;"
+
+ %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) mmaShape = [16, 8, 64] :
+ (vector<2x8xi4>, vector<2x8xi4>, vector<2x2xi32>) -> vector<2x2xi32>
+ return %d : vector<2x2xi32>
+}
+
+// CHECK-LABEL: func @mma_sp_sync_ui4_16864(
+func.func @mma_sp_sync_ui4_16864(%arg0: vector<2x8xui4>,
+ %arg1: vector<2x8xui4>,
+ %arg2: vector<2x2xi32>,
+ %arg3: vector<2xi16>) -> vector<2x2xi32> {
+
+ // CHECK: %[[d:.+]] = llvm.inline_asm has_side_effects asm_dialect = att
+ // CHECK-SAME: "mma.sp.sync.aligned.m16n8k64.row.col.satfinite.s32.u4.u4.s32 {$0,$1,$2,$3},{$4,$5},{$6,$7},{$8,$9,$10,$11},$12,0x0;"
+
+ %d = nvgpu.mma.sp.sync(%arg0, %arg1, %arg2) metadata(%arg3) mmaShape = [16, 8, 64] :
+ (vector<2x8xui4>, vector<2x8xui4>, vector<2x2xi32>) -> vector<2x2xi32>
+ return %d : vector<2x2xi32>
+}
+
!barrierType = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>
!tokenType = !nvgpu.mbarrier.token
More information about the Mlir-commits
mailing list