analyze_version.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_ANALYZE_VERSION_H
#define FREE_TENSOR_ANALYZE_VERSION_H
 
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/comp_transient_bounds.h>
#include <analyze/symbol_table.h>
#include <analyze/track_stmt.h>
#include <autograd/derivative.h>
#include <visitor.h>
 
namespace freetensor {
 
class CountScopeLen : public CompTransientBounds<SymbolTable<Visitor>> {
    typedef CompTransientBounds<SymbolTable<Visitor>> BaseClass;
 
    ID def_;
    std::string var_;
    const std::unordered_set<ID> &affectingScopes_; // For IDs
    const std::unordered_set<ID> &needVersions_;    // Store or ReduceTo IDs
    const std::unordered_set<ID> &fakeStoreIds_;
 
    // Total number of versions in a sub-tree
    std::unordered_map<Stmt, Expr> scopeLen_;
 
  public:
    CountScopeLen(const ID &def, const std::unordered_set<ID> &affectingScopes,
                  const std::unordered_set<ID> &needVersions,
                  const std::unordered_set<ID> &fakeStoreIds)
        : def_(def), affectingScopes_(affectingScopes),
          needVersions_(needVersions), fakeStoreIds_(fakeStoreIds) {}
 
    const std::unordered_map<Stmt, Expr> &scopeLen() const { return scopeLen_; }
 
  protected:
    using BaseClass::visit;
    void visit(const Store &op) override;
    void visit(const ReduceTo &op) override;
    void visit(const VarDef &op) override;
    void visit(const For &op) override;
    void visit(const StmtSeq &op) override;
    void visit(const If &op) override;
    void visit(const Assert &op) override;
};
 
class AnalyzeVersion : public TrackStmt<Visitor> {
    typedef TrackStmt<Visitor> BaseClass;
 
    ID def_;
    std::string var_;
    const std::unordered_set<ID> &affectingScopes_; // For IDs
    const std::unordered_set<ID> &needVersions_;    // Store or ReduceTo IDs
    const std::unordered_set<ID> &fakeStoreIds_;
    const std::unordered_map<Stmt, Expr> &scopeLen_;
    Expr totLen_;
    std::unordered_map<StmtOrExprID, Expr> &versions_;
    std::unordered_map<std::string, std::pair<std::string, Expr>>
        &userVersions_;
    std::string tapeName_;
    Expr offset_ = makeIntConst(0);
 
  public:
    AnalyzeVersion(const ID &def, const std::unordered_set<ID> &affectingScopes,
                   const std::unordered_set<ID> &needVersions,
                   const std::unordered_set<ID> &fakeStoreIds,
                   const std::unordered_map<Stmt, Expr> &scopeLen,
                   const Expr &totLen,
                   std::unordered_map<StmtOrExprID, Expr> &versions,
                   std::unordered_map<std::string, std::pair<std::string, Expr>>
                       &userVersions)
        : def_(def), affectingScopes_(affectingScopes),
          needVersions_(needVersions), fakeStoreIds_(fakeStoreIds),
          scopeLen_(scopeLen), totLen_(totLen), versions_(versions),
          userVersions_(userVersions) {}
 
    const std::string &tapeName() const { return tapeName_; }
 
  protected:
    void visit(const Load &op) override;
    void visit(const MarkVersion &op) override;
    void visit(const Store &op) override;
    void visit(const ReduceTo &op) override;
    void visit(const VarDef &op) override;
    void visit(const For &op) override;
    void visit(const StmtSeq &op) override;
};
 
class SetUserVersionsForInputs : public SymbolTable<Visitor> {
    typedef SymbolTable<Visitor> BaseClass;
 
    std::unordered_map<std::string, std::pair<std::string, Expr>>
        &userVersions_;
 
  public:
    SetUserVersionsForInputs(
        std::unordered_map<std::string, std::pair<std::string, Expr>>
            &userVersions)
        : userVersions_(userVersions) {}
 
  protected:
    using BaseClass::visit;
    void visit(const MarkVersion &op) override;
};
 
std::tuple<std::unordered_map<StmtOrExprID, Expr>, std::unordered_map<ID, Expr>,
           std::unordered_set<ID>,
           std::unordered_map<std::string, std::pair<std::string, Expr>>>
analyzeVersion(
    const Stmt &op,
    const std::unordered_map<ID, std::unordered_set<ID>> &needVersions,
    const std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
        &derivatives,
    bool localVersionsOnly);
 
} // namespace freetensor
 
#endif // FREE_TENSOR_ANALYZE_VERSION_H