#undef EXPR_FUNCTOR_DEFAULT
#undef STMT_FUNCTOR_DEFAULT
+/*!
+ * \brief ExprVisitor
+ */
+class TVM_DLL ExprVisitor :
+ public ExprFunctor<void(const Expr&)> {
+ public:
+ using ExprFunctor::operator();
+
+ protected:
+ using ExprFunctor::VisitExpr;
+ // list of functions to override.
+ void VisitExpr_(const Variable* op) override;
+ void VisitExpr_(const Load* op) override;
+ void VisitExpr_(const Let* op) override;
+ void VisitExpr_(const Call* op) override;
+ void VisitExpr_(const Add* op) override;
+ void VisitExpr_(const Sub* op) override;
+ void VisitExpr_(const Mul* op) override;
+ void VisitExpr_(const Div* op) override;
+ void VisitExpr_(const Mod* op) override;
+ void VisitExpr_(const FloorDiv* op) override;
+ void VisitExpr_(const FloorMod* op) override;
+ void VisitExpr_(const Min* op) override;
+ void VisitExpr_(const Max* op) override;
+ void VisitExpr_(const EQ* op) override;
+ void VisitExpr_(const NE* op) override;
+ void VisitExpr_(const LT* op) override;
+ void VisitExpr_(const LE* op) override;
+ void VisitExpr_(const GT* op) override;
+ void VisitExpr_(const GE* op) override;
+ void VisitExpr_(const And* op) override;
+ void VisitExpr_(const Or* op) override;
+ void VisitExpr_(const Reduce* op) override;
+ void VisitExpr_(const Cast* op) override;
+ void VisitExpr_(const Not* op) override;
+ void VisitExpr_(const Select* op) override;
+ void VisitExpr_(const Ramp* op) override;
+ void VisitExpr_(const Broadcast* op) override;
+ void VisitExpr_(const Shuffle* op) override;
+ void VisitExpr_(const IntImm* op) override;
+ void VisitExpr_(const UIntImm* op) override;
+ void VisitExpr_(const FloatImm* op) override;
+ void VisitExpr_(const StringImm* op) override;
+};
+
+/*!
+ * \brief ExprMutator that mutates expressions.
+ */
+class TVM_DLL ExprMutator :
+ protected ExprFunctor<Expr(const Expr&)> {
+ public:
+ using ExprFunctor::operator();
+
+ protected:
+ using ExprFunctor::VisitExpr;
+ // list of functions to override.
+ Expr VisitExpr_(const Variable* op) override;
+ Expr VisitExpr_(const Load* op) override;
+ Expr VisitExpr_(const Let* op) override;
+ Expr VisitExpr_(const Call* op) override;
+ Expr VisitExpr_(const Add* op) override;
+ Expr VisitExpr_(const Sub* op) override;
+ Expr VisitExpr_(const Mul* op) override;
+ Expr VisitExpr_(const Div* op) override;
+ Expr VisitExpr_(const Mod* op) override;
+ Expr VisitExpr_(const FloorDiv* op) override;
+ Expr VisitExpr_(const FloorMod* op) override;
+ Expr VisitExpr_(const Min* op) override;
+ Expr VisitExpr_(const Max* op) override;
+ Expr VisitExpr_(const EQ* op) override;
+ Expr VisitExpr_(const NE* op) override;
+ Expr VisitExpr_(const LT* op) override;
+ Expr VisitExpr_(const LE* op) override;
+ Expr VisitExpr_(const GT* op) override;
+ Expr VisitExpr_(const GE* op) override;
+ Expr VisitExpr_(const And* op) override;
+ Expr VisitExpr_(const Or* op) override;
+ Expr VisitExpr_(const Reduce* op) override;
+ Expr VisitExpr_(const Cast* op) override;
+ Expr VisitExpr_(const Not* op) override;
+ Expr VisitExpr_(const Select* op) override;
+ Expr VisitExpr_(const Ramp* op) override;
+ Expr VisitExpr_(const Broadcast* op) override;
+ Expr VisitExpr_(const Shuffle* op) override;
+ Expr VisitExpr_(const IntImm* op) override;
+ Expr VisitExpr_(const UIntImm* op) override;
+ Expr VisitExpr_(const FloatImm* op) override;
+ Expr VisitExpr_(const StringImm* op) override;
+};
+
+/*!
+ * \brief StmtVisitor.
+ */
+class TVM_DLL StmtVisitor :
+ protected StmtFunctor<void(const Stmt&)> {
+ public:
+ using StmtFunctor::operator();
+
+ protected:
+ using StmtFunctor::VisitStmt;
+ /*!
+ * \brief Visitor to Exprs, can be overriden
+ * to do recursive changes to Exprs.
+ * \note A common pattern is to call ExprVisitor here,
+ * or have a class sub-class both StmtVisitor and ExprVisitor
+ * and redirect Visit to ExprMutator::VisitExpr(Expr)
+ */
+ virtual void VisitExpr(const Expr& e) {}
+ // statement visitor
+ void VisitStmt_(const AttrStmt* op) override;
+ void VisitStmt_(const IfThenElse* op) override;
+ void VisitStmt_(const LetStmt* op) override;
+ void VisitStmt_(const For* op) override;
+ void VisitStmt_(const Allocate* op) override;
+ void VisitStmt_(const Store* op) override;
+ void VisitStmt_(const Free* op) override;
+ void VisitStmt_(const AssertStmt* op) override;
+ void VisitStmt_(const ProducerConsumer* op) override;
+ void VisitStmt_(const Provide* op) override;
+ void VisitStmt_(const Realize* op) override;
+ void VisitStmt_(const Prefetch* op) override;
+ void VisitStmt_(const Block* op) override;
+ void VisitStmt_(const Evaluate* op) override;
+};
+
+/*!
+ * \brief StmtMutator that mutates the statements.
+ */
+class TVM_DLL StmtMutator :
+ protected StmtFunctor<Stmt(const Stmt&)> {
+ public:
+ /*!
+ * \brief Mutate stmt.
+ * \param stmt The input statement to be mutated.
+ * \return The result of the call
+ * \note It is important that stmt is passed by value.
+ * so copy on write can be triggered correctly.
+ * do mutator(std::move(stmt)) or when copy elison is triggered.
+ */
+ Stmt operator()(Stmt stmt) {
+ allow_copy_on_write_ = true;
+ return VisitStmt(stmt);
+ }
+
+ protected:
+ // We perform copy on write optimizations on the StmtMutator
+ // so that an unique copy of parent can be mutated inplace
+ // when some of its children changed.
+ // We only do such optimization for Stmt nests(instead of Exprs) for now
+ // as Stmt's parent state is more likely remain unchanged when one of
+ // its child block changes.
+ /*!
+ * \brief Internal state to indicate whether copy on write is enabled.
+ * COW is enabled iff all the parents of the node are unique.
+ */
+ bool allow_copy_on_write_{false};
+ /*!
+ * \brief Perform copy on write on node.
+ *
+ * If CopyOnWrite is allowed, directly return
+ * a strong reference to the node container.
+ * Otherwise, return a copy of the node.
+ *
+ * \return The result object pointer.
+ */
+ template<typename TNode>
+ ObjectPtr<TNode> CopyOnWrite(const TNode* node) {
+ if (allow_copy_on_write_) {
+ // return the old node.
+ return runtime::GetObjectPtr<TNode>(const_cast<TNode*>(node));
+ } else {
+ // Make a new copy of the node.
+ // need to rely on the default copy constructor
+ return runtime::make_object<TNode>(*node);
+ }
+ }
+ /*!
+ * \brief Internal mutator that everyone calls.
+ * \note To override mutate's behavior, override VisitExpr instead.
+ * \param stmt The input stmt.
+ * \return The mutated results.
+ */
+ Stmt VisitStmt(const Stmt& stmt) override {
+ if (allow_copy_on_write_ && !stmt.unique()) {
+ allow_copy_on_write_ = false;
+ Stmt ret = StmtFunctor::VisitStmt(stmt);
+ allow_copy_on_write_ = true;
+ return ret;
+ } else {
+ return StmtFunctor::VisitStmt(stmt);
+ }
+ }
+ /*!
+ * \brief Visitor to Exprs, can be overriden
+ * to do recursive changes to Exprs.
+ * \note A common pattern is to call ExprMutator here,
+ * or have a class sub-class both StmtMutator and ExprMutator
+ * and redirect Mutate to ExprMutator::Mutate(Expr)
+ */
+ virtual Expr VisitExpr(const Expr& e) {
+ return e;
+ }
+ // statement visitor
+ Stmt VisitStmt_(const AttrStmt* op) override;
+ Stmt VisitStmt_(const IfThenElse* op) override;
+ Stmt VisitStmt_(const LetStmt* op) override;
+ Stmt VisitStmt_(const For* op) override;
+ Stmt VisitStmt_(const Allocate* op) override;
+ Stmt VisitStmt_(const Store* op) override;
+ Stmt VisitStmt_(const Free* op) override;
+ Stmt VisitStmt_(const AssertStmt* op) override;
+ Stmt VisitStmt_(const ProducerConsumer* op) override;
+ Stmt VisitStmt_(const Provide* op) override;
+ Stmt VisitStmt_(const Realize* op) override;
+ Stmt VisitStmt_(const Prefetch* op) override;
+ Stmt VisitStmt_(const Block* op) override;
+ Stmt VisitStmt_(const Evaluate* op) override;
+ // internal helper.
+ class Internal;
+};
+
+/*!
+ * \brief Visitor that recursively visit stmts and exprs on them.
+ */
+class StmtExprVisitor :
+ public StmtVisitor,
+ public ExprVisitor {
+ public:
+ using StmtVisitor::operator();
+ using ExprVisitor::operator();
+
+ protected:
+ using StmtVisitor::VisitStmt;
+ using ExprVisitor::VisitExpr;
+
+ void VisitExpr(const Expr& e) override {
+ return ExprVisitor::VisitExpr(e);
+ }
+};
+
+/*!
+ * \brief Mutator that recursively mutates stmts and exprs on them.
+ */
+class StmtExprMutator :
+ public StmtMutator,
+ public ExprMutator {
+ public:
+ using StmtMutator::operator();
+ using ExprMutator::operator();
+
+ protected:
+ using StmtMutator::VisitExpr;
+ using ExprMutator::VisitExpr;
+
+ Expr VisitExpr(const Expr& e) override {
+ return ExprMutator::VisitExpr(e);
+ }
+};
+
} // namespace ir
} // namespace tvm
#endif // TVM_IR_FUNCTOR_EXT_H_
--- /dev/null
+/*
+ * 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.
+ */
+/*!
+ * \file ir_functor.cc
+ */
+#include <tvm/ir_functor_ext.h>
+#include <tvm/packed_func_ext.h>
+#include <tvm/ir_visitor.h>
+
+namespace tvm {
+namespace ir {
+
+// visitor to implement apply
+class IRApplyVisit :
+ public StmtExprVisitor {
+ public:
+ explicit IRApplyVisit(std::function<void(const ObjectRef&)> f) : f_(f) {}
+
+ void VisitExpr(const Expr& node) final {
+ if (visited_.count(node.get()) != 0) return;
+ visited_.insert(node.get());
+ ExprVisitor::VisitExpr(node);
+ f_(node);
+ }
+
+ void VisitStmt(const Stmt& node) final {
+ if (visited_.count(node.get()) != 0) return;
+ visited_.insert(node.get());
+ StmtVisitor::VisitStmt(node);
+ f_(node);
+ }
+
+ private:
+ std::function<void(const ObjectRef&)> f_;
+ std::unordered_set<const Object*> visited_;
+};
+
+void PostOrderVisit(const ObjectRef& node,
+ std::function<void(const ObjectRef&)> fvisit) {
+ if (node.as<StmtNode>()) {
+ IRApplyVisit visitor(fvisit);
+ visitor(Downcast<Stmt>(node));
+ } else {
+ IRApplyVisit visitor(fvisit);
+ visitor(Downcast<Expr>(node));
+ }
+}
+
+class IRTransformer final :
+ public StmtExprMutator {
+ public:
+ IRTransformer(const runtime::PackedFunc& f_preorder,
+ const runtime::PackedFunc& f_postorder,
+ const std::unordered_set<uint32_t>& only_enable)
+ : f_preorder_(f_preorder),
+ f_postorder_(f_postorder),
+ only_enable_(only_enable) {
+ }
+ Stmt VisitStmt(const Stmt& stmt) final {
+ return MutateInternal<Stmt>(stmt, [this](const Stmt& s) {
+ return StmtMutator::VisitStmt(s);
+ });
+ }
+ Expr VisitExpr(const Expr& expr) final {
+ return MutateInternal<Expr>(expr, [this](const Expr& e) {
+ return ExprMutator::VisitExpr(e);
+ });
+ }
+
+ private:
+ template <typename T, typename F>
+ T MutateInternal(const T& node, F fmutate) {
+ if (only_enable_.size() &&
+ !only_enable_.count(node->type_index())) {
+ return fmutate(node);
+ }
+ if (f_preorder_ != nullptr) {
+ T pre = f_preorder_(node);
+ if (pre.defined()) return pre;
+ }
+ T new_node = fmutate(node);
+ if (f_postorder_ != nullptr) {
+ T post = f_postorder_(new_node);
+ if (post.defined()) return post;
+ }
+ return new_node;
+ }
+ // The functions
+ const runtime::PackedFunc& f_preorder_;
+ const runtime::PackedFunc& f_postorder_;
+ // type indices enabled.
+ const std::unordered_set<uint32_t>& only_enable_;
+};
+
+Stmt IRTransform(Stmt ir_node,
+ const runtime::PackedFunc& f_preorder,
+ const runtime::PackedFunc& f_postorder,
+ const Array<Expr>& only_enable) {
+ std::unordered_set<uint32_t> only_type_index;
+ for (Expr s : only_enable) {
+ only_type_index.insert(Object::TypeKey2Index(s.as<StringImm>()->value.c_str()));
+ }
+ IRTransformer transform(f_preorder, f_postorder, only_type_index);
+ return transform(std::move(ir_node));
+}
+
+// Implementation of Visitors
+template<typename T, typename F>
+inline void VisitArray(const Array<T>& arr, F fvisit) {
+ for (size_t i = 0; i < arr.size(); i++) {
+ fvisit(arr[i]);
+ }
+}
+
+void StmtVisitor::VisitStmt_(const LetStmt* op) {
+ this->VisitExpr(op->value);
+ this->VisitStmt(op->body);
+}
+
+void StmtVisitor::VisitStmt_(const AttrStmt* op) {
+ this->VisitExpr(op->value);
+ this->VisitStmt(op->body);
+}
+
+void StmtVisitor::VisitStmt_(const For* op) {
+ this->VisitExpr(op->min);
+ this->VisitExpr(op->extent);
+ this->VisitStmt(op->body);
+}
+
+void StmtVisitor::VisitStmt_(const Allocate* op) {
+ VisitArray(op->extents, [this](const Expr& e) { this->VisitExpr(e); });
+ this->VisitStmt(op->body);
+ this->VisitExpr(op->condition);
+ if (op->new_expr.defined()) {
+ this->VisitExpr(op->new_expr);
+ }
+}
+
+void StmtVisitor::VisitStmt_(const Store* op) {
+ this->VisitExpr(op->value);
+ this->VisitExpr(op->index);
+ this->VisitExpr(op->predicate);
+}
+
+void StmtVisitor::VisitStmt_(const IfThenElse* op) {
+ this->VisitExpr(op->condition);
+ this->VisitStmt(op->then_case);
+ if (op->else_case.defined()) {
+ this->VisitStmt(op->else_case);
+ }
+}
+
+void StmtVisitor::VisitStmt_(const Free* op) {}
+
+void StmtVisitor::VisitStmt_(const AssertStmt* op) {
+ this->VisitExpr(op->condition);
+ this->VisitExpr(op->message);
+ this->VisitStmt(op->body);
+}
+
+void StmtVisitor::VisitStmt_(const ProducerConsumer* op) {
+ this->VisitStmt(op->body);
+}
+
+void StmtVisitor::VisitStmt_(const Provide* op) {
+ VisitArray(op->args, [this](const Expr& e) { this->VisitExpr(e); });
+ this->VisitExpr(op->value);
+}
+
+void StmtVisitor::VisitStmt_(const Realize* op) {
+ VisitArray(op->bounds, [this](const Range& r) {
+ this->VisitExpr(r->min);
+ this->VisitExpr(r->extent);
+ });
+ this->VisitStmt(op->body);
+ this->VisitExpr(op->condition);
+}
+
+void StmtVisitor::VisitStmt_(const Prefetch* op) {
+ VisitArray(op->bounds, [this](const Range& r) {
+ this->VisitExpr(r->min);
+ this->VisitExpr(r->extent);
+ });
+}
+
+void StmtVisitor::VisitStmt_(const Block* op) {
+ this->VisitStmt(op->first);
+ this->VisitStmt(op->rest);
+}
+
+void StmtVisitor::VisitStmt_(const Evaluate* op) {
+ this->VisitExpr(op->value);
+}
+
+void ExprVisitor::VisitExpr_(const Variable* op) {}
+
+void ExprVisitor::VisitExpr_(const Load* op) {
+ this->VisitExpr(op->index);
+ this->VisitExpr(op->predicate);
+}
+
+void ExprVisitor::VisitExpr_(const Let* op) {
+ this->VisitExpr(op->value);
+ this->VisitExpr(op->body);
+}
+
+void ExprVisitor::VisitExpr_(const Call* op) {
+ VisitArray(op->args, [this](const Expr& e) { this->VisitExpr(e); });
+}
+
+#define DEFINE_BINOP_VISIT_(OP) \
+ void ExprVisitor::VisitExpr_(const OP* op) { \
+ this->VisitExpr(op->a); \
+ this->VisitExpr(op->b); \
+ }
+
+DEFINE_BINOP_VISIT_(Add);
+DEFINE_BINOP_VISIT_(Sub);
+DEFINE_BINOP_VISIT_(Mul);
+DEFINE_BINOP_VISIT_(Div);
+DEFINE_BINOP_VISIT_(Mod);
+DEFINE_BINOP_VISIT_(FloorDiv);
+DEFINE_BINOP_VISIT_(FloorMod);
+DEFINE_BINOP_VISIT_(Min);
+DEFINE_BINOP_VISIT_(Max);
+DEFINE_BINOP_VISIT_(EQ);
+DEFINE_BINOP_VISIT_(NE);
+DEFINE_BINOP_VISIT_(LT);
+DEFINE_BINOP_VISIT_(LE);
+DEFINE_BINOP_VISIT_(GT);
+DEFINE_BINOP_VISIT_(GE);
+DEFINE_BINOP_VISIT_(And);
+DEFINE_BINOP_VISIT_(Or);
+
+void ExprVisitor::VisitExpr_(const IntImm* op) {}
+void ExprVisitor::VisitExpr_(const UIntImm* op) {}
+void ExprVisitor::VisitExpr_(const FloatImm* op) {}
+void ExprVisitor::VisitExpr_(const StringImm* op) {}
+
+void ExprVisitor::VisitExpr_(const Reduce* op) {
+ VisitArray(op->axis, [this](const IterVar& r) {
+ this->VisitExpr(r->dom->min);
+ this->VisitExpr(r->dom->extent);
+ });
+ VisitArray(op->source, [this](const Expr& e) { this->VisitExpr(e); });
+ this->VisitExpr(op->condition);
+}
+
+void ExprVisitor::VisitExpr_(const Cast* op) {
+ this->VisitExpr(op->value);
+}
+
+void ExprVisitor::VisitExpr_(const Not* op) {
+ this->VisitExpr(op->a);
+}
+
+void ExprVisitor::VisitExpr_(const Select* op) {
+ this->VisitExpr(op->condition);
+ this->VisitExpr(op->true_value);
+ this->VisitExpr(op->false_value);
+}
+
+void ExprVisitor::VisitExpr_(const Ramp* op) {
+ this->VisitExpr(op->base);
+ this->VisitExpr(op->stride);
+}
+
+void ExprVisitor::VisitExpr_(const Shuffle* op) {
+ VisitArray(op->indices, [this](const Expr& e) { this->VisitExpr(e); });
+ VisitArray(op->vectors, [this](const Expr& e) { this->VisitExpr(e); });
+}
+
+void ExprVisitor::VisitExpr_(const Broadcast* op) {
+ this->VisitExpr(op->value);
+}
+
+// Implementation of mutators
+template<typename T, typename F>
+inline Array<T> MutateArray(const Array<T>& arr,
+ F fmutate,
+ bool allow_copy_on_write = false) {
+ if (allow_copy_on_write) {
+ // if we allow copy on write, we can directly
+ // call the inplace mutate function.
+ const_cast<Array<T>&>(arr).MutateByApply(fmutate);
+ return arr;
+ } else {
+ Array<T> copy = arr;
+ copy.MutateByApply(fmutate);
+ return copy;
+ }
+}
+
+class StmtMutator::Internal {
+ public:
+ static Array<Expr> Mutate(StmtMutator* self, const Array<Expr>& arr) {
+ auto fmutate = [self](const Expr& e) { return self->VisitExpr(e); };
+ return MutateArray(arr, fmutate, self->allow_copy_on_write_);
+ }
+
+ static Array<Stmt> Mutate(StmtMutator* self, const Array<Stmt>& arr) {
+ auto fmutate = [self](const Stmt& s) { return self->VisitStmt(s); };
+ return MutateArray(arr, fmutate, self->allow_copy_on_write_);
+ }
+
+ static Array<Range> Mutate(StmtMutator* self, const Array<Range>& arr) {
+ auto fmutate = [self](const Range& r) {
+ Expr min = self->VisitExpr(r->min);
+ Expr extent = self->VisitExpr(r->extent);
+ if (min.same_as(r->min) && extent.same_as(r->extent)) {
+ return r;
+ } else {
+ return Range::make_by_min_extent(min, extent);
+ }
+ };
+ return MutateArray(arr, fmutate, self->allow_copy_on_write_);
+ }
+};
+
+Stmt StmtMutator::VisitStmt_(const AttrStmt* op) {
+ Expr value = this->VisitExpr(op->value);
+ Stmt body = this->VisitStmt(op->body);
+ if (value.same_as(op->value) &&
+ body.same_as(op->body)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->value = std::move(value);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const LetStmt* op) {
+ Expr value = this->VisitExpr(op->value);
+ Stmt body = this->VisitStmt(op->body);
+ if (value.same_as(op->value) &&
+ body.same_as(op->body)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->value = std::move(value);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const For* op) {
+ Expr min = this->VisitExpr(op->min);
+ Expr extent = this->VisitExpr(op->extent);
+ Stmt body = this->VisitStmt(op->body);
+ if (min.same_as(op->min) &&
+ extent.same_as(op->extent) &&
+ body.same_as(op->body)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->min = std::move(min);
+ n->extent = std::move(extent);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Allocate* op) {
+ Array<Expr> extents = Internal::Mutate(this, op->extents);
+ Stmt body = this->VisitStmt(op->body);
+ Expr condition = this->VisitExpr(op->condition);
+ Expr new_expr;
+ if (op->new_expr.defined()) {
+ new_expr = this->VisitExpr(op->new_expr);
+ }
+ if (extents.same_as(op->extents) &&
+ body.same_as(op->body) &&
+ condition.same_as(op->condition) &&
+ new_expr.same_as(op->new_expr)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->extents = std::move(extents);
+ n->body = std::move(body);
+ n->condition = std::move(condition);
+ n->new_expr = std::move(new_expr);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const IfThenElse* op) {
+ Expr condition = this->VisitExpr(op->condition);
+ Stmt then_case = this->VisitStmt(op->then_case);
+ Stmt else_case;
+ if (op->else_case.defined()) {
+ else_case = this->VisitStmt(op->else_case);
+ }
+ if (condition.same_as(op->condition) &&
+ then_case.same_as(op->then_case) &&
+ else_case.same_as(op->else_case)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->condition = std::move(condition);
+ n->then_case = std::move(then_case);
+ n->else_case = std::move(then_case);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Store* op) {
+ Expr value = this->VisitExpr(op->value);
+ Expr index = this->VisitExpr(op->index);
+ Expr predicate = this->VisitExpr(op->predicate);
+ if (value.same_as(op->value) &&
+ index.same_as(op->index) &&
+ predicate.same_as(op->predicate)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->value = std::move(value);
+ n->index = std::move(index);
+ n->predicate = std::move(predicate);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Provide* op) {
+ Array<Expr> args = Internal::Mutate(this, op->args);
+ Expr value = this->VisitExpr(op->value);
+ if (args.same_as(op->args) &&
+ value.same_as(op->value)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->args = std::move(args);
+ n->value = std::move(value);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Realize* op) {
+ Region bounds = Internal::Mutate(this, op->bounds);
+ Stmt body = this->VisitStmt(op->body);
+ Expr condition = this->VisitExpr(op->condition);
+ if (bounds.same_as(op->bounds) &&
+ body.same_as(op->body) &&
+ condition.same_as(op->condition)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->bounds = std::move(bounds);
+ n->body = std::move(body);
+ n->condition = std::move(condition);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Prefetch* op) {
+ Region bounds = Internal::Mutate(this, op->bounds);
+ if (bounds.same_as(op->bounds)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->bounds = std::move(bounds);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Block* op) {
+ Stmt first = this->VisitStmt(op->first);
+ Stmt rest = this->VisitStmt(op->rest);
+ if (first.same_as(op->first) &&
+ rest.same_as(op->rest)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->first = std::move(first);
+ n->rest = std::move(rest);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const AssertStmt* op) {
+ Expr condition = this->VisitExpr(op->condition);
+ Expr message = this->VisitExpr(op->message);
+ Stmt body = this->VisitStmt(op->body);
+
+ if (condition.same_as(op->condition) &&
+ message.same_as(op->message) &&
+ body.same_as(op->body)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->condition = std::move(condition);
+ n->message = std::move(message);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const ProducerConsumer* op) {
+ Stmt body = this->VisitStmt(op->body);
+ if (body.same_as(op->body)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->body = std::move(body);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Evaluate* op) {
+ Expr value = this->VisitExpr(op->value);
+ if (value.same_as(op->value)) {
+ return GetRef<Stmt>(op);
+ } else {
+ auto n = CopyOnWrite(op);
+ n->value = std::move(value);
+ return Stmt(n);
+ }
+}
+
+Stmt StmtMutator::VisitStmt_(const Free* op) {
+ return GetRef<Stmt>(op);
+}
+
+
+Expr ExprMutator::VisitExpr_(const Variable* op) {
+ return GetRef<Expr>(op);
+}
+
+Expr ExprMutator::VisitExpr_(const Load* op) {
+ Expr index = this->VisitExpr(op->index);
+ Expr predicate = this->VisitExpr(op->predicate);
+ if (index.same_as(op->index) && predicate.same_as(op->predicate)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Load::make(op->dtype, op->buffer_var, index, predicate);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Let* op) {
+ Expr value = this->VisitExpr(op->value);
+ Expr body = this->VisitExpr(op->body);
+ if (value.same_as(op->value) &&
+ body.same_as(op->body)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Let::make(op->var, value, body);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Call* op) {
+ auto fmutate = [this](const Expr& e) { return this->VisitExpr(e); };
+ Array<Expr> args = MutateArray(op->args, fmutate);
+
+ if (args.same_as(op->args)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Call::make(op->dtype,
+ op->name,
+ args,
+ op->call_type,
+ op->func,
+ op->value_index);
+ }
+}
+
+#define DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(OP) \
+ Expr ExprMutator::VisitExpr_(const OP *op) { \
+ return GetRef<Expr>(op); \
+ }
+
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(IntImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(UIntImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(FloatImm)
+DEFINE_OP_RETURN_SELF_EXPR_MUTATE_(StringImm)
+
+#define DEFINE_BIOP_EXPR_MUTATE_(OP) \
+ Expr ExprMutator::VisitExpr_(const OP* op) { \
+ Expr a = this->VisitExpr(op->a); \
+ Expr b = this->VisitExpr(op->b); \
+ if (a.same_as(op->a) && \
+ b.same_as(op->b)) { \
+ return GetRef<Expr>(op); \
+ } else { \
+ return OP::make(a, b); \
+ } \
+ }
+
+DEFINE_BIOP_EXPR_MUTATE_(Add);
+DEFINE_BIOP_EXPR_MUTATE_(Sub);
+DEFINE_BIOP_EXPR_MUTATE_(Mul);
+DEFINE_BIOP_EXPR_MUTATE_(Div);
+DEFINE_BIOP_EXPR_MUTATE_(Mod);
+DEFINE_BIOP_EXPR_MUTATE_(FloorDiv);
+DEFINE_BIOP_EXPR_MUTATE_(FloorMod);
+DEFINE_BIOP_EXPR_MUTATE_(Min);
+DEFINE_BIOP_EXPR_MUTATE_(Max);
+DEFINE_BIOP_EXPR_MUTATE_(EQ);
+DEFINE_BIOP_EXPR_MUTATE_(NE);
+DEFINE_BIOP_EXPR_MUTATE_(LT);
+DEFINE_BIOP_EXPR_MUTATE_(LE);
+DEFINE_BIOP_EXPR_MUTATE_(GT);
+DEFINE_BIOP_EXPR_MUTATE_(GE);
+DEFINE_BIOP_EXPR_MUTATE_(And);
+DEFINE_BIOP_EXPR_MUTATE_(Or);
+
+Expr ExprMutator::VisitExpr_(const Reduce* op) {
+ auto fitervar = [this](const IterVar& v) {
+ Range r = v->dom;
+ Expr min = this->VisitExpr(r->min);
+ Expr extent = this->VisitExpr(r->extent);
+ if (min.same_as(r->min) &&
+ extent.same_as(r->extent)) {
+ return v;
+ } else {
+ return IterVarNode::make(
+ Range::make_by_min_extent(min, extent),
+ v->var, v->iter_type, v->thread_tag);
+ }
+ };
+ Array<IterVar> axis = MutateArray(op->axis, fitervar);
+
+ auto fexpr = [this](const Expr& e) { return this->VisitExpr(e); };
+ Array<Expr> source = MutateArray(op->source, fexpr);
+
+ Expr condition = this->VisitExpr(op->condition);
+
+ if (axis.same_as(op->axis) &&
+ source.same_as(op->source) &&
+ condition.same_as(op->condition)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Reduce::make(
+ op->combiner, source, axis, condition, op->value_index);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Cast* op) {
+ Expr value = this->VisitExpr(op->value);
+ if (value.same_as(op->value)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Cast::make(op->dtype, value);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Not* op) {
+ Expr a = this->VisitExpr(op->a);
+ if (a.same_as(op->a)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Not::make(a);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Select* op) {
+ Expr condition = this->VisitExpr(op->condition);
+ Expr true_value = this->VisitExpr(op->true_value);
+ Expr false_value = this->VisitExpr(op->false_value);
+ if (condition.same_as(op->condition) &&
+ true_value.same_as(op->true_value) &&
+ false_value.same_as(op->false_value)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Select::make(condition, true_value, false_value);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Ramp* op) {
+ Expr base = this->VisitExpr(op->base);
+ Expr stride = this->VisitExpr(op->stride);
+ if (base.same_as(op->base) &&
+ stride.same_as(op->stride)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Ramp::make(base, stride, op->lanes);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Broadcast* op) {
+ Expr value = this->VisitExpr(op->value);
+ if (value.same_as(op->value)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Broadcast::make(value, op->lanes);
+ }
+}
+
+Expr ExprMutator::VisitExpr_(const Shuffle* op) {
+ auto fexpr = [this](const Expr& e) { return this->VisitExpr(e); };
+ auto vectors = MutateArray(op->vectors, fexpr);
+ if (vectors.same_as(op->vectors)) {
+ return GetRef<Expr>(op);
+ } else {
+ return Shuffle::make(vectors, op->indices);
+ }
+}
+
+} // namespace ir
+} // namespace tvm