[llvm] [mlir] [NVPTX] Added more MMA intrinsics for F8F6F4 and FP64 types. (PR #156040)
Artem Belevich via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 2 14:57:00 PDT 2025
================
@@ -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
More information about the llvm-commits
mailing list