Function test_conv2d_offload

tests/python/relax/test_codegen_cutlass.py:252–279

(data_shape, weight_shape, dtype, epilogue, residual_block)

    ],
)
def test_conv2d_offload(data_shape, weight_shape, dtype, epilogue, residual_block):
    low, high = -1, 1
    data = np.random.randint(low, high, size=data_shape).astype(dtype)
    weight = np.random.randint(low, high, size=weight_shape).astype(dtype)
    bias = np.random.randint(low, high, size=(1, 1, 1, weight_shape[0])).astype(dtype)

    with_bias, activation = _epilogue_table[epilogue]
    residual_bin_op, residual_activation = _residual_block_table[residual_block]

    if with_bias:
        args = (data, weight, bias)
    else:
        args = (data, weight)

    mod = get_relax_conv2d_module(
        data_shape,
        weight_shape,
        dtype,
        with_bias=with_bias,
        activation=activation,
        residual_bin_op=residual_bin_op,
        residual_activation=residual_activation,
    )
    out = get_result_with_relax_cutlass_offload(mod, *args)

    ref = build_and_run(mod, args, "llvm")

    tvm.testing.assert_allclose(out, ref, rtol=1e-5, atol=1e-5)


@pytest.mark.parametrize(