# This script generates all variants of wmma builtins, verifies that clang calls
# correct LLVM intrinsics, and checks that availability of specific builtins is
# constrained by the correct PTX version and the target GPU variant.
# Dummy test run to avoid lit warnings.
# RUN: echo "This is not a real test. It's a generator for builtins-nvpts-mma.cu" >/dev/null
from __future__ import print_function
import argparse
from collections import defaultdict
from itertools import product
from string import Template
class MMAFrag:
def __init__(self, geom, frag, ptx_elt_type):
self.geom = geom
self.frag = frag
self.ptx_type = ptx_elt_type
def __repr__(self):
return "%s:%s:%s" % (self.geom, self.frag, self.ptx_type)
class MMAOp:
def __init__(self, a, b, c, d, b1op=""):
self.a = a
self.b = b
self.c = c
self.d = d
self.b1op = b1op
def __repr__(self):
return "{A:%s, B:%s, C:%s, D:%s}" % (self.a, self.b, self.c, self.d)
def make_mma_ops(geoms, types_a, types_b, types_c, types_d, b1ops=None):
ops = []
if b1ops is None:
b1ops = [""]
for geom, type_a, type_c in product(geoms, types_a, types_c):
for type_b, type_d in product(
types_b if types_b else [type_a], types_d if types_d else [type_c]
):
ops += [
MMAOp(
MMAFrag(geom, "a", type_a),
MMAFrag(geom, "b", type_b),
MMAFrag(geom, "c", type_c),
MMAFrag(geom, "d", type_d),
b1op,
)
for b1op in b1ops
]
return ops
def make_ldst_ops(geoms, frags, types):
return [
MMAFrag(geom, frag, ptx_type)
for (geom, frag, ptx_type) in product(geoms, frags, types)
]
def get_mma_ops():
return (
make_mma_ops(["m16n16k8"], ["tf32"], [], ["f32"], [])
+ make_mma_ops(["m16n16k16", "m32n8k16", "m8n32k16"], ["bf16"], [], ["f32"], [])
+ make_mma_ops(["m8n8k4"], ["f64"], [], ["f64"], [])
+ make_mma_ops(
["m16n16k16", "m32n8k16", "m8n32k16"],
["f16"],
[],
["f16", "f32"],
["f16", "f32"],
)
+ make_mma_ops(
["m16n16k16", "m32n8k16", "m8n32k16"], ["s8", "u8"], [], ["s32"], []
)
+ make_mma_ops(["m8n8k32"], ["s4", "u4"], [], ["s32"], [])
+ make_mma_ops(
["m8n8k128"], ["b1"], [], ["s32"], [], [".xor.popc", ".and.popc"]
)
)
def get_ldst_ops():
# NOTE: fragemts are from the point of view of PTX.
# fragment `d` is only for store ops, others for both loads and stores.
return (
make_ldst_ops(
["m16n16k16", "m32n8k16", "m8n32k16"],
["a", "b"],
["f16", "u8", "s8", "bf16"],
)
+ make_ldst_ops(
["m16n16k16", "m32n8k16", "m8n32k16"], ["c", "d"], ["f16", "f32", "s32"]
)
+ make_ldst_ops(["m8n8k32"], ["a", "b"], ["s4", "u4"])
+ make_ldst_ops(["m8n8k128"], ["a", "b"], ["b1"])
+ make_ldst_ops(["m8n8k32", "m8n8k128"], ["c", "d"], ["s32"])
+ make_ldst_ops(["m8n8k4"], ["a", "b", "c", "d"], ["f64"])
+
# TF32 m16n16k8 is odd.
# For fragment 'C' it uses __mma_*tf32*_m16n16k8_ld_c
# but 'D' calls __mma_m16n16k8_st_c_*f32*.
make_ldst_ops(["m16n16k8"], ["a", "b", "c"], ["tf32"])
+ make_ldst_ops(["m16n16k8"], ["d"], ["f32"])
)
def is_geom_supported(geom):
# geometries for FP and ints.
if geom in ["m8n32k16", "m32n8k16"]:
return ptx_version >= 61
# geometries for sub-ints.
if geom in ["m8n8k32", "m8n8k128"]:
return ptx_version >= 63 and gpu_arch >= 75
if geom == "m16n16k16":
return ptx_version >= 60
if geom in ["m16n16k8", "m8n8k4"]:
return ptx_version >= 70 and gpu_arch >= 80
assert False # Unexpected geometry.
def is_type_supported(ptx_type):
if ptx_type in ["s8", "u8", "s32"]:
return ptx_version >= 63 and gpu_arch >= 72
if ptx_type in ["s4", "u4", "b1"]:
return ptx_version >= 63 and gpu_arch >= 75
if ptx_type in ["bf16", "tf32", "f64"]:
return ptx_version >= 70 and gpu_arch >= 80
return ptx_version >= 60 and gpu_arch >= 70
def is_rnd_supported(op):
# rnd is only supported for FP64 WMMA
return op.a.ptx_type == "f64"
def is_mma_variant_supported(op, layout_a, layout_b, satf):
if not (is_type_supported(op.a.ptx_type) and is_geom_supported(op.a.geom)):
return False
if satf and not op.a.ptx_type in ["f16", "s8", "u8", "s4", "u4"]:
return False
# sub-integer types require row/col layout.
if op.a.ptx_type in ["s4", "u4", "b1"]:
return layout_a == "row" and layout_b == "col"
return True
def is_ldst_variant_supported(frag, layout):
if not (is_type_supported(frag.ptx_type) and is_geom_supported(frag.geom)):
return False
if frag.ptx_type in ["s4", "u4", "b1"]:
# sub-integer types require sm_75 and ptx63, row/col layout for a/b.
return (
(frag.frag == "a" and layout == "row")
or (frag.frag == "b" and layout == "col")
or frag.frag in ["c", "d"]
)
return True
def get_builtin_prefix(frag):
prefix = None
if frag.geom in ["m16n16k16", "m32n8k16", "m8n32k16"]:
if frag.ptx_type in ["f16", "f32"]:
prefix = "__hmma"
elif frag.ptx_type == "bf16":
prefix = "__mma_bf16"
else:
prefix = "__imma"
elif frag.geom == "m8n8k32":
prefix = "__imma" # sub-integers
elif frag.geom == "m8n8k128":
prefix = "__bmma"
elif frag.geom == "m8n8k4":
prefix = "__dmma"
elif frag.geom == "m16n16k8":
if frag.ptx_type == "f32":
prefix = "__mma"
else:
prefix = "__mma_tf32"
assert prefix
return prefix
def get_ldst_builtin_name(frag):
prefix = get_builtin_prefix(frag)
if prefix == "__hmma":
suffix = "" if frag.frag in ["a", "b"] else frag.ptx_type
elif prefix in ["__dmma", "__mma_bf16", "__mma_tf32"]:
suffix = "" if frag.frag in ["a", "b", "c"] else frag.ptx_type
else:
suffix = "" if frag.frag == "c" else frag.ptx_type
if suffix == "s32":
suffix = "i32"
if frag.frag == "d":
ifrag = "c"
op = "st"
else:
ifrag = frag.frag
op = "ld"
name = "%s_%s_%s_%s%s" % (
prefix,
frag.geom,
op,
ifrag,
"_" + suffix if suffix else "",
)
return name
def get_mma_builtin_name(op):
prefix = get_builtin_prefix(op.a)
if prefix == "__hmma":
suffix = op.d.ptx_type + op.c.ptx_type
elif prefix in ["__mma_bf16", "__mma_tf32"]:
suffix = op.d.ptx_type
else:
suffix = op.a.ptx_type
name = "{prefix}_{geom}_mma{b1op}_{suffix}".format(
prefix=prefix, geom=op.a.geom, b1op=op.b1op.replace(".", "_"), suffix=suffix
)
return name
def get_required_sm(frag, b1op=""):
if frag.ptx_type in ["f64", "bf16", "tf32"]:
return 80
if frag.ptx_type in ["u4", "s4", "b1"]:
if b1op == ".and.popc":
return 80
return 75
if frag.ptx_type in ["s8", "u8"]:
return 72
if frag.ptx_type == "s32":
if frag.geom in ["m8n8k32", "m8n8k128"]: # s4/u4/b1
return 75
else: # s8/u8
return 72
if frag.ptx_type in ["f16", "f32"]:
if frag.geom == "m16n16k8":
return 80
else:
return 70
assert False
def get_required_ptx(frag, b1op=""):
if frag.ptx_type == "b1" and b1op == ".and.popc":
return 71
if frag.ptx_type in ["f64", "bf16", "tf32"]:
return 70
if frag.ptx_type in ["f16", "f32"]:
if frag.geom == "m16n16k16":
return 60
if frag.geom == "m16n16k8":
return 70
return 61
return 63
def get_src_dst_prefix(frag):
if frag.ptx_type == "f32":
return "f"
if frag.ptx_type == "f64":
return "d"
if frag.ptx_type == "tf32" and frag.frag in ["c", "d"]:
return "f"
return ""
def gen_wmma_ldst_tests(results):
load_template = """
// CHECK${check_suffix}: call {{.*}} @${intrinsic}
// expected-error-re@+1 {{'${builtin}' needs target feature (sm_${min_sm}{{.*}},(ptx${min_ptx}{{.*}}}}
${builtin}(${dst}, ${src}, ldm, ${blayout});
""".rstrip()
intrinsic_template = (
"llvm.nvvm.wmma.${geom}.${op}.${frag}.${ilayout}.stride.${itype}"
)
for frag, layout in sorted(product(get_ldst_ops(), ["row", "col"]), key=str):
if not is_ldst_variant_supported(frag, layout):
continue
src_dst_prefix = get_src_dst_prefix(frag)
min_sm = get_required_sm(frag)
min_ptx = get_required_ptx(frag)
# TF32 uses f32 for accumulator loads.
if frag.geom == "m16n16k8" and frag.frag == "c":
assert frag.ptx_type == "tf32"
itype = "f32"
else:
itype = frag.ptx_type
params = {
"check_suffix": "_PTX%d_SM%d" % (min_ptx, min_sm),
"builtin": get_ldst_builtin_name(frag),
"min_ptx": min_ptx,
"min_sm": min_sm,
"dst": src_dst_prefix + "dst",
"src": src_dst_prefix + "src",
"blayout": 0 if layout == "row" else 1,
"intrinsic": Template(intrinsic_template).substitute(
{
"frag": frag.frag,
"geom": frag.geom,
"ilayout": layout,
"itype": itype,
"op": "store" if frag.frag == "d" else "load",
}
),
}
results[(min_ptx, min_sm)] += Template(load_template).substitute(params)
return results
def mma_signature(op):
if op.a.ptx_type == "f16":
# FP16 ops identified by accumulator & result type.
return "%s.%s" % (op.d.ptx_type, op.c.ptx_type)
else:
# other ops are identified by input type.
return op.a.ptx_type
# Get numeric value for rowcol parameter of the builtin
# AFAICT it uses the encoding accepted by NVVM intrinsics:
# https://docs.nvidia.com/cuda/nvvm-ir-spec/index.html#nvvm-intrin-warp-level-matrix-mma
def get_ilayout(a, b):
return {"row.row": 0, "row.col": 1, "col.row": 2, "col.col": 3}[a + "." + b]
def gen_wmma_mma_tests(results):
mma_template = """
// CHECK${check_suffix}: call {{.*}} @${intrinsic}
// expected-error-re@+1 {{'${builtin}' needs target feature (sm_${min_sm}{{.*}},(ptx${min_ptx}{{.*}}}}
${builtin}(${dst}, ${asrc}, ${asrc}, ${csrc}, ${ilayout}${maybe_satf});
""".rstrip()
intrinsic_template = "llvm.nvvm.wmma.${geom}.mma${b1op}.${alayout}.${blayout}.${intrinsic_signature}${satf}"
for op, alayout, blayout, satf in sorted(
product(get_mma_ops(), ["row", "col"], ["row", "col"], [".satfinite", ""]),
key=str,
):
if not is_mma_variant_supported(op, alayout, blayout, satf):
continue
asrc_prefix = get_src_dst_prefix(op.a)
csrc_prefix = get_src_dst_prefix(op.c)
ddst_prefix = get_src_dst_prefix(op.d)
if op.a.ptx_type == "b1": # .b1 MMA has no satf argument.
isatf_arg = ""
else:
isatf_arg = ", 1" if satf else ", 0"
min_sm = get_required_sm(op.a, op.b1op)
min_ptx = get_required_ptx(op.a, op.b1op)
params = {
"check_suffix": "_PTX%d_SM%d" % (min_ptx, min_sm),
"builtin": get_mma_builtin_name(op),
"min_ptx": min_ptx,
"min_sm": min_sm,
"dst": ddst_prefix + "dst",
"asrc": asrc_prefix + "src",
"csrc": csrc_prefix + "src",
"ilayout": get_ilayout(alayout, blayout),
"maybe_satf": isatf_arg,
"intrinsic": Template(intrinsic_template).substitute(
{
"geom": op.a.geom,
"alayout": alayout,
"blayout": blayout,
"intrinsic_signature": mma_signature(op),
"satf": satf,
"b1op": op.b1op,
}
),
}
results[(min_ptx, min_sm)] += Template(mma_template).substitute(params)
return results
def gen_tests():
results = gen_wmma_ldst_tests(defaultdict(str))
results = gen_wmma_mma_tests(results)
run_template = r"""
//
// *** DO NOT EDIT ***
//
// This test has been automatically generated by
// builtins-nvtx-mma.py --ptx=${ptx} --gpu-arch=${sm}
//
// Make sure we can handle all builtins available on sm_${sm} with PTX${ptx}
// ${run}: %clang_cc1 -triple nvptx64-unknown-unknown -target-cpu sm_${sm} \
// ${run}: -fcuda-is-device -target-feature +ptx${ptx} \
// ${run}: -DPTX=${ptx} -DSM=${sm} \
// ${run}: -S -emit-llvm -o - -x cuda %s \
// ${run}: | FileCheck -check-prefixes=${check_labels} %s
// Verify that all builtins have correct constraints.
// ${run}: %clang_cc1 -triple nvptx-unknown-unknown \
// ${run}: -target-cpu sm_60 -target-feature +ptx42 \
// ${run}: -DPTX=${ptx} -DSM=${sm} -fcuda-is-device -S -o /dev/null -x cuda \
// ${run}: -verify %s
"""
def supported_variants(ptx, sm, results):
return [(ptx_, sm_) for ptx_, sm_ in results if ptx_ <= ptx and sm_ <= sm]
print(
Template(run_template).substitute(
{
"run": "RUN", # To avoid lit misinterpreting the template
"ptx": ptx_version,
"sm": gpu_arch,
"check_labels": ",".join(
[
"CHECK_PTX%d_SM%d" % (ptx_, sm_)
for ptx_, sm_ in supported_variants(
ptx_version, gpu_arch, results
)
]
),
}
)
)
print(
"""
#if !defined(CUDA_VERSION)
#define __device__ __attribute__((device))
#define __global__ __attribute__((global))
#define __shared__ __attribute__((shared))
#define __constant__ __attribute__((constant))
typedef unsigned long long uint64_t;
#endif
// CHECK-LABEL: test_wmma_buitins
__device__ void test_wmma_buitins(int *src, int *dst,
float *fsrc, float *fdst,
double *dsrc, double *ddst, int ldm) {
"""
)
for (ptx, sm), tests in sorted(results.items()):
print()
print("#if (PTX >= %d) && (SM >= %d)" % (ptx, sm))
print(tests)
print("#endif // (PTX >= %d) && (SM >= %d)" % (ptx, sm))
print("}")
parser = argparse.ArgumentParser()
parser.add_argument("--ptx", type=int, default=60)
parser.add_argument("--gpu-arch", type=int, default=70)
args = parser.parse_args()
ptx_version = args.ptx
gpu_arch = args.gpu_arch
gen_tests()
Codebase Browser
llvm