wehub-resource-sync
666 lines
21 KiB
Python
666 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
|
|
"""Test eliminate common subexpr pass"""
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm.relax.transform import EliminateCommonSubexpr
|
|
from tvm.script.parser import ir as I
|
|
from tvm.script.parser import relax as R
|
|
from tvm.script.parser import tirx as T
|
|
|
|
|
|
def verify(input, expected, call_only=False):
|
|
tvm.ir.assert_structural_equal(EliminateCommonSubexpr(call_only)(input), expected)
|
|
|
|
|
|
def test_simple():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv0 = R.add(x, y)
|
|
lv1 = R.add(x, y)
|
|
gv = R.multiply(lv0, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv0 = R.add(x, y)
|
|
lv1 = lv0
|
|
gv = R.multiply(lv0, lv0)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_constants():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo() -> R.Tuple(R.Tensor((), dtype="int32"), R.Tensor((2, 2), dtype="int32")):
|
|
with R.dataflow():
|
|
# we are not going to bind the constant 1 to a var
|
|
lv0 = R.add(R.const(1, dtype="int32"), R.const(1, dtype="int32"))
|
|
# we expect to bind the repeated large constants
|
|
lv1 = R.add(
|
|
R.const(tvm.runtime.tensor(np.zeros((2, 2), dtype="int32"))),
|
|
R.const(tvm.runtime.tensor(np.zeros((2, 2), dtype="int32"))),
|
|
)
|
|
gv = (lv0, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo() -> R.Tuple(R.Tensor((), dtype="int32"), R.Tensor((2, 2), dtype="int32")):
|
|
with R.dataflow():
|
|
lv0 = R.add(R.const(1, dtype="int32"), R.const(1, dtype="int32"))
|
|
lv1 = R.add(
|
|
R.const(tvm.runtime.tensor(np.zeros((2, 2), dtype="int32"))),
|
|
R.const(tvm.runtime.tensor(np.zeros((2, 2), dtype="int32"))),
|
|
)
|
|
gv = (lv0, lv1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_repeated_inner_tuples():
|
|
"""CSE is only applied at variable bindings
|
|
|
|
To remain consistent with the behavior of the normalizer, tuples
|
|
are kept as-is, even if they contain repeated sub-tuples.
|
|
"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
|
|
with R.dataflow():
|
|
# repeated units: (x, x), (x, (x, x)), ((x, x), (x, (x, x)))
|
|
tup = (((x, x), (x, (x, x))), ((x, x), (x, (x, x))), (x, (x, x)))
|
|
gv = tup[0][0][1]
|
|
R.output(gv)
|
|
return gv
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_inner_function():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
|
|
with R.dataflow():
|
|
# we are going to do CSE inside the local function
|
|
@R.function
|
|
def bar(y: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
|
|
with R.dataflow():
|
|
# writing this out in ANF to illustrate why CSE behaves as it does
|
|
# result of ANF transforming R.add(R.add(y, y), R.add(y, y))
|
|
lv0 = R.add(y, y)
|
|
lv1 = R.add(y, y)
|
|
lv2 = R.add(lv0, lv1)
|
|
gv = lv2
|
|
R.output(gv)
|
|
return R.add(gv, gv)
|
|
|
|
# also making the ANF explicit to better illustrate the result of CSE
|
|
# result of ANF transforming R.add(R.add(bar(x), bar(x)), R.add(bar(x), bar(x)))
|
|
lv0 = bar(x)
|
|
lv1 = bar(x)
|
|
lv2 = R.add(lv0, lv1)
|
|
lv3 = bar(x)
|
|
lv4 = bar(x)
|
|
lv5 = R.add(lv3, lv4)
|
|
lv6 = R.add(lv2, lv5)
|
|
gv = lv6
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
|
|
with R.dataflow():
|
|
|
|
@R.function
|
|
def bar(y: R.Tensor((), dtype="int32")) -> R.Tensor((), dtype="int32"):
|
|
with R.dataflow():
|
|
lv0 = R.add(y, y)
|
|
lv1 = lv0
|
|
lv2 = R.add(lv0, lv0)
|
|
gv = lv2
|
|
R.output(gv)
|
|
return R.add(gv, gv)
|
|
|
|
# can further clean this up
|
|
# using canonicalize bindings, eliminate unused bindings, and CSE again
|
|
lv0 = bar(x)
|
|
lv1 = lv0
|
|
lv2 = R.add(lv0, lv0)
|
|
lv3 = lv0
|
|
lv4 = lv0
|
|
lv5 = lv2
|
|
lv6 = R.add(lv2, lv2)
|
|
gv = lv6
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_call_only():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((160,), dtype="float32")):
|
|
with R.dataflow():
|
|
lv1 = R.arange(R.prim_value(0), R.prim_value(160), R.prim_value(1), dtype="float32")
|
|
lv2 = R.arange(R.prim_value(0), R.prim_value(160), R.prim_value(1), dtype="float32")
|
|
lv3 = R.add(x, lv1)
|
|
out = R.add(lv3, lv2)
|
|
R.output(out)
|
|
return out
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor((160,), dtype="float32")) -> R.Tensor((160,), dtype="float32"):
|
|
with R.dataflow():
|
|
lv1 = R.arange(R.prim_value(0), R.prim_value(160), R.prim_value(1), dtype="float32")
|
|
lv2 = lv1
|
|
lv3 = R.add(x, lv1)
|
|
out = R.add(lv3, lv1)
|
|
R.output(out)
|
|
return out
|
|
|
|
verify(Before, Expected, call_only=True)
|
|
|
|
|
|
def test_cse_outside_dataflow():
|
|
# same example as previously but it will work without a dataflow wrapper
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
lv0 = R.add(x, y)
|
|
lv1 = R.add(x, y)
|
|
gv = R.multiply(lv0, lv1)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
lv0 = R.add(x, y)
|
|
lv1 = lv0
|
|
gv = R.multiply(lv0, lv0)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_no_cse_across_dataflow():
|
|
# same example as previously but it will work without a dataflow wrapper
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
lv0 = R.add(x, y)
|
|
lv1 = R.add(x, y)
|
|
gv1 = R.multiply(lv0, lv1)
|
|
R.output(gv1)
|
|
|
|
_ = R.print(format="Prevent dataflow block merging")
|
|
|
|
with R.dataflow():
|
|
lv2 = R.add(x, y)
|
|
lv3 = R.add(x, y)
|
|
gv2 = R.multiply(lv2, lv3)
|
|
R.output(gv2)
|
|
|
|
gv3 = R.add(x, y)
|
|
gv4 = R.add(x, y)
|
|
gv5 = R.multiply(gv3, gv4)
|
|
|
|
output = R.add(R.add(gv1, gv2), gv5)
|
|
return output
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
# The R.add(x,y) may be de-duplicated within a dataflow block
|
|
lv0 = R.add(x, y)
|
|
lv1 = lv0
|
|
gv1 = R.multiply(lv0, lv0)
|
|
R.output(gv1)
|
|
|
|
_ = R.print(format="Prevent dataflow block merging")
|
|
|
|
with R.dataflow():
|
|
# However, the later dataflow block may not be
|
|
# de-duplicated using variables in the earlier block.
|
|
lv2 = R.add(x, y)
|
|
lv3 = lv2
|
|
gv2 = R.multiply(lv2, lv2)
|
|
R.output(gv2)
|
|
|
|
# And while non-dataflow bindings can be de-duplicated,
|
|
# they cannot be de-duplicated using bindings that were
|
|
# valid in either of the earlier dataflow blocks.
|
|
gv3 = R.add(x, y)
|
|
gv4 = gv3
|
|
gv5 = R.multiply(gv3, gv3)
|
|
|
|
output = R.add(R.add(gv1, gv2), gv5)
|
|
return output
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_no_replacement_across_dataflow_boundary():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
A = R.add(x, y)
|
|
# B has the same value as A, and so instances of B can be replaced with A.
|
|
B = R.add(x, y)
|
|
C = R.multiply(A, B)
|
|
|
|
# However, B is exposed for use outside of the
|
|
# DataflowBlock, while A is not. Therefore, any
|
|
# additional uses of `B` must NOT be replaced with
|
|
# A.
|
|
R.output(B, C)
|
|
|
|
# In addition, because `A` is only valid within the
|
|
# dataflow block, the `R.add(x,y)` cannot be de-duplicated
|
|
# as another usage of `A`.
|
|
D = R.add(x, y)
|
|
return (B, C, D)
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def main(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
A = R.add(x, y)
|
|
B = A
|
|
C = R.multiply(A, A)
|
|
R.output(B, C)
|
|
|
|
D = B
|
|
return (B, C, B)
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_do_not_eliminate_impure():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
# it's a repeated subexpression but it would be wrong to deduplicate it
|
|
p1 = R.print(format="Message")
|
|
p2 = R.print(format="Message")
|
|
a1 = R.assert_op(R.const(False), format="Always fails")
|
|
lv0 = R.add(x, y)
|
|
lv1 = R.add(x, y)
|
|
gv = R.multiply(lv0, lv1)
|
|
a2 = R.assert_op(R.const(False), format="Always fails")
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
p1 = R.print(format="Message")
|
|
p2 = R.print(format="Message")
|
|
a1 = R.assert_op(R.const(False), format="Always fails")
|
|
lv0 = R.add(x, y)
|
|
lv1 = lv0
|
|
gv = R.multiply(lv0, lv0)
|
|
a2 = R.assert_op(R.const(False), format="Always fails")
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_do_not_eliminate_shape_expr():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
x = R.reshape(x, [6])
|
|
y = R.reshape(y, [6])
|
|
z = R.add(x, y)
|
|
return z
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_do_not_eliminate_extern_func():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def foo(x: R.Tensor((2, 3), dtype="float32")):
|
|
y = R.call_packed("extern_func_name", x, ty_args=R.Tensor([2, 3]))
|
|
z = R.call_packed("extern_func_name", y, ty_args=R.Tensor([2, 3]))
|
|
return z
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_call_tir_tuple_arg():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def main(A: R.Tensor([16, 16], "int32"), B: R.Tensor([16, 16], "int32")):
|
|
cls = Before
|
|
Prod = R.call_tir(cls.product, [A, B], out_ty=R.Tensor([16, 16], "int32"))
|
|
Sum = R.call_tir(cls.sum, [A, B], out_ty=R.Tensor([16, 16], "int32"))
|
|
return (Prod, Sum)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def product(
|
|
A: T.Buffer([16, 16], "int32"),
|
|
B: T.Buffer([16, 16], "int32"),
|
|
C: T.Buffer([16, 16], "int32"),
|
|
):
|
|
for iters in T.grid(*A.shape):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
C[i, j] = A[i, j] * B[i, j]
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def sum(
|
|
A: T.Buffer([16, 16], "int32"),
|
|
B: T.Buffer([16, 16], "int32"),
|
|
C: T.Buffer([16, 16], "int32"),
|
|
):
|
|
for iters in T.grid(*A.shape):
|
|
with T.sblock("compute"):
|
|
i, j = T.axis.remap("SS", iters)
|
|
C[i, j] = A[i, j] + B[i, j]
|
|
|
|
Expected = Before
|
|
|
|
# If EliminateCommonSubexpr produces unnormalized expressions,
|
|
# normalization of those expressions may produce additional
|
|
# variables bindings. This test case should be agnostic to those
|
|
# additional bindings, so DCE is applied after CSE.
|
|
After = tvm.ir.transform.Sequential(
|
|
[
|
|
EliminateCommonSubexpr(),
|
|
tvm.relax.transform.DeadCodeElimination(),
|
|
]
|
|
)(Before)
|
|
tvm.ir.assert_structural_equal(Expected, After)
|
|
|
|
|
|
def test_do_not_eliminate_dtype():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function(pure=False)
|
|
def foo() -> R.Tensor((32, 64), "int32"):
|
|
obj: R.Any = R.vm.alloc_storage(R.shape([24576]), runtime_device_index=0, dtype="uint8")
|
|
a: R.Tensor([32, 64], dtype="int32") = R.vm.alloc_tensor(
|
|
obj, offset=0, shape=R.shape([32, 64]), dtype="int32"
|
|
)
|
|
ret_val: R.Tensor([32, 64], dtype="int32") = R.builtin.alloc_tensor(
|
|
R.shape([32, 64]), R.dtype("int32"), R.prim_value(0)
|
|
)
|
|
_t1: R.Tuple = R.vm.kill_object(a)
|
|
_t3: R.Tuple = R.vm.kill_object(obj)
|
|
lv: R.Tensor([32, 64], dtype="int32") = ret_val
|
|
return lv
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_match_cast():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
A1 = R.add(x, y)
|
|
B1 = R.match_cast(A1, R.Tensor([2, 3], "float32"))
|
|
|
|
A2 = R.add(x, y)
|
|
B2 = R.match_cast(A2, R.Tensor([2, 3], "float32"))
|
|
|
|
gv = R.multiply(B1, B2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32")):
|
|
with R.dataflow():
|
|
A1 = R.add(x, y)
|
|
B1 = R.match_cast(A1, R.Tensor([2, 3], "float32"))
|
|
|
|
A2 = A1
|
|
B2 = B1
|
|
gv = R.multiply(B1, B1)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_match_cast_with_symbolic_vars():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor(dtype="float32"), y: R.Tensor(dtype="float32")):
|
|
with R.dataflow():
|
|
A1 = R.add(x, y)
|
|
|
|
n = T.int64()
|
|
m = T.int64()
|
|
B1 = R.match_cast(A1, R.Tensor([n, m], "float32"))
|
|
|
|
A2 = R.add(x, y)
|
|
p = T.int64()
|
|
q = T.int64()
|
|
B2 = R.match_cast(A2, R.Tensor([p, q], "float32"))
|
|
|
|
gv = R.multiply(B1, B2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(x: R.Tensor(dtype="float32"), y: R.Tensor(dtype="float32")):
|
|
with R.dataflow():
|
|
A1 = R.add(x, y)
|
|
n = T.int64()
|
|
m = T.int64()
|
|
B1 = R.match_cast(A1, R.Tensor([n, m], "float32"))
|
|
|
|
A2 = A1
|
|
p = T.int64()
|
|
q = T.int64()
|
|
B2 = R.match_cast(A1, R.Tensor([p, q], "float32"))
|
|
|
|
gv = R.multiply(B1, B2)
|
|
R.output(gv)
|
|
return gv
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_replace_binding_within_branch_with_duplicate_before_branch():
|
|
"""Bindings before a branch may be used within the branch"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
condition: R.Prim("bool"),
|
|
):
|
|
A = R.add(x, y)
|
|
if condition:
|
|
B = R.add(x, y)
|
|
C = R.multiply(x, B)
|
|
D = R.multiply(A, C)
|
|
else:
|
|
B = R.add(x, y)
|
|
C = R.multiply(y, B)
|
|
D = R.multiply(A, C)
|
|
return D
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Expected:
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
condition: R.Prim("bool"),
|
|
):
|
|
A = R.add(x, y)
|
|
if condition:
|
|
B = A
|
|
C = R.multiply(x, A)
|
|
D = R.multiply(A, C)
|
|
else:
|
|
B = A
|
|
C = R.multiply(y, A)
|
|
D = R.multiply(A, C)
|
|
return D
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_keep_duplicate_across_if_and_then():
|
|
"""Bindings in `if` are not valid within `else`"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
condition: R.Prim("bool"),
|
|
):
|
|
if condition:
|
|
A = R.add(x, y)
|
|
B = R.multiply(x, A)
|
|
else:
|
|
A = R.add(x, y)
|
|
B = R.multiply(y, A)
|
|
return B
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_keep_duplicate_after_branch():
|
|
"""Only the final binding is valid after a if/else branch"""
|
|
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(
|
|
x: R.Tensor((2, 3), dtype="float32"),
|
|
y: R.Tensor((2, 3), dtype="float32"),
|
|
condition: R.Prim("bool"),
|
|
):
|
|
if condition:
|
|
A = R.add(x, y)
|
|
B = R.multiply(x, A)
|
|
else:
|
|
A = R.add(x, y)
|
|
B = R.multiply(y, A)
|
|
|
|
C = R.add(x, y)
|
|
D = R.multiply(B, C)
|
|
return D
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_keep_alloc_tensor():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32")):
|
|
tmp_buf1 = R.builtin.alloc_tensor(R.shape([64]), R.dtype("int32"), R.prim_value(0))
|
|
tmp_buf2 = R.builtin.alloc_tensor(R.shape([64]), R.dtype("int32"), R.prim_value(0))
|
|
out = R.add(tmp_buf1, tmp_buf2)
|
|
return out
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
def test_keep_alloc_storage():
|
|
@I.ir_module(s_tir=True)
|
|
class Before:
|
|
@R.function
|
|
def foo(x: R.Tensor((2, 3), dtype="float32")):
|
|
tmp_storage1 = R.vm.alloc_storage(R.shape([64]), runtime_device_index=0, dtype="uint8")
|
|
tmp_buf1 = R.vm.alloc_tensor(tmp_storage1, offset=0, shape=R.shape([64]), dtype="int32")
|
|
tmp_storage2 = R.vm.alloc_storage(R.shape([64]), runtime_device_index=0, dtype="uint8")
|
|
tmp_buf2 = R.vm.alloc_tensor(tmp_storage2, offset=0, shape=R.shape([64]), dtype="int32")
|
|
out = R.add(tmp_buf1, tmp_buf2)
|
|
return out
|
|
|
|
Expected = Before
|
|
|
|
verify(Before, Expected)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|