mirror of
https://github.com/Gericom/teak-llvm.git
synced 2025-06-19 11:35:51 -04:00
5fe85a003f
These builtins provide access to the new integer and sub-integer variants of MMA (matrix multiply-accumulate) instructions provided by CUDA-10.x on sm_75 (AKA Turing) GPUs. Also added a feature for PTX 6.4. While Clang/LLVM does not generate any PTX instructions that need it, we still need to pass it through to ptxas in order to be able to compile code that uses the new 'mma' instruction as inline assembly (e.g used by NVIDIA's CUTLASS library https://github.com/NVIDIA/cutlass/blob/master/cutlass/arch/mma.h#L101) Differential Revision: https://reviews.llvm.org/D60279 llvm-svn: 359248
344 lines
12 KiB
Python
344 lines
12 KiB
Python
# 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):
|
|
self.a = a
|
|
self.b = b
|
|
self.c = c
|
|
self.d = d
|
|
|
|
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):
|
|
ops = []
|
|
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.append(MMAOp(MMAFrag(geom, "a", type_a),
|
|
MMAFrag(geom, "b", type_b),
|
|
MMAFrag(geom, "c", type_c),
|
|
MMAFrag(geom, "d", type_d)))
|
|
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(["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"], []))
|
|
def get_ldst_ops():
|
|
return (make_ldst_ops(["m16n16k16", "m32n8k16", "m8n32k16"],
|
|
["a", "b"], ["f16", "u8", "s8"]) +
|
|
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"]))
|
|
|
|
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
|
|
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
|
|
return ptx_version >= 60 and gpu_arch >= 70
|
|
|
|
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
|
|
# sub-integer require row/col layout, and no satf.
|
|
if op.a.ptx_type in ["s4", "u4", "b1"]:
|
|
if op.a.ptx_type == "b1" and satf:
|
|
return False
|
|
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 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"
|
|
else:
|
|
prefix = "__imma"
|
|
elif frag.geom == "m8n8k32":
|
|
prefix = "__imma" # sub-integers
|
|
elif frag.geom == "m8n8k128":
|
|
prefix = "__bmma"
|
|
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 ["__imma", "__bmma"]:
|
|
suffix = "" if frag.frag in ["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
|
|
else:
|
|
suffix = op.a.ptx_type
|
|
|
|
name = "%s_%s_mma%s_%s" % (prefix, op.a.geom,
|
|
"_xor_popc" if op.a.ptx_type == "b1" else "",
|
|
suffix)
|
|
return name
|
|
|
|
|
|
def get_required_sm(frag):
|
|
if frag.ptx_type in ["u4", "s4", "b1"]:
|
|
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"]:
|
|
return 70
|
|
assert(False)
|
|
|
|
def get_required_ptx(frag):
|
|
if frag.ptx_type in ["f16", "f32"]:
|
|
return 60 if frag.geom == "m16n16k16" else 61
|
|
return 63
|
|
|
|
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
|
|
|
|
is_fp = frag.ptx_type == "f32"
|
|
min_sm = get_required_sm(frag)
|
|
min_ptx = get_required_ptx(frag)
|
|
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": "fdst" if is_fp else "dst",
|
|
"src": "fsrc" if is_fp else "src",
|
|
"blayout" : 0 if layout == "row" else 1,
|
|
"intrinsic" : Template(intrinsic_template).substitute({
|
|
"frag" : frag.frag,
|
|
"geom" : frag.geom,
|
|
"ilayout" : layout,
|
|
"itype" : frag.ptx_type,
|
|
"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 in ["s8", "u8", "s4", "u4", "b1"]:
|
|
# int and sub-int ops are identified by input type.
|
|
return op.a.ptx_type
|
|
else:
|
|
# the rest are FP ops identified by accumulator & result type.
|
|
return "%s.%s" % (op.d.ptx_type, op.c.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_isatf});
|
|
""".rstrip()
|
|
intrinsic_template = "llvm.nvvm.wmma.${geom}.mma.${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
|
|
|
|
a_is_fp = op.a.ptx_type == "f32"
|
|
c_is_fp = op.c.ptx_type == "f32"
|
|
d_is_fp = op.d.ptx_type == "f32"
|
|
min_sm = get_required_sm(op.a)
|
|
min_ptx = get_required_ptx(op.a)
|
|
if op.a.ptx_type == "b1": # .b1 MMA has no satf argument.
|
|
isatf_arg = ""
|
|
else:
|
|
isatf_arg = ", 1" if satf else ", 0"
|
|
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": "fdst" if d_is_fp else "dst",
|
|
"asrc": "fsrc" if a_is_fp else "src",
|
|
"csrc": "fsrc" if c_is_fp else "src",
|
|
"ilayout" : get_ilayout(alayout, blayout),
|
|
"maybe_isatf" : isatf_arg,
|
|
"intrinsic" : Template(intrinsic_template).substitute({
|
|
"geom" : op.a.geom,
|
|
"alayout" : alayout,
|
|
"blayout" : blayout,
|
|
"intrinsic_signature" : mma_signature(op),
|
|
"satf" : satf,
|
|
})
|
|
}
|
|
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, 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()
|