Skip to content

T
tic
Unverified Commit 203ca7a0 by Tianqi Chen Committed by GitHub
Options

[REFACTOR] Migrate Low-level IR Passes into the New Stmt/Expr Mutator (#4607) 

* CombineContextCall

* Migrate BoundChecker

* Migrate CoprocSync

* Migrate detect_device

* Migrate loop_partition

* Migrate infer_fragement

* Migrate inject_copy_intrin

* Migrate inject double buffer

* Migrate lower_intrin and simplify

* Migrate storage flatten

* Migrate inject prefetch

* Migrate inject_virtual_thread

* migrate inline

* Migrate lift attr scope

* Migrate custom datatypes

* migrate lower_thread_all_reduce

* Migrate lower_tvm_builtin

* migrate lower_warp memory

* Migrate make_api.cc

* Migrate remap_thread_axis

* Migrate remove_no_op

* migrate rewrite_unsafe_select

* Migrate skip_assert simple_passes

* Migrate split_host_device

* Migrate ssa

* Migrate storage_access

* Migrate storage_rewrite

* Migrate tensor_core

* Migrate unroll_loop

* Migrate vectorize

* Migrate verify compact_buffer gpu_code

* Migrate verify_memory

* Migrate storage_sync

* Remove unused refs to mutator

* Migrate hybrid_op

* Migrate tensorize

* Migrate schedule ops

* Migrate schedule_dataflow_rewrite

* Migrate auto_inline_elemwise

* Remove unecessary ref to visitor

* remove unecessary ref

* Migrate bound_deducer

* Migrate domain_touched

* Migrate autotvm feature touch extractor

* Add annotations
parent 3f43bee0
Showing with 873 additions and 868 deletions
include/tvm/ir_functor_ext.h
......@@ -546,6 +546,35 @@ class StmtExprMutator :
}
};
/*!
* \brief recursively visit the ir in post DFS order node, and transform it
*
* \param node The ir to be transformed.
* \param preorder The function called in before recursive mutation
* If preorder returns None, then the transform will proceed to recursive call.
* If preorder returns a not None Stmt/Expr, the transformer will simply return it and
* won't do further recursion.
* \param postorder The function called after recursive mutation.
* The recursive mutation result is passed to postorder for further mutation.
* \param only_enable List of StringImm.
* If it is empty, all IRNode will call preorder/postorder
* If it is not empty, preorder/postorder will only be called
* when the IRNode's type key is in the list.
*/
TVM_DLL Stmt IRTransform(Stmt node,
const runtime::PackedFunc& preorder,
const runtime::PackedFunc& postorder,
const Array<Expr>& only_enable = {});
/*!
* \brief recursively visit the ir in post DFS order node, apply fvisit
* Each node is guaranteed to be visited only once.
* \param node The ir to be visited.
* \param fvisit The visitor function to be applied.
*/
TVM_DLL void PostOrderVisit(const ObjectRef& node, std::function<void(const ObjectRef&)> fvisit);
} // namespace ir
} // namespace tvm
#endif // TVM_IR_FUNCTOR_EXT_H_
include/tvm/ir_mutator.h
......@@ -122,27 +122,6 @@ class TVM_DLL IRMutator {
virtual Expr Mutate_(const StringImm* op, const Expr& e);
virtual Expr Mutate_(const Shuffle* op, const Expr& e);
};
/*!
* \brief recursively visit the ir in post DFS order node, and transform it
*
* \param node The ir to be transformed.
* \param preorder The function called in before recursive mutation
* If preorder returns None, then the transform will proceed to recursive call.
* If preorder returns a not None Stmt/Expr, the transformer will simply return it and
* won't do further recursion.
* \param postorder The function called after recursive mutation.
* The recursive mutation result is passed to postorder for further mutation.
* \param only_enable List of StringImm.
* If it is empty, all IRNode will call preorder/postorder
* If it is not empty, preorder/postorder will only be called
* when the IRNode's type key is in the list.
*/
Stmt IRTransform(Stmt node,
const runtime::PackedFunc& preorder,
const runtime::PackedFunc& postorder,
const Array<Expr>& only_enable = {});
} // namespace ir
} // namespace tvm
#endif // TVM_IR_MUTATOR_H_
include/tvm/ir_visitor.h
......@@ -145,15 +145,6 @@ class TVM_DLL IRVisitor {
virtual void Visit_(const FloatImm* op);
virtual void Visit_(const StringImm* op);
};
/*!
* \brief recursively visit the ir in post DFS order node, apply fvisit
* Each node is guaranteed to be visited only once.
* \param node The ir to be visited.
* \param fvisit The visitor function to be applied.
*/
TVM_DLL void PostOrderVisit(const ObjectRef& node, std::function<void(const ObjectRef&)> fvisit);
} // namespace ir
} // namespace tvm
......
src/api/api_pass.cc
......@@ -25,8 +25,7 @@
#include <tvm/ir.h>
#include <tvm/attrs.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/api_registry.h>
namespace tvm {
......
src/arithmetic/bound_deducer.cc
......@@ -23,7 +23,7 @@
*/
#include <tvm/expr.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/arithmetic.h>
#include <tvm/api_registry.h>
......@@ -38,17 +38,17 @@ using namespace ir;
// a visitor to find the path to the target variable
// from a expression.
class VariablePathFinder: public IRVisitor {
class VariablePathFinder: public ExprVisitor {
public:
explicit VariablePathFinder(Expr target) : target_(target) {}
void Visit(const ObjectRef& node) final {
void VisitExpr(const Expr& node) final {
if (visited_.count(node.get()) != 0) return;
visited_.insert(node.get());
if (!found_) path_.push_back(node.get());
if (node.same_as(target_)) found_ = true;
IRVisitor::Visit(node);
ExprVisitor::VisitExpr(node);
if (!found_) path_.pop_back();
}
......@@ -64,14 +64,14 @@ class VariablePathFinder: public IRVisitor {
// return empty vector to represent failure
std::vector<const Object*> GetPath(Expr target, Expr expr) {
VariablePathFinder v(target);
v.Visit(expr);
v(expr);
return v.path_;
}
enum CompareOp {kGreater, kLess, kEqual};
// a visitor to deduce the bound of a variable from a expression
class BoundDeducer: public IRVisitor {
class BoundDeducer: public ExprVisitor {
public:
friend class BoundDeduceInputChecker;
friend class Converter;
......@@ -82,39 +82,39 @@ class BoundDeducer: public IRVisitor {
void Deduce();
void Visit(const ObjectRef& e) final {
void VisitExpr(const Expr& e) final {
if (!success_) return;
if (e.get() == path_[iter_++]) {
IRVisitor::Visit(e);
ExprVisitor::VisitExpr(e);
} else {
success_ = false;
return;
}
}
void Visit_(const LT* op) final {
void VisitExpr_(const LT* op) final {
LOG(FATAL) << "unable to deduce due to multiple comparison operator";
}
void Visit_(const LE* op) final {
void VisitExpr_(const LE* op) final {
LOG(FATAL) << "unable to deduce due to multiple comparison operator";
}
void Visit_(const GT* op) final {
void VisitExpr_(const GT* op) final {
LOG(FATAL) << "unable to deduce due to multiple comparison operator";
}
void Visit_(const GE* op) final {
void VisitExpr_(const GE* op) final {
LOG(FATAL) << "unable to deduce due to multiple comparison operator";
}
void Visit_(const Add* op) final {
void VisitExpr_(const Add* op) final {
bool left = op->a.get() == path_[iter_];
result_ -= left ? op->b : op->a;
Visit(left ? op->a : op->b);
this->VisitExpr(left ? op->a : op->b);
}
void Visit_(const Sub* op) final {
void VisitExpr_(const Sub* op) final {
bool left = op->a.get() == path_[iter_];
if (left) {
result_ += op->b;
......@@ -123,10 +123,10 @@ class BoundDeducer: public IRVisitor {
result_ = - result_;
comp_op = ReverseOp(comp_op);
}
Visit(left ? op->a : op->b);
this->VisitExpr(left ? op->a : op->b);
}
void Visit_(const Mul* op) final {
void VisitExpr_(const Mul* op) final {
bool left = op->a.get() == path_[iter_];
Expr operand = left ? op->b : op->a;
Expr target_var = left ? op->a : op->b;
......@@ -171,7 +171,7 @@ class BoundDeducer: public IRVisitor {
// ( x <= -3/-2 --> x <= 1)
}
}
Visit(left ? op->a : op->b);
this->VisitExpr(left ? op->a : op->b);
}
Expr result_;
......@@ -194,17 +194,17 @@ class BoundDeducer: public IRVisitor {
Analyzer analyzer_;
};
class BoundDeduceInputChecker: public IRVisitor {
class BoundDeduceInputChecker: public ExprVisitor {
public:
bool Check(BoundDeducer* deducer) {
deducer_ = deducer;
Visit(deducer_->expr_);
this->VisitExpr(deducer_->expr_);
return target_count == 1;
}
void Visit(const ObjectRef& e) final {
void VisitExpr(const Expr& e) final {
if (e.same_as(deducer_->target_)) ++target_count;
IRVisitor::Visit(e);
ExprVisitor::VisitExpr(e);
}
private:
......@@ -305,7 +305,7 @@ void BoundDeducer::Deduce() {
}
expr_map_ = EvalSetForEachSubExpr(expr_, hint_map_);
Visit(expr_);
this->VisitExpr(expr_);
}
void BoundDeducer::Relax() {
......
src/arithmetic/canonical_simplify.cc
......@@ -23,7 +23,6 @@
*/
#include <tvm/arithmetic.h>
#include <tvm/expr_operator.h>
#include <tvm/ir_mutator.h>
#include "const_fold.h"
#include "pattern_match.h"
#include "rewrite_simplify.h"
......@@ -435,30 +434,30 @@ class CanonicalSimplifier::Impl : public RewriteSimplifier::Impl {
Expr CanonicalSimplify(Expr expr) {
expr = Mutate(expr);
expr = operator()(expr);
return expr;
}
// override the original mutate function.
Expr Mutate(Expr expr) final {
expr = IRMutator::Mutate(expr);
Expr VisitExpr(const Expr& input_expr) final {
auto expr = Rewriter::VisitExpr(input_expr);
return Normalize(expr);
}
// Normal mutation without normalization.
Expr CanonicalMutate(Expr expr) {
return IRMutator::Mutate(expr);
return Rewriter::VisitExpr(expr);
}
using Rewriter::Mutate_;
Expr Mutate_(const Add* op, const Expr& self) final;
Expr Mutate_(const Sub* op, const Expr& self) final;
Expr Mutate_(const Mul* op, const Expr& self) final;
Expr Mutate_(const Div* op, const Expr& self) final;
Expr Mutate_(const Mod* op, const Expr& self) final;
Expr Mutate_(const FloorDiv* op, const Expr& self) final;
Expr Mutate_(const FloorMod* op, const Expr& self) final;
Expr Mutate_(const Reduce* op, const Expr& self) final;
using Rewriter::VisitExpr_;
Expr VisitExpr_(const Add* op) final;
Expr VisitExpr_(const Sub* op) final;
Expr VisitExpr_(const Mul* op) final;
Expr VisitExpr_(const Div* op) final;
Expr VisitExpr_(const Mod* op) final;
Expr VisitExpr_(const FloorDiv* op) final;
Expr VisitExpr_(const FloorMod* op) final;
Expr VisitExpr_(const Reduce* op) final;
private:
/*!
......@@ -567,9 +566,9 @@ class CanonicalSimplifier::Impl : public RewriteSimplifier::Impl {
};
Expr CanonicalSimplifier::Impl::
Mutate_(const Add* op, const Expr& self) {
VisitExpr_(const Add* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
// normalize
Expr a = this->CanonicalMutate(op->a);
......@@ -593,9 +592,9 @@ Mutate_(const Add* op, const Expr& self) {
}
Expr CanonicalSimplifier::Impl::
Mutate_(const Sub* op, const Expr& self) {
VisitExpr_(const Sub* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
// normalize
Expr a = this->CanonicalMutate(op->a);
......@@ -620,9 +619,9 @@ Mutate_(const Sub* op, const Expr& self) {
Expr CanonicalSimplifier::Impl::
Mutate_(const Mul* op, const Expr& self) {
VisitExpr_(const Mul* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
// normalize
Expr a = this->CanonicalMutate(op->a);
......@@ -652,7 +651,7 @@ Mutate_(const Mul* op, const Expr& self) {
a = Normalize(a);
b = Normalize(b);
if (op->a.same_as(a) && op->b.same_as(b)) {
return self;
return GetRef<Expr>(op);
} else {
return Mul::make(a, b);
}
......@@ -727,9 +726,9 @@ SplitDivConst(SplitExpr lhs, int64_t cval, DivMode div_mode) {
}
Expr CanonicalSimplifier::Impl::
Mutate_(const Div* op, const Expr& self) {
VisitExpr_(const Div* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
Expr a = this->CanonicalMutate(op->a);
......@@ -781,16 +780,16 @@ Mutate_(const Div* op, const Expr& self) {
a = Normalize(a);
b = Normalize(b);
if (op->a.same_as(a) && op->b.same_as(b)) {
return self;
return GetRef<Expr>(op);
} else {
return Div::make(a, b);
}
}
Expr CanonicalSimplifier::Impl::
Mutate_(const FloorDiv* op, const Expr& self) {
VisitExpr_(const FloorDiv* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
Expr a = this->CanonicalMutate(op->a);
Expr b = this->CanonicalMutate(op->b);
......@@ -837,7 +836,7 @@ Mutate_(const FloorDiv* op, const Expr& self) {
a = Normalize(a);
b = Normalize(b);
if (op->a.same_as(a) && op->b.same_as(b)) {
return self;
return GetRef<Expr>(op);
} else {
return FloorDiv::make(a, b);
}
......@@ -866,7 +865,7 @@ SplitModConst(SplitExpr lhs, int64_t cval, DivMode div_mode) {
// Do a recursive call to simplify the mod with the new factor.
if (new_upper_factor < lhs->upper_factor &&
lhs->upper_factor != SplitExprNode::kPosInf) {
auto updated = ToSplitExpr(Mutate(ModImpl(
auto updated = ToSplitExpr(this->VisitExpr(ModImpl(
lhs->index, make_const(lhs.dtype(), new_upper_factor), div_mode)));
// re-apply the lower_factor
if (lhs->lower_factor != 1) {
......@@ -894,9 +893,9 @@ SplitModConst(SplitExpr lhs, int64_t cval, DivMode div_mode) {
}
Expr CanonicalSimplifier::Impl::
Mutate_(const Mod* op, const Expr& self) {
VisitExpr_(const Mod* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
// normalize
Expr a = this->CanonicalMutate(op->a);
......@@ -957,16 +956,16 @@ Mutate_(const Mod* op, const Expr& self) {
a = Normalize(a);
b = Normalize(b);
if (op->a.same_as(a) && op->b.same_as(b)) {
return self;
return GetRef<Expr>(op);
} else {
return Mod::make(a, b);
}
}
Expr CanonicalSimplifier::Impl::
Mutate_(const FloorMod* op, const Expr& self) {
VisitExpr_(const FloorMod* op) {
if (!IsIndexType(op->dtype)) {
return Rewriter::Mutate_(op, self);
return Rewriter::VisitExpr_(op);
}
// normalize
Expr a = this->CanonicalMutate(op->a);
......@@ -1017,7 +1016,7 @@ Mutate_(const FloorMod* op, const Expr& self) {
a = Normalize(a);
b = Normalize(b);
if (op->a.same_as(a) && op->b.same_as(b)) {
return self;
return GetRef<Expr>(op);
} else {
return FloorMod::make(a, b);
}
......@@ -1029,7 +1028,7 @@ SimplifyReduceCombiner(const Reduce* op) {
// First simplify the results
Array<Expr> simplified_result;
for (const auto& res : op->combiner->result) {
Expr new_res = Mutate(res);
Expr new_res = this->VisitExpr(res);
simplified_result.push_back(new_res);
}
......@@ -1078,7 +1077,7 @@ SimplifyReduceCombiner(const Reduce* op) {
if (used[i]) {
// We simplify the result and identity, but not the source
new_result.push_back(simplified_result[i]);
new_identity.push_back(Mutate(op->combiner->identity_element[i]));
new_identity.push_back(this->VisitExpr(op->combiner->identity_element[i]));
new_lhs.push_back(op->combiner->lhs[i]);
new_rhs.push_back(op->combiner->rhs[i]);
new_source.push_back(op->source[i]);
......@@ -1095,9 +1094,9 @@ SimplifyReduceCombiner(const Reduce* op) {
}
Expr CanonicalSimplifier::Impl::
Mutate_(const Reduce* op, const Expr& self) {
VisitExpr_(const Reduce* op) {
// Recursively call simplification when necessary.
Expr ret = RewriteSimplifier::Impl::Mutate_(op, self);
Expr ret = RewriteSimplifier::Impl::VisitExpr_(op);
op = ret.as<Reduce>();
// already been simplified by const reduction axis removal
if (op == nullptr) return ret;
......@@ -1106,7 +1105,7 @@ Mutate_(const Reduce* op, const Expr& self) {
// assumption we would have to perform a single iteration of the loop, i.e. use
// `(*op->combiner.get())(op->combineop->identity_element, op->source)[op->value_index]`
// instead of `op->source[op->value_index]`. The former may be more difficult to simplify.
return Mutate(
return this->VisitExpr(
Select::make(op->condition,
op->source[op->value_index],
op->combiner->identity_element[op->value_index]));
......
src/arithmetic/const_fold.h
......@@ -25,7 +25,6 @@
#define TVM_ARITHMETIC_CONST_FOLD_H_
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/expr_operator.h>
#include <algorithm>
#include <cmath>
......
src/arithmetic/detect_linear_equation.cc
......@@ -23,7 +23,6 @@
*/
#include <tvm/expr.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/arithmetic.h>
......
src/arithmetic/domain_touched.cc
......@@ -6,9 +6,9 @@
* 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
......@@ -23,7 +23,7 @@
*/
#include <tvm/expr.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/tensor.h>
#include <tvm/api_registry.h>
......@@ -36,13 +36,13 @@ namespace arith {
using namespace ir;
// Find Read region of the tensor in the stmt.
class FuncTouchedDomain final : public IRVisitor {
class FuncTouchedDomain final : public StmtExprVisitor {
public:
FuncTouchedDomain(const Tensor &tensor, bool consider_calls, bool consider_provides)
: tensor_(tensor), consider_calls_(consider_calls), consider_provides_(consider_provides) {}
Domain Find(const Stmt& stmt) {
this->Visit(stmt);
operator()(stmt);
Domain ret;
Range none;
for (size_t i = 0; i < bounds_.size(); ++i) {
......@@ -51,49 +51,49 @@ class FuncTouchedDomain final : public IRVisitor {
return ret;
}
void Visit_(const For *op) final {
void VisitStmt_(const For *op) final {
const Variable* var = op->loop_var.get();
dom_map_[var] = IntSet::range(
Range::make_by_min_extent(op->min, op->extent));
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
dom_map_.erase(var);
}
void Visit_(const LetStmt* op) final {
void VisitStmt_(const LetStmt* op) final {
dom_map_[op->var.get()] =
arith::EvalSet(op->value, dom_map_);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
dom_map_.erase(op->var.get());
}
/* TODO: Thread extent unitest not generated.*/
void Visit_(const AttrStmt* op) final {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::thread_extent) {
const IterVarNode* thread_axis = op->node.as<IterVarNode>();
CHECK(thread_axis);
const Variable* var = thread_axis->var.get();
dom_map_[var] = IntSet::range(Range(make_zero(op->value.dtype()), op->value));
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
dom_map_.erase(var);
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
void Visit_(const Call* op) final {
void VisitExpr_(const Call* op) final {
if (consider_calls_ && tensor_->op.same_as(op->func)
&& tensor_->value_index == op->value_index) {
Touch(op->args);
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
void Visit_(const Provide* op) final {
void VisitStmt_(const Provide* op) final {
if (consider_provides_ && tensor_->op.same_as(op->func)
&& tensor_->value_index == op->value_index) {
Touch(op->args);
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
private:
......
src/arithmetic/ir_mutator_with_analyzer.cc
......@@ -30,41 +30,44 @@ namespace arith {
using namespace ir;
Stmt IRMutatorWithAnalyzer::
Mutate_(const For* op, const Stmt& s) {
VisitStmt_(const For* op) {
analyzer_->Bind(op->loop_var,
Range::make_by_min_extent(op->min, op->extent));
return IRMutator::Mutate_(op, s);
Range::make_by_min_extent(op->min, op->extent));
return StmtExprMutator::VisitStmt_(op);
}
Stmt IRMutatorWithAnalyzer::
Mutate_(const LetStmt* op, const Stmt& s) {
Expr value = this->Mutate(op->value);
VisitStmt_(const LetStmt* op) {
Expr value = this->VisitExpr(op->value);
if (!ir::HasSideEffect(value)) {
analyzer_->Bind(op->var, value);
}
// We keep the let-binding here
// as sub-class may or maynot choose to replace it.
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
if (value.same_as(op->value) &&
body.same_as(op->body)) {
return s;
return GetRef<Stmt>(op);
} else {
return LetStmt::make(op->var, value, body);
auto n = this->CopyOnWrite(op);
n->value = std::move(value);
n->body = std::move(body);
return Stmt(n);
}
}
Stmt IRMutatorWithAnalyzer::
Mutate_(const IfThenElse* op, const Stmt& s) {
Expr condition = this->Mutate(op->condition);
VisitStmt_(const IfThenElse* op) {
Expr condition = this->VisitExpr(op->condition);
Stmt then_case, else_case;
{
With<ConstraintContext> ctx(analyzer_, condition);
then_case = this->Mutate(op->then_case);
then_case = this->VisitStmt(op->then_case);
}
if (op->else_case.defined()) {
With<ConstraintContext> ctx(analyzer_,
analyzer_->rewrite_simplify(Not::make(condition)));
else_case = this->Mutate(op->else_case);
else_case = this->VisitStmt(op->else_case);
}
if (is_one(condition)) return then_case;
if (is_zero(condition)) {
......@@ -77,57 +80,65 @@ Mutate_(const IfThenElse* op, const Stmt& s) {
if (condition.same_as(op->condition) &&
then_case.same_as(op->then_case) &&
else_case.same_as(op->else_case)) {
return s;
return GetRef<Stmt>(op);
} else {
return IfThenElse::make(condition, then_case, else_case);
auto n = this->CopyOnWrite(op);
n->condition = std::move(condition);
n->then_case = std::move(then_case);
n->else_case = std::move(else_case);
return Stmt(n);
}
}
Stmt IRMutatorWithAnalyzer::
Mutate_(const AttrStmt* op, const Stmt& s) {
VisitStmt_(const AttrStmt* op) {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::virtual_thread) {
IterVar iv = Downcast<IterVar>(op->node);
CHECK_NE(iv->thread_tag.length(), 0U);
analyzer_->Bind(iv->var,
Range::make_by_min_extent(0, op->value));
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
return stmt;
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt IRMutatorWithAnalyzer::
Mutate_(const AssertStmt* op, const Stmt& s) {
Expr condition = this->Mutate(op->condition);
Expr message = this->Mutate(op->message);
VisitStmt_(const AssertStmt* op) {
Expr condition = this->VisitExpr(op->condition);
Expr message = this->VisitExpr(op->message);
With<ConstraintContext> ctx(analyzer_, condition);
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
if (condition.same_as(op->condition) &&
message.same_as(op->message) &&
body.same_as(op->body)) {
return s;
return GetRef<Stmt>(op);
} else {
return AssertStmt::make(condition, message, body);
auto n = this->CopyOnWrite(op);
n->condition = std::move(condition);
n->message = std::move(message);
n->body = std::move(body);
return Stmt(n);
}
}
Expr IRMutatorWithAnalyzer::
Mutate_(const Call* op, const Expr& self) {
VisitExpr_(const Call* op) {
// add condition context to if_then_else
if (op->is_intrinsic(ir::intrinsic::tvm_if_then_else)) {
Expr cond = Mutate(op->args[0]);
Expr cond = this->VisitExpr(op->args[0]);
Expr true_value, false_value;
{
With<ConstraintContext> constraint(analyzer_, cond);
true_value = Mutate(op->args[1]);
true_value = this->VisitExpr(op->args[1]);
}
{
With<ConstraintContext> constraint(analyzer_,
analyzer_->rewrite_simplify(Not::make(cond)));
false_value = Mutate(op->args[2]);
false_value = this->VisitExpr(op->args[2]);
}
if (is_zero(cond)) {
return false_value;
......@@ -138,45 +149,45 @@ Mutate_(const Call* op, const Expr& self) {
if (cond.same_as(op->args[0]) &&
true_value.same_as(op->args[1]) &&
false_value.same_as(op->args[2])) {
return self;
return GetRef<Expr>(op);
} else {
return Call::make(op->dtype, op->name,
{cond, true_value, false_value},
op->call_type);
}
}
return IRMutator::Mutate_(op, self);
return StmtExprMutator::VisitExpr_(op);
}
Expr IRMutatorWithAnalyzer::
Mutate_(const Let* op, const Expr& self) {
Expr value = this->Mutate(op->value);
VisitExpr_(const Let* op) {
Expr value = this->VisitExpr(op->value);
if (!ir::HasSideEffect(value)) {
analyzer_->Bind(op->var, value);
}
// We keep the let-binding here
// as sub-class may or maynot choose to replace it.
Expr body = this->Mutate(op->body);
Expr body = this->VisitExpr(op->body);
if (value.same_as(op->value) &&
body.same_as(op->body)) {
return self;
return GetRef<Expr>(op);
} else {
return Let::make(op->var, value, body);
}
}
Expr IRMutatorWithAnalyzer::
Mutate_(const Select* op, const Expr& self) {
Expr cond = Mutate(op->condition);
VisitExpr_(const Select* op) {
Expr cond = this->VisitExpr(op->condition);
Expr true_value, false_value;
{
With<ConstraintContext> constraint(analyzer_, cond);
true_value = Mutate(op->true_value);
true_value = VisitExpr(op->true_value);
}
{
With<ConstraintContext> constraint(analyzer_,
analyzer_->rewrite_simplify(Not::make(cond)));
false_value = Mutate(op->false_value);
false_value = VisitExpr(op->false_value);
}
if (is_zero(cond)) {
return false_value;
......@@ -188,20 +199,20 @@ Mutate_(const Select* op, const Expr& self) {
if (cond.same_as(op->condition) &&
true_value.same_as(op->true_value) &&
false_value.same_as(op->false_value)) {
return self;
return GetRef<Expr>(op);
} else {
return Select::make(cond, true_value, false_value);
}
}
Expr IRMutatorWithAnalyzer::
Mutate_(const Reduce* op, const Expr& self) {
VisitExpr_(const Reduce* op) {
// Setup the domain information before simplification.
for (const IterVar& iv : op->axis) {
analyzer_->Bind(iv->var, iv->dom);
}
// Recursively call simplification when necessary.
return IRMutator::Mutate_(op, self);
return StmtExprMutator::VisitExpr_(op);
}
} // namespace arith
......
src/arithmetic/ir_mutator_with_analyzer.h
......@@ -24,9 +24,9 @@
#ifndef TVM_ARITHMETIC_IR_MUTATOR_WITH_ANALYZER_H_
#define TVM_ARITHMETIC_IR_MUTATOR_WITH_ANALYZER_H_
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/arithmetic.h>
#include <utility>
namespace tvm {
namespace arith {
......@@ -40,23 +40,24 @@ namespace arith {
*
* \sa src/arithmetic/ir_mutator_with_analyzer.cc
*/
class IRMutatorWithAnalyzer : public ir::IRMutator {
class IRMutatorWithAnalyzer : public ir::StmtExprMutator {
public:
explicit IRMutatorWithAnalyzer(Analyzer* analyzer)
: analyzer_(analyzer) {}
using IRMutator::Mutate_;
using StmtExprMutator::VisitStmt_;
using StmtExprMutator::VisitExpr_;
// override functions that need to populate the context information.
Stmt Mutate_(const ir::For* op, const Stmt& self) override;
Stmt Mutate_(const ir::LetStmt* op, const Stmt& self) override;
Stmt Mutate_(const ir::IfThenElse* op, const Stmt& self) override;
Stmt Mutate_(const ir::AttrStmt* op, const Stmt& self) override;
Stmt Mutate_(const ir::AssertStmt* op, const Stmt& self) override;
Expr Mutate_(const ir::Let* op, const Expr& self) override;
Expr Mutate_(const ir::Select* op, const Expr& self) override;
Expr Mutate_(const ir::Call* op, const Expr& self) override;
Expr Mutate_(const ir::Reduce* op, const Expr& self) override;
Stmt VisitStmt_(const ir::For* op) override;
Stmt VisitStmt_(const ir::LetStmt* op) override;
Stmt VisitStmt_(const ir::IfThenElse* op) override;
Stmt VisitStmt_(const ir::AttrStmt* op) override;
Stmt VisitStmt_(const ir::AssertStmt* op) override;
Expr VisitExpr_(const ir::Let* op) override;
Expr VisitExpr_(const ir::Select* op) override;
Expr VisitExpr_(const ir::Call* op) override;
Expr VisitExpr_(const ir::Reduce* op) override;
protected:
/*! \brief internal analyzer field. */
......
src/arithmetic/ir_visitor_with_analyzer.h
......@@ -27,43 +27,43 @@
#include <tvm/arithmetic.h>
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
namespace tvm {
namespace ir {
class IRVisitorWithAnalyzer final : public IRVisitor {
class IRVisitorWithAnalyzer final : public StmtExprVisitor {
public:
Expr Simplify(const Expr& expr) {
return analyzer_.Simplify(expr);
}
void Visit_(const For* op) {
void VisitStmt_(const For* op) {
analyzer_.Bind(op->loop_var,
Range::make_by_min_extent(op->min, op->extent));
return IRVisitor::Visit_(op);
return StmtExprVisitor::VisitStmt_(op);
}
void Visit_(const AttrStmt* op) {
void VisitStmt_(const AttrStmt* op) {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::virtual_thread) {
IterVar iv = Downcast<IterVar>(op->node);
CHECK_NE(iv->thread_tag.length(), 0U);
analyzer_.Bind(iv->var,
Range::make_by_min_extent(0, op->value));
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
void Visit_(const Reduce* op) {
void VisitExpr_(const Reduce* op) {
// Setup the domain information before simplification.
for (const IterVar& iv : op->axis) {
analyzer_.Bind(iv->var, iv->dom);
}
// Recursively call simplification when necessary.
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
protected:
......
src/arithmetic/rewrite_simplify.cc
......@@ -24,7 +24,6 @@
// Acknowledgement: Most rewrite-rules are from Halide.
#include <tvm/arithmetic.h>
#include <tvm/expr_operator.h>
#include <tvm/ir_mutator.h>
#include <algorithm>
#include "const_fold.h"
#include "pattern_match.h"
......@@ -69,7 +68,7 @@ using namespace ir;
// try to prove x equals val
RewriteSimplifier::Impl::CompareResult RewriteSimplifier::Impl::
TryCompare(const Expr& x, int64_t val) {
Expr diff = Mutate(x);
Expr diff = this->VisitExpr(x);
if (const auto* ptr = diff.as<IntImm>()) {
if (ptr->value == val) {
return kEQ;
......@@ -117,8 +116,8 @@ Update(const Var& var, const Expr& info, bool override) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Add* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Add* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Add>();
Expr const_res = TryConstFold<Add>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -232,8 +231,8 @@ std::function<void()> RewriteSimplifier::Impl::EnterConstraint(const Expr& const
}
Expr RewriteSimplifier::Impl::
Mutate_(const Sub* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Sub* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Sub>();
Expr const_res = TryConstFold<Sub>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -431,8 +430,8 @@ Mutate_(const Sub* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Mul* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Mul* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Mul>();
Expr const_res = TryConstFold<Mul>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -470,8 +469,8 @@ Mutate_(const Mul* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Div* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Div* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Div>();
Expr const_res = TryConstFold<Div>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -692,8 +691,8 @@ Mutate_(const Div* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Mod* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Mod* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Mod>();
Expr const_res = TryConstFold<Mod>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -782,8 +781,8 @@ Mutate_(const Mod* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const FloorDiv* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const FloorDiv* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<FloorDiv>();
Expr const_res = TryConstFold<FloorDiv>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -926,8 +925,8 @@ Mutate_(const FloorDiv* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const FloorMod* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const FloorMod* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<FloorMod>();
Expr const_res = TryConstFold<FloorMod>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -996,8 +995,8 @@ Mutate_(const FloorMod* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Min* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Min* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Min>();
Expr const_res = TryConstFold<Min>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1181,8 +1180,8 @@ Mutate_(const Min* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Max* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Max* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Max>();
Expr const_res = TryConstFold<Max>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1354,8 +1353,8 @@ Mutate_(const Max* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const EQ* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const EQ* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<EQ>();
Expr const_res = TryConstFold<EQ>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1388,28 +1387,28 @@ Mutate_(const EQ* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const NE* op, const Expr& self) {
return Mutate(Not::make(op->a == op->b));
VisitExpr_(const NE* op) {
return this->VisitExpr(Not::make(op->a == op->b));
}
Expr RewriteSimplifier::Impl::
Mutate_(const LE* op, const Expr& self) {
return Mutate(Not::make(op->b < op->a));
VisitExpr_(const LE* op) {
return this->VisitExpr(Not::make(op->b < op->a));
}
Expr RewriteSimplifier::Impl::
Mutate_(const GT* op, const Expr& self) {
return Mutate(op->b < op->a);
VisitExpr_(const GT* op) {
return this->VisitExpr(op->b < op->a);
}
Expr RewriteSimplifier::Impl::
Mutate_(const GE* op, const Expr& self) {
return Mutate(Not::make(op->a < op->b));
VisitExpr_(const GE* op) {
return this->VisitExpr(Not::make(op->a < op->b));
}
Expr RewriteSimplifier::Impl::
Mutate_(const LT* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const LT* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<LT>();
Expr const_res = TryConstFold<LT>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1564,8 +1563,8 @@ Mutate_(const LT* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Not* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Not* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Not>();
Expr const_res = TryConstFold<Not>(op->a);
if (const_res.defined()) return const_res;
......@@ -1589,8 +1588,8 @@ Mutate_(const Not* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const And* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const And* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<And>();
Expr const_res = TryConstFold<And>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1638,8 +1637,8 @@ Mutate_(const And* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Or* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Or* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Or>();
Expr const_res = TryConstFold<Or>(op->a, op->b);
if (const_res.defined()) return const_res;
......@@ -1688,8 +1687,8 @@ Mutate_(const Or* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Select* op, const Expr& self) {
Expr ret = IRMutatorWithAnalyzer::Mutate_(op, self);
VisitExpr_(const Select* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Select>();
if (op == nullptr) return ret;
// Pattern var to match any expression
......@@ -1699,9 +1698,9 @@ Mutate_(const Select* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Call* op, const Expr& self) {
VisitExpr_(const Call* op) {
// add condition context to if_then_else
Expr ret = IRMutatorWithAnalyzer::Mutate_(op, self);
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Call>();
if (op == nullptr) return ret;
if (op->is_intrinsic(Call::likely) && is_const(op->args[0])) {
......@@ -1729,35 +1728,35 @@ Mutate_(const Call* op, const Expr& self) {
}
Expr RewriteSimplifier::Impl::
Mutate_(const Variable* op, const Expr& self) {
VisitExpr_(const Variable* op) {
Var var = GetRef<Var>(op);
auto it = var_map_.find(var);
if (it != var_map_.end()) {
return it->second;
}
return self;
return GetRef<Expr>(op);
}
Expr RewriteSimplifier::Impl::
Mutate_(const Cast* op, const Expr& self) {
Expr ret = IRMutator::Mutate_(op, self);
VisitExpr_(const Cast* op) {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<Cast>();
return cast(op->dtype, op->value);
}
Expr RewriteSimplifier::Impl::
Mutate_(const Let* op, const Expr& self) {
Expr value = this->Mutate(op->value);
VisitExpr_(const Let* op) {
Expr value = this->VisitExpr(op->value);
if (!ir::HasSideEffect(value)) {
// it is fine to discard the let binding
// because the value will always be inlined in the simplifier.
analyzer_->Bind(op->var, value);
return this->Mutate(op->body);
return this->VisitExpr(op->body);
}
Expr body = this->Mutate(op->body);
Expr body = this->VisitExpr(op->body);
if (value.same_as(op->value) &&
body.same_as(op->body)) {
return self;
return GetRef<Expr>(op);
} else {
return Let::make(op->var, value, body);
}
......@@ -1768,7 +1767,7 @@ Expr RewriteSimplifier::operator()(const Expr& expr) {
Expr res = expr;
int max_iter = 2;
for (int i = 0; i < max_iter; ++i) {
Expr new_expr = impl_->Mutate(res);
Expr new_expr = impl_->operator()(res);
if (new_expr.same_as(res)) return res;
res = new_expr;
}
......
src/arithmetic/rewrite_simplify.h
......@@ -26,7 +26,6 @@
#include <tvm/arithmetic.h>
#include <tvm/expr_operator.h>
#include <tvm/ir_mutator.h>
#include <unordered_map>
#include <vector>
#include "const_fold.h"
......@@ -45,35 +44,35 @@ using namespace ir;
*/
class RewriteSimplifier::Impl : public IRMutatorWithAnalyzer {
public:
using IRMutatorWithAnalyzer::Mutate_;
using IRMutatorWithAnalyzer::VisitExpr_;
explicit Impl(Analyzer* parent)
: IRMutatorWithAnalyzer(parent) {}
void Update(const Var& var, const Expr& info, bool override);
Expr Mutate_(const Add* op, const Expr& self) override;
Expr Mutate_(const Sub* op, const Expr& self) override;
Expr Mutate_(const Mul* op, const Expr& self) override;
Expr Mutate_(const Div* op, const Expr& self) override;
Expr Mutate_(const Mod* op, const Expr& self) override;
Expr Mutate_(const FloorDiv* op, const Expr& self) override;
Expr Mutate_(const FloorMod* op, const Expr& self) override;
Expr Mutate_(const Min* op, const Expr& self) override;
Expr Mutate_(const Max* op, const Expr& self) override;
Expr Mutate_(const EQ* op, const Expr& self) override;
Expr Mutate_(const NE* op, const Expr& self) override;
Expr Mutate_(const LT* op, const Expr& self) override;
Expr Mutate_(const LE* op, const Expr& self) override;
Expr Mutate_(const GT* op, const Expr& self) override;
Expr Mutate_(const GE* op, const Expr& self) override;
Expr Mutate_(const And* op, const Expr& self) override;
Expr Mutate_(const Or* op, const Expr& self) override;
Expr Mutate_(const Not* op, const Expr& self) override;
Expr Mutate_(const Select* op, const Expr& self) override;
Expr Mutate_(const Call* op, const Expr& self) override;
Expr Mutate_(const Variable* op, const Expr& self) override;
Expr Mutate_(const Cast* op, const Expr& self) override;
Expr Mutate_(const Let* op, const Expr& self) override;
void Update(const Var& var, const Expr& info, bool override_info);
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 Not* op) override;
Expr VisitExpr_(const Select* op) override;
Expr VisitExpr_(const Call* op) override;
Expr VisitExpr_(const Variable* op) override;
Expr VisitExpr_(const Cast* op) override;
Expr VisitExpr_(const Let* op) override;
std::function<void()> EnterConstraint(const Expr& constraint);
......@@ -123,7 +122,7 @@ class RewriteSimplifier::Impl : public IRMutatorWithAnalyzer {
Expr RecursiveRewrite(const Expr& x) {
if (recur_depth_ >= kMaxRecurDepth) return x;
++recur_depth_;
Expr res = Mutate(x);
Expr res = this->VisitExpr(x);
--recur_depth_;
return res;
}
......
src/arithmetic/stmt_simplify.cc
......@@ -24,7 +24,6 @@
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/arithmetic.h>
#include <tvm/ir_mutator.h>
#include <tvm/expr_operator.h>
#include <tvm/arithmetic.h>
#include "ir_mutator_with_analyzer.h"
......@@ -40,44 +39,47 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
: IRMutatorWithAnalyzer(analyzer) {}
using Parent = IRMutatorWithAnalyzer;
using Parent::Mutate;
using Parent::Mutate_;
using Parent::VisitStmt;
using Parent::VisitStmt_;
Expr Mutate(Expr expr) final {
Expr VisitExpr(const Expr& expr) final {
return analyzer_->Simplify(expr);
}
Stmt Simplify(Stmt stmt) {
return Mutate(stmt);
return operator()(std::move(stmt));
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
analyzer_->Bind(op->loop_var, Range::make_by_min_extent(op->min, op->extent));
With<ConstraintContext> ctx1(analyzer_, op->loop_var >= op->min);
With<ConstraintContext> ctx2(analyzer_, op->loop_var < op->min + op->extent);
return IRMutator::Mutate_(op, s);
return Parent::VisitStmt_(op);
}
Stmt Mutate_(const LetStmt* op, const Stmt& s) {
Expr value = this->Mutate(op->value);
Stmt VisitStmt_(const LetStmt* op) {
Expr value = this->VisitExpr(op->value);
if (!ir::HasSideEffect(value)) {
// it is fine to discard the let binding
// because the call to simplify will always inline the var.
analyzer_->Bind(op->var, value);
return Mutate(op->body);
return this->VisitStmt(op->body);
}
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
if (value.same_as(op->value) &&
body.same_as(op->body)) {
return s;
return GetRef<Stmt>(op);
} else {
return LetStmt::make(op->var, value, body);
auto n = this->CopyOnWrite(op);
n->value = std::move(value);
n->body = std::move(body);
return Stmt(n);
}
}
// eliminate useless stores
Stmt Mutate_(const Store* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Store* op) final {
Stmt stmt = Parent::VisitStmt_(op);
op = stmt.as<Store>();
if (const Load* load = op->value.as<Load>()) {
if (load->buffer_var.same_as(op->buffer_var) &&
......@@ -85,7 +87,7 @@ class StmtSimplifier : public IRMutatorWithAnalyzer {
return Evaluate::make(0);
}
}
return stmt;
return GetRef<Stmt>(op);
}
};
......@@ -98,7 +100,7 @@ Stmt CanonicalSimplify(Stmt stmt, Map<Var, Range> vrange) {
for (auto kv : vrange) {
analyzer.Bind(kv.first, kv.second);
}
return arith::StmtSimplifier(&analyzer).Simplify(stmt);
return arith::StmtSimplifier(&analyzer).Simplify(std::move(stmt));
}
Expr CanonicalSimplify(Expr expr, Map<Var, Range> vrange) {
......@@ -119,7 +121,7 @@ Expr Simplify(Expr expr, Map<Var, Range> vrange) {
}
Stmt Simplify(Stmt stmt, Map<Var, Range> vrange) {
return CanonicalSimplify(stmt, vrange);
return CanonicalSimplify(std::move(stmt), vrange);
}
} // namespace ir
} // namespace tvm
src/autotvm/feature_visitor.cc
......@@ -6,9 +6,9 @@
* 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
......@@ -29,7 +29,7 @@ namespace tvm {
namespace autotvm {
// for loop
void FeatureVisitor::Visit_(const For *op) {
void FeatureVisitor::VisitStmt_(const For* op) {
const auto *extent = op->extent.as<IntImm>();
int64_t loop_extent = -1;
if (extent != nullptr)
......@@ -51,13 +51,13 @@ void FeatureVisitor::Visit_(const For *op) {
}
if (EnterItervar_(op->loop_var, loop_extent, ann)) {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
ExitItervar_();
}
}
// parallel axis, virtual thread
void FeatureVisitor::Visit_(const AttrStmt *op) {
void FeatureVisitor::VisitStmt_(const AttrStmt* op) {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::virtual_thread) {
VarExpr var = op->node.as<tvm::IterVarNode>()->var;
......@@ -86,24 +86,24 @@ void FeatureVisitor::Visit_(const AttrStmt *op) {
}
if (EnterItervar_(var, extent->value, ann)) {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
ExitItervar_();
}
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
// memory access
void FeatureVisitor::Visit_(const Load *op) {
void FeatureVisitor::VisitExpr_(const Load* op) {
EnterMem_(op->buffer_var, op->index);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
ExitMem_();
}
void FeatureVisitor::Visit_(const Store *op) {
void FeatureVisitor::VisitStmt_(const Store* op) {
EnterMem_(op->buffer_var, op->index);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
ExitMem_();
}
......
src/autotvm/feature_visitor.h
......@@ -6,9 +6,9 @@
* 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
......@@ -27,7 +27,7 @@
#define TVM_AUTOTVM_FEATURE_VISITOR_H_
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <string>
namespace tvm {
......@@ -48,15 +48,18 @@ enum AnnotationType {
* \brief A base class for feature extractor, used for processing
* for loop and memory access in the IR
*/
class FeatureVisitor : public IRVisitor {
class FeatureVisitor : public StmtExprVisitor {
public:
// for loop
void Visit_(const For *op);
void Visit_(const AttrStmt *op);
void VisitStmt_(const For *op);
void VisitStmt_(const AttrStmt *op);
// memory access
void Visit_(const Load *op);
void Visit_(const Store *op);
void VisitExpr_(const Load *op);
void VisitStmt_(const Store *op);
using StmtExprVisitor::VisitStmt_;
using StmtExprVisitor::VisitExpr_;
protected:
/*!
......
src/autotvm/touch_extractor.cc
......@@ -44,14 +44,14 @@ int ParallelLevel(AnnotationType ann) {
}
// get touch pattern from index expression
class IndexParser: public IRVisitor {
class IndexParser: public ExprVisitor {
public:
void Parse(Expr expr) {
pattern_map.clear();
this->Visit(expr);
this->VisitExpr(expr);
}
void Visit_(const Variable *op) {
void VisitExpr_(const Variable *op) {
// TODO(lmzheng): handle more index types (multiple occurrence)
if (pattern_map.count(op) == 0) {
pattern_map[op] = TouchPattern();
......@@ -60,13 +60,13 @@ class IndexParser: public IRVisitor {
}
}
void Visit_(const Mul *op) {
void VisitExpr_(const Mul *op) {
if (op->a.as<Variable>()) {
if (const auto stride = op->b.as<IntImm>()) {
next_stride_ = stride->value;
}
}
IRVisitor::Visit_(op);
ExprVisitor::VisitExpr_(op);
}
std::unordered_map<const Variable*, TouchPattern> pattern_map;
......
src/autotvm/touch_extractor.h
......@@ -6,9 +6,9 @@
* 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
......@@ -26,7 +26,7 @@
#define TVM_AUTOTVM_TOUCH_EXTRACTOR_H_
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/api_registry.h>
#include <stack>
#include <vector>
......@@ -85,39 +85,39 @@ struct ItervarFeature {
// extract iter vars and their touch pattern from ir
class TouchExtractor : public FeatureVisitor {
public:
void Analyze(Stmt stmt) {
this->Visit(stmt);
void Analyze(const Stmt& stmt) {
operator()(stmt);
}
// arithmetic stats
void Visit_(const Add *op) {
void VisitExpr_(const Add *op) {
if (op->dtype.is_float())
itervar_map[itervar_stack_.back()].add_ct++;
IRVisitor::Visit_(op);
FeatureVisitor::VisitExpr_(op);
}
void Visit_(const Sub *op) {
void VisitExpr_(const Sub *op) {
if (op->dtype.is_float())
itervar_map[itervar_stack_.back()].add_ct++;
IRVisitor::Visit_(op);
FeatureVisitor::VisitExpr_(op);
}
void Visit_(const Mul *op) {
void VisitExpr_(const Mul *op) {
if (op->dtype.is_float())
itervar_map[itervar_stack_.back()].mul_ct++;
IRVisitor::Visit_(op);
FeatureVisitor::VisitExpr_(op);
}
void Visit_(const Div *op) {
void VisitExpr_(const Div *op) {
if (op->dtype.is_float())
itervar_map[itervar_stack_.back()].div_ct++;
IRVisitor::Visit_(op);
FeatureVisitor::VisitExpr_(op);
}
void Visit_(const Mod *op) {
void VisitExpr_(const Mod *op) {
if (op->dtype.is_float())
itervar_map[itervar_stack_.back()].div_ct++;
IRVisitor::Visit_(op);
FeatureVisitor::VisitExpr_(op);
}
std::unordered_map<VarExpr, ItervarFeature, tvm::ExprHash, tvm::ExprEqual> itervar_map;
......@@ -134,7 +134,7 @@ class TouchExtractor : public FeatureVisitor {
std::deque<VarExpr> itervar_stack_; // use deque instead of stack for indexing
std::deque<size_t> skip_stack_size_;
using IRVisitor::Visit_;
using FeatureVisitor::VisitExpr_;
};
} // namespace autotvm
......
src/op/compute_op.cc
......@@ -24,8 +24,8 @@
#include <tvm/operation.h>
#include <tvm/arithmetic.h>
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_functor_ext.h>
#include <unordered_set>
#include <string>
#include <utility>
......@@ -538,7 +538,7 @@ namespace {
* must be Reduce as well; and their inputs should have the
* same attribute except value_index.
*/
class ComputeVerifier final : protected ir::IRVisitor {
class ComputeVerifier final : protected ir::ExprVisitor {
public:
/// Special member functions
//@{
......@@ -567,20 +567,20 @@ class ComputeVerifier final : protected ir::IRVisitor {
}
level_ = 0;
ir::IRVisitor::Visit(e);
ExprVisitor::VisitExpr(e);
}
}
protected:
/// Visitor implementation
//@{
void Visit(const ObjectRef& n) final {
void VisitExpr(const Expr& n) final {
++level_;
ir::IRVisitor::Visit(n);
ExprVisitor::VisitExpr(n);
--level_;
}
void Visit_(const ir::Reduce* op) final {
void VisitExpr_(const ir::Reduce* op) final {
// Check for non top level reductions
CHECK(0 == level_)
<< "Reductions are only allowed at the top level of compute. "
......
src/op/hybrid_op.cc
......@@ -24,7 +24,7 @@
#include <tvm/operation.h>
#include <tvm/arithmetic.h>
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <tvm/expr_operator.h>
#include <unordered_set>
......@@ -221,7 +221,7 @@ namespace op {
Stmt ApplyLoopShapes(const Stage &stage,
const std::unordered_map<IterVar, Range> &dom_map, Stmt stmt) {
class LoopSpliter : public IRMutator {
class LoopSpliter : public StmtExprMutator {
Expr factor;
const Variable *parent;
IterVar inner, outer;
......@@ -247,7 +247,7 @@ Stmt ApplyLoopShapes(const Stage &stage,
outer = IterVarNode::make(outer_dom, outer_->var, outer_->iter_type);
}
Stmt Mutate_(const For *op, const Stmt &stmt) {
Stmt VisitStmt_(const For *op) final {
if (op->loop_var.get() == parent) {
std::unordered_map<const Variable *, Expr> rmap;
rmap[op->loop_var.get()] = inner + outer * factor;
......@@ -261,11 +261,11 @@ Stmt ApplyLoopShapes(const Stage &stage,
splitted = true;
return ret;
}
return IRMutator::Mutate_(op, stmt);
return StmtExprMutator::VisitStmt_(op);
}
};
class LoopFuser : public IRMutator {
class LoopFuser : public StmtExprMutator {
const IterVar &parent;
const Variable *inner;
const Variable *outer;
......@@ -280,8 +280,7 @@ Stmt ApplyLoopShapes(const Stage &stage,
extent(0), fused(false) {}
// TODO(@were): Handle imperfect loops
Stmt Mutate_(const For *op, const Stmt &stmt) {
Stmt VisitStmt_(const For* op) final {
if (op->loop_var.get() == inner) {
CHECK(under_outer);
std::unordered_map<const Variable *, Expr> rmap;
......@@ -291,7 +290,7 @@ Stmt ApplyLoopShapes(const Stage &stage,
return ir::Substitute(op->body, rmap);
} else if (op->loop_var.get() == outer) {
under_outer = true;
Stmt body = IRMutator::Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
std::unordered_map<const Variable *, Expr> rmap;
rmap[op->loop_var.get()] = indexdiv(parent, extent);
body = ir::Substitute(body, rmap);
......@@ -299,25 +298,25 @@ Stmt ApplyLoopShapes(const Stage &stage,
return For::make(parent->var, Expr(0), extent * op->extent,
op->for_type, op->device_api, body);
} else if (under_outer) {
Stmt body = IRMutator::Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
std::unordered_map<const Variable *, Expr> rmap;
rmap[op->loop_var.get()] = indexmod(indexdiv(parent, extent), op->extent);
body = ir::Substitute(body, rmap);
extent = extent * op->extent;
return body;
}
return IRMutator::Mutate(stmt);
return StmtExprMutator::VisitStmt_(op);
}
};
for (auto &rel : stage->relations) {
if (const SplitNode *split = rel.as<SplitNode>()) {
LoopSpliter Spliter(split, dom_map);
stmt = Spliter.Mutate(stmt);
stmt = Spliter(stmt);
CHECK(Spliter.splitted);
} else if (const FuseNode *fuse = rel.as<FuseNode>()) {
LoopFuser Fuser(fuse);
stmt = Fuser.Mutate(stmt);
stmt = Fuser(stmt);
CHECK(Fuser.fused);
}
}
......@@ -327,14 +326,14 @@ Stmt ApplyLoopShapes(const Stage &stage,
Stmt ApplyLoopAnnotations(const Stage &stage,
const std::unordered_map<IterVar, IterVar> &rebased, Stmt stmt) {
class LoopAnnotator : public IRMutator {
class LoopAnnotator : public StmtMutator {
const Variable *var;
const IterVarAttr &attr;
public:
LoopAnnotator(const Variable *var_, const IterVarAttr &attr_) : var(var_), attr(attr_) {}
Stmt Mutate_(const For *op, const Stmt &stmt) {
Stmt VisitStmt_(const For *op) final {
if (op->loop_var.get() == var) {
if (attr->bind_thread.defined()) {
const auto &iter_var = attr->bind_thread;
......@@ -352,7 +351,7 @@ Stmt ApplyLoopAnnotations(const Stage &stage,
IterVarTypeToForType(attr->iter_type), op->device_api, op->body);
}
}
return IRMutator::Mutate_(op, stmt);
return StmtMutator::VisitStmt_(op);
}
};
......@@ -381,7 +380,7 @@ Stmt ApplyLoopAnnotations(const Stage &stage,
CHECK_EQ(found, 1) << " iter var should be found exactly once!";
if (need_change) {
stmt = LoopAnnotator(var, attr).Mutate(stmt);
stmt = LoopAnnotator(var, attr)(std::move(stmt));
}
}
return stmt;
......@@ -411,7 +410,7 @@ Stmt ApplyLoopOrder(const Stage &stage,
}
}
class LoopReorder : public IRMutator {
class LoopReorder : public StmtMutator {
const Stage &stage;
const std::unordered_map<IterVar, Range> &dom_map;
const std::unordered_map<const Variable *, IterVar> &reorder;
......@@ -422,13 +421,13 @@ Stmt ApplyLoopOrder(const Stage &stage,
const std::unordered_map<const Variable*, IterVar> &reorder)
: stage(stage), dom_map(dom_map), reorder(reorder) {}
Stmt Mutate_(const For *op, const Stmt &stmt) {
Stmt VisitStmt_(const For* op) final {
// Reorder from in to out
Stmt body_ = IRMutator::Mutate(op->body);
Stmt body_ = this->VisitStmt(op->body);
CHECK(reorder.count(op->loop_var.get()));
auto target = reorder.find(op->loop_var.get())->second;
if (body_.same_as(op->body) && op->loop_var.get() == target->var.get())
return stmt;
return GetRef<Stmt>(op);
const Stmt &body = op->body.same_as(body_) ? op->body : body_;
ForType for_type = IterVarTypeToForType(target->iter_type);
if (stage->iter_var_attrs.count(target)) {
......@@ -441,7 +440,7 @@ Stmt ApplyLoopOrder(const Stage &stage,
};
if (need_reorder)
return LoopReorder(stage, dom_map, reorder).Mutate(stmt);
return LoopReorder(stage, dom_map, reorder)(stmt);
return stmt;
}
......@@ -479,21 +478,21 @@ std::vector<IterVar> GatherLoopVars(Stmt stmt) {
}
// replacer to replace tensors' usage in Provide
class ProviderReplacer : public ir::IRMutator {
class ProviderReplacer : public ir::StmtMutator {
public:
explicit ProviderReplacer(const std::unordered_map<Tensor, Tensor> &vmap)
: vmap_(vmap) {}
Stmt Mutate_(const ir::Provide* op, const Stmt &s) {
Stmt VisitStmt_(const ir::Provide* op) final {
Tensor t = Downcast<Operation>(op->func).output(op->value_index);
auto it = vmap_.find(t);
if (it != vmap_.end()) {
Stmt ret = ir::Provide::make(
it->second->op, it->second->value_index, op->value, op->args);
found = true;
return IRMutator::Mutate_(ret.as<ir::Provide>(), ret);
return this->VisitStmt(ret);
}
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
// whether it is found.
......@@ -506,7 +505,7 @@ class ProviderReplacer : public ir::IRMutator {
Stmt ReplaceProvideTensor(Stmt stmt,
const std::unordered_map<Tensor, Tensor> &replace) {
ProviderReplacer repl(replace);
Stmt ret = repl.Mutate(stmt);
Stmt ret = repl(stmt);
return repl.found ? ret : stmt;
}
} // namespace op
......
src/op/hybrid_op.h
......@@ -6,9 +6,9 @@
* 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
......@@ -25,8 +25,6 @@
#define TVM_OP_HYBRID_OP_H_
#include <tvm/expr.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/schedule.h>
#include <unordered_map>
#include <unordered_set>
......
src/op/op_util.cc
......@@ -23,8 +23,8 @@
*/
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/operation.h>
#include <tvm/ir_mutator.h>
#include <string>
#include "op_util.h"
#include "../schedule/message_passing.h"
......@@ -186,12 +186,12 @@ std::vector<Stmt> MakeIfNest(const std::vector<Expr>& predicates) {
}
// replacer to replace tensors
class TensorReplacer : public ir::IRMutator {
class TensorReplacer : public ir::StmtExprMutator {
public:
explicit TensorReplacer(const std::unordered_map<Tensor, Tensor>& vmap)
: vmap_(vmap) {}
Expr Mutate_(const ir::Call* op, const Expr& e) {
Expr VisitExpr_(const ir::Call* op) final {
if (op->call_type == ir::Call::Halide) {
Tensor t = Downcast<Operation>(op->func).output(op->value_index);
auto it = vmap_.find(t);
......@@ -200,10 +200,10 @@ class TensorReplacer : public ir::IRMutator {
op->dtype, it->second->op->name, op->args,
op->call_type, it->second->op, it->second->value_index);
found = true;
return IRMutator::Mutate_(ret.as<ir::Call>(), ret);
return this->VisitExpr(ret);
}
}
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
// whether it is found.
......@@ -216,13 +216,13 @@ class TensorReplacer : public ir::IRMutator {
Stmt ReplaceTensor(Stmt stmt,
const std::unordered_map<Tensor, Tensor>& replace) {
TensorReplacer repl(replace);
Stmt ret = repl.Mutate(stmt);
Stmt ret = repl(stmt);
return repl.found ? ret : stmt;
}
Expr ReplaceTensor(Expr expr,
const std::unordered_map<Tensor, Tensor>& replace) {
TensorReplacer repl(replace);
Expr ret = repl.Mutate(expr);
Expr ret = repl(expr);
return repl.found ? ret : expr;
}
......
src/op/tensor_compute_op.cc
......@@ -24,7 +24,6 @@
#include <tvm/operation.h>
#include <tvm/arithmetic.h>
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_pass.h>
#include <unordered_set>
#include "./op_util.h"
......
src/op/tensorize.cc
......@@ -22,7 +22,7 @@
* \file tensorize.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <tvm/api_registry.h>
#include "op_util.h"
......@@ -157,10 +157,10 @@ void VerifyTensorizeLoopNest(const ComputeOpNode* self,
}
// Remap the tensor placeholder, index and inline things.
class TensorIntrinMatcher final : public IRMutator {
class TensorIntrinMatcher final : public StmtExprMutator {
public:
Expr Mutate_(const Call* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Call* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
if (op->call_type == Call::Halide) {
Tensor t = Downcast<Operation>(op->func).output(op->value_index);
......@@ -180,17 +180,17 @@ class TensorIntrinMatcher final : public IRMutator {
return expr;
}
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
auto it = var_remap_.find(op);
if (it != var_remap_.end()) {
return it->second;
} else {
return e;
return GetRef<Expr>(op);
}
}
Expr Mutate_(const Reduce* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Reduce* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Reduce>();
Array<IterVar> axis;
for (size_t i = 0; i < op->axis.size(); ++i) {
......@@ -317,7 +317,7 @@ Array<Expr> MatchTensorizeBody(
matcher.Init(self, stage, dom_map, out_dom, in_region, intrin, compute_intrin_iter_space);
Array<Expr> ret;
for (Expr expr : self->body) {
ret.push_back(matcher.Mutate(expr));
ret.push_back(matcher(expr));
}
return ret;
}
......
src/pass/bound_checker.cc
......@@ -23,9 +23,8 @@
// Instrument checkers for out of the bounds access.
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <vector>
#include <unordered_map>
#include <utility>
......@@ -33,48 +32,48 @@
namespace tvm {
namespace ir {
class BoundCollector : public IRVisitor {
class BoundCollector : public StmtVisitor {
public:
BoundCollector() {}
void Visit_(const AttrStmt *op) {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == ir::attr::buffer_bound) {
if (const Variable *key = op->node.as<Variable>()) {
mem_to_shape[key] = op->value;
}
}
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
// Hashtable which maps buffer_var to shape.
std::unordered_map<const Variable *, Expr> mem_to_shape;
};
class BoundChecker : public IRMutator {
class BoundChecker : public StmtExprMutator {
public:
explicit BoundChecker(
const std::unordered_map<const Variable *, Expr> &mem_to_shape)
: mem_to_shape_(mem_to_shape) {}
Stmt Mutate_(const Allocate *op, const Stmt &s) final {
Stmt VisitStmt_(const Allocate* op) final {
// If the shape was updated we should update the hashtable.
if (UpdateIsNeeded(op->buffer_var)) {
Update(op->buffer_var, op->extents, op->dtype);
}
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Expr Mutate_(const Call *op, const Expr &ex) final {
Expr VisitExpr_(const Call* op) final {
if (process_store_ && op->is_intrinsic(intrinsic::tvm_if_then_else)) {
unsafe_rewritten_ = true;
}
return IRMutator::Mutate_(op, ex);
return StmtExprMutator::VisitExpr_(op);
}
Stmt Mutate_(const Store *op, const Stmt &s) final {
Stmt VisitStmt_(const Store* op) final {
store_scope_bound_collector_.clear();
process_store_ = true;
unsafe_rewritten_ = false;
IRMutator::Mutate_(op, s);
StmtExprMutator::VisitStmt_(op);
process_store_ = false;
if (CanInstrument(op->index, op->buffer_var)) {
Collect(op->index, op->buffer_var);
......@@ -92,23 +91,24 @@ class BoundChecker : public IRMutator {
return body;
}
}
return s;
return GetRef<Stmt>(op);
}
Expr Mutate_(const Load *op, const Expr &ex) final {
Expr VisitExpr_(const Load* op) final {
if (CanInstrument(op->index, op->buffer_var)) {
Collect(op->index, op->buffer_var);
}
return IRMutator::Mutate_(op, ex);
return StmtExprMutator::VisitExpr_(op);
}
private:
bool UpdateIsNeeded(const VarExpr &buffer_var) const {
bool UpdateIsNeeded(const VarExpr& buffer_var) const {
return (buffer_var.defined() && mem_to_shape_.count(buffer_var.get()));
}
void Update(const VarExpr &buffer_var, const Array<Expr> &new_shape,
const DataType &type) {
void Update(const VarExpr& buffer_var,
const Array<Expr>& new_shape,
const DataType& type) {
// Sanity check at first.
if (!new_shape.size()) {
return;
......@@ -132,7 +132,7 @@ class BoundChecker : public IRMutator {
mem_to_shape_[buffer_var.get()] = shape;
}
bool IndexIsValid(const Expr &index) const {
bool IndexIsValid(const Expr& index) const {
if (!index.defined()) {
return false;
}
......@@ -146,7 +146,7 @@ class BoundChecker : public IRMutator {
return true;
}
bool CanInstrument(const Expr &index, const VarExpr &buffer_var) const {
bool CanInstrument(const Expr& index, const VarExpr& buffer_var) const {
return buffer_var.defined() && mem_to_shape_.count(buffer_var.get()) &&
IndexIsValid(index) && !unsafe_rewritten_;
}
......@@ -206,8 +206,8 @@ class BoundChecker : public IRMutator {
Stmt InstrumentBoundCheckers(Stmt stmt) {
BoundCollector bound_collector;
// At first walk recursively and collect bound attributes.
bound_collector.Visit(stmt);
return BoundChecker(bound_collector.mem_to_shape).Mutate(stmt);
bound_collector(stmt);
return BoundChecker(bound_collector.mem_to_shape)(std::move(stmt));
}
} // namespace ir
} // namespace tvm
src/pass/combine_context_call.cc
......@@ -23,7 +23,7 @@
* \file combine_context_call.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <map>
......@@ -32,7 +32,7 @@ namespace ir {
// Calculate the statistics of packed function.
// These information are needed during codegen.
class ContextCallCombiner final : public IRMutator {
class ContextCallCombiner final : public StmtExprMutator {
public:
struct CompareExpr {
bool operator()(const Expr& lhs, const Expr& rhs) const {
......@@ -40,7 +40,7 @@ class ContextCallCombiner final : public IRMutator {
}
};
Expr Mutate_(const Call* op, const Expr& e) final {
Expr VisitExpr_(const Call* op) final {
if (op->is_intrinsic(intrinsic::tvm_thread_context)) {
CHECK_EQ(op->args.size(), 1U);
Expr ctx = op->args[0];
......@@ -60,39 +60,39 @@ class ContextCallCombiner final : public IRMutator {
return std::move(ctx_var);
}
} else {
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::coproc_uop_scope) {
// Map of comparison expression to variable
std::map<Expr, Var, CompareExpr> temp;
std::swap(temp, ctx_map_);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
std::swap(temp, ctx_map_);
return BuildContext(temp, stmt);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
if (op->for_type == ForType::Parallel) {
// Map of comparison expression to variable
std::map<Expr, Var, CompareExpr> temp;
std::swap(temp, ctx_map_);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
std::swap(temp, ctx_map_);
return BuildContext(temp, stmt);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Combine(Stmt stmt) {
return BuildContext(ctx_map_, this->Mutate(stmt));
return BuildContext(ctx_map_, this->VisitStmt(stmt));
}
private:
......
src/pass/coproc_sync.cc
......@@ -22,8 +22,7 @@
*/
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <unordered_map>
#include <unordered_set>
#include "ir_util.h"
......@@ -33,25 +32,25 @@ namespace tvm {
namespace ir {
// Visitor to find touched set by co-processor scope.
class CoProcTouchedBuffer : public IRVisitor {
class CoProcTouchedBuffer : public StmtExprVisitor {
public:
void Visit_(const Load* op) final {
void VisitExpr_(const Load* op) final {
if (in_scope_) {
touched_[op->buffer_var.get()].coproc = true;
} else {
touched_[op->buffer_var.get()].normal = true;
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
void Visit_(const Store* op) final {
void VisitStmt_(const Store* op) final {
if (in_scope_) {
touched_[op->buffer_var.get()].coproc = true;
} else {
touched_[op->buffer_var.get()].normal = true;
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
void Visit_(const Call* op) final {
void VisitExpr_(const Call* op) final {
if (op->is_intrinsic(intrinsic::tvm_access_ptr)) {
const Variable* buffer = op->args[1].as<Variable>();
if (in_scope_) {
......@@ -60,17 +59,17 @@ class CoProcTouchedBuffer : public IRVisitor {
touched_[buffer].normal = true;
}
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
void Visit_(const AttrStmt* op) final {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::coproc_scope && !in_scope_) {
in_scope_ = true;
IterVar iv = Downcast<IterVar>(op->node);
coproc_.insert(iv);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
in_scope_ = false;
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
......@@ -96,7 +95,7 @@ class CoProcSyncPlanner : public StorageAccessVisitor {
}
void Plan(const Stmt& stmt) {
this->Visit(stmt);
this->VisitStmt(stmt);
PlanSync(scope_.back(), nullptr, true);
if (sync_.size() == 0) {
sync_[stmt.get()] = GetSync(coproc_name_ + ".coproc_sync");
......@@ -218,14 +217,14 @@ class CoProcBarrierDetector : public StorageAccessVisitor {
write_barrier_name_ = coproc_name + ".coproc_write_barrier";
}
void PlanReadBarrier(Stmt stmt) {
void PlanReadBarrier(const Stmt& stmt) {
read_barrier_ = true;
this->Visit(stmt);
this->VisitStmt(stmt);
PlanReadBarrier(scope_.back(), nullptr);
}
void PlanWriteBarrier(Stmt stmt) {
void PlanWriteBarrier(const Stmt& stmt) {
read_barrier_ = false;
this->Visit(stmt);
this->VisitStmt(stmt);
PlanWriteBarrier(scope_.back(), nullptr);
}
......@@ -356,7 +355,7 @@ class CoProcBarrierDetector : public StorageAccessVisitor {
};
class CoProcInstDepDetector : public IRVisitor {
class CoProcInstDepDetector : public StmtVisitor {
public:
explicit CoProcInstDepDetector(
const IterVar& coproc_axis,
......@@ -366,15 +365,15 @@ class CoProcInstDepDetector : public IRVisitor {
sync_pop_name_ = coproc_name + ".coproc_dep_pop";
}
void Plan(Stmt stmt) {
this->Visit(stmt);
void Plan(const Stmt& stmt) {
this->VisitStmt(stmt);
if (last_state_.node != nullptr) {
MatchFixEnterPop(first_state_);
MatchFixExitPush(last_state_);
}
}
void Visit_(const AttrStmt* op) final {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::coproc_scope &&
op->node.same_as(coproc_axis_)) {
const IntImm* ctx_id = op->value.as<IntImm>();
......@@ -385,15 +384,15 @@ class CoProcInstDepDetector : public IRVisitor {
curr_state_.exit_ctx.insert(ctx_id->value);
UpdateState();
} else {
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
}
void Visit_(const For* op) final {
void VisitStmt_(const For* op) final {
SyncState temp_first, temp_last;
std::swap(first_state_, temp_first);
std::swap(last_state_, temp_last);
this->Visit(op->body);
this->VisitStmt(op->body);
curr_state_.clear();
if (last_state_.node != nullptr) {
curr_state_.node = op;
......@@ -412,13 +411,13 @@ class CoProcInstDepDetector : public IRVisitor {
}
}
void Visit_(const IfThenElse* op) final {
void VisitStmt_(const IfThenElse* op) final {
SyncState temp_first, temp_last, curr_state;
std::swap(first_state_, temp_first);
std::swap(last_state_, temp_last);
{
// then stmt
this->Visit(op->then_case);
this->VisitStmt(op->then_case);
if (last_state_.node != nullptr) {
curr_state.node = op;
MatchFixEnterPop(first_state_);
......@@ -434,7 +433,7 @@ class CoProcInstDepDetector : public IRVisitor {
last_state_.clear();
}
if (op->else_case.defined()) {
this->Visit(op->else_case);
this->VisitStmt(op->else_case);
if (last_state_.node != nullptr) {
curr_state.node = op;
MatchFixEnterPop(first_state_);
......@@ -606,11 +605,11 @@ class CoProcInstDepDetector : public IRVisitor {
};
class CoProcSyncInserter : public IRMutator {
class CoProcSyncInserter : public StmtMutator {
public:
Stmt Insert(Stmt stmt) {
CoProcTouchedBuffer visitor;
visitor.Visit(stmt);
visitor(stmt);
if (visitor.coproc_.size() == 0) return stmt;
std::unordered_set<const Variable*> touched;
......@@ -652,10 +651,10 @@ class CoProcSyncInserter : public IRMutator {
auto& vec = insert_after_[kv.first];
vec.insert(vec.end(), kv.second.begin(), kv.second.end());
}
return Mutate(stmt);
return operator()(std::move(stmt));
}
Stmt Mutate(Stmt stmt) final {
Stmt VisitStmt(const Stmt& stmt) final {
Stmt before, after;
auto it = insert_before_.find(stmt.get());
if (it != insert_before_.end()) {
......@@ -666,14 +665,14 @@ class CoProcSyncInserter : public IRMutator {
if (it != insert_after_.end()) {
after = MergeSeq(it->second);
}
stmt = IRMutator::Mutate(stmt);
Stmt new_stmt = StmtMutator::VisitStmt(stmt);
if (before.defined()) {
stmt = Block::make(before, stmt);
new_stmt = Block::make(before, new_stmt);
}
if (after.defined()) {
stmt = Block::make(stmt, after);
new_stmt = Block::make(new_stmt, after);
}
return stmt;
return new_stmt;
}
private:
......@@ -685,7 +684,7 @@ class CoProcSyncInserter : public IRMutator {
Stmt CoProcSync(Stmt stmt) {
return CoProcSyncInserter().Insert(stmt);
return CoProcSyncInserter().Insert(std::move(stmt));
}
} // namespace ir
......
src/pass/detect_device.cc
......@@ -6,9 +6,9 @@
* 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
......@@ -22,7 +22,6 @@
*/
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include "../pass/ir_util.h"
namespace tvm {
......
src/pass/hoist_if_then_else.cc
......@@ -21,9 +21,7 @@
* \file hoist_if_then_else.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/arithmetic.h>
#include <tvm/api_registry.h>
#include <unordered_map>
......
src/pass/infer_fragment.cc
......@@ -23,8 +23,7 @@
*/
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <unordered_map>
#include <unordered_set>
#include "ir_util.h"
......@@ -35,7 +34,7 @@ namespace tvm {
namespace ir {
// Get fragment information from tensor intrinsics
class FragmentGetter : public IRVisitor {
class FragmentGetter : public StmtExprVisitor {
public:
// fragment metadata
struct FragmentInfo {
......@@ -48,8 +47,8 @@ class FragmentGetter : public IRVisitor {
: m(_m), n(_n), k(_k), layout(_layout) {}
};
void Visit_(const Call* op) final {
IRVisitor::Visit_(op);
void VisitExpr_(const Call* op) final {
StmtExprVisitor::VisitExpr_(op);
if (op->is_intrinsic(intrinsic::tvm_load_matrix_sync) ||
op->is_intrinsic(intrinsic::tvm_store_matrix_sync)) {
......@@ -116,13 +115,13 @@ class FragmentGetter : public IRVisitor {
}
// Get memory scope
void Visit_(const AttrStmt* op) final {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::storage_scope) {
const Variable* buffer = op->node.as<Variable>();
CHECK(buffer);
scopes[buffer] = op->value.as<StringImm>()->value;
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
// Memory scope for allocations
......@@ -132,11 +131,12 @@ class FragmentGetter : public IRVisitor {
};
// Check shape of fragment making sure it is a valid shape for tvm_mma_sync
class FragmentChecker : public IRVisitor {
class FragmentChecker : public StmtExprVisitor {
public:
explicit FragmentChecker(const FragmentGetter &getter) : fragment_getter(getter) {}
void Visit_(const Call* op) final {
void VisitExpr_(const Call* op) final {
StmtExprVisitor::VisitExpr_(op);
// Check shape when calling tvm_mma_sync
if (op->is_intrinsic(intrinsic::tvm_mma_sync)) {
CHECK_EQ(op->args.size(), 8U);
......@@ -170,12 +170,12 @@ class FragmentChecker : public IRVisitor {
};
// Store the metadata into attributes
class InferFragmenter : public IRMutator {
class InferFragmenter : public StmtMutator {
public:
explicit InferFragmenter(const FragmentGetter &getter) : fragment_getter(getter) {}
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Allocate* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
const Variable* buffer = op->buffer_var.get();
if (fragment_getter.fragments.count(buffer)) {
// Add attribute to fragments allocation
......@@ -206,9 +206,10 @@ class InferFragmenter : public IRMutator {
Stmt InferFragment(Stmt stmt) {
FragmentGetter getter;
getter.Visit(stmt);
FragmentChecker(getter).Visit(stmt);
stmt = InferFragmenter(getter).Mutate(stmt);
getter(stmt);
FragmentChecker checker(getter);
checker(stmt);
stmt = InferFragmenter(getter)(std::move(stmt));
return stmt;
}
......
src/pass/inject_copy_intrin.cc
......@@ -23,7 +23,7 @@
*/
#include <tvm/ir.h>
#include <tvm/packed_func_ext.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include "../arithmetic/pattern_match.h"
......@@ -32,7 +32,7 @@ namespace ir {
using runtime::PackedFunc;
class CopyIntrinInjector : public IRMutator {
class CopyIntrinInjector : public StmtMutator {
public:
CopyIntrinInjector(const std::string& pragma_key,
const PackedFunc& flower_copy_fromto)
......@@ -40,7 +40,7 @@ class CopyIntrinInjector : public IRMutator {
flower_copy_fromto_(flower_copy_fromto) {
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::storage_scope) {
const Variable* buf = op->node.as<Variable>();
storage_scope_[buf] = op->value.as<StringImm>()->value;
......@@ -50,7 +50,7 @@ class CopyIntrinInjector : public IRMutator {
<< "Cannot match copy pattern of " << op->body;
return ret;
}
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
private:
......@@ -193,8 +193,7 @@ class CopyIntrinInjector : public IRMutator {
Stmt InjectCopyIntrin(Stmt stmt,
const std::string& pragma_key,
const PackedFunc& flower_copy_fromto) {
return CopyIntrinInjector(pragma_key, flower_copy_fromto)
.Mutate(stmt);
return CopyIntrinInjector(pragma_key, flower_copy_fromto)(std::move(stmt));
}
} // namespace ir
......
src/pass/inject_double_buffer.cc
......@@ -22,8 +22,7 @@
* \file inject_double_buffer.cc
*/
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/expr_operator.h>
#include "ir_util.h"
#include "../arithmetic/compute_expr.h"
......@@ -32,18 +31,18 @@ namespace tvm {
namespace ir {
// Detect double buffer variables.
class DoubleBufferDetector : public IRVisitor {
class DoubleBufferDetector : public StmtExprVisitor {
public:
void Visit_(const AttrStmt* op) final {
void VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::double_buffer_scope) {
touched_.insert(op->node.as<Variable>());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
void Visit_(const Variable* op) final {
void VisitExpr_(const Variable* op) final {
if (touched_.count(op)) {
touched_.erase(op);
}
......@@ -53,55 +52,55 @@ class DoubleBufferDetector : public IRVisitor {
};
class StripDoubleBufferWrite : public IRMutator {
class StripDoubleBufferWrite : public StmtMutator {
public:
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::double_buffer_write) {
return Mutate(op->body);
return VisitStmt(op->body);
} else {
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
}
};
class DoubleBufferInjector : public IRMutator {
class DoubleBufferInjector : public StmtExprMutator {
public:
explicit DoubleBufferInjector(int split_loop)
: split_loop_(split_loop) {}
Stmt Inject(const Stmt& stmt) {
Stmt Inject(Stmt stmt) {
DoubleBufferDetector detector;
detector.Visit(stmt);
detector(stmt);
if (detector.touched_.empty()) return stmt;
for (const Variable* v : detector.touched_) {
dbuffer_info_[v] = StorageEntry();
}
return ConvertSSA(this->Mutate(stmt));
return ConvertSSA(operator()(std::move(stmt)));
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::storage_scope) {
const Variable* buf = op->node.as<Variable>();
auto it = dbuffer_info_.find(buf);
if (it != dbuffer_info_.end()) {
it->second.scope = op->value.as<StringImm>()->value;
return Mutate(op->body);
return this->VisitStmt(op->body);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
} else if (op->attr_key == attr::double_buffer_scope) {
return MakeProducer(op, s);
return MakeProducer(op);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt VisitStmt_(const Allocate* op) final {
auto it = dbuffer_info_.find(op->buffer_var.get());
if (it != dbuffer_info_.end()) {
it->second.stride = arith::ComputeReduce<Mul>(
op->extents, Expr()) * op->dtype.lanes();
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
Array<Expr> new_extents{make_const(op->extents[0].dtype(), 2)};
for (Expr e : op->extents) {
......@@ -118,13 +117,13 @@ class DoubleBufferInjector : public IRMutator {
Evaluate::make(0)));
return op->body;
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
loop_nest_.push_back(op);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
auto it = loop_pre_.find(op);
if (it != loop_pre_.end()) {
const For* old_loop = stmt.as<For>();
......@@ -151,7 +150,7 @@ class DoubleBufferInjector : public IRMutator {
MergeSeq(loop_seq));
// tail
std::vector<Stmt> tail_seq;
Stmt tail_body = StripDoubleBufferWrite().Mutate(old_loop->body);
Stmt tail_body = StripDoubleBufferWrite()(old_loop->body);
for (int32_t i = 0; i < split_loop_; ++i) {
Expr idx = tail_base + make_const(tail_base.dtype(), i);
vmap[old_loop->loop_var.get()] = idx;
......@@ -171,8 +170,8 @@ class DoubleBufferInjector : public IRMutator {
return stmt;
}
Stmt Mutate_(const Store* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Store* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Store>();
auto it = dbuffer_info_.find(op->buffer_var.get());
if (it != dbuffer_info_.end()) {
......@@ -188,8 +187,8 @@ class DoubleBufferInjector : public IRMutator {
}
}
Expr Mutate_(const Load* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Load* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Load>();
auto it = dbuffer_info_.find(op->buffer_var.get());
if (it != dbuffer_info_.end()) {
......@@ -205,20 +204,20 @@ class DoubleBufferInjector : public IRMutator {
}
}
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
CHECK(!dbuffer_info_.count(op));
return e;
return GetRef<Expr>(op);
}
private:
Stmt MakeProducer(const AttrStmt* op, const Stmt& s) {
Stmt MakeProducer(const AttrStmt* op) {
const VarExpr buffer = Downcast<VarExpr>(op->node);
CHECK_NE(loop_nest_.size(), 0U)
<< "Double buffer scope must be inside a loop";
auto it = dbuffer_info_.find(buffer.get());
if (it == dbuffer_info_.end()) {
LOG(WARNING) << "Skip double buffer scope " << op->node;
return Mutate(op->body);
return this->VisitStmt(op->body);
}
StorageEntry& e = it->second;
e.loop = loop_nest_.back();
......@@ -230,7 +229,7 @@ class DoubleBufferInjector : public IRMutator {
e.loop->loop_var.dtype());
e.switch_read_var = indexmod(e.loop->loop_var, two);
in_double_buffer_scope_ = true;
Stmt body = Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
in_double_buffer_scope_ = false;
std::unordered_map<const Variable*, Expr> vmap;
vmap[e.switch_write_var.get()] = zero;
......
src/pass/inject_prefetch.cc
......@@ -22,8 +22,7 @@
*/
// Inject prefetch op in HalideIR
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <tvm/arithmetic.h>
#include <unordered_set>
......@@ -34,10 +33,10 @@ namespace ir {
using arith::IntSet;
using arith::DomainTouched;
class PrefetchInjector : public IRMutator {
class PrefetchInjector : public StmtMutator {
public:
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt ret = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const AttrStmt* op) final {
Stmt ret = StmtMutator::VisitStmt_(op);
op = ret.as<AttrStmt>();
if (op && op->attr_key == attr::prefetch_scope) {
Tensor ts = Downcast<Tensor>(op->node);
......@@ -65,13 +64,13 @@ class PrefetchInjector : public IRMutator {
return ret;
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
auto &var = op->loop_var;
loop_nest_.push_back(var);
if (op->for_type == ForType::Vectorized) {
vectorized_[var.get()] = IntSet::interval(op->min, (op->min + op->extent) - 1);
}
Stmt ret = IRMutator::Mutate_(op, s);
Stmt ret = StmtMutator::VisitStmt_(op);
if (op->for_type == ForType::Vectorized) {
vectorized_.erase(var.get());
}
......@@ -88,7 +87,7 @@ class PrefetchInjector : public IRMutator {
const Range PrefetchInjector::none;
Stmt InjectPrefetch(Stmt stmt) {
return PrefetchInjector().Mutate(stmt);
return PrefetchInjector()(std::move(stmt));
}
} // namespace ir
......
src/pass/inline.cc
......@@ -6,9 +6,9 @@
* 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
......@@ -21,8 +21,8 @@
* \file inline.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_functor_ext.h>
namespace tvm {
namespace ir {
......@@ -30,13 +30,13 @@ namespace ir {
// inliner to inline a function
// the result may not be SSA,
// ConvertSSA need to be applied after this pass
class IRInline final : public IRMutator {
class IRInline final : public StmtExprMutator {
public:
IRInline(FunctionRef f, Array<Var> args, Expr body)
: f_(f), args_(args), body_(body) {}
Expr Mutate_(const Call* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Call* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
if (op->func == f_) {
......@@ -78,7 +78,7 @@ Stmt Inline(Stmt stmt,
Expr body) {
CHECK_EQ(f->num_outputs(), 1)
<< "can only inline output single value operation";
Stmt ret = IRInline(f, args, body).Mutate(stmt);
Stmt ret = IRInline(f, args, body)(std::move(stmt));
if (ret.same_as(stmt)) return ret;
return ConvertSSA(ret);
}
......
src/pass/ir_functor.cc
......@@ -21,7 +21,6 @@
*/
#include <tvm/ir_functor_ext.h>
#include <tvm/packed_func_ext.h>
#include <tvm/ir_visitor.h>
namespace tvm {
namespace ir {
......
src/pass/lift_attr_scope.cc
......@@ -24,7 +24,7 @@
* \file lift_attr_scope.cc
*/
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include "ir_util.h"
namespace tvm {
......@@ -32,13 +32,13 @@ namespace ir {
// NOTE: this optimization can only be applied
// to a few specified attr keys
class AttrScopeLifter : public IRMutator {
class AttrScopeLifter : public StmtMutator {
public:
explicit AttrScopeLifter(std::string attr_key)
: attr_key_(attr_key) {}
Stmt Lift(Stmt stmt) {
stmt = Mutate(stmt);
stmt = operator()(std::move(stmt));
if (attr_node_.defined()) {
stmt = AttrStmt::make(
attr_node_, attr_key_, attr_value_, stmt);
......@@ -47,8 +47,8 @@ class AttrScopeLifter : public IRMutator {
}
// do not go beyond
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Allocate* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
if (attr_node_.defined()) {
Stmt body = AttrStmt::make(
......@@ -65,17 +65,17 @@ class AttrScopeLifter : public IRMutator {
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr_key_) {
attr_node_ = op->node;
attr_value_ = op->value;
return op->body;
} else {
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const Block* op, const Stmt& s) final {
Stmt VisitStmt_(const Block* op) final {
std::vector<Stmt> seq;
FlattenSeq(op->first, &seq);
FlattenSeq(op->rest, &seq);
......@@ -83,21 +83,21 @@ class AttrScopeLifter : public IRMutator {
if (seq.size() == 2 &&
seq[0].same_as(op->first) &&
seq[1].same_as(op->rest)) {
return s;
return GetRef<Stmt>(op);
}
return MergeSeq(seq);
}
Stmt Mutate_(const IfThenElse* op, const Stmt& s) final {
Stmt VisitStmt_(const IfThenElse* op) final {
if (!op->else_case.defined()) {
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
Stmt then_case = this->Mutate(op->then_case);
Stmt then_case = this->VisitStmt(op->then_case);
ObjectRef first_node;
Expr first_value;
std::swap(first_node, attr_node_);
std::swap(first_value, attr_value_);
Stmt else_case = this->Mutate(op->else_case);
Stmt else_case = this->VisitStmt(op->else_case);
if (attr_node_.defined() &&
attr_value_.defined() &&
first_node.defined() &&
......@@ -106,7 +106,7 @@ class AttrScopeLifter : public IRMutator {
ValueSame(attr_value_, first_value)) {
if (then_case.same_as(op->then_case) &&
else_case.same_as(op->else_case)) {
return s;
return GetRef<Stmt>(op);
} else {
return IfThenElse::make(op->condition, then_case, else_case);
}
......@@ -124,7 +124,7 @@ class AttrScopeLifter : public IRMutator {
}
if (then_case.same_as(op->then_case) &&
else_case.same_as(op->else_case)) {
return s;
return GetRef<Stmt>(op);
} else {
return IfThenElse::make(op->condition, then_case, else_case);
}
......@@ -155,7 +155,7 @@ class AttrScopeLifter : public IRMutator {
for (const Stmt & stmt : seq) {
attr_node_ = ObjectRef();
attr_value_ = Expr();
Stmt rest = this->Mutate(stmt);
Stmt rest = this->VisitStmt(stmt);
if (attr_node_.defined() &&
attr_value_.defined() &&
curr_node.defined() &&
......@@ -214,7 +214,7 @@ class AttrScopeLifter : public IRMutator {
};
Stmt LiftAttrScope(Stmt stmt, std::string attr_key) {
return AttrScopeLifter(attr_key).Lift(stmt);
return AttrScopeLifter(attr_key).Lift(std::move(stmt));
}
} // namespace ir
......
src/pass/lower_custom_datatypes.cc
......@@ -21,7 +21,7 @@
* \brief Pass for lowering custom datatypes
*/
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <tvm/packed_func_ext.h>
#include "../codegen/datatype/registry.h"
......@@ -37,17 +37,17 @@ namespace ir {
* datatype) for lowering this type of expression, and uses it to lower the
* expression.
*/
class CustomDatatypesLowerer : public IRMutator {
class CustomDatatypesLowerer : public StmtExprMutator {
public:
explicit CustomDatatypesLowerer(const std::string& target) : target_(target) {}
inline Expr Mutate_(const Cast* op, const Expr& e) final {
inline Expr VisitExpr_(const Cast* op) final {
auto type_code = op->dtype.code();
auto src_type_code = op->value.dtype().code();
// If either datatype is a registered custom datatype, we must lower.
bool toBeLowered = datatype::Registry::Global()->GetTypeRegistered(type_code) ||
datatype::Registry::Global()->GetTypeRegistered(src_type_code);
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Cast>();
if (toBeLowered) {
auto lower = datatype::GetCastLowerFunc(target_, type_code, src_type_code);
......@@ -59,8 +59,9 @@ class CustomDatatypesLowerer : public IRMutator {
return expr;
}
inline Expr Mutate_(const FloatImm* imm, const Expr& e) final {
inline Expr VisitExpr_(const FloatImm* imm) final {
auto type_code = imm->dtype.code();
auto e = GetRef<Expr>(imm);
if (datatype::Registry::Global()->GetTypeRegistered(type_code)) {
auto lower = datatype::GetFloatImmLowerFunc(target_, type_code);
CHECK(lower) << "FloatImm lowering function for target " << target_ << " type "
......@@ -70,9 +71,9 @@ class CustomDatatypesLowerer : public IRMutator {
return e;
}
inline Stmt Mutate_(const Allocate* allocate, const Stmt& s) final {
inline Stmt VisitStmt_(const Allocate* allocate) final {
bool toBeLowered = datatype::Registry::Global()->GetTypeRegistered(allocate->dtype.code());
Stmt stmt = IRMutator::Mutate_(allocate, s);
Stmt stmt = StmtExprMutator::VisitStmt_(allocate);
allocate = stmt.as<Allocate>();
if (toBeLowered) {
......@@ -84,9 +85,9 @@ class CustomDatatypesLowerer : public IRMutator {
return stmt;
}
inline Expr Mutate_(const Load* load, const Expr& e) final {
inline Expr VisitExpr_(const Load* load) final {
bool toBeLowered = datatype::Registry::Global()->GetTypeRegistered(load->dtype.code());
Expr expr = IRMutator::Mutate_(load, e);
Expr expr = StmtExprMutator::VisitExpr_(load);
load = expr.as<Load>();
if (toBeLowered) {
auto new_load_type = DataType::UInt(load->dtype.bits());
......@@ -96,10 +97,10 @@ class CustomDatatypesLowerer : public IRMutator {
}
#define DEFINE_MUTATE__(OP) \
inline Expr Mutate_(const OP* op, const Expr& e) final { \
inline Expr VisitExpr_(const OP* op) final { \
auto type_code = op->dtype.code(); \
bool toBeLowered = datatype::Registry::Global()->GetTypeRegistered(type_code); \
Expr expr = IRMutator::Mutate_(op, e); \
Expr expr = StmtExprMutator::VisitExpr_(op); \
op = expr.as<OP>(); \
if (toBeLowered) { \
auto lower = datatype::Get##OP##LowerFunc(target_, type_code); \
......@@ -131,7 +132,7 @@ class CustomDatatypesLowerer : public IRMutator {
LoweredFunc LowerCustomDatatypes(LoweredFunc f, const std::string& target) {
auto n = make_object<LoweredFuncNode>(*f.operator->());
n->body = CustomDatatypesLowerer(target).Mutate(n->body);
n->body = CustomDatatypesLowerer(target)(n->body);
return LoweredFunc(n);
}
......
src/pass/lower_intrin.cc
......@@ -22,7 +22,6 @@
* \file lower_intrin.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
#include <tvm/api_registry.h>
#include <tvm/expr_operator.h>
......@@ -34,9 +33,10 @@
namespace tvm {
namespace ir {
class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
class IntrinInjecter : public tvm::arith::IRMutatorWithAnalyzer {
public:
using IRMutatorWithAnalyzer::Mutate_;
using IRMutatorWithAnalyzer::VisitStmt_;
using IRMutatorWithAnalyzer::VisitExpr_;
IntrinInjecter(arith::Analyzer* analyzer, std::string target)
: IRMutatorWithAnalyzer(analyzer) {
......@@ -51,28 +51,29 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
}
}
Expr Mutate_(const Call* op, const Expr& e) final {
Expr VisitExpr_(const Call* op) final {
if (op->call_type == Call::Intrinsic ||
op->call_type == Call::PureIntrinsic) {
Expr r = ApplyPattern(op->name, e);
Expr r = ApplyPattern(op->name, GetRef<Expr>(op));
if (r.defined()) return r;
}
return IRMutator::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
Expr Mutate_(const Add* op, const Expr& e) final {
Expr VisitExpr_(const Add* op) final {
if (const Mul* mb = op->b.as<Mul>()) {
return MakeFMA(mb->a, mb->b, op->a, op, e);
return MakeFMA(mb->a, mb->b, op->a, op);
} else if (const Mul* ma = op->a.as<Mul>()) {
return MakeFMA(ma->a, ma->b, op->b, op, e);
return MakeFMA(ma->a, ma->b, op->b, op);
}
return IRMutator::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
// We use floordiv for integer analysis,
// but will need to lower them to native truncdiv instructions
Expr Mutate_(const FloorDiv* op, const Expr& e) final {
Expr ret = IRMutatorWithAnalyzer::Mutate_(op, e);
Expr VisitExpr_(const FloorDiv* op) final {
auto e = GetRef<Expr>(op);
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<FloorDiv>();
if (op == nullptr) return ret;
int shift;
......@@ -117,8 +118,8 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
}
}
Expr Mutate_(const FloorMod* op, const Expr& e) final {
Expr ret = IRMutatorWithAnalyzer::Mutate_(op, e);
Expr VisitExpr_(const FloorMod* op) final {
Expr ret = IRMutatorWithAnalyzer::VisitExpr_(op);
op = ret.as<FloorMod>();
if (op == nullptr) return ret;
// Lower floordiv to native truncdiv.
......@@ -167,34 +168,37 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
}
}
Expr Mutate_(const Max* op, const Expr& e) final {
Expr VisitExpr_(const Max* op) final {
using namespace arith;
PVar<Expr> x, y;
PVar<Integer> c;
auto e = GetRef<Expr>(op);
if (max(floordiv(x, y), c).Match(e) &&
c.Eval()->value >= 0 &&
analyzer_->CanProveGreaterEqual(y.Eval(), 0)) {
return max(Mutate(truncdiv(x, y).Eval()), c.Eval());
return max(VisitExpr(truncdiv(x, y).Eval()), c.Eval());
}
return IRMutatorWithAnalyzer::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
Expr Mutate_(const EQ* op, const Expr& e) final {
Expr VisitExpr_(const EQ* op) final {
using namespace arith;
PVar<Expr> x, y;
auto e = GetRef<Expr>(op);
if ((floormod(x, y) == 0).Match(e)) {
return Mutate((truncmod(x, y) == 0).Eval());
return VisitExpr((truncmod(x, y) == 0).Eval());
}
return IRMutatorWithAnalyzer::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
Expr Mutate_(const NE* op, const Expr& e) final {
Expr VisitExpr_(const NE* op) final {
using namespace arith;
PVar<Expr> x, y;
auto e = GetRef<Expr>(op);
if ((floormod(x, y) != 0).Match(e)) {
return Mutate((truncmod(x, y) != 0).Eval());
return VisitExpr((truncmod(x, y) != 0).Eval());
}
return IRMutatorWithAnalyzer::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
private:
......@@ -231,7 +235,7 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
}
Expr MakeFMA(const Expr& a, const Expr& b, const Expr& c,
const Add* op, const Expr& e) {
const Add* op) {
// emit fma instruction: a * b + c
Expr lhs = SwapBroadcastCast(a);
Expr rhs = SwapBroadcastCast(b);
......@@ -239,14 +243,14 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
if (fma_ != nullptr && op->dtype.is_float()) {
Expr r = (*fma_)(Call::make(
op->dtype, "fma", {lhs, rhs, c}, Call::PureIntrinsic));
if (r.defined()) return this->Mutate(r);
if (r.defined()) return this->VisitExpr(r);
} else {
if (!lhs.same_as(a) || !rhs.same_as(b)) {
Expr mul = this->Mutate(Mul::make(lhs, rhs));
return Add::make(mul, this->Mutate(c));
Expr mul = this->VisitExpr(Mul::make(lhs, rhs));
return Add::make(mul, this->VisitExpr(c));
}
}
return IRMutator::Mutate_(op, e);
return IRMutatorWithAnalyzer::VisitExpr_(op);
}
Expr ApplyPattern(const std::string& name, const Expr& e) {
......@@ -262,7 +266,7 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
Expr r = (*f)(e);
CHECK(r.defined()) << "intrinsic rule must always return valid Expr";
if (!r.same_as(e)) {
return this->Mutate(r);
return this->VisitExpr(r);
}
}
}
......@@ -277,7 +281,7 @@ class IntrinInjecter : public arith::IRMutatorWithAnalyzer {
Stmt LowerIntrinStmt(Stmt stmt, const std::string& target) {
arith::Analyzer analyzer;
return IntrinInjecter(&analyzer, target).Mutate(stmt);
return IntrinInjecter(&analyzer, target)(std::move(stmt));
}
LoweredFunc
......
src/pass/lower_thread_allreduce.cc
......@@ -22,7 +22,7 @@
* \file lower_thread_allreduce.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <unordered_set>
#include "ir_util.h"
......@@ -32,19 +32,19 @@
namespace tvm {
namespace ir {
class ThreadAllreduceBuilder final : public IRMutator {
class ThreadAllreduceBuilder final : public StmtExprMutator {
public:
explicit ThreadAllreduceBuilder(int warp_size)
: warp_size_(warp_size) {}
Stmt Mutate_(const AttrStmt *op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt *op) final {
if (op->attr_key == attr::thread_extent) {
thread_extents_.push_back(op);
Stmt ret = IRMutator::Mutate_(op, s);
Stmt ret = StmtExprMutator::VisitStmt_(op);
thread_extents_.pop_back();
return ret;
} else if (op->attr_key == attr::storage_scope) {
Stmt ret = IRMutator::Mutate_(op, s);
Stmt ret = StmtExprMutator::VisitStmt_(op);
op = ret.as<AttrStmt>();
const Variable* v = op->node.as<Variable>();
if (alloc_remap_.count(v)) {
......@@ -56,15 +56,15 @@ class ThreadAllreduceBuilder final : public IRMutator {
const CommReducerNode *combiner = op->node.as<CommReducerNode>();
CHECK(combiner);
reduce_combiner_.push_back(combiner);
Stmt ret = IRMutator::Mutate_(op, s);
Stmt ret = StmtExprMutator::VisitStmt_(op);
reduce_combiner_.pop_back();
return ret;
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const Evaluate* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Evaluate* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Evaluate>();
const Call* call = op->value.as<Call>();
if (call && call->is_intrinsic(intrinsic::tvm_thread_allreduce)) {
......@@ -73,8 +73,8 @@ class ThreadAllreduceBuilder final : public IRMutator {
return stmt;
}
}
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Allocate* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
auto it = alloc_remap_.find(op->buffer_var.get());
if (it != alloc_remap_.end()) {
......@@ -93,13 +93,13 @@ class ThreadAllreduceBuilder final : public IRMutator {
return stmt;
}
}
Expr Mutate_(const Load* op, const Expr& e) final {
Expr VisitExpr_(const Load* op) final {
auto it = load_remap_.find(op->buffer_var.get());
if (it != load_remap_.end()) {
CHECK(is_zero(op->index)) << e;
CHECK(is_zero(op->index));
return it->second;
} else {
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
}
......@@ -339,7 +339,7 @@ LoweredFunc
LowerThreadAllreduce(LoweredFunc f, int warp_size) {
CHECK_NE(f->func_type, kHostFunc);
auto n = make_object<LoweredFuncNode>(*f.operator->());
n->body = ThreadAllreduceBuilder(warp_size).Mutate(n->body);
n->body = ThreadAllreduceBuilder(warp_size)(n->body);
return LoweredFunc(n);
}
} // namespace ir
......
src/pass/lower_tvm_builtin.cc
......@@ -22,7 +22,7 @@
* \file lower_tvm_buildin.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <unordered_set>
#include "ir_util.h"
......@@ -43,14 +43,14 @@ inline Expr StackAlloca(std::string type, size_t num) {
// Calculate the statistics of packed function.
// These information are needed during codegen.
class BuiltinLower : public IRMutator {
class BuiltinLower : public StmtExprMutator {
public:
Stmt Build(Stmt stmt) {
stack_shape_ = Var("stack_shape", DataType::Handle());
stack_array_ = Var("stack_array", DataType::Handle());
stack_value_ = Var("stack_value", DataType::Handle());
stack_tcode_ = Var("stack_tcode", DataType::Handle());
stmt = this->Mutate(stmt);
stmt = this->VisitStmt(stmt);
if (max_shape_stack_ != 0) {
stmt = LetStmt::make(
stack_shape_, StackAlloca("shape", max_shape_stack_), stmt);
......@@ -68,8 +68,8 @@ class BuiltinLower : public IRMutator {
return stmt;
}
Stmt Mutate(Stmt stmt) final {
stmt = IRMutator::Mutate(stmt);
Stmt VisitStmt(const Stmt& s) final {
auto stmt = StmtExprMutator::VisitStmt(s);
CHECK_EQ(run_shape_stack_, 0);
CHECK_EQ(run_array_stack_, 0);
while (prep_seq_.size() != 0) {
......@@ -79,9 +79,9 @@ class BuiltinLower : public IRMutator {
return stmt;
}
Stmt Mutate_(const Allocate* op, const Stmt& s) {
Stmt VisitStmt_(const Allocate* op) {
// Lower allocate to device allocate when needed.
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
if (op->new_expr.defined()) return stmt;
// Get constant allocation bound.
......@@ -141,39 +141,39 @@ class BuiltinLower : public IRMutator {
return body;
}
Stmt Mutate_(const AttrStmt* op, const Stmt &s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::device_context_id) {
CHECK(!device_id_.defined());
device_id_ = op->value;
return Mutate(op->body);
return this->VisitStmt(op->body);
} else if (op->attr_key == attr::device_context_type) {
CHECK(!device_type_.defined());
device_type_ = op->value;
return Mutate(op->body);
return this->VisitStmt(op->body);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Expr Mutate_(const Call* op, const Expr &e) final {
Expr VisitExpr_(const Call* op) final {
if (op->is_intrinsic(intrinsic::tvm_call_packed)) {
return MakeCallPacked(op, e);
return MakeCallPacked(op);
} else if (op->is_intrinsic(intrinsic::tvm_call_trace_packed)) {
return MakeCallTracePacked(op, e);
return MakeCallTracePacked(op);
} else if (op->is_intrinsic(intrinsic::tvm_stack_make_shape)) {
return MakeShape(op, e);
return MakeShape(op);
} else if (op->is_intrinsic(intrinsic::tvm_stack_make_array)) {
return MakeArray(op, e);
return MakeArray(op);
} else if (op->is_intrinsic(intrinsic::tvm_context_id)) {
return make_zero(op->dtype);
} else {
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
}
// call shape
Expr MakeShape(const Call* op, const Expr& e) {
Expr MakeShape(const Call* op) {
size_t stack_begin = run_shape_stack_;
run_shape_stack_ += op->args.size();
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
for (size_t i = 0; i < op->args.size(); ++i) {
prep_seq_.emplace_back(
......@@ -183,10 +183,10 @@ class BuiltinLower : public IRMutator {
return AddressOffset(stack_shape_, DataType::Int(64), stack_begin);
}
// make array
Expr MakeArray(const Call* op, const Expr& e) {
Expr MakeArray(const Call* op) {
size_t idx = run_array_stack_;
run_array_stack_ += 1;
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
prep_seq_.emplace_back(
TVMStructSet(stack_array_, idx, intrinsic::kArrData, op->args[0]));
......@@ -230,13 +230,13 @@ class BuiltinLower : public IRMutator {
return TVMStructGet(DataType::Handle(), stack_array_, idx, intrinsic::kArrAddr);
}
// call packed.
Expr MakeCallPacked(const Call* op, const Expr& e) {
Expr MakeCallPacked(const Call* op) {
size_t restore_shape_stack = run_shape_stack_;
size_t restore_array_stack = run_array_stack_;
size_t arg_stack_begin = run_arg_stack_;
run_arg_stack_ += op->args.size();
// Specially handle the buffer packed intrinsic
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
for (size_t i = 1; i < op->args.size(); ++i) {
Expr stack_index = ConstInt32(arg_stack_begin + i - 1);
......@@ -278,14 +278,14 @@ class BuiltinLower : public IRMutator {
packed_args, Call::Intrinsic);
}
Expr MakeCallTracePacked(const Call *op, const Expr &e) {
Expr MakeCallTracePacked(const Call *op) {
size_t restore_shape_stack = run_shape_stack_;
size_t restore_array_stack = run_array_stack_;
size_t arg_stack_begin = run_arg_stack_;
run_arg_stack_ += op->args.size();
size_t args_size = op->args.size();
CHECK_GT(args_size, 0);
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
for (size_t i = 1; i < op->args.size(); ++i) {
Expr stack_index = ConstInt32(arg_stack_begin + i - 1);
......
src/pass/lower_warp_memory.cc
......@@ -26,8 +26,7 @@
// Thanks to Andrew Adams and Vinod Grover for
// explaining the concept of warp shuffle.
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <unordered_set>
#include "ir_util.h"
......@@ -75,7 +74,7 @@ namespace ir {
// Visitor to find m in pattern
// store warp_mem[m * warp_index + (warp_size * m) * y + x]
class WarpStoreCoeffFinder : private IRVisitor {
class WarpStoreCoeffFinder : private StmtVisitor {
public:
WarpStoreCoeffFinder(const Variable* buffer,
Var warp_index,
......@@ -86,13 +85,13 @@ class WarpStoreCoeffFinder : private IRVisitor {
}
// find the warp co-efficient in the statement given the warp size
int Find(const Stmt& stmt) {
this->Visit(stmt);
this->VisitStmt(stmt);
return warp_coeff_;
}
private:
/// Visitor implementation
void Visit_(const Store *op) final {
void VisitStmt_(const Store *op) final {
if (op->buffer_var.get() == buffer_) {
if (op->value.dtype().lanes() == 1) {
UpdatePattern(op->index);
......@@ -104,7 +103,7 @@ class WarpStoreCoeffFinder : private IRVisitor {
UpdatePattern(base);
}
} else {
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
}
......@@ -141,14 +140,14 @@ class WarpStoreCoeffFinder : private IRVisitor {
// Visitor to find the warp index
class WarpIndexFinder : private IRVisitor {
class WarpIndexFinder : private StmtVisitor {
public:
explicit WarpIndexFinder(int warp_size)
: warp_size_(warp_size) {
}
// find the warp co-efficient in the statement given the warp size
IterVar Find(const Stmt& stmt) {
this->Visit(stmt);
this->VisitStmt(stmt);
CHECK(warp_index_.defined())
<< "Cannot find warp index(threadIdx.x) within the scope of warp memory";
return warp_index_;
......@@ -156,7 +155,7 @@ class WarpIndexFinder : private IRVisitor {
private:
/// Visitor implementation
void Visit_(const AttrStmt *op) final {
void VisitStmt_(const AttrStmt *op) final {
if (op->attr_key == attr::thread_extent) {
IterVar iv = Downcast<IterVar>(op->node);
if (iv->thread_tag == "threadIdx.x") {
......@@ -177,7 +176,7 @@ class WarpIndexFinder : private IRVisitor {
}
}
}
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
// warp size
int warp_size_{0};
......@@ -185,13 +184,13 @@ class WarpIndexFinder : private IRVisitor {
IterVar warp_index_{nullptr};
};
// Mutator to change the read pattern
class WarpAccessRewriter : protected IRMutator {
class WarpAccessRewriter : protected StmtExprMutator {
public:
explicit WarpAccessRewriter(int warp_size, arith::Analyzer* analyzer)
: warp_size_(warp_size), analyzer_(analyzer) {}
// Rewrite the allocate statement which transforms
// warp memory to local memory.
Stmt Rewrite(const Allocate* op, const Stmt& stmt) {
Stmt Rewrite(const Allocate* op) {
buffer_ = op->buffer_var.get();
int alloc_size = op->constant_allocation_size();
CHECK_GT(alloc_size, 0)
......@@ -208,27 +207,27 @@ class WarpAccessRewriter : protected IRMutator {
op->dtype,
{make_const(DataType::Int(32), alloc_size / warp_size_)},
op->condition,
this->Mutate(op->body));
this->VisitStmt(op->body));
}
protected:
Expr Mutate_(const Variable* op, const Expr& expr) {
Expr Mutate_(const Variable* op) {
CHECK(op != buffer_)
<< "Cannot access address of warp memory directly";
return IRMutator::Mutate_(op, expr);
return StmtExprMutator::VisitExpr_(op);
}
Stmt Mutate_(const Store* op, const Stmt& stmt) {
Stmt VisitStmt_(const Store* op) {
if (op->buffer_var.get() == buffer_) {
Expr local_index, group;
std::tie(local_index, group) = SplitIndexByGroup(op->index);
return Store::make(op->buffer_var, op->value, local_index, op->predicate);
} else {
return IRMutator::Mutate_(op, stmt);
return StmtExprMutator::VisitStmt_(op);
}
}
Expr Mutate_(const Load* op, const Expr& expr) {
Expr Mutate_(const Load* op) {
if (op->buffer_var.get() == buffer_) {
Expr local_index, group;
std::tie(local_index, group) = SplitIndexByGroup(op->index);
......@@ -243,7 +242,7 @@ class WarpAccessRewriter : protected IRMutator {
{load_value, group},
Call::Intrinsic);
} else {
return IRMutator::Mutate_(op, expr);
return StmtExprMutator::VisitExpr_(op);
}
}
// Split the index to the two component
......@@ -297,18 +296,18 @@ class WarpAccessRewriter : protected IRMutator {
// Bind bound information of variables to make analyzer more effective
// TODO(tqchen): consider a pass to inline the bound info into the expr
// so analysis can be context independent.
class BindVarBoundInfo : public IRVisitor {
class BindVarBoundInfo : public StmtVisitor {
public:
explicit BindVarBoundInfo(arith::Analyzer* analyzer)
: analyzer_(analyzer) {}
void Visit_(const For* op) final {
void VisitStmt_(const For* op) final {
const Var& loop_var = op->loop_var;
analyzer_->Bind(loop_var, Range::make_by_min_extent(op->min, op->extent));
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
void Visit_(const AttrStmt* op) {
void VisitStmt_(const AttrStmt* op) {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::virtual_thread) {
IterVar iv = Downcast<IterVar>(op->node);
......@@ -319,7 +318,7 @@ class BindVarBoundInfo : public IRVisitor {
analyzer_->Bind(iv->var, dom);
}
}
IRVisitor::Visit_(op);
StmtVisitor::VisitStmt_(op);
}
protected:
......@@ -330,7 +329,7 @@ class BindVarBoundInfo : public IRVisitor {
};
// Mutator to change the read pattern
class WarpMemoryRewriter : private IRMutator {
class WarpMemoryRewriter : private StmtMutator {
public:
explicit WarpMemoryRewriter(int warp_size)
: warp_size_(warp_size) {
......@@ -338,36 +337,37 @@ class WarpMemoryRewriter : private IRMutator {
Stmt Rewrite(Stmt stmt) {
if (warp_size_ == 1) return stmt;
BindVarBoundInfo(&analyzer_).Visit(stmt);
stmt = this->Mutate(stmt);
BindVarBoundInfo binder(&analyzer_);
binder(stmt);
stmt = operator()(std::move(stmt));
stmt = CanonicalSimplify(stmt);
return stmt;
}
private:
Stmt Mutate_(const Allocate* op, const Stmt& stmt) {
Stmt VisitStmt_(const Allocate* op) {
if (warp_buffer_.count(op->buffer_var.get())) {
WarpAccessRewriter rewriter(warp_size_, &analyzer_);
return rewriter.Rewrite(op, stmt);
return rewriter.Rewrite(op);
} else {
return IRMutator::Mutate_(op, stmt);
return StmtMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& stmt) {
Stmt VisitStmt_(const AttrStmt* op) {
using runtime::StorageScope;
if (op->attr_key == attr::storage_scope) {
const Variable* buf = op->node.as<Variable>();
StorageScope scope = StorageScope::make(op->value.as<StringImm>()->value);
if (scope.rank == runtime::StorageRank::kWarp) {
warp_buffer_.insert(buf);
Stmt ret = IRMutator::Mutate_(op, stmt);
Stmt ret = StmtMutator::VisitStmt_(op);
op = ret.as<AttrStmt>();
return AttrStmt::make(
op->node, op->attr_key, StringImm::make("local"), op->body);
}
}
return IRMutator::Mutate_(op, stmt);
return StmtMutator::VisitStmt_(op);
}
int warp_size_{0};
......
src/pass/make_api.cc
......@@ -22,8 +22,7 @@
*/
#include <tvm/ir_pass.h>
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/buffer.h>
#include <tvm/runtime/device_api.h>
#include <vector>
......@@ -207,29 +206,29 @@ LoweredFunc MakeAPI(Stmt body,
return f;
}
class DeviceTypeBinder: public IRMutator {
class DeviceTypeBinder: public StmtExprMutator {
public:
explicit DeviceTypeBinder(int device_type)
: device_type_(device_type) {}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::device_context_type) {
if (const Variable* var = op->value.as<Variable>()) {
var_ = var;
Expr value = make_const(op->value.dtype(), device_type_);
Stmt body = IRMutator::Mutate_(op, s);
Stmt body = StmtExprMutator::VisitStmt_(op);
var_ = nullptr;
std::ostringstream os;
os << "device_type need to be " << device_type_;
return AssertStmt::make(op->value == value, os.str(), body);
}
}
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Stmt Mutate_(const IfThenElse* op, const Stmt& s) final {
Stmt VisitStmt_(const IfThenElse* op) final {
// eager simplify if guard.
Stmt res = IRMutator::Mutate_(op, s);
Stmt res = StmtExprMutator::VisitStmt_(op);
op = res.as<IfThenElse>();
if (is_zero(op->condition)) {
if (op->else_case.defined()) return op->else_case;
......@@ -241,9 +240,9 @@ class DeviceTypeBinder: public IRMutator {
return res;
}
Expr Mutate_(const NE* op, const Expr& e) final {
Expr VisitExpr_(const NE* op) final {
// eager check NE for device check
Expr res = IRMutator::Mutate_(op, e);
Expr res = StmtExprMutator::VisitExpr_(op);
op = res.as<NE>();
if (ir::Equal(op->a, op->b)) {
return make_const(op->dtype, false);
......@@ -251,11 +250,11 @@ class DeviceTypeBinder: public IRMutator {
return res;
}
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
if (op == var_) {
return make_const(op->dtype, device_type_);
} else {
return e;
return GetRef<Expr>(op);
}
}
......@@ -267,7 +266,7 @@ class DeviceTypeBinder: public IRMutator {
LoweredFunc BindDeviceType(LoweredFunc f,
int device_type) {
auto n = make_object<LoweredFuncNode>(*f.operator->());
n->body = DeviceTypeBinder(device_type).Mutate(n->body);
n->body = DeviceTypeBinder(device_type)(n->body);
return LoweredFunc(n);
}
......
src/pass/remap_thread_axis.cc
......@@ -21,8 +21,7 @@
* \file remap_thread_axis.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <unordered_map>
......@@ -31,7 +30,7 @@ namespace tvm {
namespace ir {
// Mutator to change the read pattern
class ThreadAxisRewriter : private IRMutator {
class ThreadAxisRewriter : private StmtExprMutator {
public:
explicit ThreadAxisRewriter(
const std::unordered_map<std::string, IterVar>& tmap)
......@@ -39,11 +38,11 @@ class ThreadAxisRewriter : private IRMutator {
}
Stmt Rewrite(Stmt stmt) {
return Mutate(stmt);
return operator()(std::move(stmt));
}
private:
Stmt Mutate_(const AttrStmt* op, const Stmt& stmt) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::thread_extent) {
IterVar iv = Downcast<IterVar>(op->node);
CHECK_NE(iv->thread_tag.length(), 0U);
......@@ -56,18 +55,18 @@ class ThreadAxisRewriter : private IRMutator {
} else {
CHECK(vmap_[v].same_as(new_iv->var));
}
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
return AttrStmt::make(
new_iv, op->attr_key, op->value, body);
}
}
return IRMutator::Mutate_(op, stmt);
return StmtExprMutator::VisitStmt_(op);
}
Expr Mutate_(const Variable* op, const Expr& expr) final {
Expr VisitExpr_(const Variable* op) final {
auto it = vmap_.find(op);
if (it != vmap_.end()) return it->second;
return IRMutator::Mutate_(op, expr);
return StmtExprMutator::VisitExpr_(op);
}
// The thread map
const std::unordered_map<std::string, IterVar>& tmap_;
......
src/pass/remove_no_op.cc
......@@ -6,9 +6,9 @@
* 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
......@@ -23,30 +23,30 @@
*/
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <unordered_map>
namespace tvm {
namespace ir {
// Mark the statment of each stage.
class NoOpRemover : public IRMutator {
class NoOpRemover : public StmtMutator {
public:
Stmt Mutate_(const LetStmt* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const LetStmt* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<LetStmt>();
return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == "pragma_debug_skip_region") {
return MakeEvaluate(0);
}
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<AttrStmt>();
return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
}
Stmt Mutate_(const IfThenElse* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const IfThenElse* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<IfThenElse>();
if (op->else_case.defined()) {
if (is_no_op(op->else_case)) {
......@@ -66,35 +66,35 @@ class NoOpRemover : public IRMutator {
}
}
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const For* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<For>();
if (is_zero(op->extent)) {
return Evaluate::make(0);
}
return is_no_op(op->body) ? MakeEvaluate({op->min, op->extent}) : stmt;
}
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Allocate* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
return is_no_op(op->body) ? MakeEvaluate(op->extents) : stmt;
}
Stmt Mutate_(const ProducerConsumer* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const ProducerConsumer* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<ProducerConsumer>();
return is_no_op(op->body) ? op->body : stmt;
}
Stmt Mutate_(const Realize* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Realize* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<Realize>();
return is_no_op(op->body) ? op->body : stmt;
}
Stmt Mutate_(const Evaluate* op, const Stmt& s) final {
if (HasSideEffect(op->value)) return s;
Stmt VisitStmt_(const Evaluate* op) final {
if (HasSideEffect(op->value)) return GetRef<Stmt>(op);
return Evaluate::make(0);
}
Stmt Mutate_(const Block* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Block* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<Block>();
if (is_no_op(op->first)) {
return op->rest;
......@@ -129,7 +129,7 @@ class NoOpRemover : public IRMutator {
};
Stmt RemoveNoOp(Stmt stmt) {
return NoOpRemover().Mutate(stmt);
return NoOpRemover()(std::move(stmt));
}
} // namespace ir
} // namespace tvm
src/pass/rewrite_unsafe_select.cc
......@@ -23,7 +23,6 @@
*/
#include <tvm/ir.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_pass.h>
namespace tvm {
......@@ -109,10 +108,10 @@ class UnsafeExprDetector : public ExprFunctor<bool(const Expr& n)> {
}
};
class UnsafeSelectRewriter : public IRMutator {
class UnsafeSelectRewriter : public StmtExprMutator {
public:
Expr Mutate_(const Select* op, const Expr& e) {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Select* op) {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Select>();
UnsafeExprDetector unsafe;
bool cond_is_scalar_bool = op->condition.dtype().is_bool() && op->condition.dtype().is_scalar();
......@@ -131,7 +130,7 @@ class UnsafeSelectRewriter : public IRMutator {
};
Stmt RewriteUnsafeSelect(Stmt stmt) {
return UnsafeSelectRewriter().Mutate(stmt);
return UnsafeSelectRewriter()(std::move(stmt));
}
} // namespace ir
......
src/pass/simple_passes.cc
......@@ -22,25 +22,24 @@
* \brief Implementation of simple passes
*/
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
namespace tvm {
namespace ir {
class IRSideEffect : public IRVisitor {
class IRSideEffect : public ExprVisitor {
public:
void Visit(const ObjectRef& e) final {
void VisitExpr(const Expr& e) final {
if (has_side_effect_) return;
IRVisitor::Visit(e);
ExprVisitor::VisitExpr(e);
}
void Visit_(const Call* op) final {
void VisitExpr_(const Call* op) final {
if (!op->is_pure()) {
has_side_effect_ = true; return;
} else {
IRVisitor::Visit_(op);
ExprVisitor::VisitExpr_(op);
}
}
......@@ -49,23 +48,23 @@ class IRSideEffect : public IRVisitor {
bool HasSideEffect(const Expr& e) {
IRSideEffect v;
v.Visit(e);
v(e);
return v.has_side_effect_;
}
class IRSubstitue : public IRMutator {
class IRSubstitue : public StmtExprMutator {
public:
explicit IRSubstitue(
const std::unordered_map<const Variable*, Expr>& smap)
: smap_(smap) {
}
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
auto it = smap_.find(op);
if (it != smap_.end()) {
return it->second;
} else {
return e;
return GetRef<Expr>(op);
}
}
......@@ -76,13 +75,13 @@ class IRSubstitue : public IRMutator {
Stmt Substitute(Stmt stmt,
const std::unordered_map<const Variable*, Expr>& value_map) {
if (value_map.size() == 0) return stmt;
return IRSubstitue(value_map).Mutate(stmt);
return IRSubstitue(value_map)(std::move(stmt));
}
Expr Substitute(Expr expr,
const std::unordered_map<const Variable*, Expr>& value_map) {
if (value_map.size() == 0) return expr;
return IRSubstitue(value_map).Mutate(expr);
return IRSubstitue(value_map)(std::move(expr));
}
Stmt Substitute(Stmt stmt, const Map<Var, Expr>& value_map) {
......@@ -101,20 +100,20 @@ Expr Substitute(Expr expr, const Map<Var, Expr>& value_map) {
return Substitute(expr, vmap);
}
class VarTouchVisitor : public IRVisitor {
class VarTouchVisitor : public ExprVisitor {
public:
void Visit(const ObjectRef& e) final {
void VisitExpr(const Expr& e) final {
if (use_var_) return;
IRVisitor::Visit(e);
ExprVisitor::VisitExpr(e);
}
void Visit_(const Variable* op) final {
void VisitExpr_(const Variable* op) final {
Handle(op);
}
void Visit_(const Load* op) final {
void VisitExpr_(const Load* op) final {
Handle(op->buffer_var.get());
IRVisitor::Visit_(op);
ExprVisitor::VisitExpr_(op);
}
virtual void Handle(const Variable* var) = 0;
......@@ -149,14 +148,14 @@ class ExprUseVSetVisitor : public VarTouchVisitor {
bool ExprUseVar(const Expr& e, const Var& v) {
ExprUseVarVisitor visitor(v.get());
visitor.Visit(e);
visitor(e);
return visitor.use_var_;
}
bool ExprUseVar(const Expr& e,
const std::unordered_set<const Variable*>& vset) {
ExprUseVSetVisitor visitor(vset);
visitor.Visit(e);
visitor(e);
return visitor.use_var_;
}
......
src/pass/skip_assert.cc
......@@ -19,22 +19,22 @@
#include <tvm/ir.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
namespace tvm {
namespace ir {
class AssertSkipper : public IRMutator {
class AssertSkipper : public StmtMutator {
public:
Stmt Mutate_(const AssertStmt* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const AssertStmt* op) final {
Stmt stmt = StmtMutator::VisitStmt_(op);
op = stmt.as<AssertStmt>();
return op->body;
}
};
Stmt SkipAssert(Stmt stmt) {
return AssertSkipper().Mutate(stmt);
return AssertSkipper()(std::move(stmt));
}
LoweredFunc SkipAssert(LoweredFunc f) {
......
src/pass/split_host_device.cc
......@@ -24,7 +24,7 @@
#include <tvm/ir.h>
#include <tvm/lowered_func.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/runtime/module.h>
#include <unordered_map>
......@@ -32,9 +32,9 @@ namespace tvm {
namespace ir {
// use/def analysis, also delete unreferenced lets
class IRUseDefAnalysis : public IRMutator {
class IRUseDefAnalysis : public StmtExprMutator {
public:
Stmt Mutate_(const AttrStmt *op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::thread_extent) {
IterVar iv = Downcast<IterVar>(op->node);
CHECK_NE(iv->thread_tag.length(), 0U);
......@@ -48,75 +48,77 @@ class IRUseDefAnalysis : public IRMutator {
Expr value = op->value;
if (visit_thread_extent_) {
value = this->Mutate(value);
value = this->VisitExpr(value);
}
Stmt body = this->VisitStmt(op->body);
if (value.same_as(op->value) && body.same_as(op->body)) {
return GetRef<Stmt>(op);
}
Stmt body = this->Mutate(op->body);
if (value.same_as(op->value) && body.same_as(op->body)) return s;
return AttrStmt::make(op->node, op->attr_key, value, body);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const LetStmt *op, const Stmt& s) final {
Stmt VisitStmt_(const LetStmt* op) final {
this->HandleDef(op->var.get());
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
// eliminate unreferenced let
if (use_count_.at(op->var.get()) == 0 &&
!HasSideEffect(op->value)) {
return body;
} else {
Expr value = this->Mutate(op->value);
Expr value = this->VisitExpr(op->value);
if (body.same_as(op->body) &&
value.same_as(op->value)) {
return s;
return GetRef<Stmt>(op);
} else {
return LetStmt::make(op->var, value, body);
}
}
}
Stmt Mutate_(const For *op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
this->HandleDef(op->loop_var.get());
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Stmt Mutate_(const Allocate *op, const Stmt& s) final {
Stmt VisitStmt_(const Allocate* op) final {
this->HandleDef(op->buffer_var.get());
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Stmt Mutate_(const Store *op, const Stmt& s) final {
Stmt VisitStmt_(const Store* op) final {
this->HandleUse(op->buffer_var);
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Expr Mutate_(const Let *op, const Expr& e) final {
Expr VisitExpr_(const Let* op) final {
this->HandleDef(op->var.get());
Expr body = this->Mutate(op->body);
Expr body = this->VisitExpr(op->body);
// eliminate unreferenced let
if (use_count_.at(op->var.get()) == 0 &&
!HasSideEffect(op->value)) {
return body;
} else {
Expr value = this->Mutate(op->value);
Expr value = this->VisitExpr(op->value);
if (body.same_as(op->body) &&
value.same_as(op->value)) {
return e;
return GetRef<Expr>(op);
} else {
return Let::make(op->var, value, body);
}
}
}
Expr Mutate_(const Variable *op, const Expr& e) final {
this->HandleUse(e);
return IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Variable* op) final {
this->HandleUse(GetRef<Expr>(op));
return StmtExprMutator::VisitExpr_(op);
}
Expr Mutate_(const Load *op, const Expr& e) final {
Expr VisitExpr_(const Load* op) final {
this->HandleUse(op->buffer_var);
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
void HandleDef(const Variable* v) {
......@@ -154,20 +156,20 @@ class IRUseDefAnalysis : public IRMutator {
std::unordered_map<const Variable*, int> def_count_;
};
class HostDeviceSplitter : public IRMutator {
class HostDeviceSplitter : public StmtMutator {
public:
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt VisitStmt_(const Allocate* op) final {
handle_data_type_[op->buffer_var.get()] = make_const(op->dtype, 0);
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
Stmt Mutate_(const AttrStmt *op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::thread_extent ||
op->attr_key == attr::pipeline_exec_scope ||
op->attr_key == attr::device_scope) {
return SplitDeviceFunc(s);
return SplitDeviceFunc(GetRef<Stmt>(op));
}
return IRMutator::Mutate_(op, s);
return StmtMutator::VisitStmt_(op);
}
Array<LoweredFunc> Split(LoweredFunc f) {
......@@ -178,7 +180,7 @@ class HostDeviceSplitter : public IRMutator {
name_ = f->name;
ObjectPtr<LoweredFuncNode> n =
make_object<LoweredFuncNode>(*f.operator->());
n->body = this->Mutate(f->body);
n->body = operator()(f->body);
n->func_type = kHostFunc;
Array<LoweredFunc> ret{LoweredFunc(n)};
for (LoweredFunc x : device_funcs_) {
......@@ -195,7 +197,7 @@ class HostDeviceSplitter : public IRMutator {
// isolate the device function.
IRUseDefAnalysis m;
m.visit_thread_extent_ = false;
n->body = m.Mutate(body);
n->body = m(std::move(body));
n->name = os.str();
n->func_type = kDeviceFunc;
n->thread_axis = m.thread_axis_;
......@@ -243,7 +245,7 @@ Array<Var> UndefinedVars(const Stmt& stmt, const Array<Var>& args) {
for (Var arg : args) {
m.use_count_[arg.get()] = 0;
}
m.Mutate(stmt);
m(stmt);
return m.undefined_;
}
......
src/pass/ssa.cc
......@@ -24,8 +24,7 @@
* \file ssa.cc
*/
#include <tvm/ir.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/ir_pass.h>
#include <unordered_set>
#include <unordered_map>
......@@ -34,29 +33,33 @@
namespace tvm {
namespace ir {
namespace {
class IRVerifySSA final : public IRVisitor {
class IRVerifySSA final : public StmtExprVisitor {
public:
bool is_ssa{true};
void Visit(const ObjectRef& n) final {
void VisitExpr(const Expr& n) final {
if (!is_ssa) return;
IRVisitor::Visit(n);
StmtExprVisitor::VisitExpr(n);
}
void Visit_(const Let* op) final {
void VisitStmt(const Stmt& n) final {
if (!is_ssa) return;
StmtExprVisitor::VisitStmt(n);
}
void VisitExpr_(const Let* op) final {
MarkDef(op->var.get());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
void Visit_(const LetStmt* op) final {
void VisitStmt_(const LetStmt* op) final {
MarkDef(op->var.get());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
void Visit_(const For* op) final {
void VisitStmt_(const For* op) final {
MarkDef(op->loop_var.get());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
void Visit_(const Allocate* op) final {
void VisitStmt_(const Allocate* op) final {
MarkDef(op->buffer_var.get());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
private:
......@@ -70,31 +73,32 @@ class IRVerifySSA final : public IRVisitor {
std::unordered_map<const Variable*, int> defined_;
};
class IRConvertSSA final : public IRMutator {
class IRConvertSSA final : public StmtExprMutator {
public:
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
if (scope_.count(op)) {
return scope_[op].back();
} else {
return e;
return GetRef<Expr>(op);
}
}
Expr Mutate_(const Let* op, const Expr& e) final {
Expr VisitExpr_(const Let* op) final {
const VarExpr& v = op->var;
if (defined_.count(v.get())) {
Expr value = IRMutator::Mutate(op->value);
Expr value = this->VisitExpr(op->value);
VarExpr new_var = Variable::make(v.dtype(), v->name_hint);
scope_[v.get()].push_back(new_var);
Expr body = IRMutator::Mutate(op->body);
Expr body = this->VisitExpr(op->body);
scope_[v.get()].pop_back();
return Let::make(new_var, value, body);
} else {
defined_.insert(v.get());
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
}
Expr Mutate_(const Load* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Load* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Load>();
if (scope_.count(op->buffer_var.get())) {
return Load::make(
......@@ -104,8 +108,8 @@ class IRConvertSSA final : public IRMutator {
return expr;
}
}
Stmt Mutate_(const Store* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Store* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Store>();
if (scope_.count(op->buffer_var.get())) {
return Store::make(
......@@ -115,41 +119,41 @@ class IRConvertSSA final : public IRMutator {
return stmt;
}
}
Stmt Mutate_(const LetStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const LetStmt* op) final {
const VarExpr& v = op->var;
if (defined_.count(v.get())) {
Expr value = IRMutator::Mutate(op->value);
Expr value = this->VisitExpr(op->value);
VarExpr new_var = Variable::make(v.dtype(), v->name_hint);
scope_[v.get()].push_back(new_var);
Stmt body = IRMutator::Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
scope_[v.get()].pop_back();
return LetStmt::make(new_var, value, body);
} else {
defined_.insert(v.get());
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const For* op, const Stmt& s) final {
Stmt VisitStmt_(const For* op) final {
const VarExpr& v = op->loop_var;
if (defined_.count(v.get())) {
VarExpr new_var = Variable::make(v.dtype(), v->name_hint);
scope_[v.get()].push_back(new_var);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
scope_[v.get()].pop_back();
op = stmt.as<For>();
return For::make(
new_var, op->min, op->extent, op->for_type, op->device_api, op->body);
} else {
defined_.insert(v.get());
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt VisitStmt_(const Allocate* op) final {
const VarExpr& v = op->buffer_var;
if (defined_.count(v.get())) {
VarExpr new_var = Variable::make(v.dtype(), v->name_hint);
scope_[v.get()].push_back(new_var);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
scope_[v.get()].pop_back();
op = stmt.as<Allocate>();
return Allocate::make(
......@@ -157,23 +161,23 @@ class IRConvertSSA final : public IRMutator {
op->body, op->new_expr, op->free_function);
} else {
defined_.insert(v.get());
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (const Variable* v = op->node.as<Variable>()) {
if (op->attr_key == attr::storage_scope) {
const Allocate* alloc = op->body.as<Allocate>();
if (alloc && op->node.same_as(alloc->buffer_var)) {
Stmt new_alloc = Mutate(op->body);
if (new_alloc.same_as(op->body)) return s;
Stmt new_alloc = this->VisitStmt(op->body);
if (new_alloc.same_as(op->body)) return GetRef<Stmt>(op);
alloc = new_alloc.as<Allocate>();
CHECK(alloc);
return AttrStmt::make(
alloc->buffer_var, op->attr_key, op->value, new_alloc);
}
}
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<AttrStmt>();
if (scope_.count(v) && scope_[v].size() != 0) {
return AttrStmt::make(
......@@ -182,7 +186,7 @@ class IRConvertSSA final : public IRMutator {
return stmt;
}
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
......@@ -194,13 +198,13 @@ class IRConvertSSA final : public IRMutator {
} // namespace
bool VerifySSA(const Stmt& ir) {
IRVerifySSA v;
v.Visit(ir);
return v.is_ssa;
IRVerifySSA visitor;
visitor(ir);
return visitor.is_ssa;
}
Stmt ConvertSSA(Stmt stmt) {
return IRConvertSSA().Mutate(stmt);
return IRConvertSSA()(std::move(stmt));
}
} // namespace ir
......
src/pass/storage_access.cc
......@@ -21,7 +21,6 @@
* \file storage_access.cc
*/
#include <tvm/ir_pass.h>
#include <tvm/ir_mutator.h>
#include <tvm/target_info.h>
#include <string>
#include <utility>
......@@ -32,7 +31,7 @@
namespace tvm {
namespace ir {
void StorageAccessVisitor::Visit_(const Load* op) {
void StorageAccessVisitor::VisitExpr_(const Load* op) {
const Variable* buf = op->buffer_var.as<Variable>();
StorageScope scope = GetScope(buf);
if (Enabled(buf, scope)) {
......@@ -47,10 +46,10 @@ void StorageAccessVisitor::Visit_(const Load* op) {
curr_stmt_.access.emplace_back(std::move(e));
}
// traverse child
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
void StorageAccessVisitor::Visit_(const Store* op) {
void StorageAccessVisitor::VisitStmt_(const Store* op) {
allow_append_ = true;
CHECK_EQ(curr_stmt_.access.size(), 0U);
curr_stmt_.stmt = op;
......@@ -67,7 +66,7 @@ void StorageAccessVisitor::Visit_(const Store* op) {
curr_stmt_.access.emplace_back(std::move(e));
}
// traverse child
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
// push to the scope
scope_.back().push_back(curr_stmt_);
// clear access entry.
......@@ -75,11 +74,11 @@ void StorageAccessVisitor::Visit_(const Store* op) {
allow_append_ = false;
}
void StorageAccessVisitor::Visit_(const Evaluate* op) {
void StorageAccessVisitor::VisitStmt_(const Evaluate* op) {
allow_append_ = true;
CHECK_EQ(curr_stmt_.access.size(), 0U);
curr_stmt_.stmt = op;
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
// push to the scope
if (curr_stmt_.access.size() != 0) {
scope_.back().push_back(curr_stmt_);
......@@ -88,17 +87,17 @@ void StorageAccessVisitor::Visit_(const Evaluate* op) {
allow_append_ = false;
}
void StorageAccessVisitor::Visit_(const AttrStmt* op) {
void StorageAccessVisitor::VisitStmt_(const AttrStmt* op) {
if (op->attr_key == attr::storage_scope) {
const Variable* buf = op->node.as<Variable>();
storage_scope_[buf] =
StorageScope::make(op->value.as<StringImm>()->value);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
} else if (op->attr_key == attr::double_buffer_write) {
CHECK(double_buffer_write_ == nullptr);
double_buffer_write_ = op->node.as<Variable>();
scope_.push_back(std::vector<StmtEntry>());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
StmtEntry s;
s.stmt = op;
s.access = Summarize(std::move(scope_.back()), nullptr);
......@@ -115,7 +114,7 @@ void StorageAccessVisitor::Visit_(const AttrStmt* op) {
} else if (op->attr_key == attr::coproc_scope) {
IterVar iv = Downcast<IterVar>(op->node);
env_threads_.push_back(iv);
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
env_threads_.CopyOnWrite()->data.pop_back();
} else if (op->attr_key == attr::thread_extent) {
IterVar iv = Downcast<IterVar>(op->node);
......@@ -123,23 +122,23 @@ void StorageAccessVisitor::Visit_(const AttrStmt* op) {
if (!in_device_env_) {
in_device_env_ = true;
scope_.push_back(std::vector<StmtEntry>());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
// no need to take the result as the thread barrier automatically syncs.
Summarize(std::move(scope_.back()), nullptr);
in_device_env_ = false;
scope_.pop_back();
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
env_threads_.CopyOnWrite()->data.pop_back();
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
}
}
void StorageAccessVisitor::Visit_(const For* op) {
void StorageAccessVisitor::VisitStmt_(const For* op) {
scope_.push_back(std::vector<StmtEntry>());
IRVisitor::Visit_(op);
StmtExprVisitor::VisitStmt_(op);
StmtEntry s;
s.stmt = op;
s.access = Summarize(std::move(scope_.back()), op);
......@@ -161,11 +160,11 @@ void StorageAccessVisitor::Visit_(const For* op) {
}
}
void StorageAccessVisitor::Visit_(const IfThenElse* op) {
void StorageAccessVisitor::VisitStmt_(const IfThenElse* op) {
++condition_counter_;
this->Visit(op->condition);
this->VisitExpr(op->condition);
scope_.push_back(std::vector<StmtEntry>());
this->Visit(op->then_case);
this->VisitStmt(op->then_case);
StmtEntry s;
s.stmt = op;
s.access = Summarize(std::move(scope_.back()), nullptr);
......@@ -180,10 +179,10 @@ void StorageAccessVisitor::Visit_(const IfThenElse* op) {
--condition_counter_;
}
void StorageAccessVisitor::Visit_(const Call* op) {
void StorageAccessVisitor::VisitExpr_(const Call* op) {
if (op->is_intrinsic(intrinsic::tvm_address_of)) {
const Load *l = op->args[0].as<Load>();
IRVisitor::Visit_(l);
StmtExprVisitor::VisitExpr_(l);
} else if (op->is_intrinsic(intrinsic::tvm_access_ptr)) {
CHECK_EQ(op->args.size(), 5U);
DataType dtype = op->args[0].dtype();
......@@ -211,7 +210,7 @@ void StorageAccessVisitor::Visit_(const Call* op) {
curr_stmt_.access.emplace_back(e);
}
}
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
} else if (op->is_intrinsic(intrinsic::tvm_storage_sync)) {
CHECK(allow_append_);
const std::string& s = op->args[0].as<StringImm>()->value;
......@@ -224,7 +223,7 @@ void StorageAccessVisitor::Visit_(const Call* op) {
curr_stmt_.access.emplace_back(std::move(e));
}
} else {
IRVisitor::Visit_(op);
StmtExprVisitor::VisitExpr_(op);
}
}
......@@ -236,11 +235,12 @@ StorageScope StorageAccessVisitor::GetScope(const Variable* buf) const {
return it->second;
}
class StorageAccessInfoLower : public IRMutator {
class StorageAccessInfoLower : public StmtExprMutator {
public:
Stmt Mutate_(const Allocate* op, const Stmt& s) final {
Stmt VisitStmt_(const Allocate* op) final {
// Lower allocate to device allocate when needed.
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Allocate>();
// For special memory, remove allocate, or use head expr
auto it = storage_info_.find(op->buffer_var.get());
......@@ -259,7 +259,7 @@ class StorageAccessInfoLower : public IRMutator {
return stmt;
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::storage_scope) {
const Variable* buf = op->node.as<Variable>();
StorageScope scope = StorageScope::make(op->value.as<StringImm>()->value);
......@@ -270,26 +270,26 @@ class StorageAccessInfoLower : public IRMutator {
CHECK(e.info.defined()) << "Cannot find memory info of " << scope.to_string();
}
storage_info_[buf] = e;
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
} else {
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
}
Expr Mutate_(const Call* op, const Expr &e) final {
Expr VisitExpr_(const Call* op) final {
if (op->is_intrinsic(intrinsic::tvm_access_ptr)) {
return MakeAccessPtr(op, e);
return MakeAccessPtr(op);
} else {
return IRMutator::Mutate_(op, e);
return StmtExprMutator::VisitExpr_(op);
}
}
private:
// tvm_access_ptr
Expr MakeAccessPtr(const Call* op, const Expr& e) {
Expr MakeAccessPtr(const Call* op) {
// Specially handle the buffer packed intrinsic
Expr expr = IRMutator::Mutate_(op, e);
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
CHECK_EQ(op->args.size(), 5U);
DataType dtype = op->args[0].dtype();
......@@ -337,7 +337,7 @@ class StorageAccessInfoLower : public IRMutator {
};
Stmt LowerStorageAccessInfo(Stmt stmt) {
return StorageAccessInfoLower().Mutate(stmt);
return StorageAccessInfoLower()(std::move(stmt));
}
LoweredFunc LowerDeviceStorageAccessInfo(LoweredFunc f) {
......
src/pass/storage_access.h
......@@ -27,7 +27,7 @@
#include <tvm/ir.h>
#include <tvm/attrs.h>
#include <tvm/ir_pass.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <vector>
#include <unordered_map>
#include "../runtime/thread_storage_scope.h"
......@@ -40,7 +40,7 @@ using runtime::StorageRank;
/*!
* \brief Base class of storage access analysis
*/
class StorageAccessVisitor : public IRVisitor {
class StorageAccessVisitor : public StmtExprVisitor {
public:
/*! \brief Storage access type */
enum AccessType {
......@@ -76,13 +76,13 @@ class StorageAccessVisitor : public IRVisitor {
std::vector<AccessEntry> access;
};
// override visitor pattern
void Visit_(const Load* op) final;
void Visit_(const Store* op) final;
void Visit_(const Evaluate* op) final;
void Visit_(const AttrStmt* op) final;
void Visit_(const For* op) final;
void Visit_(const IfThenElse* op) final;
void Visit_(const Call* op) final;
void VisitExpr_(const Load* op) final;
void VisitStmt_(const Store* op) final;
void VisitStmt_(const Evaluate* op) final;
void VisitStmt_(const AttrStmt* op) final;
void VisitStmt_(const For* op) final;
void VisitStmt_(const IfThenElse* op) final;
void VisitExpr_(const Call* op) final;
protected:
StorageAccessVisitor() {
......
src/pass/storage_flatten.cc
......@@ -26,8 +26,7 @@
#include <tvm/ir.h>
#include <tvm/expr.h>
#include <tvm/operation.h>
#include <tvm/ir_mutator.h>
#include <tvm/ir_visitor.h>
#include <tvm/ir_functor_ext.h>
#include <tvm/expr_operator.h>
#include <tvm/ir_pass.h>
#include <tvm/buffer.h>
......@@ -48,7 +47,7 @@ using runtime::StorageScope;
using runtime::ThreadScope;
using intrinsic::tvm_address_of;
class StorageFlattener : public IRMutator {
class StorageFlattener : public StmtExprMutator {
public:
explicit StorageFlattener(Map<Tensor, Buffer> extern_buffer,
int cache_line_size, bool create_bound_attributes,
......@@ -64,8 +63,8 @@ class StorageFlattener : public IRMutator {
cache_line_size_ = cache_line_size;
}
Stmt Mutate_(const Store* op, const Stmt& s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Store* op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Store>();
auto it = var_remap_.find(op->buffer_var.get());
if (it != var_remap_.end() &&
......@@ -78,14 +77,14 @@ class StorageFlattener : public IRMutator {
}
}
Stmt Mutate_(const AttrStmt* op, const Stmt& s) final {
Stmt VisitStmt_(const AttrStmt* op) final {
if (op->attr_key == attr::realize_scope) {
storage_scope_[op->node.get()] = op->value.as<StringImm>()->value;
return this->Mutate(op->body);
return this->VisitStmt(op->body);
} else if (op->attr_key == attr::double_buffer_scope &&
op->node->IsInstance<OperationNode>()) {
Operation func = Downcast<Operation>(op->node);
Stmt body = Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
for (int i = 0; i < func->num_outputs(); ++i) {
TensorKey key{func, i};
auto it = buf_map_.find(key);
......@@ -99,7 +98,7 @@ class StorageFlattener : public IRMutator {
IterVar iv = Downcast<IterVar>(op->node);
ThreadScope ts = ThreadScope::make(iv->thread_tag);
curr_thread_scope_.push_back(ts);
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
curr_thread_scope_.pop_back();
return stmt;
} else if (op->attr_key == attr::buffer_bind_scope) {
......@@ -116,17 +115,17 @@ class StorageFlattener : public IRMutator {
}
vinfo[dim].align_factor = tuple->args[1].as<IntImm>()->value;
vinfo[dim].align_offset = tuple->args[2].as<IntImm>()->value;
return this->Mutate(op->body);
return this->VisitStmt(op->body);
} else if (op->attr_key == attr::opengl_stage_scope) {
is_opengl_ = true;
}
return IRMutator::Mutate_(op, s);
return StmtExprMutator::VisitStmt_(op);
}
Stmt Mutate_(const Provide* op, const Stmt& s) final {
Stmt VisitStmt_(const Provide* op) final {
if (create_bound_attributes_)
shape_collector_.clear();
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Provide>();
TensorKey key{op->func, op->value_index};
auto it = buf_map_.find(key);
......@@ -159,11 +158,11 @@ class StorageFlattener : public IRMutator {
}
}
Stmt Mutate_(const Realize* op, const Stmt& s) final {
Stmt VisitStmt_(const Realize* op) final {
TensorKey key{op->func, op->value_index};
if (buf_map_.count(key)) {
CHECK(buf_map_.at(key).external);
return this->Mutate(op->body);
return this->VisitStmt(op->body);
} else {
// create a buffer entry
BufferEntry e;
......@@ -226,7 +225,7 @@ class StorageFlattener : public IRMutator {
align, 0, kDefault);
buf_map_[key] = e;
Stmt body = this->Mutate(op->body);
Stmt body = this->VisitStmt(op->body);
buf_map_[key].released = true;
Stmt ret;
......@@ -263,8 +262,8 @@ class StorageFlattener : public IRMutator {
}
}
Expr Mutate_(const Load* op, const Expr& e) final {
Expr expr = IRMutator::Mutate_(op, e);
Expr VisitExpr_(const Load* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Load>();
auto it = var_remap_.find(op->buffer_var.get());
if (it != var_remap_.end() &&
......@@ -277,17 +276,17 @@ class StorageFlattener : public IRMutator {
}
}
Expr Mutate_(const Variable* op, const Expr& e) final {
Expr VisitExpr_(const Variable* op) final {
auto it = var_remap_.find(op);
if (it != var_remap_.end()) {
return it->second;
} else {
return e;
return GetRef<Expr>(op);
}
}
Expr Mutate_(const Call* op, const Expr& olde) final {
Expr expr = IRMutator::Mutate_(op, olde);
Expr VisitExpr_(const Call* op) final {
Expr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<Call>();
if (op != nullptr && op->call_type == Call::Halide) {
TensorKey key{op->func, op->value_index};
......@@ -308,8 +307,8 @@ class StorageFlattener : public IRMutator {
}
}
Stmt Mutate_(const Prefetch *op, const Stmt &s) final {
Stmt stmt = IRMutator::Mutate_(op, s);
Stmt VisitStmt_(const Prefetch *op) final {
Stmt stmt = StmtExprMutator::VisitStmt_(op);
op = stmt.as<Prefetch>();
CHECK(op != nullptr);
TensorKey key{op->func, op->value_index};
......@@ -443,7 +442,7 @@ class StorageFlattener : public IRMutator {
// Apply the remaps
Stmt body = MergeNest(binder.asserts(), op->body);
body = MergeNest(binder.init_nest(), body);
body = this->Mutate(body);
body = this->VisitStmt(body);
// remove the binds
for (const Var& v : binder.defs()) {
var_remap_.erase(v.get());
......@@ -531,10 +530,10 @@ class StorageFlattener : public IRMutator {
Stmt StorageFlatten(Stmt stmt, Map<Tensor, Buffer> extern_buffer,
int cache_line_size, bool create_bound_attributes) {
IRVisitorWithAnalyzer bounded_analyzer;
bounded_analyzer.Visit(stmt);
bounded_analyzer(stmt);
stmt =
StorageFlattener(extern_buffer, cache_line_size,
create_bound_attributes, &bounded_analyzer).Mutate(stmt);
create_bound_attributes, &bounded_analyzer)(std::move(stmt));
return stmt;
}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment