chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,574 @@
|
||||
# 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.
|
||||
# ruff: noqa: RUF005
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.testing import env
|
||||
|
||||
pytest.importorskip("scipy") # tvm.topi.testing imports scipy
|
||||
|
||||
import tvm.topi.testing
|
||||
from tvm import relax
|
||||
from tvm.relax.backend.cuda.cublas import partition_for_cublas
|
||||
from tvm.relax.testing import get_relax_matmul_module
|
||||
from tvm.script import relax as R
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder import relax as relax_builder
|
||||
|
||||
try:
|
||||
import ml_dtypes
|
||||
except ImportError:
|
||||
ml_dtypes = None
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_seed():
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.gpu,
|
||||
pytest.mark.skipif(not env.has_cublas(), reason="need cublas"),
|
||||
]
|
||||
|
||||
|
||||
def build_and_run(mod, inputs_np, target, legalize=False, cuda_graph=False):
|
||||
with tvm.transform.PassContext(
|
||||
config={
|
||||
"relax.backend.use_cuda_graph": cuda_graph,
|
||||
"relax.transform.apply_legalize_ops": legalize,
|
||||
}
|
||||
):
|
||||
ex = tvm.compile(mod, target)
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.device(target, 0)
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
f = vm["main"]
|
||||
inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
|
||||
|
||||
# For cuda graph, run the compiled function twice to make sure that we can launch the
|
||||
# cached graph on the second run.
|
||||
if cuda_graph:
|
||||
f(*inputs)
|
||||
|
||||
return f(*inputs).numpy()
|
||||
|
||||
if tvm.target.Target(target).kind.name == "cuda":
|
||||
return tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
return run_and_check()
|
||||
|
||||
|
||||
def get_result_with_relax_cublas_offload(mod, np_inputs, cuda_graph=False, bind_constants=False):
|
||||
mod = partition_for_cublas(mod, bind_constants=bind_constants)
|
||||
mod = relax.transform.RunCodegen()(mod)
|
||||
|
||||
return build_and_run(mod, np_inputs, "cuda", cuda_graph)
|
||||
|
||||
|
||||
def _to_concrete_shape(symbolic_shape, var_table):
|
||||
result = []
|
||||
for dim in symbolic_shape:
|
||||
if not isinstance(dim, tvm.tirx.expr.Var):
|
||||
result.append(dim)
|
||||
continue
|
||||
|
||||
if dim not in var_table:
|
||||
var_table[dim] = np.random.randint(10, 50)
|
||||
result.append(var_table[dim])
|
||||
|
||||
return tuple(result)
|
||||
|
||||
|
||||
_vars = {
|
||||
"a": tvm.tirx.expr.Var("a", "int64"),
|
||||
"b": tvm.tirx.expr.Var("b", "int64"),
|
||||
}
|
||||
|
||||
|
||||
_epilogue_table = {
|
||||
"none": (False, None),
|
||||
"bias": (True, None),
|
||||
"relu": (True, R.nn.relu),
|
||||
"gelu": (True, R.nn.gelu),
|
||||
}
|
||||
|
||||
|
||||
def get_relax_matmul_dequantize_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
transposed_y=False,
|
||||
scale_const=1.0,
|
||||
zero_point_const=0.0,
|
||||
):
|
||||
"""Create a matmul op followd by dequantize operations."""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
x = R.arg("x", R.Tensor(x_shape, in_dtype))
|
||||
y = R.arg("y", R.Tensor(y_shape, in_dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
if transposed_y:
|
||||
axes = list(range(len(y_shape) - 2)) + [-1, -2]
|
||||
y = R.emit(R.permute_dims(y, axes=axes))
|
||||
result = R.emit(R.matmul(x, y, out_dtype="float32"))
|
||||
result = R.emit(
|
||||
R.dequantize(
|
||||
result,
|
||||
scale=R.const(scale_const, "float16"),
|
||||
zero_point=R.const(zero_point_const, "float16"),
|
||||
axis=-1,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
)
|
||||
R.output(result)
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
def get_relax_matmul_multiply_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
z_shape,
|
||||
in_dtype,
|
||||
acc_dtype,
|
||||
out_dtype,
|
||||
transposed_y=False,
|
||||
):
|
||||
"""Create a matmul op followd by multiply operations."""
|
||||
with IRBuilder() as builder:
|
||||
with relax_builder.function():
|
||||
R.func_name("main")
|
||||
x = R.arg("x", R.Tensor(x_shape, in_dtype))
|
||||
y = R.arg("y", R.Tensor(y_shape, in_dtype))
|
||||
scaleA = R.arg("scaleA", R.Tensor(z_shape, acc_dtype))
|
||||
scaleB = R.arg("scaleB", R.Tensor(z_shape, acc_dtype))
|
||||
|
||||
with R.dataflow() as frame:
|
||||
if transposed_y:
|
||||
axes = list(range(len(y_shape) - 2)) + [-1, -2]
|
||||
y = R.emit(R.permute_dims(y, axes=axes))
|
||||
result = R.emit(R.matmul(x, y, out_dtype=acc_dtype))
|
||||
z = R.emit(R.multiply(scaleA, scaleB))
|
||||
result = R.emit(R.multiply(result, z))
|
||||
if acc_dtype != out_dtype:
|
||||
result = R.emit(R.astype(result, out_dtype))
|
||||
R.output(result)
|
||||
R.func_ret_value(frame.output_vars[0])
|
||||
|
||||
func = builder.get()
|
||||
return tvm.IRModule({"main": func})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"x_shape, y_shape, transpose_y, epilogue",
|
||||
[
|
||||
# Regular
|
||||
((8, 8), (8, 8), False, "none"),
|
||||
((_vars["a"], 6), (6, 16), False, "bias"),
|
||||
# Transposed
|
||||
((4, 16), (16, 128), True, "relu"),
|
||||
((35, 8), (8, 8), True, "gelu"),
|
||||
# # 3D x 3D
|
||||
((6, 32, 8), (6, 8, 10), False, "bias"),
|
||||
((6, 32, 8), (6, 8, 10), True, "none"),
|
||||
((_vars["a"], 32, 8), (_vars["a"], 8, 10), True, "gelu"),
|
||||
# ND x ND
|
||||
((5, 3, 32, 8), (5, 3, 8, 10), True, "relu"),
|
||||
((_vars["a"], 3, 32, 8), (_vars["a"], 3, 8, 10), True, "relu"),
|
||||
((_vars["a"], _vars["b"], 32, 8), (_vars["a"], _vars["b"], 8, 10), True, "relu"),
|
||||
# ND x 2D
|
||||
((5, 3, 32, 8), (8, 10), False, "none"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"in_dtype, out_dtype",
|
||||
[
|
||||
("float16", "float16"),
|
||||
("float32", "float32"),
|
||||
],
|
||||
)
|
||||
def test_matmul_offload(
|
||||
x_shape,
|
||||
y_shape,
|
||||
transpose_y,
|
||||
epilogue,
|
||||
in_dtype,
|
||||
out_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(in_dtype)
|
||||
y = np.random.randn(*concrete_y_shape).astype(in_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(out_dtype)
|
||||
args = (x, y, bias)
|
||||
else:
|
||||
bias = None
|
||||
args = (x, y)
|
||||
|
||||
mod = get_relax_matmul_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
bias_shape=bias.shape if with_bias else None,
|
||||
transposed_y=transpose_y,
|
||||
activation=activation,
|
||||
)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args)
|
||||
ref = build_and_run(mod, args, "llvm", legalize=True)
|
||||
|
||||
tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"x_shape, y_shape, transpose_y, epilogue",
|
||||
[
|
||||
# Regular
|
||||
((8, 8), (8, 8), False, "none"),
|
||||
((_vars["a"], 8), (8, 16), False, "none"),
|
||||
# Transposed
|
||||
((4, 16), (16, 128), True, "none"),
|
||||
((35, 16), (16, 128), False, "none"),
|
||||
# # 3D x 3D
|
||||
((6, 32, 8), (6, 8, 12), False, "none"),
|
||||
((6, 32, 8), (6, 8, 10), True, "none"),
|
||||
((_vars["a"], 32, 8), (_vars["a"], 8, 10), True, "none"),
|
||||
# ND x ND
|
||||
((5, 3, 32, 8), (5, 3, 8, 12), False, "none"),
|
||||
# ND x 2D
|
||||
((5, 3, 32, 8), (8, 12), False, "none"),
|
||||
],
|
||||
)
|
||||
def test_matmul_igemm_offload(
|
||||
x_shape,
|
||||
y_shape,
|
||||
transpose_y,
|
||||
epilogue,
|
||||
):
|
||||
in_dtype = "int8"
|
||||
out_dtype = "int32"
|
||||
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(in_dtype)
|
||||
y = np.random.randn(*concrete_y_shape).astype(in_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(out_dtype)
|
||||
args = (x, y, bias)
|
||||
else:
|
||||
bias = None
|
||||
args = (x, y)
|
||||
|
||||
mod = get_relax_matmul_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
bias_shape=bias.shape if with_bias else None,
|
||||
transposed_y=transpose_y,
|
||||
activation=activation,
|
||||
)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args)
|
||||
ref = build_and_run(mod, args, "llvm", legalize=True)
|
||||
|
||||
tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.skipif(ml_dtypes is None, reason="requires ml_dtypes to be installed")
|
||||
@pytest.mark.parametrize(
|
||||
"x_shape, y_shape, transpose_y, out_dtype",
|
||||
[
|
||||
((10, 32), (64, 32), True, "float32"),
|
||||
((32, 16), (32, 16), True, "float16"),
|
||||
((2, 10, 32), (2, 64, 32), True, "float32"),
|
||||
],
|
||||
)
|
||||
def test_matmul_fp8_offload(
|
||||
x_shape,
|
||||
y_shape,
|
||||
transpose_y,
|
||||
out_dtype,
|
||||
):
|
||||
in_dtype = "float8_e4m3fn"
|
||||
mod = get_relax_matmul_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
bias_shape=None,
|
||||
transposed_y=transpose_y,
|
||||
activation=None,
|
||||
)
|
||||
numpytype = "float8_e4m3fn"
|
||||
x = np.random.uniform(low=0, high=5, size=x_shape).astype(numpytype)
|
||||
y = np.random.uniform(low=0, high=5, size=y_shape).astype(numpytype)
|
||||
z = np.swapaxes(y, -2, -1) if transpose_y else y
|
||||
args = (x, y)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args)
|
||||
ref_out = np.matmul(x, z).astype(out_dtype)
|
||||
|
||||
tvm.testing.assert_allclose(out, ref_out, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.skipif(ml_dtypes is None, reason="requires ml_dtypes to be installed")
|
||||
def test_matmul_fp8_dequantize_offload():
|
||||
x_shape = (10, 32)
|
||||
y_shape = (64, 32)
|
||||
in_dtype = "float8_e4m3fn"
|
||||
mod = get_relax_matmul_dequantize_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
"float16",
|
||||
transposed_y=True,
|
||||
scale_const=0.34786,
|
||||
zero_point_const=0.0,
|
||||
)
|
||||
|
||||
numpytype = "float8_e4m3fn"
|
||||
x = np.random.uniform(low=0, high=5, size=x_shape).astype(numpytype)
|
||||
y = np.random.uniform(low=0, high=5, size=y_shape).astype(numpytype)
|
||||
args = (x, y)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args, bind_constants=True)
|
||||
ref = build_and_run(mod, args, "llvm", legalize=True)
|
||||
tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0")
|
||||
@pytest.mark.skipif(ml_dtypes is None, reason="requires ml_dtypes to be installed")
|
||||
def test_matmul_fp8_multiply_offload():
|
||||
x_shape = (10, 32)
|
||||
y_shape = (64, 32)
|
||||
z_shape = (1,)
|
||||
in_dtype, acc_dtype = ("float8_e4m3fn", "float32")
|
||||
|
||||
mod = get_relax_matmul_multiply_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
z_shape,
|
||||
in_dtype,
|
||||
acc_dtype,
|
||||
"float16",
|
||||
transposed_y=True,
|
||||
)
|
||||
|
||||
numpytype = "float8_e4m3fn"
|
||||
x = np.random.uniform(low=0, high=5, size=x_shape).astype(numpytype)
|
||||
y = np.random.uniform(low=0, high=5, size=y_shape).astype(numpytype)
|
||||
scaleA = np.random.uniform(low=0, high=5, size=z_shape).astype(acc_dtype)
|
||||
scaleB = np.random.uniform(low=0, high=5, size=z_shape).astype(acc_dtype)
|
||||
args = (x, y, scaleA, scaleB)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args)
|
||||
ref = build_and_run(mod, args, "llvm", legalize=True)
|
||||
tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.skipif(ml_dtypes is None, reason="requires ml_dtypes to be installed")
|
||||
@pytest.mark.parametrize(
|
||||
"x_shape, y_shape, transpose_y, out_dtype",
|
||||
[
|
||||
((10, 32), (64, 32), True, "float32"),
|
||||
((32, 16), (32, 16), True, "float32"),
|
||||
((2, 10, 32), (2, 64, 32), True, "float32"),
|
||||
],
|
||||
)
|
||||
def test_matmul_bfloat16_offload(
|
||||
x_shape,
|
||||
y_shape,
|
||||
transpose_y,
|
||||
out_dtype,
|
||||
):
|
||||
in_dtype = "bfloat16"
|
||||
mod = get_relax_matmul_module(
|
||||
x_shape,
|
||||
y_shape,
|
||||
in_dtype,
|
||||
out_dtype,
|
||||
bias_shape=None,
|
||||
transposed_y=transpose_y,
|
||||
activation=None,
|
||||
)
|
||||
# Generate input data in float32 and then convert to bfloat16 using ml_dtypes.
|
||||
x_float32 = np.random.uniform(low=0, high=5, size=x_shape).astype("float32")
|
||||
y_float32 = np.random.uniform(low=0, high=5, size=y_shape).astype("float32")
|
||||
x_bf16 = ml_dtypes.bfloat16(x_float32)
|
||||
y_bf16 = ml_dtypes.bfloat16(y_float32)
|
||||
|
||||
# For the reference result, adjust y (if needed) in float32.
|
||||
z = np.swapaxes(y_float32, -2, -1) if transpose_y else y_float32
|
||||
args = (x_bf16, y_bf16)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(mod, args)
|
||||
ref_out = np.matmul(x_float32, z).astype(out_dtype)
|
||||
|
||||
tvm.testing.assert_allclose(out, ref_out, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, out_dtype, transposed_y, partition_done",
|
||||
[
|
||||
(15, 64, 32, "float32", True, True),
|
||||
(15, 64, 32, "float8_e4m3fn", True, True),
|
||||
(15, 64, 32, "float8_e5m2", True, False),
|
||||
(16, 32, 60, "float32", True, False),
|
||||
(16, 30, 64, "float32", True, False),
|
||||
(16, 8, 16, "float16", True, True),
|
||||
(16, 16, 16, "float16", False, False),
|
||||
],
|
||||
)
|
||||
def test_cublas_partition_fp8_matmul(M, N, K, out_dtype, transposed_y, partition_done):
|
||||
mod = get_relax_matmul_module(
|
||||
(M, K), (N, K), "float8_e4m3fn", out_dtype, transposed_y=transposed_y
|
||||
)
|
||||
mod = partition_for_cublas(mod)
|
||||
func_name = "relax_matmul_cublas" if partition_done else "R.matmul"
|
||||
assert func_name in mod["main"].script()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, scale, zp, num_bindings",
|
||||
[
|
||||
(16, 64, 32, 2.0, 0.0, 1),
|
||||
(16, 64, 32, 2.0, 1.0, 2),
|
||||
(16, 64, 32, [2.0] * 64, [2.0] * 64, 2),
|
||||
],
|
||||
)
|
||||
def test_cublas_partition_fp8_matmul_dequantize(M, N, K, scale, zp, num_bindings):
|
||||
mod = get_relax_matmul_dequantize_module(
|
||||
(M, K),
|
||||
(N, K),
|
||||
"float8_e4m3fn",
|
||||
"float16",
|
||||
transposed_y=True,
|
||||
scale_const=scale,
|
||||
zero_point_const=zp,
|
||||
)
|
||||
mod = partition_for_cublas(mod)
|
||||
# Check whether R.dequantize is still in main function or not
|
||||
assert len(mod["main"].body.blocks[0].bindings) == num_bindings
|
||||
|
||||
|
||||
def test_cublas_partition_fp8_matmul_multiply():
|
||||
M, N, K = (32, 64, 128)
|
||||
mod = get_relax_matmul_multiply_module(
|
||||
(M, K),
|
||||
(N, K),
|
||||
(1,),
|
||||
"float8_e4m3fn",
|
||||
"float32",
|
||||
"float16",
|
||||
transposed_y=True,
|
||||
)
|
||||
mod = partition_for_cublas(mod)
|
||||
assert len(mod["main"].body.blocks[0].bindings) == 1
|
||||
|
||||
|
||||
def test_cublas_partition_matmul_without_bias():
|
||||
# cuBLAS does not handle 2D bias (residual input)
|
||||
mod = get_relax_matmul_module((16, 32), (32, 32), "float16", "float16", bias_shape=(16, 32))
|
||||
mod = partition_for_cublas(mod)
|
||||
|
||||
# R.add is still in the main function
|
||||
assert len(mod["main"].body.blocks[0].bindings) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, was_partitioned", [(16, 8, 32, True), (16, 8, 33, False), (16, 9, 32, False)]
|
||||
)
|
||||
def test_cublas_partition_igemm(M, N, K, was_partitioned):
|
||||
mod = get_relax_matmul_module((M, K), (K, N), "int8", "int32")
|
||||
mod = partition_for_cublas(mod)
|
||||
func_name = "fused_relax_matmul_cublas" if was_partitioned else "R.matmul"
|
||||
assert func_name in mod["main"].script()
|
||||
|
||||
|
||||
def test_cublas_partition_igemm_with_bias():
|
||||
mod = get_relax_matmul_module((16, 32), (32, 8), "int8", "int32", bias_shape=(8,))
|
||||
mod = partition_for_cublas(mod)
|
||||
func = mod["main"].script()
|
||||
assert "fused_relax_matmul_cublas" in func and "R.add" in func
|
||||
|
||||
|
||||
def test_cublas_matmul_cuda_graph():
|
||||
@tvm.script.ir.ir_module
|
||||
class Mod:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((16, 16), "float16"),
|
||||
w0: R.Tensor((16, 16), "float16"),
|
||||
w1: R.Tensor((16, 16), "float16"),
|
||||
w2: R.Tensor((16, 16), "float16"),
|
||||
):
|
||||
R.func_attr({"num_input": 1})
|
||||
with R.dataflow():
|
||||
lv0 = R.matmul(x, w0)
|
||||
lv1 = R.matmul(lv0, w1)
|
||||
lv2 = R.matmul(lv1, w2)
|
||||
R.output(lv2)
|
||||
return lv2
|
||||
|
||||
mod = Mod
|
||||
shape = [16, 16]
|
||||
data = np.random.rand(*shape).astype(np.float16)
|
||||
w0 = np.random.rand(*shape).astype(np.float16)
|
||||
w1 = np.random.rand(*shape).astype(np.float16)
|
||||
w2 = np.random.rand(*shape).astype(np.float16)
|
||||
inputs = (data, w0, w1, w2)
|
||||
|
||||
out = get_result_with_relax_cublas_offload(Mod, inputs, cuda_graph=True)
|
||||
|
||||
with tvm.target.Target("cuda"):
|
||||
mod = tvm.s_tir.transform.DefaultGPUSchedule()(mod)
|
||||
ref = build_and_run(mod, inputs, "llvm", legalize=True)
|
||||
tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user