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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+1-2) 
- (modified) mlir/test/Dialect/LLVMIR/nvvm.mlir (+15) 


``````````diff
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
+}

``````````

</details>


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