This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new deb11d384e [TIR] Use IRModuleNode::Remove to remove None in
PrimFuncPass (#14494)
deb11d384e is described below
commit deb11d384e753907c8ce1ea93957888609622ad5
Author: Eric Lunderberg <[email protected]>
AuthorDate: Wed Apr 5 09:09:28 2023 -0500
[TIR] Use IRModuleNode::Remove to remove None in PrimFuncPass (#14494)
Prior to this commit, `PrimFuncPass` directly removed empty `PrimFunc`
objects from the module's `Map<GlobalVar, BaseFunc> functions`.
Because it didn't update the `global_var_map_` as well, these two maps
could become out of sync. Since the `global_var_map_` is checked as
part of `StructuralEqual()`, but isn't displayed when printing to
TVMScript, this can result in identical printouts being flagged as
non-identical.
This commit updates `PrimFuncPass` to call the `IRModuleNode::Remove`
method, which updates both the `functions` and `global_var_map_`
variables.
---
src/tir/ir/transform.cc | 11 +++++---
.../python/unittest/test_tir_transform_helpers.py | 30 +++++++++++++++++++++-
2 files changed, 36 insertions(+), 5 deletions(-)
diff --git a/src/tir/ir/transform.cc b/src/tir/ir/transform.cc
index 4c59a17673..38ec16ea37 100644
--- a/src/tir/ir/transform.cc
+++ b/src/tir/ir/transform.cc
@@ -88,7 +88,8 @@ PrimFuncPass::PrimFuncPass(
// Perform Module -> Module optimizations at the PrimFunc level.
IRModule PrimFuncPassNode::operator()(IRModule mod, const PassContext&
pass_ctx) const {
ICHECK(mod.defined());
- std::vector<ObjectRef> deleted_list;
+ std::vector<GlobalVar> deleted_list;
+
IRModuleNode* mod_ptr = mod.CopyOnWrite();
auto* func_dict = mod_ptr->functions.CopyOnWrite();
// directly loop over the underlying dict
@@ -101,14 +102,16 @@ IRModule PrimFuncPassNode::operator()(IRModule mod, const
PassContext& pass_ctx)
kv.second = std::move(func);
if (!kv.second.defined()) {
- deleted_list.push_back(kv.first);
+ deleted_list.push_back(Downcast<GlobalVar>(kv.first));
}
}
}
- // automatic removal of None
+ // Automatic removal of None. This uses IRModuleNode::Remove
+ // instead of manipulating func_dict directly, to ensure that both
+ // the function map and the global_var_map_ are correctly updated.
for (const auto& gv : deleted_list) {
- func_dict->erase(gv);
+ mod_ptr->Remove(gv);
}
return mod;
}
diff --git a/tests/python/unittest/test_tir_transform_helpers.py
b/tests/python/unittest/test_tir_transform_helpers.py
index f8dc0f682d..657bda591a 100644
--- a/tests/python/unittest/test_tir_transform_helpers.py
+++ b/tests/python/unittest/test_tir_transform_helpers.py
@@ -17,7 +17,7 @@
import pytest
import tvm
-from tvm.script import tir as T
+from tvm.script import tir as T, ir as I
import tvm.testing
@@ -119,5 +119,33 @@ def test_filter_primfunc():
assert len(after.functions) == 0
+class TestFilterRemovesGlobalVarMap(tvm.testing.CompareBeforeAfter):
+ """Filtering out a function should be identical to never adding it
+
+ This test is to guard against hidden state in the IRModule that
+ remains after filtering. Previously, this was observed in the
+ `IRModuleNode::global_var_map_`, which retained entries of
+ filtered-out functions.
+ """
+
+ transform = tvm.tir.transform.Filter(lambda prim_func: False)
+
+ def before(self):
+ @I.ir_module
+ class module:
+ @T.prim_func
+ def func():
+ T.evaluate(0)
+
+ return module
+
+ def expected(self):
+ @I.ir_module
+ class module:
+ pass
+
+ return module
+
+
if __name__ == "__main__":
tvm.testing.main()
