wehub-resource-sync
345 lines
10 KiB
Python
345 lines
10 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 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()
|