find_loop_variance.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_FIND_LOOP_VARIANCE_H
#define FREE_TENSOR_FIND_LOOP_VARIANCE_H
 
#include <unordered_map>
#include <vector>
 
#include <analyze/symbol_table.h>
#include <analyze/track_stmt.h>
#include <visitor.h>
 
namespace freetensor {
 
enum class LoopVariability : int {
    // Non-existence = Invariant
    Unknown,
    Variant,
};
 
typedef std::unordered_map<std::string, std::unordered_map<ID, LoopVariability>>
    LoopVariTransVarMap;
typedef std::unordered_map<ID /* of VarDef */,
                           std::unordered_map<ID, LoopVariability>>
    LoopVariUniqVarMap;
typedef std::unordered_map<StmtOrExprID /* of Expr */,
                           std::unordered_map<ID, LoopVariability>>
    LoopVariExprMap;
 
class InitExprVari : public TrackStmt<Visitor> {
    typedef TrackStmt<Visitor> BaseClass;
 
    LoopVariExprMap &exprInfo_;
    std::vector<ID> loopStack_;
 
  public:
    InitExprVari(LoopVariExprMap &exprInfo) : exprInfo_(exprInfo) {}
 
  protected:
    void visitExpr(const Expr &expr) override;
    void visit(const For &op) override;
};
 
class MarkStores : public TrackStmt<Visitor> {
    typedef TrackStmt<Visitor> BaseClass;
 
    const std::string &var_;
    std::vector<StmtOrExprID> &condStack_;
    LoopVariTransVarMap &varInfo_;
    const LoopVariExprMap &exprInfo_;
 
  public:
    MarkStores(const std::string &var, std::vector<StmtOrExprID> &condStack,
               LoopVariTransVarMap &varInfo, const LoopVariExprMap &exprInfo)
        : var_(var), condStack_(condStack), varInfo_(varInfo),
          exprInfo_(exprInfo) {
        varInfo_.erase(var_); // Invariant
    }
 
  private:
    // to = from meet to
    void meetTo(const Expr &from, const std::string &to);
    void meetTo(const StmtOrExprID &from, const std::string &to);
 
    template <class T> void visitMemWrite(const T &op) {
        BaseClass::visit(op);
        if (op->var_ == var_) {
            meetTo(op->expr_, op->var_);
            for (auto &&index : op->indices_) {
                meetTo(index, op->var_);
            }
            for (auto &&cond : condStack_) {
                meetTo(cond, op->var_);
            }
        }
    }
 
  protected:
    void visit(const For &op) override;
    void visit(const If &op) override;
    void visit(const Store &op) override { visitMemWrite(op); }
    void visit(const ReduceTo &op) override {
        visitMemWrite(op);
        for (auto p = op->parentStmt(); p.isValid(); p = p->parentStmt()) {
            if (p->nodeType() == ASTNodeType::VarDef &&
                p.as<VarDefNode>()->name_ == op->var_) {
                break;
            }
            if (p->nodeType() == ASTNodeType::For) {
                varInfo_[op->var_][p->id()] = LoopVariability::Variant;
            }
        }
    }
};
 
class FindLoopVariance : public SymbolTable<TrackStmt<Visitor>> {
    typedef SymbolTable<TrackStmt<Visitor>> BaseClass;
 
    std::vector<ID> loopStack_;
    std::vector<StmtOrExprID> condStack_;
    LoopVariTransVarMap varInfo_;
    LoopVariUniqVarMap uniqVarInfo_;
    LoopVariExprMap exprInfo_;
 
  public:
    FindLoopVariance(const Stmt &root) { InitExprVari{exprInfo_}(root); }
 
    const LoopVariExprMap &exprInfo() const { return exprInfo_; }
    const LoopVariUniqVarMap &varInfo() const { return uniqVarInfo_; }
 
    int unknownCnt() const;
 
  private:
    // to = from
    void copyInfo(const Expr &from, const Expr &to);
    // to = from meet to
    void meetTo(const Expr &from, const Expr &to);
 
    void visitConst(const Const &op);
    void visitBinOp(const BinaryExpr &op);
    void visitUnaryOp(const UnaryExpr &op);
 
  protected:
    void visit(const For &op) override;
    void visit(const If &op) override;
    void visit(const VarDef &op) override;
 
    void visitExpr(const Expr &op) override;
    void visit(const Var &op) override;
    void visit(const Load &op) override;
    void visit(const IfExpr &op) override;
    void visit(const Cast &op) override;
    void visit(const Intrinsic &op) override;
};
 
bool isVariant(const LoopVariExprMap &exprInfo, const StmtOrExprID &expr,
               const ID &loop);
bool isVariant(const LoopVariUniqVarMap &varInfo, const VarDef &def,
               const ID &loop);
bool isVariant(const LoopVariUniqVarMap &varInfo, const ID &defId,
               const ID &loop);
 
std::pair<LoopVariExprMap, LoopVariUniqVarMap> findLoopVariance(const Stmt &op);
 
} // namespace freetensor
 
#endif // FREE_TENSOR_FIND_LOOP_VARIANCE_H