all_uses.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_ALL_USES_H
#define FREE_TENSOR_ALL_USES_H
 
#include <unordered_set>
 
#include <analyze/find_stmt.h>
#include <visitor.h>
 
namespace freetensor {
 
class AllUses : public Visitor {
  public:
    typedef int AllUsesType;
    static constexpr AllUsesType CHECK_LOAD = 0x1;
    static constexpr AllUsesType CHECK_STORE = 0x2;
    static constexpr AllUsesType CHECK_REDUCE = 0x4;
    static constexpr AllUsesType CHECK_VAR = 0x8;
 
  private:
    AllUsesType type_;
    bool noRecurseIdx_, noRecurseSubStmt_;
    bool inFirstStmt_ = false;
    std::unordered_set<std::string> uses_;
 
  public:
    AllUses(AllUsesType type, bool noRecurseIdx, bool noRecurseSubStmt)
        : type_(type), noRecurseIdx_(noRecurseIdx),
          noRecurseSubStmt_(noRecurseSubStmt) {}
 
    const std::unordered_set<std::string> &uses() const { return uses_; }
 
  protected:
    // TODO: Do we need to check viewOf, reduction, Alloc or Free?
    void visitStmt(const Stmt &s) override;
    void visit(const Load &op) override;
    void visit(const Store &op) override;
    void visit(const ReduceTo &op) override;
    void visit(const Var &op) override;
};
 
std::unordered_set<std::string>
allUses(const AST &op,
        AllUses::AllUsesType type = AllUses::CHECK_LOAD | AllUses::CHECK_STORE |
                                    AllUses::CHECK_REDUCE,
        bool noRecurseIdx = false, bool noRecurseSubStmt = false);
inline std::unordered_set<std::string>
allUses(const auto &filter, const Stmt &op,
        AllUses::AllUsesType type = AllUses::CHECK_LOAD | AllUses::CHECK_STORE |
                                    AllUses::CHECK_REDUCE,
        bool noRecurseIdx = false) {
    std::unordered_set<std::string> ret;
    for (auto &&stmt : findAllStmt(op, filter)) {
        for (auto &&use : allUses(stmt, type, noRecurseIdx, true)) {
            ret.emplace(use);
        }
    }
    return ret;
}
inline std::unordered_set<std::string> allReads(const AST &op,
                                                bool noRecurseIdx = false,
                                                bool noRecurseSubStmt = false) {
    return allUses(op, AllUses::CHECK_LOAD, noRecurseIdx, noRecurseSubStmt);
}
inline std::unordered_set<std::string>
allReads(const auto &filter, const Stmt &op, bool noRecurseIdx = false) {
    return allUses(filter, op, AllUses::CHECK_LOAD, noRecurseIdx);
}
inline std::unordered_set<std::string>
allWrites(const AST &op, bool noRecurseIdx = false,
          bool noRecurseSubStmt = false) {
    return allUses(op, AllUses::CHECK_STORE | AllUses::CHECK_REDUCE,
                   noRecurseIdx, noRecurseSubStmt);
}
inline std::unordered_set<std::string>
allWrites(const auto &filter, const Stmt &op, bool noRecurseIdx = false) {
    return allUses(filter, op, AllUses::CHECK_STORE | AllUses::CHECK_REDUCE,
                   noRecurseIdx);
}
inline std::unordered_set<std::string> allIters(const AST &op,
                                                bool noRecurseIdx = false,
                                                bool noRecurseSubStmt = false) {
    return allUses(op, AllUses::CHECK_VAR, noRecurseIdx, noRecurseSubStmt);
}
inline std::unordered_set<std::string>
allIters(const auto &filter, const Stmt &op, bool noRecurseIdx = false) {
    return allUses(filter, op, AllUses::CHECK_VAR, noRecurseIdx);
}
inline std::unordered_set<std::string> allNames(const AST &op,
                                                bool noRecurseIdx = false,
                                                bool noRecurseSubStmt = false) {
    return allUses(op,
                   AllUses::CHECK_LOAD | AllUses::CHECK_STORE |
                       AllUses::CHECK_REDUCE | AllUses::CHECK_VAR,
                   noRecurseIdx, noRecurseSubStmt);
}
inline std::unordered_set<std::string>
allNames(const auto &filter, const Stmt &op, bool noRecurseIdx = false) {
    return allUses(filter, op,
                   AllUses::CHECK_LOAD | AllUses::CHECK_STORE |
                       AllUses::CHECK_REDUCE | AllUses::CHECK_VAR,
                   noRecurseIdx);
}
} // namespace freetensor
 
#endif // FREE_TENSOR_ALL_USES_H