output_intermediates.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_OUTPUT_INTERMEDIATES_H
#define FREE_TENSOR_OUTPUT_INTERMEDIATES_H
 
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/symbol_table.h>
#include <autograd/derivative.h>
#include <mutator.h>
 
namespace freetensor {
 
enum class OutputIntermediatesStage : int { Forward, Backward };
 
class OutputIntermediates : public SymbolTable<Mutator> {
    typedef SymbolTable<Mutator> BaseClass;
 
    const std::unordered_map<StmtOrExprID, Expr> &versions_;
    const std::unordered_map<ID, Expr> &totLens_;
    const std::unordered_set<ID> &trivials_;
    OutputIntermediatesStage stage_;
    std::string varSuffix_;
 
    std::unordered_map<ID, std::string> savedNames_;
    std::unordered_set<ID> insertedStmts_;
    std::unordered_map<ID, std::vector<Stmt>> toSave_;
    ID curStmt_;
 
  public:
    OutputIntermediates(const std::unordered_map<StmtOrExprID, Expr> &versions,
                        const std::unordered_map<ID, Expr> &totLens,
                        const std::unordered_set<ID> &trivials,
                        OutputIntermediatesStage stage,
                        const std::string &varSuffix)
        : versions_(versions), totLens_(totLens), trivials_(trivials),
          stage_(stage), varSuffix_(varSuffix) {}
 
    const auto &savedNames() const { return savedNames_; }
    const auto &insertedStmts() const { return insertedStmts_; }
 
  private:
    std::string savingName(const std::string &oldName) const;
 
  protected:
    using BaseClass::visit;
    Stmt visitStmt(const Stmt &stmt) override;
    Expr visit(const Load &op) override;
    Stmt visit(const Store &op) override;
    Stmt visit(const ReduceTo &op) override;
    Stmt visit(const VarDef &op) override;
};
 
std::tuple<Stmt, std::unordered_map<ID, std::string>,
           std::unordered_map<StmtOrExprID, Expr>, std::unordered_map<ID, Expr>,
           std::unordered_set<ID>,
           std::unordered_map<std::string, std::pair<std::string, Expr>>>
outputIntermediates(
    const Stmt &op,
    const std::unordered_map<ID, std::unordered_set<ID>> &needVersions,
    const std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
        &derivatives,
    OutputIntermediatesStage stage = OutputIntermediatesStage::Forward,
    const std::string &varSuffix = ".tape");
 
Stmt outputAllIntermedaites(
    const Stmt &op, const std::unordered_set<ID> &intermediates,
    OutputIntermediatesStage stage = OutputIntermediatesStage::Forward,
    const std::string &varSuffix = ".tape");
 
} // namespace freetensor
 
#endif // FREE_TENSOR_OUTPUT_INTERMEDIATES_H