wehub-resource-sync
615 lines
21 KiB
Python
615 lines
21 KiB
Python
# 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 ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
def gen_mod(mod, name, binding):
|
|
"""Select relax function with name, rename to main and bind constant.
|
|
|
|
Parameters
|
|
----------
|
|
mod: IRModule
|
|
The input module
|
|
|
|
name: str
|
|
The name of relax function to preserve and rename to main
|
|
|
|
binding: Dict[str, array]
|
|
The const parameter bindings
|
|
"""
|
|
funcs = {}
|
|
binding = {k: tvm.runtime.tensor(v) for k, v in binding.items()}
|
|
|
|
for k, v in mod.functions.items():
|
|
if isinstance(v, tvm.relax.Function):
|
|
if k.name_hint == name:
|
|
# rename to main
|
|
gv = tvm.ir.GlobalVar("main")
|
|
funcs[gv] = tvm.relax.Function(v.params, v.body, v.ret_ty).with_attr(
|
|
"global_symbol", "main"
|
|
)
|
|
else:
|
|
funcs[k] = v
|
|
mod = tvm.IRModule(funcs)
|
|
return relax.transform.BindParams("main", binding)(mod)
|
|
|
|
|
|
def test_one_fold_addone():
|
|
# put before after in a single module
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def addone(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")) -> None:
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock("addone"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + T.float32(1)
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((16, 16), dtype="float32"))
|
|
return lv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((16, 16), "float32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c1_np = c0_np + 1
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_one_fold_transpose():
|
|
# put before after in a single module
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def func(A: T.Buffer((2, 3), "float32"), B: T.Buffer((3, 2), "float32")) -> None:
|
|
for i, j in T.grid(3, 2):
|
|
with T.sblock("transpose"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vj, vi]
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((2, 3), "float32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(cls.func, (c0,), R.Tensor((3, 2), dtype="float32"))
|
|
return lv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((3, 2), "float32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(2 * 3).astype("float32").reshape(2, 3)
|
|
c1_np = c0_np.T
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_two_hop_addone():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def addone(A: T.Buffer((2, 2), "float32"), B: T.Buffer((2, 2), "float32")) -> None:
|
|
for i, j in T.grid(2, 2):
|
|
with T.sblock("addone"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + T.float32(1)
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((2, 2), "float32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((2, 2), dtype="float32"))
|
|
lv1 = relax.call_tir(cls.addone, (lv0,), R.Tensor((2, 2), dtype="float32"))
|
|
return lv1
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((2, 2), "float32"), c2: R.Tensor((2, 2), "float32")):
|
|
return c2
|
|
|
|
c0_np = np.arange(2 * 2).astype("float32").reshape(2, 2)
|
|
c1_np = c0_np + 1
|
|
c2_np = c1_np + 1
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np, "c2": c2_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_dataflow_fold():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def identity(A: T.Buffer((16, 16), "float32"), B: T.Buffer((16, 16), "float32")) -> None:
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock("identity"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj]
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32")):
|
|
cls = Module
|
|
with R.dataflow():
|
|
gv0 = relax.call_tir(cls.identity, (c0,), R.Tensor((16, 16), dtype="float32"))
|
|
R.output(gv0)
|
|
return gv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((16, 16), "float32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c1_np = c0_np
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_mixed_case():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
# TIR function can handle different cases.
|
|
@T.prim_func(s_tir=True)
|
|
def addone(a: T.handle, b: T.handle) -> None:
|
|
n = T.int32()
|
|
m = T.int32()
|
|
A = T.match_buffer(a, (n, m))
|
|
B = T.match_buffer(b, (n, m))
|
|
for i, j in T.grid(n, m):
|
|
with T.sblock("addone"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + T.float32(1)
|
|
|
|
@T.prim_func(s_tir=True)
|
|
def sub(
|
|
A: T.Buffer((16, 16), "float32"),
|
|
B: T.Buffer((16, 16), "float32"),
|
|
C: T.Buffer((16, 16), "float32"),
|
|
) -> None:
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock("sub"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = A[vi, vj] - B[vi, vj]
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32"), x: R.Tensor("float32", ndim=2)):
|
|
n, m = T.int64(), T.int64()
|
|
cls = Module
|
|
x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
|
|
# this line cannot be folded because n is unknown
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((n, 16), dtype="float32"))
|
|
# this line can be folded
|
|
lv1 = relax.call_tir(cls.addone, (c0,), R.Tensor((16, 16), dtype="float32"))
|
|
# this line can be folded because all inputs are const
|
|
lv2 = relax.call_tir(cls.sub, (c0, lv1), R.Tensor((16, 16), dtype="float32"))
|
|
# this line can not be folded because x's shape is unknown
|
|
lv3 = relax.call_tir(cls.sub, (lv2, x), R.Tensor((16, 16), dtype="float32"))
|
|
return (lv0, lv3)
|
|
|
|
@R.function
|
|
def expected(
|
|
c0: R.Tensor((16, 16), "float32"),
|
|
c1: R.Tensor((16, 16), "float32"),
|
|
c2: R.Tensor((16, 16), "float32"),
|
|
x: R.Tensor("float32", ndim=2),
|
|
):
|
|
n, m = T.int64(), T.int64()
|
|
cls = Module
|
|
x0 = R.match_cast(x, R.Tensor((n, m), "float32"))
|
|
# this line cannot be folded because n is unknown
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((n, 16), dtype="float32"))
|
|
# this line can not be folded because x's shape is unknown
|
|
lv3 = relax.call_tir(cls.sub, (c2, x), R.Tensor((16, 16), dtype="float32"))
|
|
return (lv0, lv3)
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c1_np = c0_np + 1
|
|
c2_np = c0_np - c1_np
|
|
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c0": c0_np, "c1": c1_np, "c2": c2_np})
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_int32_fold():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def addone(A: T.Buffer((16, 16), "int32"), B: T.Buffer((16, 16), "int32")) -> None:
|
|
for i, j in T.grid(16, 16):
|
|
with T.sblock("addone"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj] + T.int32(1)
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "int32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((16, 16), dtype="int32"))
|
|
return lv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((16, 16), "int32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(16 * 16).astype("int32").reshape(16, 16)
|
|
c1_np = c0_np + 1
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_single_relax_op():
|
|
# put before after in a single module
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32")):
|
|
with R.dataflow():
|
|
gv = R.add(c0, c0)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((16, 16), "float32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c1_np = c0_np + c0_np
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_multiple_relax_ops():
|
|
# put before after in a single module
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32"), c1: R.Tensor((16, 16), "float32")):
|
|
with R.dataflow():
|
|
lv0 = R.add(c0, c1)
|
|
lv1 = R.multiply(c0, lv0)
|
|
gv = R.subtract(lv1, c1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(c4: R.Tensor((16, 16), "float32")):
|
|
return c4
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c1_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
c2_np = c0_np + c1_np
|
|
c3_np = c0_np * c2_np
|
|
c4_np = c3_np - c1_np
|
|
before = gen_mod(Module, "before", {"c0": c0_np, "c1": c1_np})
|
|
expected = gen_mod(Module, "expected", {"c4": c4_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_do_not_fold_ops_outside_dataflow():
|
|
# put before after in a single module
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32")):
|
|
gv = R.add(c0, c0)
|
|
return gv
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, before)
|
|
|
|
|
|
def test_fold_multiple_relax_ops_with_data_dependent_reshape():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before(
|
|
data: R.Tensor((256,), "float32"),
|
|
c0: R.Tensor((2,), "int64"),
|
|
c1: R.Tensor((2,), "int64"),
|
|
):
|
|
with R.dataflow():
|
|
lv0 = R.add(c0, c0)
|
|
target_shape = R.multiply(lv0, c1)
|
|
lv2: R.Shape(ndim=2) = R.tensor_to_shape(target_shape)
|
|
gv: R.Tensor(ndim=2, dtype="float32") = R.reshape(data, lv2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(data: R.Tensor((256,), "float32")) -> R.Tensor((16, 16), dtype="float32"):
|
|
with R.dataflow():
|
|
gv: R.Tensor((16, 16), dtype="float32") = R.reshape(data, R.shape([16, 16]))
|
|
R.output(gv)
|
|
return gv
|
|
|
|
c0_np = [8, 8]
|
|
c1_np = [1, 1]
|
|
before = gen_mod(Module, "before", {"c0": c0_np, "c1": c1_np})
|
|
relax.analysis.well_formed(before)
|
|
|
|
expected = gen_mod(Module, "expected", {})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_unsupported_fold_ops_legalized_to_multiple_calls():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before(c0: R.Tensor((16, 16), "float32")):
|
|
with R.dataflow():
|
|
gv = R.nn.relu(c0)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
c0_np = np.arange(16 * 16).astype("float32").reshape(16, 16)
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
|
|
from tvm.relax.transform.legalize_ops.common import register_legalize
|
|
|
|
def customized_legalize_relu(bb: relax.BlockBuilder, call: relax.Call):
|
|
from tvm import topi # pylint: disable=import-outside-toplevel
|
|
|
|
x = bb.emit_te(topi.nn.relu, *call.args)
|
|
return bb.call_te(topi.identity, x)
|
|
|
|
# register custom legalization for relu that emits multiple bindings for testing
|
|
relu_legalize = tvm.ir.Op.get("relax.nn.relu").get_attr("FLegalize")
|
|
tvm.ir.Op.get("relax.nn.relu").reset_attr("FLegalize")
|
|
register_legalize("relax.nn.relu", customized_legalize_relu)
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, before)
|
|
|
|
# revert to correct legalization of relu
|
|
tvm.ir.Op.get("relax.nn.relu").reset_attr("FLegalize")
|
|
register_legalize("relax.nn.relu", relu_legalize)
|
|
|
|
|
|
def test_fold_shape_computation():
|
|
@I.ir_module(s_tir=True)
|
|
class Module:
|
|
@R.function
|
|
def before(
|
|
data: R.Tensor((5, 4, 3, 2), dtype="float32"),
|
|
indices: R.Tensor((1,), dtype="int64"),
|
|
) -> R.Tensor((1, 1), dtype="int64"):
|
|
with R.dataflow():
|
|
lv: R.Tensor((4,), dtype="int64") = R.shape_to_tensor(R.shape([5, 4, 3, 2]))
|
|
lv1: R.Tensor((1,), dtype="int64") = R.take(lv, indices, axis=0)
|
|
lv2: R.Tensor((1, 1), dtype="int64") = R.expand_dims(lv1, axis=[0])
|
|
gv: R.Tensor((1, 1), dtype="int64") = R.concat((lv2,), axis=0)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(
|
|
data: R.Tensor((5, 4, 3, 2), dtype="float32"), new_shape: R.Tensor((1, 1), "int64")
|
|
) -> R.Tensor((1, 1), dtype="int64"):
|
|
return new_shape
|
|
|
|
before = gen_mod(
|
|
Module, "before", {"indices": tvm.runtime.tensor(np.array([0]).astype("int64"))}
|
|
)
|
|
after = relax.transform.FoldConstant()(before)
|
|
np_take = np.take([5, 4, 3, 2], [0], axis=0)
|
|
np_expand = np.expand_dims(np_take, axis=[0])
|
|
np_concat = np.concatenate([np_expand], axis=0)
|
|
expected = gen_mod(Module, "expected", {"new_shape": tvm.runtime.tensor(np_concat)})
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_tuple_output():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def split(
|
|
A: T.Buffer((4, 4), "float32"),
|
|
B: T.Buffer((2, 4), "float32"),
|
|
C: T.Buffer((2, 4), "float32"),
|
|
) -> None:
|
|
for i, j in T.grid(2, 4):
|
|
with T.sblock("upper"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
B[vi, vj] = A[vi, vj]
|
|
for i, j in T.grid(2, 4):
|
|
with T.sblock("lower"):
|
|
vi, vj = T.axis.remap("SS", [i, j])
|
|
C[vi, vj] = A[vi + 2, vj]
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((4, 4), "float32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(
|
|
cls.split,
|
|
(c0,),
|
|
out_ty=[
|
|
R.Tensor((2, 4), dtype="float32"),
|
|
R.Tensor((2, 4), dtype="float32"),
|
|
],
|
|
)
|
|
return lv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((2, 4), "float32"), c2: R.Tensor((2, 4), "float32")) -> R.Tuple(
|
|
R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")
|
|
):
|
|
lv0: R.Tuple(R.Tensor((2, 4), dtype="float32"), R.Tensor((2, 4), dtype="float32")) = (
|
|
c1,
|
|
c2,
|
|
)
|
|
return lv0
|
|
|
|
c0_np = np.arange(16).astype("float32").reshape(4, 4)
|
|
c1_np = c0_np[:2]
|
|
c2_np = c0_np[2:]
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np, "c2": c2_np})
|
|
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_skip_folding_large_creation_op():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before():
|
|
with R.dataflow():
|
|
# 2048 elements > 1024 threshold, no tensor input
|
|
gv = R.zeros((2048,), "float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
before = Module
|
|
after = relax.transform.FoldConstant()(before)
|
|
# The zeros op should NOT be folded because the output is large
|
|
tvm.ir.assert_structural_equal(after, before)
|
|
|
|
|
|
def test_fold_small_creation_op():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before():
|
|
with R.dataflow():
|
|
# 16 elements <= 1024 threshold
|
|
gv = R.zeros((4, 4), "float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(c0: R.Tensor((4, 4), "float32")):
|
|
return c0
|
|
|
|
before = gen_mod(Module, "before", {})
|
|
expected = gen_mod(Module, "expected", {"c0": np.zeros((4, 4), dtype="float32")})
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_boundary_creation_op():
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@R.function
|
|
def before():
|
|
with R.dataflow():
|
|
# Exactly 1024 elements == threshold, should fold
|
|
gv = R.zeros((1024,), "float32")
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@R.function
|
|
def expected(c0: R.Tensor((1024,), "float32")):
|
|
return c0
|
|
|
|
before = gen_mod(Module, "before", {})
|
|
expected = gen_mod(Module, "expected", {"c0": np.zeros((1024,), dtype="float32")})
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_fold_large_op_with_tensor_input():
|
|
"""Ops with tensor inputs should be folded even if output is large."""
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(s_tir=True)
|
|
def addone(A: T.Buffer((2048,), "float32"), B: T.Buffer((2048,), "float32")) -> None:
|
|
for i in range(2048):
|
|
with T.sblock("addone"):
|
|
vi = T.axis.remap("S", [i])
|
|
B[vi] = A[vi] + T.float32(1)
|
|
|
|
@R.function
|
|
def before(c0: R.Tensor((2048,), "float32")):
|
|
cls = Module
|
|
lv0 = relax.call_tir(cls.addone, (c0,), R.Tensor((2048,), dtype="float32"))
|
|
return lv0
|
|
|
|
@R.function
|
|
def expected(c1: R.Tensor((2048,), "float32")):
|
|
return c1
|
|
|
|
c0_np = np.arange(2048).astype("float32")
|
|
c1_np = c0_np + 1
|
|
before = gen_mod(Module, "before", {"c0": c0_np})
|
|
expected = gen_mod(Module, "expected", {"c1": c1_np})
|
|
after = relax.transform.FoldConstant()(before)
|
|
tvm.ir.assert_structural_equal(after, expected)
|
|
|
|
|
|
def test_call_tir_with_tir_vars_not_folded():
|
|
"""call_tir with symbolic tir_vars cannot be const-evaluated."""
|
|
|
|
@tvm.script.ir_module
|
|
class Module:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def shape_to_tensor(out: T.Buffer((T.int64(1),), "int64"), m: T.int64):
|
|
for i in range(T.int64(1)):
|
|
with T.sblock("out"):
|
|
vi = T.axis.remap("S", [i])
|
|
out[vi] = m
|
|
|
|
@R.function
|
|
def main(x: R.Tensor(("m",), "float32")):
|
|
m = T.int64()
|
|
cls = Module
|
|
gv = relax.call_tir(
|
|
cls.shape_to_tensor, R.tuple(), R.Tensor((1,), "int64"), tir_vars=R.shape([m])
|
|
)
|
|
return gv
|
|
|
|
after = relax.transform.FoldConstant()(Module)
|
|
tvm.ir.assert_structural_equal(after, Module)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|