junrushao1994 commented on a change in pull request #9871:
URL: https://github.com/apache/tvm/pull/9871#discussion_r780639158
##########
File path: src/tir/schedule/transform.cc
##########
@@ -136,5 +136,52 @@ void LeafBlockRemovalPlan(const ScheduleState& self, const
StmtSRef& leaf_block_
throw OnlyLeafError(self->mod, GetRef<Block>(leaf_block),
GetRef<Block>(scope_block));
}
+/******** IR Substitution ********/
+class IRSubstituteInScope : public StmtExprMutator {
+ public:
+ explicit IRSubstituteInScope(std::function<PrimExpr(const VarNode*)> fmap)
+ : fmap_(std::move(fmap)) {}
+
+ PrimExpr VisitExpr_(const VarNode* op) final {
+ auto it = fmap_(op);
+ if (it.defined()) {
+ return it;
+ } else {
+ return GetRef<PrimExpr>(op);
+ }
+ }
+
+ Stmt VisitStmt_(const BlockRealizeNode* op) final {
+ auto fmutate = [&](const PrimExpr& e) { return this->VisitExpr(e); };
+ Array<PrimExpr> v = op->iter_values;
+ v.MutateByApply(fmutate);
+ PrimExpr pred = this->VisitExpr(op->predicate);
+ if (v.same_as(op->iter_values) && pred.same_as(op->predicate)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->iter_values = std::move(v);
+ n->predicate = std::move(analyzer.Simplify(pred));
+ return Stmt(n);
+ }
+ }
+
+ private:
+ const std::function<PrimExpr(const VarNode*)> fmap_;
+ arith::Analyzer analyzer;
+};
+
+Stmt SubstituteInScope(const Stmt& stmt, const Map<Var, PrimExpr>& subst_map) {
Review comment:
I'm reading the implementation. Looks like what `SubstituteInScope`
really does is that we only substitute block bindings and predicates, stops at
the `BlockRealize` boundary and does nothing else. Is that accurate? If so,
shall we rename this to `SubstituteBlockBinding`?
--
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]
