grad.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_GRAD_H
#define FREE_TENSOR_GRAD_H
 
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/symbol_table.h>
#include <autograd/derivative.h>
#include <autograd/invert_stmts.h>
#include <autograd/replace_by_saved.h>
#include <autograd/tape_strategy.h>
#include <autograd/user_grad.h>
#include <func.h>
#include <mutator.h>
#include <visitor.h>
 
namespace freetensor {
 
class InsertUserGrad : public Mutator {
    const SymbolTableInterface &symbolTable_;
    const std::unordered_map<ID, std::string> &intermediatesMap_;
    const std::unordered_map<std::string, std::pair<std::string, Expr>>
        &userVersions_;
    const std::unordered_map<std::string, std::string>
        &gradNames_; // x -> dy/dx
 
    std::unordered_set<std::string> localVarDefNames_;
 
  public:
    InsertUserGrad(
        const SymbolTableInterface &symbolTable,
        const std::unordered_map<ID, std::string> &intermediatesMap,
        const std::unordered_map<std::string, std::pair<std::string, Expr>>
            &userVersions,
        const std::unordered_map<std::string, std::string> &gradNames)
        : symbolTable_(symbolTable), intermediatesMap_(intermediatesMap),
          userVersions_(userVersions), gradNames_(gradNames) {}
 
  protected:
    Expr visit(const LoadAtVersion &op) override;
    Stmt visit(const Store &op) override;
    Stmt visit(const ReduceTo &op) override;
    Stmt visit(const VarDef &op) override;
};
 
template <class BaseClass> class RenewIDs : public BaseClass {
  protected:
    Stmt visitStmt(const Stmt &s) override {
        auto ret = BaseClass::visitStmt(s);
        ret->setId();
        return ret;
    }
};
 
class Grad : public RenewIDs<SymbolTable<Mutator>> {
    // Because a statement can be both recomputed and computed for gradient, and
    // we may even recompute a statement several times, all IDs must be renewed,
    // even for recomputation
    typedef RenewIDs<SymbolTable<Mutator>> BaseClass;
 
    std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
        &derivatives_; // Mutable for lazy operations
    const std::unordered_set<std::string> &requires_;
    const std::unordered_set<std::string> &provides_;
    const std::unordered_set<ID> &tapes_;
    const std::unordered_set<ID> &defsNeedGrad_;
    const std::unordered_map<ID, std::string>
        &intermediatesMap_; // All saved variables, including in forward stage
                            // (tapes) and backward stage (during recomputation)
    const std::unordered_map<StmtOrExprID, Expr> &versions_;
    const std::unordered_map<std::string, std::pair<std::string, Expr>>
        &userVersions_;
    const std::unordered_map<ID, Expr> &totLens_;
    const std::unordered_set<ID> &saveLocalStmts_;
    const std::unordered_set<Stmt> &notSingleWrite_;
    bool resetProvidedGrad_;
    const std::unordered_map<ID, InversionInfo> &inverseStmts_;
    std::vector<RangeToUserGrad> userGrads_; // mutable
 
    std::unordered_map<std::string, std::string> requireGrads_; // var name map
    std::unordered_map<std::string, std::string> provideGrads_; // var name map
 
    std::unordered_map<std::string, std::string> gradNames_; // x -> dy/dx
    std::unordered_map<Expr, Expr> equLoads_;
    std::unordered_map<std::string, std::unordered_set<Stmt>>
        recomputed_; // var name -> set{stmt}
    bool isRecompute_ = false;
 
    std::unordered_set<ID> inverselyUpdated_; // {VarDef IDs}
 
    std::optional<RangeToUserGrad> userGradOpen_;
    ID userGradInsertPos_;
 
  private:
    ReplaceBySaved getReplacer(const Stmt &stmt,
                               const Store &alreadyStored = nullptr) const;
 
    Stmt doVisitStmt(const Stmt &s);
 
  public:
    Grad(std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
             &derivatives,
         const std::unordered_set<std::string> &_requires,
         const std::unordered_set<std::string> &provides,
         const std::unordered_set<ID> &tapes,
         const std::unordered_set<ID> &defsNeedGrad,
         const std::unordered_map<ID, std::string> &intermediatesMap,
         const std::unordered_map<StmtOrExprID, Expr> &versions,
         const std::unordered_map<std::string, std::pair<std::string, Expr>>
             &userVersions,
         const std::unordered_map<ID, Expr> &totLens,
         const std::unordered_set<ID> &saveLocalStmts,
         const std::unordered_set<Stmt> &notSingleWrite, bool resetProvidedGrad,
         const std::unordered_map<ID, InversionInfo> &inverseStmts,
         const std::vector<RangeToUserGrad> &userGrads)
        : derivatives_(derivatives), requires_(_requires), provides_(provides),
          tapes_(tapes), defsNeedGrad_(defsNeedGrad),
          intermediatesMap_(intermediatesMap), versions_(versions),
          userVersions_(userVersions), totLens_(totLens),
          saveLocalStmts_(saveLocalStmts), notSingleWrite_(notSingleWrite),
          resetProvidedGrad_(resetProvidedGrad), inverseStmts_(inverseStmts),
          userGrads_(userGrads) {}
 
    const std::unordered_map<std::string, std::string> &requireGrads() const {
        return requireGrads_;
    }
    const std::unordered_map<std::string, std::string> &provideGrads() const {
        return provideGrads_;
    }
 
  protected:
    Stmt visitStmt(const Stmt &s) override;
    Stmt visit(const StmtSeq &op) override;
    Stmt visit(const For &op) override;
    Stmt visit(const If &op) override;
    Stmt visit(const Assert &op) override;
    Stmt visit(const VarDef &op) override;
    Stmt visit(const Store &op) override;
    Stmt visit(const ReduceTo &op) override;
};
 
std::tuple<Stmt, Stmt, std::unordered_map<std::string, std::string>,
           std::unordered_map<std::string, std::string>,
           std::unordered_map<ID, std::string>>
gradBody(const Stmt &op, const std::unordered_set<std::string> &_requires,
         const std::unordered_set<std::string> &provides,
         const TapeStrategy &tapes, bool resetProvidedGrad = true,
         bool invert = false,
         const std::vector<StmtSetToUserGrad> &userGrads = {});
 
std::tuple<Func, Func, std::unordered_map<std::string, std::string>,
           std::unordered_map<std::string, std::string>>
gradFuncInplace(const Func &func,
                const std::unordered_set<std::string> &_requires,
                const std::unordered_set<std::string> &provides,
                const TapeStrategy &tapes, bool tapeInClosure = true,
                bool resetProvidedGrad = true, bool invert = false,
                const std::vector<StmtSetToUserGrad> &userGrads = {});
 
std::tuple<Func, Func, std::unordered_map<std::string, std::string>,
           std::unordered_map<std::string, std::string>>
gradFuncOutOfPlace(const Func &func,
                   const std::unordered_set<std::string> &_requires,
                   const std::unordered_set<std::string> &provides,
                   const TapeStrategy &tapes, bool tapeInClosure = true,
                   bool resetProvidedGrad = true, bool invert = false,
                   const std::vector<StmtSetToUserGrad> &userGrads = {});
} // namespace freetensor
 
#endif // FREE_TENSOR_GRAD_H