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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From 8979aebd92a941f1f563976ef81fa5820658a6b6 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 2 Sep 2026 15:41:47 -0700
Subject: [PATCH] [mlir][NVVM] Reject f16 WGMMA accumulators for BF16 inputs

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. This addresses CID 1571353.

The pre-fix test failed under Valgrind with:

  error: expected error "op f16 += bf16 * bf16, it is not supported"
  was not produced
  ==610105== ERROR SUMMARY: 0 errors from 0 contexts

Valgrind reports no invalid access because this is a verifier omission; the
missing expected diagnostic exposes the defect.

Found by Coverity.

Assisted-by: Codex
---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp |  3 +--
 mlir/test/Dialect/LLVMIR/nvvm.mlir         | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 843344703d0a5..d9b60b0f4c260 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 70e7816c6c1ec..7003cfe319f8f 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
+}