chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,665 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user