chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,344 @@
|
||||
# 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 pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.relax.analysis import name_to_binding
|
||||
from tvm.relax.binding_rewrite import DataflowBlockRewrite
|
||||
from tvm.relax.expr import DataflowVar, Var
|
||||
from tvm.script import relax as R
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Identity:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
|
||||
def assert_immutability(rwt, original_dfb, original_root_fn):
|
||||
assert rwt.mutated_dfb() != original_dfb
|
||||
assert rwt.mutated_root_fn() != original_root_fn
|
||||
assert rwt.mutated_root_fn().body.blocks[0] != original_dfb
|
||||
assert rwt.mutated_root_fn().body.blocks[0] == rwt.mutated_dfb()
|
||||
|
||||
|
||||
def test_null_construct():
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
DataflowBlockRewrite(dfb, root_fn)
|
||||
|
||||
|
||||
def test_simple_add():
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.add(name="tmp", expr=Identity["main"].params[0], is_dfvar=True)
|
||||
|
||||
assert_immutability(rwt, dfb, root_fn)
|
||||
|
||||
# check "tmp" added
|
||||
assert "tmp" in name_to_binding(rwt.mutated_root_fn())
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
tmp: R.Tensor((32, 32), "float32") = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
def test_simple_auto_add_var():
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.add(root_fn.params[0], is_dfvar=False)
|
||||
|
||||
assert isinstance(rwt.mutated_dfb().bindings[-1].var, Var)
|
||||
|
||||
assert_immutability(rwt, dfb, root_fn)
|
||||
|
||||
|
||||
def test_simple_auto_add_dfvar():
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.add(root_fn.params[0], is_dfvar=True)
|
||||
|
||||
assert isinstance(rwt.mutated_dfb().bindings[-1].var, DataflowVar)
|
||||
|
||||
# immutatbility
|
||||
assert_immutability(rwt, dfb, root_fn)
|
||||
|
||||
|
||||
def test_simple_remove_unused():
|
||||
@tvm.script.ir_module
|
||||
class IdentityUnused:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
unused = lv0
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
root_fn = IdentityUnused["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
n2binding = name_to_binding(IdentityUnused["main"])
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_unused(n2binding["unused"][0].var)
|
||||
|
||||
assert_immutability(rwt, dfb, root_fn)
|
||||
|
||||
# check "unused" removed
|
||||
assert "unused" not in name_to_binding(rwt.mutated_root_fn())
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
def test_remove_unused_undef():
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_unused(Var("whatever"))
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_unused(Var("whatever"), allow_undef=True)
|
||||
|
||||
assert root_fn == rwt.mutated_root_fn()
|
||||
|
||||
|
||||
def test_simple_rm_all_unused():
|
||||
@tvm.script.ir_module
|
||||
class IdentityUnused:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
unused0 = lv0
|
||||
unused1 = lv0
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
root_fn = IdentityUnused["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_all_unused()
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DeadDFBlock:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return x
|
||||
|
||||
|
||||
def test_empty_dfb_after_removal():
|
||||
root_fn = DeadDFBlock["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_unused(DeadDFBlock["main"].body.blocks[0].bindings[0].var)
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
|
||||
return x
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
def test_empty_dfb_after_all_removal():
|
||||
dfb = DeadDFBlock["main"].body.blocks[0]
|
||||
root_fn = DeadDFBlock["main"]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_all_unused()
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
|
||||
return x
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
def test_chained_rm_all_unused():
|
||||
@tvm.script.ir_module
|
||||
class IdentityChainedUnused:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
unused0 = R.call_dps_packed("my_sigmoid", (x,), R.Tensor((32, 32), dtype="float32"))
|
||||
unused1 = R.call_dps_packed(
|
||||
"my_sigmoid", (unused0,), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
root_fn = IdentityChainedUnused["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.remove_all_unused()
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
tvm.ir.assert_structural_equal(rwt.mutated_root_fn(), GroundTruth["main"])
|
||||
|
||||
|
||||
def test_simple_replace_all_uses():
|
||||
@tvm.script.ir_module
|
||||
class Lv0To1:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor((32, 32), "float32"):
|
||||
# lv0 => lv1
|
||||
# / \
|
||||
# lv2 lv3
|
||||
# \ /
|
||||
# lv4
|
||||
with R.dataflow():
|
||||
lv0: R.Tensor((32, 32), "float32") = R.call_dps_packed(
|
||||
"my_relu", (x,), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
lv1: R.Tensor((32, 32), "float32") = R.call_dps_packed(
|
||||
"my_sigmoid", (x,), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
lv2: R.Tensor((32, 32), "float32") = R.call_dps_packed(
|
||||
"my_add", (x, lv0), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
lv3: R.Tensor((32, 32), "float32") = R.call_dps_packed(
|
||||
"my_mul", (x, lv0), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
lv4: R.Tensor((32, 32), "float32") = R.call_dps_packed(
|
||||
"my_whatever", (lv2, lv3), R.Tensor((32, 32), dtype="float32")
|
||||
)
|
||||
R.output(lv4)
|
||||
return lv4
|
||||
|
||||
root_fn = Lv0To1["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
n2binding = name_to_binding(root_fn)
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
missing = Var("missing")
|
||||
with pytest.raises(RuntimeError, match="Cannot find"):
|
||||
rwt.replace_all_uses(missing, missing)
|
||||
|
||||
rwt.replace_all_uses(n2binding["lv0"][0].var, n2binding["lv1"][0].var)
|
||||
rwt.replace_all_uses(n2binding["lv1"][0].var, n2binding["lv1"][0].var)
|
||||
with pytest.raises(RuntimeError, match=r"is used by 2 vars"):
|
||||
rwt.remove_unused(n2binding["lv1"][0].var)
|
||||
rwt.remove_unused(n2binding["lv0"][0].var)
|
||||
|
||||
assert_immutability(rwt, dfb, root_fn)
|
||||
|
||||
n2binding_after = name_to_binding(rwt.mutated_root_fn())
|
||||
assert "lv0" not in n2binding_after
|
||||
|
||||
|
||||
def test_simple_module_update():
|
||||
@tvm.script.ir_module
|
||||
class Identity:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
root_fn = Identity["main"]
|
||||
dfb = root_fn.body.blocks[0]
|
||||
|
||||
rwt = DataflowBlockRewrite(dfb, root_fn)
|
||||
rwt.add(name="tmp", expr=root_fn.params[0], is_dfvar=True)
|
||||
|
||||
new_ir = rwt.mutate_irmodule(Identity)
|
||||
|
||||
# immutatbility
|
||||
assert new_ir != Identity
|
||||
assert 2 == len(new_ir["main"].body.blocks[0].bindings)
|
||||
|
||||
@tvm.script.ir_module
|
||||
class GroundTruth:
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
with R.dataflow():
|
||||
lv0 = x
|
||||
tmp: R.Tensor((32, 32), "float32") = x
|
||||
R.output(lv0)
|
||||
return lv0
|
||||
|
||||
tvm.ir.assert_structural_equal(new_ir, GroundTruth)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user