test_cutlass.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import logging
import math
import tvm
from tvm import relay
from tvm.contrib.cudnn import conv_output_shape
import numpy as np
from tvm.runtime.vm import VirtualMachine
from tvm.relay.op.contrib.cutlass import partition_for_cutlass
from tvm.relay.transform import FirstOrderGradient, ToMixedPrecision, InferType
from tvm.contrib.cutlass import num_cutlass_partitions, finalize_modules, finalize_modules_vm

logging.basicConfig(level=logging.INFO)


def has_cublas():
    return tvm.get_global_func("tvm.contrib.cublas.matmul", True) != None


def has_cutlass():
    return tvm.get_global_func("relay.ext.cutlass", True) != None


def get_ref_rt_mod(mod, params, target="cuda"):
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod, target=target, params=params)
    dev = tvm.device(target, 0)
    rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
    return rt_mod, dev


def get_ref_vm(mod, params, target="cuda"):
    with tvm.transform.PassContext(opt_level=3):
        vm_exec = relay.vm.compile(mod, target=target, params=params)
        code, lib = vm_exec.save()
    dev = tvm.device(target, 0)
    vm_exec = tvm.runtime.vm.Executable.load_exec(code, lib)
    return VirtualMachine(vm_exec, dev), dev


def get_output(rt_mod, names, inputs):
    for name, inp in zip(names, inputs):
        rt_mod.set_input(name, inp)
    rt_mod.run()
    return rt_mod.get_output(0).asnumpy()


def get_output_vm(vm, names, inputs):
    params = dict(zip(names, inputs))
    return vm.invoke("main", **params).numpy()


def get_dense_with_shape(
        data_shape, weight_shape, out_dtype="float16", data_dtype="float16", weight_dtype="float16"
):
    data = relay.var("data", shape=data_shape, dtype=data_dtype)
    weight = relay.var("weight", shape=weight_shape, dtype=weight_dtype)
    return relay.nn.dense(data, weight, out_dtype=out_dtype)


def get_dense(M, N, K, out_dtype="float16", data_dtype="float16", weight_dtype="float16"):
    return get_dense_with_shape((M, K), (N, K), out_dtype, data_dtype, weight_dtype)


def get_dense_bias(M, N, K, out_dtype="float16"):
    dense = get_dense(M, N, K, out_dtype=out_dtype)
    bias = relay.var("bias", shape=(N,), dtype=out_dtype)
    return relay.nn.bias_add(dense, bias)


def get_dense_bias_relu(M, N, K, out_dtype="float16"):
    return relay.nn.relu(get_dense_bias(M, N, K, out_dtype=out_dtype))


def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
    bias_add = get_dense_bias(M, N, K, out_dtype)
    mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
    if out_dtype == "float16":
        erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
    else:
        erf = relay.op.erf(mul)
    mul_half = erf * relay.const(0.5, dtype=out_dtype)
    add = mul_half + relay.const(0.5, dtype=out_dtype)
    return add * bias_add


def get_batch_matmul_with_shape(x_shape, y_shape, out_dtype="float16"):
    x = relay.var("x", shape=x_shape, dtype="float16")
    y = relay.var("y", shape=y_shape, dtype="float16")
    return relay.nn.batch_matmul(x, y, out_dtype=out_dtype)


def get_batch_matmul(batch, M, N, K, out_dtype="float16"):
    return get_batch_matmul_with_shape((batch, M, K), (batch, N, K), out_dtype="float16")


def get_conv2d_nchw(
        d_shape,
        w_shape,
        padding,
        strides=(1, 1),
        out_dtype="float16",
        data_dtype="float16",
        weight_dtype="float16",
):
    data = relay.var("data", shape=d_shape, dtype=data_dtype)
    weight = relay.var("weight", shape=w_shape, dtype=weight_dtype)
    out_channel = w_shape[0]
    return relay.nn.conv2d(
        data=data,
        weight=weight,
        kernel_size=w_shape[2:],
        channels=out_channel,
        padding=padding,
        strides=strides,
        out_dtype=out_dtype,
    )


def get_conv2d_nchw_bias(d_shape, w_shape, padding, out_dtype="float16"):
    conv2d = get_conv2d_nchw(d_shape, w_shape, padding, out_dtype=out_dtype)
    bias = relay.var("bias", shape=(w_shape[0],), dtype=out_dtype)
    return relay.nn.bias_add(conv2d, bias)


def silu(x):
    return x * relay.sigmoid(x)


def hardswish(x, out_dtype="float16"):
    return x * (
            relay.clip(x + relay.const(3, dtype=out_dtype), a_min=0, a_max=6)
            / relay.const(6, dtype=out_dtype)
    )


def get_conv2d_nchw_bias_relu(d_shape, w_shape, padding, out_dtype="float16"):
    return relay.nn.relu(get_conv2d_nchw_bias(d_shape, w_shape, padding, out_dtype=out_dtype))


def get_conv2d_nchw_bias_sigmoid(d_shape, w_shape, padding, out_dtype="float16"):
    return relay.sigmoid(get_conv2d_nchw_bias(d_shape, w_shape, padding, out_dtype=out_dtype))


def get_conv2d_nchw_bias_silu(d_shape, w_shape, padding, out_dtype="float16"):
    conv_out = get_conv2d_nchw_bias(d_shape, w_shape, padding, out_dtype=out_dtype)
    return silu(conv_out)


def get_conv2d_nchw_bias_hardswish(d_shape, w_shape, padding, out_dtype="float16"):
    conv_out = get_conv2d_nchw_bias(d_shape, w_shape, padding, out_dtype=out_dtype)
    return hardswish(conv_out, out_dtype)


def get_conv2d_nchw_bias_residual(d_shape, w_shape, padding, out_dtype="float16"):
    data = relay.var("data", shape=d_shape, dtype="float16")
    weight = relay.var("weight", shape=w_shape, dtype="float16")
    bias = relay.var("bias", shape=(w_shape[0],), dtype=out_dtype)
    out_channel = w_shape[0]
    conv2d = relay.nn.conv2d(
        data=data,
        weight=weight,
        kernel_size=w_shape[2:],
        channels=out_channel,
        padding=padding,
        out_dtype=out_dtype,
    )
    bias_add = relay.nn.bias_add(conv2d, bias)
    return bias_add, data


def get_conv2d_transpose_nchw(
        d_shape,
        w_shape,
        padding,
        output_padding,
        strides,
        out_dtype="float32",
        data_dtype="float32",
        weight_dtype="float32",
):
    data = relay.var("data", shape=d_shape, dtype=data_dtype)
    weight = relay.var("weight", shape=w_shape, dtype=weight_dtype)
    out_channel = w_shape[1]
    return relay.nn.conv2d_transpose(
        data=data,
        weight=weight,
        kernel_size=w_shape[2:],
        channels=out_channel,
        padding=padding,
        output_padding=output_padding,
        strides=strides,
        out_dtype=out_dtype,
    )


def get_conv2d_backward_weight(
        d_shape,
        w_shape,
        o_shape,
        padding,
        strides,
        out_dtype="float32",
        data_dtype="float32",
        weight_dtype="float32",
):
    grad = relay.var("grad", shape=o_shape, dtype=weight_dtype)
    data = relay.var("data", shape=d_shape, dtype=data_dtype)
    out_channel = o_shape[1]
    return relay.nn.conv2d_backward_weight(
        grad=grad,
        data=data,
        kernel_size=w_shape[2:],
        channels=out_channel,
        padding=padding,
        strides=strides,
        out_dtype=out_dtype,
    )


def convert_conv2d_layout(mod, desired_layouts):
    with tvm.transform.PassContext(opt_level=3):
        seq = tvm.transform.Sequential([relay.transform.ConvertLayout(desired_layouts)])
        return seq(mod)


def get_random_ndarray(shape, dtype):
    if dtype == "int8":
        return np.random.randint(-128, 128, shape).astype(dtype)
    elif dtype == "uint8":
        return np.random.randint(0, 256, shape).astype(dtype)
    return np.random.uniform(-1, 1, shape).astype(dtype)


def profile_and_build(
        mod,
        params,
        sm,
        split_k_slices=[1],
        tmp_dir="./tmp",
        use_fast_math=False,
        use_3xtf32=True,
):
    print("------------ before partitioning ---------")
    print(mod)
    mod = partition_for_cutlass(mod)
    print("----------- after partitioning -----------")
    print(mod)

    num_cutlass_partition = num_cutlass_partitions(mod)
    host = tvm.target.Target("llvm")
    cuda = tvm.target.Target("cuda", host=host)
    cutlass = tvm.target.Target(
        {"kind": "cutlass", "sm": sm, "use_3xtf32": use_3xtf32, "split_k_slices": split_k_slices,
         "profile_all_alignments": False, "find_first_valid": True, "use_multiprocessing": True,
         "use_fast_math": use_fast_math, "tmp_dir": tmp_dir}, host=host)
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod, target=[cuda, cutlass], params=params)
    lib = finalize_modules(lib, "compile.so", tmp_dir)
    dev = tvm.device("cuda", 0)
    rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
    return rt_mod, dev, num_cutlass_partition


def profile_and_build_vm(
        mod,
        params,
        sm,
        split_k_slices=[1],
        tmp_dir="./tmp",
        use_fast_math=False,
        use_3xtf32=True,
):
    mod = partition_for_cutlass(mod)
    num_cutlass_partition = num_cutlass_partitions(mod)
    host = tvm.target.Target("llvm")
    cuda = tvm.target.Target("cuda", host=host)
    cutlass = tvm.target.Target(
        {"kind": "cutlass", "sm": sm, "use_3xtf32": use_3xtf32, "split_k_slices": split_k_slices,
         "profile_all_alignments": False, "find_first_valid": True, "use_multiprocessing": True,
         "use_fast_math": use_fast_math, "tmp_dir": tmp_dir}, host=host)
    with tvm.transform.PassContext(opt_level=3):
        vm_exec = relay.vm.compile(mod, target=[cuda, cutlass], params=params)
    vm_exec = finalize_modules_vm(vm_exec, "compile.so", tmp_dir)
    dev = tvm.device("cuda", 0)
    return VirtualMachine(vm_exec, dev), dev, num_cutlass_partition


def verify_dense(
        func,
        M,
        N,
        K,
        ref_target="cuda",
        sm=80,
        atol=1e-5,
        rtol=1e-5,
        run_benchmark=False,
        data_dtype="float16",
        weight_dtype="float16",
        use_3xtf32=True,
):
    if not has_cutlass():
        return
    if sm < 80 and data_dtype == "float32":
        return

    mod = tvm.IRModule.from_expr(func)
    typ = relay.transform.InferType()(mod)["main"].body.checked_type
    out_dtype = typ.dtype
    use_vm = any(isinstance(s, tvm.tir.Any) for s in typ.shape)
    np_data = get_random_ndarray((M, K), data_dtype)
    np_weight = get_random_ndarray((N, K), weight_dtype)
    np_bias = get_random_ndarray((N,), out_dtype)

    params = {"weight": np_weight, "bias": np_bias}

    if use_vm:
        if ref_target == "cuda" and out_dtype == "float16":
            # Uncomment "return" below to see the accuracy difference of static vs dynamic TVM native fp16 dense
            # The static one can use a tensorcore schedule, but the dynamic one cannot
            rt_mod, dev = get_ref_vm(tvm.IRModule.from_expr(get_dense(M, N, K)), params)
            num_partition = 1
            logging.warning(
                "The reference fp16 dense with dynamic shape using fp16 accumulation has accuracy issues."
            )
            return
        else:
            rt_mod, dev, num_partition = profile_and_build_vm(
                mod, params, sm, use_3xtf32=use_3xtf32
            )

        rt_mod_ref, dev = get_ref_vm(mod, params, target=ref_target)
        x = tvm.nd.array(np_data, device=dev)
        out = get_output_vm(rt_mod, ["data"], [x])
        ref_out = get_output_vm(rt_mod_ref, ["data"], [x])
    else:
        rt_mod_ref, dev = get_ref_rt_mod(mod, params, target=ref_target)
        rt_mod, dev, num_partition = profile_and_build(mod, params, sm, use_3xtf32=use_3xtf32)
        x = tvm.nd.array(np_data, device=dev)
        out = get_output(rt_mod, ["data"], [x])
        ref_out = get_output(rt_mod_ref, ["data"], [x])

    assert num_partition > 0
    np.testing.assert_allclose(out, ref_out, atol=atol, rtol=rtol)

    if run_benchmark:
        print("CUTLASS:", rt_mod.benchmark(dev, number=1, repeat=600))
        print("TVM with target %s:" % ref_target, rt_mod_ref.benchmark(dev, number=1, repeat=600))


def verify_batch_matmul(
        func, batch, M, N, K, ref_target="cuda", sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
):
    if not has_cutlass():
        return
    mod = tvm.IRModule.from_expr(func)
    typ = relay.transform.InferType()(mod)["main"].body.checked_type
    use_vm = any(isinstance(s, tvm.tir.Any) for s in typ.shape)
    x_np = np.random.uniform(-1, 1, (batch, M, K)).astype("float16")
    y_np = np.random.uniform(-1, 1, (batch, N, K)).astype("float16")

    if use_vm:
        rt_mod, dev, num_partition = profile_and_build_vm(mod, {}, sm)
        rt_mod_ref, dev = get_ref_vm(mod, {}, target=ref_target)
        assert num_partition > 0
        x = tvm.nd.array(x_np, device=dev)
        y = tvm.nd.array(y_np, device=dev)
        out = get_output_vm(rt_mod, ["x", "y"], [x, y])
        ref_out = get_output_vm(rt_mod_ref, ["x", "y"], [x, y])
    else:
        rt_mod, dev, num_partition = profile_and_build(mod, {}, sm)
        rt_mod_ref, dev = get_ref_rt_mod(mod, {})
        assert num_partition > 0

        x = tvm.nd.array(x_np, device=dev)
        y = tvm.nd.array(y_np, device=dev)
        out = get_output(rt_mod, ["x", "y"], [x, y])
        ref_out = get_output(rt_mod_ref, ["x", "y"], [x, y])

    np.testing.assert_allclose(out, ref_out, atol=atol, rtol=rtol)

    if run_benchmark:
        print("CUTLASS:", rt_mod.benchmark(dev, number=1, repeat=600))
        print("TVM Tensorcore (no tuning):", rt_mod_ref.benchmark(dev, number=1, repeat=600))


M = 96
N = 64
K = 64


def test_dense():
    verify_dense(get_dense(M, N, K), M, N, K)
    verify_dense(get_dense(M, N, K, out_dtype="float32"), M, N, K)
    # Test align1 case
    verify_dense(get_dense_bias(M, N + 1, K), M, N + 1, K)
    # int8
    verify_dense(
        get_dense(M, N, K, "int32", "int8", "int8"), M, N, K, data_dtype="int8", weight_dtype="int8"
    )

    dense_fp32 = get_dense(M, N, K, "float32", "float32", "float32")
    # tf32
    verify_dense(
        dense_fp32,
        M,
        N,
        K,
        data_dtype="float32",
        weight_dtype="float32",
        use_3xtf32=False,
        atol=1e-2,
        rtol=1e-2,
    )
    # 3xtf32
    verify_dense(
        dense_fp32,
        M,
        N,
        K,
        data_dtype="float32",
        weight_dtype="float32",
    )


def test_dense_bias():
    verify_dense(get_dense_bias(M, N, K), M, N, K)
    verify_dense(get_dense_bias(M, N, K, out_dtype="float32"), M, N, K)


def test_dense_bias_relu():
    verify_dense(get_dense_bias_relu(M, N, K), M, N, K)
    verify_dense(get_dense_bias_relu(M, N, K, out_dtype="float32"), M, N, K)


def test_dense_bias_gelu():
    verify_dense(get_dense_bias_gelu(M, N, K), M, N, K, atol=1e-3, rtol=1e-3)
    verify_dense(get_dense_bias_gelu(M, N, K, out_dtype="float32"), M, N, K, atol=1e-3, rtol=1e-3)


def test_dense_dynamic():
    data_shape = (relay.Any(), K)
    weight_shape = (relay.Any(), K)

    if has_cublas():
        # TVM native fp16 dense (without tensorcore), using fp16 accum, seems to have accuracy issues
        # Use cublas as a reference

        verify_dense(
            get_dense_with_shape(data_shape, weight_shape),
            M,
            N,
            K,
            ref_target="cuda -libs=cublas",
        )

    verify_dense(
        get_dense_with_shape(data_shape, weight_shape, out_dtype="float32"),
        M,
        N,
        K,
        atol=1e-4,
        rtol=1e-4,
    )


def test_batch_matmul():
    batch = 8
    verify_batch_matmul(get_batch_matmul(batch, M, N, K), batch, M, N, K)
    verify_batch_matmul(get_batch_matmul(batch, M, N, K, out_dtype="float32"), batch, M, N, K)

    if has_cublas():
        # Test dynamic shape batch_matmul
        # AutoTVM does not seem to support it
        x_shape = (relay.Any(), relay.Any(), K)
        y_shape = (relay.Any(), relay.Any(), K)

        verify_batch_matmul(
            get_batch_matmul_with_shape(x_shape, y_shape),
            batch,
            M,
            N,
            K,
            ref_target="cuda -libs=cublas",
        )


def verify_conv2d_common(
        expr_nchw,  # can be dynamic batch
        expr_ref,  # always static batch
        input_names,
        inputs,
        params,
        sm=80,
        split_k_slices=[1],
        atol=1e-5,
        rtol=1e-5,
        use_cudnn_ref=False,
        run_benchmark=False,
        use_fast_math=False,
        ref_target="cuda",
        use_vm=False,
):
    if not has_cutlass():
        return
    if sm < 80 and inputs[0].dtype == "float32":
        return

    mod_nchw = tvm.IRModule.from_expr(expr_nchw)
    mod_ref = tvm.IRModule.from_expr(expr_ref)

    if use_vm:
        profile_and_build_func = profile_and_build_vm
        get_output_func = get_output_vm
        ref_build_func = get_ref_vm
    else:
        profile_and_build_func = profile_and_build
        get_output_func = get_output
        ref_build_func = get_ref_rt_mod

    mod_weight_ohwi = convert_conv2d_layout(
        mod_nchw,
        {
            "nn.conv2d": ["NHWC", "OHWI"],
            "nn.conv2d_transpose": ["NHWC", "IHWO"],
            "nn.conv2d_backward_weight": ["NHWC", "OHWI"],
        },
    )

    rt_mod, _, num_cutlass_partition = profile_and_build_func(
        mod_weight_ohwi, params, sm, split_k_slices, use_fast_math=use_fast_math
    )
    out = get_output_func(rt_mod, input_names, inputs)

    assert num_cutlass_partition > 0

    if use_cudnn_ref:
        rt_mod_ref, dev = ref_build_func(
            convert_conv2d_layout(mod_ref, {"nn.conv2d": ["NHWC", "OHWI"]}),
            params,
            target="cuda -libs=cudnn",
        )
    else:
        rt_mod_ref, dev = ref_build_func(
            convert_conv2d_layout(mod_ref, {"nn.conv2d": ["NHWC", "HWIO"]}),
            params,
            target=ref_target,
        )

    ref_out = get_output_func(rt_mod_ref, input_names, inputs)

    if run_benchmark:
        print("CUTLASS:", rt_mod.benchmark(dev, number=1, repeat=600))
        print("TVM Tensorcore (no tuning):", rt_mod_ref.benchmark(dev, number=1, repeat=600))

    np.testing.assert_allclose(out, ref_out, atol=atol, rtol=rtol)


def verify_conv2d(
        expr_nchw,  # can be dynamic batch
        expr_ref,  # always static batch
        d_shape,
        w_shape,
        sm=80,
        atol=1e-5,
        rtol=1e-5,
        use_cudnn_ref=False,
        run_benchmark=False,
        use_fast_math=False,
        data_dtype="float16",
        weight_dtype="float16",
        ref_target="cuda",
        use_vm=False,
):
    mod_nchw = tvm.IRModule.from_expr(expr_nchw)
    typ = relay.transform.InferType()(mod_nchw)["main"].body.checked_type

    use_vm = use_vm or any(isinstance(s, tvm.tir.Any) for s in typ.shape)

    np_data = get_random_ndarray(d_shape, data_dtype)
    np_weight = get_random_ndarray(w_shape, weight_dtype)
    np_bias = get_random_ndarray((w_shape[0],), typ.dtype)
    params = {"weight": np_weight, "bias": np_bias}

    split_k_slices = [1]

    return verify_conv2d_common(
        expr_nchw,
        expr_ref,
        ["data"],
        [np_data],
        params,
        sm,
        split_k_slices,
        atol,
        rtol,
        use_cudnn_ref,
        run_benchmark,
        use_fast_math,
        ref_target,
        use_vm,
    )


def verify_conv2d_backward_weight(
        expr_nchw,  # can be dynamic batch
        expr_ref,  # always static batch
        grad_shape,
        data_shape,
        sm=80,
        split_k_slices=[1],
        atol=1e-5,
        rtol=1e-5,
        use_cudnn_ref=False,
        use_fast_math=False,
        grad_dtype="float16",
        data_dtype="float16",
        ref_target="cuda",
        use_vm=False,
):
    np_grad = get_random_ndarray(grad_shape, grad_dtype)
    np_data = get_random_ndarray(data_shape, data_dtype)
    params = {}
    input_names = ["grad", "data"]
    return verify_conv2d_common(
        expr_nchw,
        expr_ref,
        input_names,
        [np_grad, np_data],
        params,
        sm,
        split_k_slices,
        atol,
        rtol,
        use_cudnn_ref,
        False,
        use_fast_math,
        ref_target,
        use_vm,
    )


def test_conv2d():
    padding = (1, 1)
    for IC in [3, 16]:
        d_shape = (16, IC, 32, 32)
        w_shape = (32, IC, 3, 3)
        mod_nchw = get_conv2d_nchw(d_shape, w_shape, padding)

        verify_conv2d(
            mod_nchw,
            mod_nchw,
            d_shape,
            w_shape,
            sm=80,
            atol=1e-5,
            rtol=1e-5,
            use_cudnn_ref=(IC == 3),  # The autotvm kernel has an accuracy issue with IC == 3 case
            run_benchmark=False,
        )

    d_shape = (16, 16, 32, 32)
    w_shape = (32, 16, 3, 3)
    padding = (1, 1)
    dyn_batch_shape = (relay.Any(),) + d_shape[1:]

    mod_nchw = get_conv2d_nchw(d_shape, w_shape, padding)
    mod_dyn = get_conv2d_nchw(dyn_batch_shape, w_shape, padding)

    verify_conv2d(
        mod_dyn, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )

    for data_dtype, weight_dtype, out_dtype in [
        ("float32", "float32", "float32"),  # 3xtf32
        ("int8", "int8", "int32"),
        ("uint8", "int8", "int32"),
    ]:
        expr = get_conv2d_nchw(
            d_shape,
            w_shape,
            padding,
            out_dtype=out_dtype,
            data_dtype=data_dtype,
            weight_dtype=weight_dtype,
        )

        verify_conv2d(
            expr,
            expr,
            d_shape,
            w_shape,
            sm=80,
            atol=1e-5,
            rtol=1e-5,
            run_benchmark=False,
            data_dtype=data_dtype,
            weight_dtype=weight_dtype,
            ref_target="llvm",
        )

    # align1 + int8 case
    d_shape = (16, 3, 32, 32)
    w_shape = (32, 3, 3, 3)
    mod_nchw = get_conv2d_nchw(
        d_shape, w_shape, padding, out_dtype="int32", data_dtype="uint8", weight_dtype="int8"
    )

    verify_conv2d(
        mod_nchw,
        mod_nchw,
        d_shape,
        w_shape,
        sm=80,
        atol=1e-5,
        rtol=1e-5,
        ref_target="llvm",
        data_dtype="uint8",
        weight_dtype="int8",
    )


def test_conv2d_fusion():
    d_shape = (16, 16, 32, 32)
    w_shape = (32, 16, 3, 3)
    padding = (1, 1)

    mod_nchw = get_conv2d_nchw_bias(d_shape, w_shape, padding)
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )

    mod_nchw = get_conv2d_nchw_bias_relu(d_shape, w_shape, padding)
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )

    mod_nchw = get_conv2d_nchw_bias_sigmoid(d_shape, w_shape, padding, out_dtype="float16")
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )
    verify_conv2d(
        mod_nchw,
        mod_nchw,
        d_shape,
        w_shape,
        sm=80,
        atol=1e-3,
        rtol=1e-3,
        run_benchmark=False,
        use_fast_math=True,
    )

    mod_nchw = get_conv2d_nchw_bias_sigmoid(d_shape, w_shape, padding, out_dtype="float32")
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )

    mod_nchw = get_conv2d_nchw_bias_silu(d_shape, w_shape, padding, out_dtype="float32")
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )

    mod_nchw = get_conv2d_nchw_bias_hardswish(d_shape, w_shape, padding, out_dtype="float16")
    verify_conv2d(
        mod_nchw, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False
    )


def test_conv2d_residual_block():
    d_shape = (16, 16, 32, 32)
    w_shape = (16, 16, 3, 3)
    padding = (1, 1)

    bias_add, residual_input = get_conv2d_nchw_bias_residual(d_shape, w_shape, padding)

    for func, tol in [
        (relay.nn.relu(bias_add + residual_input), 1e-5),
        (relay.nn.relu(bias_add) + residual_input, 1e-5),
        (relay.sigmoid(bias_add) * residual_input, 1e-5),
        (relay.nn.relu(silu(bias_add) * residual_input), 1e-5),
        # HardSwish requires higher tolerance since vectoring the residual block epilogue
        # in cutlass.
        # TODO(masahi): Invesitigate this issue
        (relay.nn.relu(hardswish(bias_add) + residual_input), 1e-3),
    ]:
        verify_conv2d(func, func, d_shape, w_shape, sm=80, atol=tol, rtol=tol, run_benchmark=False)


def test_conv2d_transpose():
    OC = 8
    IC = 16
    d_shape = (16, IC, 32, 32)
    w_shape = (OC, IC, 3, 3)
    padding = (1, 1)
    dtype = "float32"

    for strides in [(1, 1), (2, 2)]:
        o_shape = conv_output_shape(
            0, padding, strides, (1, 1), d_shape, (OC, IC, 3, 3), "float32", "float32"
        )
        output_padding = (1, 1) if strides[0] > 1 else (0, 0)
        mod_nchw = get_conv2d_transpose_nchw(
            o_shape,
            w_shape,
            padding,
            output_padding,
            strides,
            out_dtype=dtype,
            data_dtype=dtype,
            weight_dtype=dtype,
        )

        verify_conv2d(
            mod_nchw,
            mod_nchw,
            o_shape,
            w_shape,
            sm=80,
            atol=1e-3,
            rtol=1e-3,
            use_cudnn_ref=False,
            run_benchmark=False,
            data_dtype=dtype,
            weight_dtype=dtype,
        )


def test_conv2d_backward_weight():
    OC = 8
    IC = 16
    d_shape = (16, IC, 32, 32)
    w_shape = (OC, IC, 3, 3)
    dtype = "float16"

    for strides in [(1, 1), (2, 2)]:
        o_shape = (16, OC, 32 // strides[0], 32 // strides[1])
        padding = (1, 1)

        mod_nchw = get_conv2d_backward_weight(
            d_shape,
            w_shape,
            o_shape,
            padding,
            strides,
            out_dtype="float32",
            data_dtype=dtype,
            weight_dtype=dtype,
        )

        for split_k_slices in [1, 8]:
            verify_conv2d_backward_weight(
                mod_nchw,
                mod_nchw,
                o_shape,
                d_shape,
                sm=80,
                split_k_slices=[split_k_slices],
                atol=1e-3,
                rtol=1e-3,
                use_cudnn_ref=False,
                grad_dtype=dtype,
                data_dtype=dtype,
            )


def test_conv2d_bwd():
    IC = 16
    OC = 8
    dshape = (16, IC, 32, 32)
    wshape = (OC, IC, 3, 3)
    padding = (0, 0)
    strides = (1, 1)

    conv = get_conv2d_nchw(
        dshape,
        wshape,
        padding,
        strides=strides,
        out_dtype="float32",
        data_dtype="float32",
        weight_dtype="float32",
    )
    fwd_mod = InferType()(tvm.IRModule.from_expr(conv))

    # Note: large difference in tvm and cutlass Wgrad results if use fp16.
    # Cutlass wgrad uses fp32 accumulation even if the output is fp16.
    use_fp16 = False
    verify_dgrad = False  # False to verify wgrad
    tol = 1e-5 if verify_dgrad else 1e-4  # Wgrad slightly less accurate

    if use_fp16:
        fwd_mod = ToMixedPrecision("float16")(fwd_mod)

    fwd_bwd_func = FirstOrderGradient()(fwd_mod)["main"]

    bwd_func = relay.Function(
        fwd_bwd_func.params,
        relay.TupleGetItem(relay.TupleGetItem(fwd_bwd_func.body, 1), 0 if verify_dgrad else 1),
    )

    verify_conv2d(
        bwd_func,
        bwd_func,
        dshape,
        wshape,
        sm=80,
        atol=1e-2 if use_fp16 else tol,
        rtol=1e-2 if use_fp16 else tol,
        use_cudnn_ref=False,
        data_dtype="float32",
        weight_dtype="float32",
        use_vm=True,
    )


if __name__ == "__main__":
    import sys
    import pytest

    sys.exit(pytest.main([__file__] + sys.argv[1:]))