[Mlir-commits] [mlir] [mlir][nvgpu] Preserve unsigned types in MMA lowering to NVVM (PR #224837)

[llvmlistbot at llvm.org]

https://github.com/Vaisman created https://github.com/llvm/llvm-project/pull/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 MMA lowering.

>From 0cb3a63493b8677e857d227dd53609a5ae2b9a2a Mon Sep 17 00:00:00 2001
From: Vasili Svirski <vasili.svirski at gmail.com>
Date: Sat, 19 Sep 2026 11:50:21 +0200
Subject: [PATCH] [mlir][nvgpu] Preserve unsigned types in MMA lowering to NVVM

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 MMA lowering.
---
 .../Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp    |  4 +--
 .../Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir | 36 +++++++++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 7b8303d87db684..d90492db742295 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..e3150c9f954cba 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,20 @@ 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>
+}
+
 !barrierType = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>
 !tokenType = !nvgpu.mbarrier.token