| # This script generates all variants of wmma builtins, verifies that clang calls |
| # correct LLVM instrinsics, 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() |
