wmma.py

# This test generates all variants of wmma intrinsics and verifies that LLVM
# generates correct instructions for them. This is the test generator only.  The
# test scripts themselves are in wmma-ptx*-sm*.py files.

# RUN: true

from __future__ import print_function

import argparse
from itertools import product
from string import Template

class MMAType:
    def __init__(self, ptx_type):
        self.ptx_type = ptx_type
        self.llvm_type = {
            "f16": "<2 x half>",
            "f32": "float",
            "f64": "double",
            "s32": "i32",
            "b16": "i32",
            "s8": "i32",
            "u8": "i32",
            "s4": "i32",
            "u4": "i32",
            "b1": "i32",
            "bf16": "i32",
            "tf32": "i32",
        }[ptx_type]

        self.ptx_reg_pattern = {
            "f16": "%r[0-9]+",
            "f32": "%f[0-9]+",
            "f64": "%fd[0-9]+",
        }.get(ptx_type, "%r[0-9]+")

    def __repr__(self):
        return "%s/%s" % (self.ptx_type, self.llvm_type)


class MMAFrag:
    def __init__(self, geom, frag, ptx_elt_type):
        self.geom = geom
        self.frag = frag
        self.mma_type = MMAType(ptx_elt_type)
        self.nregs = {
            # u8/s8 -> s32 @ m16n16k16/m8n32k16/m32n8k16
            "m16n16k16:a:u8": 2,
            "m16n16k16:a:s8": 2,
            "m16n16k16:b:u8": 2,
            "m16n16k16:b:s8": 2,
            "m16n16k16:c:s32": 8,
            "m16n16k16:d:s32": 8,
            "m8n32k16:a:u8": 1,
            "m8n32k16:a:s8": 1,
            "m8n32k16:b:u8": 4,
            "m8n32k16:b:s8": 4,
            "m8n32k16:c:s32": 8,
            "m8n32k16:d:s32": 8,
            "m32n8k16:a:u8": 4,
            "m32n8k16:a:s8": 4,
            "m32n8k16:b:u8": 1,
            "m32n8k16:b:s8": 1,
            "m32n8k16:c:s32": 8,
            "m32n8k16:d:s32": 8,
            "m8n8k16:a:u8": 1,
            "m8n8k16:a:s8": 1,
            "m8n8k16:b:u8": 1,
            "m8n8k16:b:s8": 1,
            "m8n8k16:c:s32": 2,
            "m8n8k16:d:s32": 2,
            "m16n8k16:a:u8": 2,
            "m16n8k16:a:s8": 2,
            "m16n8k16:b:u8": 1,
            "m16n8k16:b:s8": 1,
            "m16n8k16:c:s32": 4,
            "m16n8k16:d:s32": 4,
            "m16n8k32:a:u8": 4,
            "m16n8k32:a:s8": 4,
            "m16n8k32:b:u8": 2,
            "m16n8k32:b:s8": 2,
            "m16n8k32:c:s32": 4,
            "m16n8k32:d:s32": 4,
            # u4/s4 -> s32 @ m8n8k32 (u4/s4)
            "m8n8k32:a:u4": 1,
            "m8n8k32:a:s4": 1,
            "m8n8k32:b:u4": 1,
            "m8n8k32:b:s4": 1,
            "m8n8k32:c:s32": 2,
            "m8n8k32:d:s32": 2,
            "m16n8k32:a:u4": 2,
            "m16n8k32:a:s4": 2,
            "m16n8k32:b:u4": 1,
            "m16n8k32:b:s4": 1,
            "m16n8k32:c:s32": 4,
            "m16n8k32:d:s32": 4,
            "m16n8k64:a:u4": 4,
            "m16n8k64:a:s4": 4,
            "m16n8k64:b:u4": 2,
            "m16n8k64:b:s4": 2,
            "m16n8k64:c:s32": 4,
            "m16n8k64:d:s32": 4,
            # b1 -> s32 @ m8n8k128(b1)
            "m8n8k128:a:b1": 1,
            "m8n8k128:b:b1": 1,
            "m8n8k128:c:s32": 2,
            "m8n8k128:d:s32": 2,
            "m16n8k128:a:b1": 2,
            "m16n8k128:b:b1": 1,
            "m16n8k128:c:s32": 4,
            "m16n8k128:d:s32": 4,
            "m16n8k256:a:b1": 4,
            "m16n8k256:b:b1": 2,
            "m16n8k256:c:s32": 4,
            "m16n8k256:d:s32": 4,
            # bf16 -> s32 @ m16n16k16/m8n32k16/m32n8k16
            "m16n16k16:a:bf16": 4,
            "m16n16k16:b:bf16": 4,
            "m8n32k16:a:bf16": 2,
            "m8n32k16:b:bf16": 8,
            "m32n8k16:a:bf16": 8,
            "m32n8k16:b:bf16": 2,
            "m16n8k16:a:bf16": 4,
            "m16n8k16:b:bf16": 2,
            "m16n8k16:c:f32": 4,
            "m16n8k16:d:f32": 4,
            "m16n8k8:a:bf16": 2,
            "m16n8k8:b:bf16": 1,
            "m16n8k8:c:f32": 4,
            "m16n8k8:d:f32": 4,
            "m8n8k4:a:f64": 1,
            "m8n8k4:b:f64": 1,
            "m8n8k4:c:f64": 2,
            "m8n8k4:d:f64": 2,
            # tf32 -> s32 @ m16n16k8
            "m16n16k8:a:tf32": 4,
            "m16n16k8:b:tf32": 4,
            "m16n8k4:a:tf32": 2,
            "m16n8k4:b:tf32": 1,
            "m16n8k4:c:f32": 4,
            "m16n8k4:d:f32": 4,
            "m16n8k8:a:tf32": 4,
            "m16n8k8:b:tf32": 2,
            "m16n8k8:c:f32": 4,
            "m16n8k8:d:f32": 4,
            "m8n8k4:a:f16": 2,
            "m8n8k4:b:f16": 2,
            "m16n8k8:a:f16": 2,
            "m16n8k8:b:f16": 1,
            "m16n8k8:c:f16": 2,
            "m16n8k8:d:f16": 2,
            "m16n8k8:c:f32": 4,
            "m16n8k8:d:f32": 4,
            "m16n8k16:a:f16": 4,
            "m16n8k16:b:f16": 2,
            "m16n8k16:c:f16": 2,
            "m16n8k16:d:f16": 2,
            "m16n8k16:c:f32": 4,
            "m16n8k16:d:f32": 4,
            # ldmatrix
            "m8n8:x1:b16": 1,
            "m8n8:x2:b16": 2,
            "m8n8:x4:b16": 4,
        }.get(
            "%s:%s:%s" % (geom, frag, ptx_elt_type),
            {
                # All other FP shape/fragment/type combinations have the same size
                "a:f16": 8,
                "b:f16": 8,
                "c:f16": 4,
                "d:f16": 4,
                "c:f32": 8,
                "d:f32": 8,
            }.get("%s:%s" % (frag, ptx_elt_type), None),
        )
        assert self.nregs

    def __repr__(self):
        return "%s:%s:%s%s" % (
            self.geom,
            self.frag,
            self.mma_type,
            "" if self.nregs == 1 else ("*%d" % self.nregs),
        )


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 make_ldmatrix_ops(geoms, frags, types):
    return [
        MMAFrag(geom, frag, ptx_type)
        for (geom, frag, ptx_type) in product(geoms, frags, types)
    ]


def get_wmma_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"], [])
    )


def get_mma_ops():
    return (
        make_mma_ops(["m8n8k4"], ["f64"], [], ["f64"], [])
        + make_mma_ops(["m16n8k4", "m16n8k8"], ["tf32"], [], ["f32"], [])
        + make_mma_ops(["m16n8k16", "m16n8k8"], ["bf16"], [], ["f32"], [])
        + make_mma_ops(
            ["m8n8k4", "m16n8k8", "m16n8k16"],
            ["f16"],
            [],
            ["f16", "f32"],
            ["f16", "f32"],
        )
        + make_mma_ops(
            ["m8n8k16", "m16n8k16", "m16n8k32"], ["s8", "u8"], ["s8", "u8"], ["s32"], []
        )
        + make_mma_ops(
            ["m8n8k32", "m16n8k32", "m16n8k64"], ["s4", "u4"], ["s4", "u4"], ["s32"], []
        )
        + make_mma_ops(["m8n8k128", "m16n8k128", "m16n8k256"], ["b1"], [], ["s32"], [])
    )


def get_ldst_ops(kind):
    ldst_ops = (
        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"])
        + make_ldst_ops(["m16n16k8"], ["a", "b"], ["tf32"])
        + make_ldst_ops(["m16n16k8"], ["c", "d"], ["f32"])
    )
    return [x for x in ldst_ops if (x.frag == "d") == (kind == "store")]


def get_ldmatrix_ops():
    return make_ldmatrix_ops(["m8n8"], ["x1", "x2", "x4"], ["b16"])


def is_wmma_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 == "m16n8k8":
        return ptx_version >= 65
    if geom in ["m16n16k8", "m8n8k4"]:
        return ptx_version >= 70
    assert False  # Unexpected geometry.


def is_mma_geom_supported(geom):
    # geometries for FP and ints.
    if geom == "m8n8k4":
        return ptx_version >= 64
    if geom in ["m16n8k8", "m8n8k16", "m8n8k32"]:
        return ptx_version >= 65
    if geom in [
        "m16n8k16",
        "m16n8k4",
        "m16n8k32",
        "m16n8k64",
        "m8n8k128",
        "m16n8k128",
        "m16n8k256",
    ]:
        return ptx_version >= 70
    assert False  # Unexpected geometry.


def is_ldmatrix_geom_supported(geom):
    if geom in ["m8n8"]:
        return ptx_version >= 65 and gpu_arch >= 75
    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 == "b16":
        return ptx_version >= 65 and gpu_arch >= 75
    if ptx_type in ["bf16", "tf32", "f64"]:
        return ptx_version >= 70
    return ptx_version >= 60 and gpu_arch >= 70


def is_wmma_variant_supported(op, layout_a, layout_b, rnd, satf):
    if not (
        is_type_supported(op.a.mma_type.ptx_type) and is_wmma_geom_supported(op.a.geom)
    ):
        return False

    # rnd is only supported for FP64 WMMA
    if rnd and op.a.mma_type.ptx_type != "f64":
        return False

    if satf:
        # satfinite for floating points was removed in PTX 6.5
        if op.a.mma_type.ptx_type == "f16" and ptx_version >= 65:
            return False
        if not op.a.mma_type.ptx_type in ["f16", "s8", "u8", "s4", "u4"]:
            return False

    # sub-integer require row/col layout.
    if op.a.mma_type.ptx_type in ["s4", "u4", "b1"]:
        return layout_a == "row" and layout_b == "col"
    return True


def is_mma_variant_supported(op, layout_a, layout_b, satf):
    if not (
        is_type_supported(op.a.mma_type.ptx_type) and is_mma_geom_supported(op.a.geom)
    ):
        return False

    if satf and not op.a.mma_type.ptx_type in ["s8", "u8", "s4", "u4"]:
        return False

    # If the type of C is f32 then so must the type of D
    if (
        op.a.geom == "m8n8k4"
        and op.c.mma_type.ptx_type == "f32"
        and op.d.mma_type.ptx_type != "f32"
    ):
        return False

    # A and B type must be the same. C and D type must be the same
    if op.a.geom == "m16n8k8" and (
        op.a.mma_type.ptx_type != op.b.mma_type.ptx_type
        or op.c.mma_type.ptx_type != op.d.mma_type.ptx_type
    ):
        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:
        return False

    # Require row/col layout for all MMA except m8n8k4 on FP16
    if not (op.a.geom == "m8n8k4" and op.a.mma_type.ptx_type == "f16"):
        return layout_a == "row" and layout_b == "col"
    return True


def is_ldst_variant_supported(frag, layout):
    if not (
        is_type_supported(frag.mma_type.ptx_type) and is_wmma_geom_supported(frag.geom)
    ):
        return False
    if frag.mma_type.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 is_ldmatrix_variant_supported(frag):
    if not (
        is_type_supported(frag.mma_type.ptx_type)
        and is_ldmatrix_geom_supported(frag.geom)
    ):
        return False
    return frag.frag in ["x1", "x2", "x4"]


def make_wmma_slice_ty(frag):
    return [frag.mma_type.llvm_type] * frag.nregs


def make_wmma_ld_ret_ty(frag):
    results = make_wmma_slice_ty(frag)
    if len(results) == 1:
        return "%s" % results[0]
    return "{%s}" % ", ".join(results)


# returns address space
def get_aspace(space):
    space_map = {
        ".global": 1,
        ".shared": 3,
        ".const": 4,
        ".local": 5,
        ".param": 101,
        "": 0,
        ".generic": 0,
    }
    return space_map[space]


def get_pspace(space):
    return "p%di8" % get_aspace(space)


def check_pattern(frag):
    return "{{%s}}" % ", *".join([frag.mma_type.ptx_reg_pattern] * frag.nregs)


def gen_wmma_load_tests():
    load_template = """
declare ${ret_ty} @${intrinsic}(i8 ${as}* %src ${extra_args});

; CHECK-LABEL: .func {{.*}}test_${function}(
define ${ret_ty} @test_${function}(i8 ${as}* %src ${extra_args}) {
; CHECK: ${instruction}
; CHECK: {${check_result}}
; CHECK: [%rd{{[0-9]+}}]${stride_pattern}
  %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src ${extra_args});
  ret ${ret_ty} %v0;
}

; CHECK-LABEL: .func{{.*}}test_${function}_o(
define ${ret_ty} @test_${function}_o(i8 ${as}* %src ${extra_args}) {
; CHECK: ${instruction}
; CHECK: {${check_result}}
; CHECK: [%rd{{[0-9]+}}+128]${stride_pattern}
  %src1 = getelementptr i8, i8 ${as}* %src, i32 128;
  %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src1 ${extra_args});
  ret ${ret_ty} %v0;
}
"""
    intrinsic_template = (
        "llvm.nvvm.wmma.${geom}.load.${abc}.${layout}${stride}.${itype}.${pspace}"
    )
    instruction_template = (
        "wmma.load.${abc}.sync${aligned}.${layout}.${geom}${space}.${itype}"
    )

    generated_items = []

    for frag, layout, space, stride in product(
        get_ldst_ops("load"),
        ["row", "col"],
        ["", ".shared", ".global"],
        ["", ".stride"],
    ):
        if not is_ldst_variant_supported(frag, layout):
            continue

        params = {
            "abc": frag.frag,
            "aligned": ".aligned" if ptx_version >= 63 else "",
            "layout": layout,
            "space": space,
            "stride": stride,
            "itype": frag.mma_type.ptx_type,
            "pspace": get_pspace(space),
            "as": "addrspace(%d)" % get_aspace(space),
            "geom": frag.geom,
        }

        test_params = params
        test_params["intrinsic"] = Template(intrinsic_template).substitute(params)
        test_params["function"] = test_params["intrinsic"].replace(".", "_")
        test_params["instruction"] = Template(instruction_template).substitute(params)
        test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)
        test_params["check_result"] = check_pattern(frag)

        if stride:
            test_params["extra_args"] = ", i32 %stride"
            test_params["stride_pattern"] = ", %r{{[0-9]+}}"
        else:
            test_params["extra_args"] = ""
            test_params["stride_pattern"] = ""

        print(Template(load_template).substitute(test_params))

        generated_items.append((test_params["intrinsic"], test_params["instruction"]))

    return generated_items


def make_wmma_slice_args(frag):
    return ", ".join(
        [
            "%s %%%s%d" % (t, frag.frag, i)
            for i, t in enumerate(make_wmma_slice_ty(frag))
        ]
    )


def gen_wmma_store_tests():
    store_template = """
declare void @${intrinsic}(i8 ${as}* %src, ${args}${extra_args});

; CHECK-LABEL: .func {{.*}}test_${function}(
define void @test_${function}(i8 ${as}* %src, ${args}${extra_args}) {
; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}
; CHECK: {${check_args}}
; CHECK: ${stride_pattern}
  call void @${intrinsic}(i8 ${as}* %src, ${args} ${extra_args});
  ret void
}

; CHECK-LABEL: .func{{.*}}test_${function}_o(
define void @test_${function}_o(i8 ${as}* %src, ${args}${extra_args}) {
; CHECK: ${instruction} {{.*}}[%rd{{[0-9+]}}+128]
; CHECK: ${check_args}
; CHECK: ${stride_pattern}
  %src1 = getelementptr i8, i8 ${as}* %src, i32 128;
  call void @${intrinsic}(i8 ${as}* %src1, ${args}${extra_args});
  ret void
}
"""
    intrinsic_template = (
        "llvm.nvvm.wmma.${geom}.store.${abc}.${layout}${stride}.${itype}.${pspace}"
    )
    instruction_template = (
        "wmma.store.${abc}.sync${aligned}.${layout}.${geom}${space}.${itype}"
    )

    generated_items = []

    for frag, layout, space, stride in product(
        get_ldst_ops("store"),
        ["row", "col"],
        ["", ".shared", ".global"],
        ["", ".stride"],
    ):

        if not is_ldst_variant_supported(frag, layout):
            continue

        params = {
            "abc": frag.frag,
            "aligned": ".aligned" if ptx_version >= 63 else "",
            "layout": layout,
            "space": space,
            "stride": stride,
            "itype": frag.mma_type.ptx_type,
            "pspace": get_pspace(space),
            "as": "addrspace(%d)" % get_aspace(space),
            "geom": frag.geom,
        }

        test_params = params
        test_params["intrinsic"] = Template(intrinsic_template).substitute(params)
        test_params["function"] = test_params["intrinsic"].replace(".", "_")
        test_params["instruction"] = Template(instruction_template).substitute(params)
        test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)
        test_params["check_args"] = check_pattern(frag)
        if stride:
            test_params["extra_args"] = ", i32 %stride"
            test_params["stride_pattern"] = ", %r{{[0-9]+}};"
        else:
            test_params["extra_args"] = ""
            test_params["stride_pattern"] = ";"
        test_params["args"] = make_wmma_slice_args(frag)

        print(Template(store_template).substitute(test_params))
        generated_items.append((test_params["intrinsic"], test_params["instruction"]))

    return generated_items


def gen_ldmatrix_tests():
    ldmatrix_template = """
declare ${ret_ty} @${intrinsic}(i8 ${as}* %src);

; CHECK-LABEL: .func {{.*}}test_${function}(
define ${ret_ty} @test_${function}(i8 ${as}* %src) {
; CHECK: ${instruction}
; CHECK: {${check_result}}
; CHECK: [%rd{{[0-9]+}}]
  %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src);
  ret ${ret_ty} %v0;
}

; CHECK-LABEL: .func{{.*}}test_${function}_o(
define ${ret_ty} @test_${function}_o(i8 ${as}* %src) {
; CHECK: ${instruction}
; CHECK: {${check_result}}
; CHECK: [%rd{{[0-9]+}}+128]
  %src1 = getelementptr i8, i8 ${as}* %src, i32 128;
  %v0 = call ${ret_ty} @${intrinsic}(i8 ${as}* %src1);
  ret ${ret_ty} %v0;
}
"""
    intrinsic_template = (
        "llvm.nvvm.ldmatrix.sync.aligned.${geom}.${frag}${trans}.${itype}.${pspace}"
    )
    instruction_template = (
        "ldmatrix.sync.aligned.${geom}.${frag}${trans}${space}.${itype}"
    )

    generated_items = []

    for frag, space, trans in product(
        get_ldmatrix_ops(),
        ["", ".shared"],
        ["", ".trans"],
    ):
        if not is_ldmatrix_variant_supported(frag):
            continue

        params = {
            "frag": frag.frag,
            "space": space,
            "trans": trans,
            "itype": frag.mma_type.ptx_type,
            "pspace": get_pspace(space),
            "as": "addrspace(%d)" % get_aspace(space),
            "geom": frag.geom,
        }

        test_params = params
        test_params["intrinsic"] = Template(intrinsic_template).substitute(params)
        test_params["function"] = test_params["intrinsic"].replace(".", "_")
        test_params["instruction"] = Template(instruction_template).substitute(params)
        test_params["ret_ty"] = make_wmma_ld_ret_ty(frag)
        test_params["check_result"] = check_pattern(frag)

        print(Template(ldmatrix_template).substitute(test_params))

        generated_items.append((test_params["intrinsic"], test_params["instruction"]))

    return generated_items


def mma_signature(op):
    if op.a.mma_type.ptx_type == "f16":
        # FP16 ops identified by accumulator & result type.
        return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)
    elif op.a.mma_type.ptx_type != op.b.mma_type.ptx_type:
        # other ops are identified by input types.
        return "%s.%s" % (op.a.mma_type.ptx_type, op.b.mma_type.ptx_type)
    else:
        # if input types are the same, it only appears once.
        return op.a.mma_type.ptx_type


def mma_ptx_signature(op):
    # Encode all four types as D.A.B.C
    return ".".join(x.mma_type.ptx_type for x in (op.d, op.a, op.b, op.c))


def wmma_signature(op):
    if op.a.mma_type.ptx_type == "f16":
        # FP16 ops identified by accumulator & result type.
        return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)
    else:
        # other ops are identified by input type.
        return op.a.mma_type.ptx_type


def wmma_ptx_signature(op):
    if op.a.mma_type.ptx_type == "f16":
        # FP16 instructions use D.C
        return "%s.%s" % (op.d.mma_type.ptx_type, op.c.mma_type.ptx_type)
    else:
        # other instructions encode all four types as D.A.B.C
        return ".".join(x.mma_type.ptx_type for x in (op.d, op.a, op.b, op.c))


def common_mma_test_gen(params, op, intrinsic_template, instruction_template):
    mma_template = """
declare ${ret_ty} @${intrinsic}(
        ${args});

; CHECK-LABEL: .func {{.*}}test_${function}(
define ${ret_ty} @test_${function}(
        ${args}) {
; CHECK: ${instruction}
; CHECK-NEXT: ${check_d}
; CHECK-NEXT: ${check_a}
; CHECK-NEXT: ${check_b}
; CHECK-NEXT: ${check_c}
  %r = call ${ret_ty} @${intrinsic}(
        ${args});
  ret ${ret_ty} %r;
}
"""

    test_params = params
    test_params["intrinsic"] = Template(intrinsic_template).substitute(params)
    test_params["function"] = test_params["intrinsic"].replace(".", "_")
    test_params["instruction"] = Template(instruction_template).substitute(params)
    test_params["ret_ty"] = make_wmma_ld_ret_ty(op.d)
    test_params["check_a"] = check_pattern(op.a)
    test_params["check_b"] = check_pattern(op.b)
    test_params["check_c"] = check_pattern(op.c)
    test_params["check_d"] = check_pattern(op.d)
    args = ",\n        ".join(make_wmma_slice_args(frag) for frag in (op.a, op.b, op.c))
    test_params["args"] = args
    print(Template(mma_template).substitute(test_params))
    return (test_params["intrinsic"], test_params["instruction"])


def get_b1_ops(ptx_type):
    if ptx_type != "b1":
        return [""]
    if ptx_version >= 71:
        return [".xor.popc", ".and.popc"]
    return [".xor.popc"]


def gen_wmma_mma_tests():
    wmma_intrinsic_template = "llvm.nvvm.wmma.${geom}.mma${b1op}.${alayout}.${blayout}${rnd}.${intrinsic_signature}${satf}"
    wmma_instruction_template = "wmma.mma${b1op}.sync${aligned}.${alayout}.${blayout}.${geom}${rnd}.${ptx_signature}${satf}"

    generated_items = []

    for op, alayout, blayout, rnd, satf in product(
        get_wmma_ops(),
        ["row", "col"],
        ["row", "col"],
        [".rn", ".rz", ".rm", ".rp", ""],
        [".satfinite", ""],
    ):

        if not is_wmma_variant_supported(op, alayout, blayout, rnd, satf):
            continue

        for b1op in get_b1_ops(op.a.mma_type.ptx_type):
            params = {
                "aligned": ".aligned" if ptx_version >= 63 else "",
                "alayout": alayout,
                "blayout": blayout,
                "intrinsic_signature": wmma_signature(op),
                "ptx_signature": wmma_ptx_signature(op),
                "satf": satf,
                "rnd": rnd,
                "geom": op.a.geom,
                "b1op": b1op,
            }

            intrinsic_template = wmma_intrinsic_template
            instruction_template = wmma_instruction_template

            generated_items.append(
                common_mma_test_gen(
                    params, op, intrinsic_template, instruction_template
                )
            )

    return generated_items


def gen_mma_tests():
    mma_intrinsic_template = "llvm.nvvm.mma${b1op}.${geom}.${alayout}.${blayout}${satf}.${intrinsic_signature}"
    mma_instruction_template = "mma.sync${aligned}.${geom}.${alayout}.${blayout}${satf}.${ptx_signature}${b1op}"

    generated_items = []

    for op, alayout, blayout, satf in product(
        get_mma_ops(), ["row", "col"], ["row", "col"], [".satfinite", ""]
    ):

        if not is_mma_variant_supported(op, alayout, blayout, satf):
            continue

        for b1op in get_b1_ops(op.a.mma_type.ptx_type):
            params = {
                "aligned": ".aligned" if ptx_version >= 63 else "",
                "alayout": alayout,
                "blayout": blayout,
                "intrinsic_signature": mma_signature(op),
                "ptx_signature": mma_ptx_signature(op),
                "satf": satf,
                "geom": op.a.geom,
                "b1op": b1op,
            }

            intrinsic_template = mma_intrinsic_template
            instruction_template = mma_instruction_template

            generated_items.append(
                common_mma_test_gen(
                    params, op, intrinsic_template, instruction_template
                )
            )

    return generated_items


# Append complete list of intrinsics and instructions we've generated tests for.
# Generate set of checks to verify that that we did generate sensible set of
# tests for the given combination of PTX and SM variants.
#
def gen_check_unsupported_ops(items):
    print(
        "; Complete list of intrinsics supported by PTX%d on sm_%d"
        % (ptx_version, gpu_arch)
    )
    print("; INTRINSICS: {{^; INTRINSICS_LIST_BEGIN}}")
    print(
        """

; NOEXTGEOM-NOT: {{m8n32|m32n8}}
; NOINT-NOT: .{{s32|s8}}
; NOSUBINT-NOT: {{s4|u4|b1}}
; NOMMA-NOT: .m8n8k4.
; NOALTFLOAT-NOT: .{{bf16|tf32}}
; NODOUBLE-NOT: .f64
; NOLDMATRIX-NOT: ldmatrix.sync.aligned

; M16N16-DAG: m16n16k16.load.{{[ab].*}}.f16.p
; M16N16-DAG: m16n16k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p
; M16N16-DAG: m16n16k16.mma.{{.*}}.f16.f32
; M16N16-DAG: m16n16k16.mma.{{.*}}.f32.f16
; M16N16-DAG: m16n16k16.mma.{{.*}}.f16.f16
; M16N16-DAG: m16n16k16.mma.{{.*}}.f32.f32

; PTX60 adds support for m32n8k16/m8n32k16 geometries.
; EXTGEOM-DAG: m32n8k16.load.{{[ab].*}}.f16.p
; EXTGEOM-DAG: m32n8k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p
; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f16.f32
; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f32.f16
; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f16.f16
; EXTGEOM-DAG: m32n8k16.mma.{{.*}}.f32.f32

; EXTGEOM-DAG: m8n32k16.load.{{[ab].*}}.f16.p
; EXTGEOM-DAG: m8n32k16.{{load|store}}.{{[cd].*\.(f16|f32)}}.p
; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f16.f32
; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f32.f16
; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f16.f16
; EXTGEOM-DAG: m8n32k16.mma.{{.*}}.f32.f32

; INT-DAG: m16n16k16.load.{{[ab].*}}.s8.p
; INT-DAG: m8n32k16.load.{{[ab].*}}.s8.p
; INT-DAG: m32n8k16.load.{{[ab].*}}.s8.p
; INT-DAG: m16n16k16.load.{{[ab].*}}.u8.p
; INT-DAG: m8n32k16.load.{{[ab].*}}.u8.p
; INT-DAG: m32n8k16.load.{{[ab].*}}.u8.p
; INT-DAG: m32n8k16.{{load|store}}.{{[cd].*\.s32}}.p
; INT-DAG: m16n16k16.mma.{{.*}}.u8
; INT-DAG: m16n16k16.mma.{{.*}}.s8
; INT-DAG: m8n32k16.mma.{{.*}}.u8
; INT-DAG: m8n32k16.mma.{{.*}}.s8
; INT-DAG: m32n8k16.mma.{{.*}}.u8
; INT-DAG: m32n8k16.mma.{{.*}}.s8

; SUBINT-DAG: m8n8k128.load.{{[ab].*}}.b1.p
; SUBINT-DAG: m8n8k32.load.{{[ab].*}}.s4.p
; SUBINT-DAG: m8n8k32.load.{{[ab].*}}.u4.p
; SUBINT-DAG: m8n8k128.{{load|store}}.{{[cd].*\.s32}}.p
; SUBINT-DAG: m8n8k32.{{load|store}}.{{[cd].*\.s32}}.p
; SUBINT-DAG: m8n8k32.mma.{{.*}}.u4
; SUBINT-DAG: m8n8k32.mma.{{.*}}.s4
; SUBINT-DAG: m8n8k128.mma.{{.*}}.b1

; ALTFLOAT-DAG: m16n16k16.load.{{[ab].*}}.bf16.p
; ALTFLOAT-DAG: m8n32k16.load.{{[ab].*}}.bf16.p
; ALTFLOAT-DAG: m32n8k16.load.{{[ab].*}}.bf16.p
; ALTFLOAT-DAG: m16n16k8.load.{{[ab].*}}.tf32.p
; ALTFLOAT-DAG: m16n16k16.mma.{{.*}}.bf16
; ALTFLOAT-DAG: m8n32k16.mma.{{.*}}.bf16
; ALTFLOAT-DAG: m32n8k16.mma.{{.*}}.bf16
; ALTFLOAT-DAG: m16n16k8.mma.{{.*}}.tf32

; DOUBLE-DAG: m8n8k4.load.{{[abc].*}}.f64.p
; DOUBLE-DAG: m8n8k4.store.d.{{.*}}.f64.p
; DOUBLE-DAG: m8n8k4.mma.{{.*}}.f64

; MMA-DAG: mma.m8n8k4.{{.*}}.f16.f32
; MMA-DAG: mma.m8n8k4.{{.*}}.f32.f16
; MMA-DAG: mma.m8n8k4.{{.*}}.f16.f16
; MMA-DAG: mma.m8n8k4.{{.*}}.f32.f32

; PTX65MMA-DAG: mma.m16n8k8.row.col.f16.f16
; PTX65MMA-DAG: mma.m16n8k8.row.col.f32.f32
; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.u8.u8
; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.s8
; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.s8.u8
; PTX65MMA-DAG: mma.m8n8k16.row.col{{.*}}.u8.s8
; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.u4
; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.s4
; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.s4.u4
; PTX65MMA-DAG: mma.m8n8k32.row.col{{.*}}.u4.s4

; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.trans.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.trans.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.trans.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.shared.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.shared.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.shared.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16
; PTX65LDMATRIX-DAG: ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16

; PTX71MMA-DAG: mma.m8n8k4.row.col.f64
; PTX71MMA-DAG: mma.m16n8k4.row.col.tf32
; PTX71MMA-DAG: mma.m16n8k8.row.col.tf32
; PTX71MMA-DAG: mma.m16n8k16.row.col.bf16
; PTX71MMA-DAG: mma.m16n8k8.row.col.bf16
; PTX71MMA-DAG: mma.m16n8k16.row.col.f16.f16
; PTX71MMA-DAG: mma.m16n8k16.row.col.f32.f32
; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.u8
; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.s8
; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.s8.u8
; PTX71MMA-DAG: mma.m16n8k16.row.col{{.*}}.u8.s8
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.u8
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.s8
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s8.u8
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u8.s8
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.u4
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.s4
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.s4.u4
; PTX71MMA-DAG: mma.m16n8k32.row.col{{.*}}.u4.s4
; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.u4
; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.s4
; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.s4.u4
; PTX71MMA-DAG: mma.m16n8k64.row.col{{.*}}.u4.s4
; PTX71MMA-DAG: mma.and.popc.m8n8k128.row.col.b1
; PTX71MMA-DAG: mma.xor.popc.m8n8k128.row.col.b1
; PTX71MMA-DAG: mma.and.popc.m16n8k128.row.col.b1
; PTX71MMA-DAG: mma.xor.popc.m16n8k128.row.col.b1
; PTX71MMA-DAG: mma.and.popc.m16n8k256.row.col.b1
; PTX71MMA-DAG: mma.xor.popc.m16n8k256.row.col.b1
;

"""
    )

    print("; INTRINSICS_LIST_BEGIN")
    for intrinsic, instruction in sorted(items):
        print("; ", intrinsic, " -> ", instruction, "")
    print("; INTRINSICS_LIST_END")
    print("; INTRINSICS: ; INTRINSICS_LIST_END")


def gen_tests():
    items = gen_wmma_load_tests()
    items += gen_wmma_store_tests()
    items += gen_ldmatrix_tests()
    items += gen_wmma_mma_tests()
    items += gen_mma_tests()
    gen_check_unsupported_ops(items)


def main():
    global ptx_version
    global gpu_arch
    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()


if __name__ == "__main__":
    main()