[Mlir-commits] [mlir] 2431992 - [mlir][NVVM] Reject f16 WGMMA accumulators for BF16 inputs (#220869)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-09-07T13:54:28+02:00
New Revision: 24319927cb0c4acbe9fb3f2bd042b07dad2a4f9f

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

LOG: [mlir][NVVM] Reject f16 WGMMA accumulators for BF16 inputs (#220869)

WGMMA supports an f32 accumulator for BF16 multiplicands. The verifier
also accepted an f16 accumulator, allowing an unsupported operation
through dialect verification.

Restrict the BF16 case to f32 and add verifier coverage to the existing
NVVM dialect test.

Found by Coverity.

Assisted-by: Codex

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
    mlir/test/Dialect/LLVMIR/nvvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 43280c4a6f00af..6bf52e218aae70 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3018,8 +3018,7 @@ static LogicalResult isAllowedWGMMADataType(NVVM::WGMMATypes typeD,
       return success();
     break;
   case NVVM::WGMMATypes::bf16:
-    if ((typeD == NVVM::WGMMATypes::f32 || typeD == NVVM::WGMMATypes::f16) &&
-        typeB == NVVM::WGMMATypes::bf16)
+    if (typeD == NVVM::WGMMATypes::f32 && typeB == NVVM::WGMMATypes::bf16)
       return success();
     break;
   case NVVM::WGMMATypes::e4m3:

diff  --git a/mlir/test/Dialect/LLVMIR/nvvm.mlir b/mlir/test/Dialect/LLVMIR/nvvm.mlir
index 70e7816c6c1ec7..7003cfe319f8f1 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm.mlir
@@ -661,3 +661,18 @@ llvm.func @kernel_func(%arg0: !llvm.ptr {nvvm.grid_constant}) attributes {nvvm.k
 llvm.func @kernel_func(%arg0: !llvm.ptr {llvm.byval = i32, nvvm.grid_constant = true}) attributes {nvvm.kernel} {
   llvm.return
 }
+
+// -----
+
+func.func @wgmma_f16_bf16_bf16(%descA : i64, %descB : i64) {
+  %result = llvm.mlir.undef : !llvm.struct<(f16, f16, f16, f16)>
+  // expected-error @+1 {{op f16 += bf16 * bf16, it is not supported}}
+  %res = nvvm.wgmma.mma_async %descA, %descB, %result,
+      #nvvm.shape<m = 64, n = 16, k = 16>,
+      D [<f16>, <zero>],
+      A [<bf16>, #nvvm.wgmma_scale_in<neg>, <col>],
+      B [<bf16>, #nvvm.wgmma_scale_in<neg>, <col>]
+      : !llvm.struct<(f16, f16, f16, f16)>
+      -> !llvm.struct<(f16, f16, f16, f16)>
+  return
+}