Function test_matmul_offload

tests/python/relax/test_codegen_cutlass.py:387–428

(
    x_shape,
    y_shape,
    transpose_y,
    epilogue,
    residual_block,
    dtype,
)

    ],
)
def test_matmul_offload(
    x_shape,
    y_shape,
    transpose_y,
    epilogue,
    residual_block,
    dtype,
):
    with_bias, activation = _epilogue_table[epilogue]
    var_table = {}
    concrete_x_shape = _to_concrete_shape(x_shape, var_table)
    concrete_y_shape = _to_concrete_shape(y_shape, var_table)
    x = np.random.randn(*concrete_x_shape).astype(dtype)
    y = np.random.randn(*concrete_y_shape).astype(dtype)

    if transpose_y:
        y = np.swapaxes(y, -2, -1)
        y_shape = (*y_shape[:-2], y_shape[-1], y_shape[-2])

    if with_bias:
        bias = np.random.randn(concrete_y_shape[-1]).astype(dtype)
        args = (x, y, bias)
    else:
        bias = None
        args = (x, y)

    residual_bin_op, residual_activation = _residual_block_table[residual_block]

    mod = get_relax_matmul_module(
        x_shape,
        y_shape,
        dtype,
        bias_shape=bias.shape if with_bias else None,
        transposed_y=transpose_y,
        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-2, atol=1e-2)


def test_matmul_with_3d_bias_offload():