[Mlir-commits] [mlir] [mlir][AMDGPU] Add scaled wmma ops for gfx1250 (PR #169854)

[Krzysztof Drewniak]

================
@@ -1225,24 +1226,19 @@ def AMDGPU_ScaledMFMAOp :
   let hasCanonicalizer = 1;
 }
 
-def AMDGPU_ScaledWMMAOp :
-    AMDGPU_Op<"scaled_wmma", [AllTypesMatch<["destC", "destD"]>,
-                              Pure]>,
-    Arguments<(ins
-                   ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$m,
-                   ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$n,
-                   ConfinedAttr<I32Attr, [IntIsOneOf<[128]>]>:$k,
-                   ScaledWMMAInTypes:$sourceA,
-                   ScaledWMMAInTypes:$sourceB,
-                   ScaledWMMAOutTypes:$destC,
-                   AnyTypeOf<[I32, I64]>:$scaleA,
-                   AnyTypeOf<[I32, I64]>:$scaleB,
-                   DefaultValuedAttr<I32Attr, "0">:$scaleAType,
-                   DefaultValuedAttr<I32Attr, "0">:$fmtScaleA,
-                   DefaultValuedAttr<I32Attr, "0">:$scaleBType,
-                   DefaultValuedAttr<I32Attr, "0">:$fmtScaleB
-                   )>,
-    Results<(outs ScaledWMMAOutTypes: $destD)> {
+def AMDGPU_ScaledWMMAOp
+    : AMDGPU_Op<"scaled_wmma", [AllTypesMatch<["destC", "destD"]>, Pure]>,
+      Arguments<(ins ConfinedAttr<I32Attr, [IntIsOneOf<[16, 32]>]>:$m,
+          ConfinedAttr<I32Attr, [IntIsOneOf<[16]>]>:$n,
+          ConfinedAttr<I32Attr, [IntIsOneOf<[128]>]>:$k,
+          ScaledWMMAInTypes:$sourceA, ScaledWMMAInTypes:$sourceB,
+          ScaledWMMAOutTypes:$destC, AnyTypeOf<[I32, I64]>:$scaleA,
+          AnyTypeOf<[I32, I64]>:$scaleB,
----------------
krzysz00 wrote:

.. Wait, hold on, I don't think an E4M3 scale is allowed? Unless I missed something?

But yeah, E5M3FNU scales are a future sequence of PRs (you'll need to add them to APFloat, plumb them though to MLIR, probably add software fallbacks, etc.) ... I'd just leave a note saying that E5M3 scales will be added here once the type is plumbed through

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