This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 4380b206bc [Unity][Bugfix] Do not include `PrimFunc`s in the
dependency graph when checking for recursion (#14228)
4380b206bc is described below
commit 4380b206bc0ba97d945b427a88fe3361fce99d16
Author: Steven S. Lyubomirsky <[email protected]>
AuthorDate: Wed Mar 8 01:46:37 2023 -0500
[Unity][Bugfix] Do not include `PrimFunc`s in the dependency graph when
checking for recursion (#14228)
This PR fixes a small oversight in #14149. Namely, the analysis for
detecting recursion in Relax adds all reference to global vars into the
dependency graph, but some global vars might refer to TIR PrimFuncs. Those
should not be in the dependency graph, and this PR ignores any that it detects.
---
src/relax/analysis/detect_recursion.cc | 13 +++++++--
.../python/relax/test_analysis_detect_recursion.py | 33 +++++++++++++++++++++-
2 files changed, 43 insertions(+), 3 deletions(-)
diff --git a/src/relax/analysis/detect_recursion.cc
b/src/relax/analysis/detect_recursion.cc
index 1c2fc69971..9c150fed8b 100644
--- a/src/relax/analysis/detect_recursion.cc
+++ b/src/relax/analysis/detect_recursion.cc
@@ -77,15 +77,24 @@ namespace relax {
class DependencyGatherer : public ExprVisitor {
public:
+ explicit DependencyGatherer(const IRModule& m) : m_(m) {}
+
std::unordered_set<std::string> Track(const Function& func) {
this->VisitExpr(func);
return deps_;
}
- void VisitExpr_(const GlobalVarNode* gv) override {
deps_.insert(gv->name_hint); }
+ void VisitExpr_(const GlobalVarNode* gv) override {
+ // disregard PrimFuncs
+ if (!m_->Lookup(GetRef<GlobalVar>(gv)).as<relax::FunctionNode>()) {
+ return;
+ }
+ deps_.insert(gv->name_hint);
+ }
private:
std::unordered_set<std::string> deps_;
+ const IRModule& m_;
};
using adjacency_map = std::unordered_map<std::string,
std::unordered_set<std::string>>;
@@ -101,7 +110,7 @@ adjacency_map GatherDependencyGraph(const IRModule& m) {
continue;
}
std::string name = gv_func.first->name_hint;
- auto deps = DependencyGatherer().Track(GetRef<relax::Function>(func));
+ auto deps = DependencyGatherer(m).Track(GetRef<relax::Function>(func));
ret.insert({name, deps});
}
return ret;
diff --git a/tests/python/relax/test_analysis_detect_recursion.py
b/tests/python/relax/test_analysis_detect_recursion.py
index 43227c7a58..c1afcd3aad 100644
--- a/tests/python/relax/test_analysis_detect_recursion.py
+++ b/tests/python/relax/test_analysis_detect_recursion.py
@@ -18,7 +18,7 @@ from typing import List
import tvm
import tvm.testing
from tvm import relax as rx
-from tvm.script import relax as R
+from tvm.script import relax as R, tir as T
from tvm.relax.analysis import detect_recursion
@@ -401,5 +401,36 @@ def test_mutual_recursion_via_references():
assert_groups(groups, [["a", "b", "c"]])
+def test_disregard_primfuncs():
+ @tvm.script.ir_module
+ class CallPrimFunc:
+ # copied from test_analysis.py
+ @T.prim_func
+ def identity_identity(A: T.Buffer((4, 4), "float32"), B: T.Buffer((4,
4), "float32")):
+ C = T.alloc_buffer((128, 128), "float32")
+ for i0, i1 in T.grid(4, 4):
+ with T.block("identity"):
+ vi0, vi1 = T.axis.remap("SS", [i0, i1])
+ C[vi0, vi1] = A[vi0, vi1]
+ for i0, i1 in T.grid(4, 4):
+ with T.block("identity"):
+ vi0, vi1 = T.axis.remap("SS", [i0, i1])
+ B[vi0, vi1] = C[vi0, vi1]
+
+ @R.function
+ def a(x: R.Tensor((4, 4), "float32")) -> R.Object:
+ y = R.call_tir(identity_identity, x, R.Tensor((4, 4), "float32"))
+ return b(y)
+
+ @R.function
+ def b(x: R.Tensor((4, 4), "float32")) -> R.Object:
+ y = R.call_tir(identity_identity, x, R.Tensor((4, 4), "float32"))
+ return a(y)
+
+ groups = detect_recursion(CallPrimFunc)
+ # the prim func should not be listed here
+ assert_groups(groups, [["a", "b"]])
+
+
if __name__ == "__main__":
tvm.testing.main()
