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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Vaisman

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/224837.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp (+2-2) 
- (modified) mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir (+36) 


``````````diff
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 7b8303d87db68..d90492db74229 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 fe11c4f038af8..e3150c9f954cb 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
 

``````````

</details>


https://github.com/llvm/llvm-project/pull/224837