replace_by_saved.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_REPLACE_BY_SAVED_H
#define FREE_TENSOR_REPLACE_BY_SAVED_H
 
#include <optional>
 
#include <analyze/symbol_table.h>
#include <autograd/invert_from_store.h>
#include <mutator.h>
 
namespace freetensor {
 
class ReplaceBySaved : public Mutator {
    const SymbolTableInterface &symbolTable_;
    const std::unordered_map<ID, std::string> &intermediatesMap_;
    const std::unordered_map<StmtOrExprID, Expr> &versions_;
    ID rootStmtID_;
    std::optional<InvertFromStore> invertFromStore_;
    bool isGrad_ = false;
 
  public:
    ReplaceBySaved(
        const SymbolTableInterface &symbolTable,
        const std::unordered_map<ID, std::string> &intermediatesMap,
        const std::unordered_map<StmtOrExprID, Expr> &versions,
        const ID &rootStmtID,
        const std::optional<InvertFromStore> &invertFromStore = std::nullopt)
        : symbolTable_(symbolTable), intermediatesMap_(intermediatesMap),
          versions_(versions), rootStmtID_(rootStmtID),
          invertFromStore_(invertFromStore) {}
 
    // Replace recomputing expressions
    auto recomp(const auto &op) {
        isGrad_ = false;
        return (*this)(op);
    }
 
    // Replace gradient expressions
    auto grad(const auto &op) {
        isGrad_ = true;
        return (*this)(op);
    }
 
  private:
    // Disabled. Use `ReplcaeBySaved::recomp` or `RepalceBySaved::grad` instaed
    using Mutator::operator();
 
  protected:
    Expr visitExpr(const Expr &expr) override;
    Expr visit(const Load &op) override;
};
 
} // namespace freetensor
 
#endif // FREE_TENSOR_REPLACE_BY_SAVED_H