[Mlir-commits] [llvm] [mlir] [NVPTX] Added more MMA intrinsics for F8F6F4 and FP64 types. (PR #156040)

[Artem Belevich]

================
@@ -516,13 +559,53 @@ def is_mma_variant_supported(op, layout_a, layout_b, satf):
     ):
         return False
 
+    if (
+        op.a.geom != "m8n8k4"
+        and op.a.mma_type.ptx_type == "f64"
+        and (ptx_version < 78 or gpu_arch < 90)
+    ):
+        return False
+
     # C and D type must be the same
-    if op.a.geom == "m16n8k16" and op.c.mma_type.ptx_type != op.d.mma_type.ptx_type:
+    if (
+        op.a.geom in ["m16n8k16", "m16n8k32"]
+        and op.c.mma_type.ptx_type != op.d.mma_type.ptx_type
+    ):
+        return False
+
+    if (
+        op.a.geom in ["m16n8k16", "m16n8k32"]
+        and any(
+            x in ["e4m3", "e5m2"]
+            for x in (op.a.mma_type.ptx_type, op.b.mma_type.ptx_type)
+        )
+        and ptx_version < 87
+    ):
+        return False
+
+    if kind != "" and (ptx_version < 87 or gpu_arch < 120 or not aa):
----------------
Artem-B wrote:

Nit: I'd write it as `and not (ptx_version >=87 and gpu_arch >= 120 and aa)`

Easier to reason about narrowly specific subset of what we want, rather than a much larger set of things we do not want, along multiple dimansions.

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