[Mlir-commits] [mlir] [mlir][nvgpu] Align integer WGMMA type checks with i8 operands (PR #212215)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu

@llvm/pr-subscribers-mlir

Author: Vaisman

<details>
<summary>Changes</summary>

Integer WGMMA uses i8 operands with an i32 accumulator, but the NVGPU verifier and K-shape selection still checked for i16.

Changes both checks to i8 and adds tests for rejected i16 inputs and the existing i8 "not supported yet" limitation.

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


3 Files Affected:

- (modified) mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp (+1-1) 
- (modified) mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp (+1-1) 
- (modified) mlir/test/Dialect/NVGPU/invalid.mlir (+22) 


``````````diff
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index b155505a2dae1..f37a2af13abf9 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -1248,7 +1248,7 @@ struct NVGPUWarpgroupMmaOpLowering
       } else if (inputElemType.isF16() || inputElemType.isBF16()) {
         wgmmaK = 16;
       } else if (isa<Float8E4M3FNType, Float8E5M2Type>(inputElemType) ||
-                 inputElemType.isInteger(16)) {
+                 inputElemType.isInteger(8)) {
         wgmmaK = 32;
       } else if (inputElemType.isInteger(1)) {
         wgmmaK = 256;
diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index 642e588488641..edf5e8b49e0da 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -540,7 +540,7 @@ LogicalResult isAllowedWGMMADataType(Type typeD, Type typeA, Type typeB) {
   if (typeA.isTF32() && typeD.isF32() && typeB.isTF32())
     return success();
   // s32 += i8 + i8
-  if (typeA.isInteger(16) && typeB.isInteger(16) && typeD.isInteger(32))
+  if (typeA.isInteger(8) && typeB.isInteger(8) && typeD.isInteger(32))
     return success();
   // s32 += i1 + i1
   if (typeA.isInteger(1) && typeB.isInteger(1) && typeD.isInteger(32))
diff --git a/mlir/test/Dialect/NVGPU/invalid.mlir b/mlir/test/Dialect/NVGPU/invalid.mlir
index 0341a08a2c778..22c84992cd630 100644
--- a/mlir/test/Dialect/NVGPU/invalid.mlir
+++ b/mlir/test/Dialect/NVGPU/invalid.mlir
@@ -276,6 +276,28 @@ func.func @warpgroup_mma_wrong_input(%descA: !tDescA, %descB: !tDescB, %acc: !tR
 
 // -----
 
+!tResult = !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xi32>>
+!tDescA  = !nvgpu.warpgroup.descriptor<tensor = memref<64x32xi16, 3>>
+!tDescB  = !nvgpu.warpgroup.descriptor<tensor = memref<32x128xi16, 3>>
+func.func @warpgroup_mma_unsupported_i16_inputs(%descA: !tDescA, %descB: !tDescB, %acc: !tResult) {
+  // expected-error @+1 {{'nvgpu.warpgroup.mma' op 'i32' += 'i16' * 'i16', it is not supported.}}
+  %0 = nvgpu.warpgroup.mma %descA, %descB, %acc: !tDescA, !tDescB, !tResult -> !tResult
+  return
+}
+
+// -----
+
+!tResult = !nvgpu.warpgroup.accumulator<fragmented = vector<64x128xi32>>
+!tDescA  = !nvgpu.warpgroup.descriptor<tensor = memref<64x32xi8, 3>>
+!tDescB  = !nvgpu.warpgroup.descriptor<tensor = memref<32x128xi8, 3>>
+func.func @warpgroup_mma_i8_inputs_not_enabled(%descA: !tDescA, %descB: !tDescB, %acc: !tResult) {
+  // expected-error @+1 {{'nvgpu.warpgroup.mma' op hit a limitation: 'i32' += 'i8' * 'i8', it is not supported yet}}
+  %0 = nvgpu.warpgroup.mma %descA, %descB, %acc: !tDescA, !tDescB, !tResult -> !tResult
+  return
+}
+
+// -----
+
 !desc = !nvgpu.tensormap.descriptor<tensor = memref<32x8xf32,3>, swizzle=swizzle_32b, l2promo = none, oob = zero, interleave = none>
 !mbarrier = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>
 func.func @tma_load_1(%desc: !desc, %buffer1: memref<128xf32,3>, %buffer2: memref<32x8xf32,3>, %buffer3: memref<32x32xf32>, %mbarrier: !mbarrier) {

``````````

</details>


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