chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,173 @@
|
||||
# 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: F841
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
import tvm.script
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
use_np_array = tvm.testing.parameter(False, True)
|
||||
|
||||
|
||||
def test_bind_params(use_np_array):
|
||||
@tvm.script.ir_module
|
||||
class InputModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "tir_matmul"})
|
||||
A = T.match_buffer(x, (16, 16))
|
||||
B = T.match_buffer(y, (16, 16))
|
||||
C = T.match_buffer(z, (16, 16))
|
||||
for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi = T.axis.S(16, i0 * 4 + i1)
|
||||
vj = T.axis.S(16, j)
|
||||
vk = T.axis.R(16, k0 * 4 + k1)
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((16, 16), "float32"), w: R.Tensor((16, 16), "float32")) -> R.Tensor(
|
||||
(16, 16), "float32"
|
||||
):
|
||||
gv0 = R.call_tir(InputModule.tir_matmul, (x, w), R.Tensor((16, 16), dtype="float32"))
|
||||
return gv0
|
||||
|
||||
x_np = np.random.rand(16, 16).astype(np.float32)
|
||||
w_np = np.random.rand(16, 16).astype(np.float32)
|
||||
x_tvm = tvm.runtime.tensor(x_np)
|
||||
w_tvm = tvm.runtime.tensor(w_np)
|
||||
params_dict = {"w": w_np if use_np_array else w_tvm}
|
||||
mod = relax.transform.BindParams("main", params_dict)(InputModule)
|
||||
assert len(mod["main"].params) == 1
|
||||
|
||||
target = tvm.target.Target("llvm")
|
||||
ex_after = tvm.compile(mod, target)
|
||||
vm_after = relax.VirtualMachine(ex_after, tvm.cpu())
|
||||
res_after = vm_after["main"](x_tvm)
|
||||
|
||||
ex_before = tvm.compile(InputModule, target)
|
||||
vm_before = relax.VirtualMachine(ex_before, tvm.cpu())
|
||||
res_before = vm_before["main"](x_tvm, w_tvm)
|
||||
|
||||
tvm.testing.assert_allclose(res_before.numpy(), res_after.numpy())
|
||||
|
||||
|
||||
def test_bind_params_symbolic_vars():
|
||||
@tvm.script.ir_module
|
||||
class Before:
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor(("batch", "m"), dtype="float32"),
|
||||
w0: R.Tensor(("n", "m"), dtype="float32"),
|
||||
b0: R.Tensor(("n",), dtype="float32"),
|
||||
w1: R.Tensor(("k", "n"), dtype="float32"),
|
||||
b1: R.Tensor(("k",), dtype="float32"),
|
||||
) -> R.Tensor(("batch", "k"), dtype="float32"):
|
||||
batch = T.Var("batch", "int64")
|
||||
k = T.Var("k", "int64")
|
||||
m = T.Var("m", "int64")
|
||||
n = T.Var("n", "int64")
|
||||
with R.dataflow():
|
||||
lv0 = R.call_dps_packed(
|
||||
"linear0", (x, w0, b0), out_ty=R.Tensor((batch, n), dtype="float32")
|
||||
)
|
||||
out = R.call_dps_packed(
|
||||
"linear1", (lv0, w1, b1), out_ty=R.Tensor((batch, k), dtype="float32")
|
||||
)
|
||||
R.output(out)
|
||||
return out
|
||||
|
||||
m, n, k = 4, 6, 8
|
||||
w0_tvm = tvm.runtime.tensor(np.random.rand(n, m).astype(np.float32))
|
||||
b0_tvm = tvm.runtime.tensor(np.random.rand(n).astype(np.float32))
|
||||
w1_tvm = tvm.runtime.tensor(np.random.rand(k, n).astype(np.float32))
|
||||
b1_tvm = tvm.runtime.tensor(np.random.rand(k).astype(np.float32))
|
||||
params_dict = {"w0": w0_tvm, "b0": b0_tvm, "w1": w1_tvm, "b1": b1_tvm}
|
||||
mod = relax.transform.BindParams("main", params_dict)(Before)
|
||||
|
||||
# Since it contains ConstantNode, it's hard to check with structural equality.
|
||||
func = mod["main"]
|
||||
assert len(func.params) == 1
|
||||
batch = func.params[0].ty.shape[0]
|
||||
tvm.ir.assert_structural_equal(func.params[0].ty, relax.TensorType((batch, 4), "float32"))
|
||||
tvm.ir.assert_structural_equal(func.ret_ty, relax.TensorType((batch, 8), "float32"))
|
||||
bindings = func.body.blocks[0].bindings
|
||||
tvm.ir.assert_structural_equal(bindings[0].var.ty, relax.TensorType((batch, 6), "float32"))
|
||||
tvm.ir.assert_structural_equal(bindings[1].var.ty, relax.TensorType((batch, 8), "float32"))
|
||||
|
||||
|
||||
param_specification = tvm.testing.parameter("by_string", "by_var")
|
||||
|
||||
|
||||
def test_bind_params_by_var_obj(param_specification):
|
||||
@tvm.script.ir_module
|
||||
class Before:
|
||||
@R.function
|
||||
def main(A: R.Tensor([16], "float32")):
|
||||
return A
|
||||
|
||||
np_data = np.arange(16).astype("float32")
|
||||
inlined_relax_const = relax.const(np_data)
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Expected:
|
||||
@R.function
|
||||
def main():
|
||||
return inlined_relax_const
|
||||
|
||||
if param_specification == "by_string":
|
||||
var = "A"
|
||||
elif param_specification == "by_var":
|
||||
var = Before["main"].params[0]
|
||||
else:
|
||||
raise ValueError("Unknown param_specification: {param_specification}")
|
||||
|
||||
After = relax.transform.BindParams("main", {var: np_data})(Before)
|
||||
|
||||
tvm.ir.assert_structural_equal(Expected, After)
|
||||
|
||||
|
||||
def test_bind_params_by_var_name():
|
||||
@tvm.script.ir_module
|
||||
class Before:
|
||||
@R.function
|
||||
def main(A: R.Tensor([16], "float32")):
|
||||
return A
|
||||
|
||||
np_data = np.arange(16).astype("float32")
|
||||
inlined_relax_const = relax.const(np_data)
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Expected:
|
||||
@R.function
|
||||
def main():
|
||||
return inlined_relax_const
|
||||
|
||||
After = relax.transform.BindParams("main", {"A": np_data})(Before)
|
||||
|
||||
tvm.ir.assert_structural_equal(Expected, After)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user