Fix not enabling fp8 WMMA on Navi4 by default (#1677)

Fix not enabling WMMA on Navi4 by default getDefaultFeatures() in the architecture database hadn't been told that sometimes there's a fp8 WMMA, so it wasn't sending those operations down that path. This commit fixes the issue and tests that enablement.
ROCm · Oct 9, 2024 · 82378ac · 82378ac
1 parent a1f9fe1
commit 82378ac
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/mlir/lib/Dialect/Rock/utility/AmdArchDb.cpp b/mlir/lib/Dialect/Rock/utility/AmdArchDb.cpp
@@ -111,8 +111,10 @@ GemmFeatures mlir::rock::AmdArchInfo::getDefaultFeatures(Type dataType) {
   bool isWmma = bitEnumContainsAll(theseFeatures, GemmFeatures::wmma);
   Type elementType = getElementTypeOrSelf(dataType);
   if (isWmma) {
-    if (!elementType.isF16() && !elementType.isBF16() &&
-        !elementType.isInteger(8)) {
+    if (!(isa<Float16Type, BFloat16Type>(elementType) ||
+          elementType.isInteger(8) ||
+          (hasFp8ConversionInstrs &&
+           isa<Float8E5M2Type, Float8E4M3FNType>(elementType)))) {
       theseFeatures = bitEnumClear(theseFeatures, GemmFeatures::wmma);
     }
   }

diff --git a/mlir/test/rocmlir-gen/wmma-enablement.mlir b/mlir/test/rocmlir-gen/wmma-enablement.mlir
@@ -0,0 +1,12 @@
+// RUN: rocmlir-gen --arch gfx1201 --operation gemm --operation gemm -wmma infer -t f16 -p | grep '|wmma' | count 1
+// RUN: rocmlir-gen --arch gfx1201 --operation gemm -wmma infer -t fp8_fp8 -p | grep '|wmma' | count 1
+// RUN: rocmlir-gen --arch gfx1201 --operation gemm -wmma infer -t bf8_bf8 -p | grep '|wmma' | count 1
+// RUN: rocmlir-gen --arch gfx1201 --operation gemm -wmma infer -t fp8_fp8 -force-f8-types=fnuz -p | not grep '|wmma'
+
+// RUN: rocmlir-gen --arch gfx1100 --operation gemm -wmma infer -t f16 -p | grep '|wmma' | count 1
+// RUN: rocmlir-gen --arch gfx1100 --operation gemm -wmma infer -t fp8_fp8 -p | not grep '|wmma'
+
+// YES: rock.gemm
+// YES-SAME: features = {{[^ ]*}}wmma
+// NO: rock.gemm
+// NO-NOT: wmma