remove_writes.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_REMOVE_WRITES_H
#define FREE_TENSOR_REMOVE_WRITES_H
 
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/find_loop_variance.h>
#include <analyze/symbol_table.h>
#include <func.h>
#include <mutator.h>
#include <visitor.h>
 
namespace freetensor {
 
class FindLoopInvariantWrites : public SymbolTable<Visitor> {
    typedef SymbolTable<Visitor> BaseClass;
 
    std::vector<For> loopStack_;
    std::vector<If> ifStack_;
    std::unordered_map<Store, std::tuple<VarDef, Expr, For>>
        results_; 
    std::unordered_map<std::string, int> defDepth_;
    const LoopVariExprMap &variantExpr_;
    ID singleDefId_;
 
  public:
    FindLoopInvariantWrites(const LoopVariExprMap &variantExpr,
                            const ID &singleDefId)
        : variantExpr_(variantExpr), singleDefId_(singleDefId) {}
 
    const std::unordered_map<Store, std::tuple<VarDef, Expr, For>> &
    results() const {
        return results_;
    }
 
  protected:
    using BaseClass::visit;
    void visit(const For &op) override;
    void visit(const If &op) override;
    void visit(const VarDef &op) override;
    void visit(const Store &op) override;
};
 
class RemoveWrites : public Mutator {
    const std::unordered_set<Stmt> &redundant_;
    const std::unordered_map<Stmt, Stmt> &replacement_;
 
  public:
    RemoveWrites(const std::unordered_set<Stmt> &redundant,
                 const std::unordered_map<Stmt, Stmt> &replacement)
        : redundant_(redundant), replacement_(replacement) {}
 
    template <class T> Stmt doVisit(const T &op) {
        if (redundant_.count(op)) {
            return makeStmtSeq({}, op->metadata(), op->id());
        } else if (replacement_.count(op)) {
            return replacement_.at(op);
        } else {
            return Mutator::visit(op);
        }
    }
 
  protected:
    Stmt visit(const Store &op) override { return doVisit(op); }
    Stmt visit(const ReduceTo &op) override { return doVisit(op); }
    Stmt visit(const StmtSeq &op) override;
    Stmt visit(const For &op) override;
    Stmt visit(const If &op) override;
};
 
Stmt removeWrites(const Stmt &op, const ID &singleDefId = {});
 
DEFINE_PASS_FOR_FUNC(removeWrites)
 
} // namespace freetensor
 
#endif // FREE_TENSOR_REMOVE_WRITES_H