slyubomirsky commented on code in PR #16204:
URL: https://github.com/apache/tvm/pull/16204#discussion_r1416181388
##########
src/relax/transform/canonicalize_bindings.cc:
##########
@@ -170,17 +191,26 @@ class CanonicalizePlanner : public ExprVisitor {
}
void VisitExpr_(const VarNode* var) override {
- if (!inside_dataflow_) {
- used_outside_dataflow_.insert(GetRef<Var>(var));
+ auto var_ref = GetRef<Var>(var);
+ // if a var is used in a dataflow block but *not* the one
+ // where it was defined, it also needs to be exposed, so also we treat
that as
+ // used outside of a dataflow block
+ if (!inside_dataflow_ ||
+ (def_blocks_.count(var_ref) &&
!current_block_.same_as(def_blocks_.at(var_ref)))) {
+ used_outside_home_dataflow_.insert(GetRef<Var>(var));
}
}
bool inside_dataflow_{false};
+ BindingBlock current_block_;
Review Comment:
Yeah, I was going to do that for a second reason too: There is no "current
block" when you're looking at the `body` field of a SeqExpr. This could lead to
weird results.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
