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

[llvmlistbot at llvm.org]

Author: Vaisman
Date: 2026-10-04T02:09:35Z
New Revision: ed17053cb1381a578ce3243a0712ffe5aaaad35d

URL: https://github.com/llvm/llvm-project/commit/ed17053cb1381a578ce3243a0712ffe5aaaad35d
DIFF: https://github.com/llvm/llvm-project/commit/ed17053cb1381a578ce3243a0712ffe5aaaad35d.diff

LOG: [mlir][nvgpu] Align integer WGMMA type checks with i8 operands (#212215)

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.

Added: 
    

Modified: 
    mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
    mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
    mlir/test/Dialect/NVGPU/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 7b8303d87db68..c2f74baf1ff89 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -1249,7 +1249,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 2649564a7828e..de712e68cfefe 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -539,7 +539,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 15249577241cf..121ed43a29b70 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) {