derivative.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_DERIVATIVE_H
#define FREE_TENSOR_DERIVATIVE_H
 
#include <optional>
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/symbol_table.h>
#include <autograd/invert_from_store.h>
#include <mutator.h>
#include <visitor.h>
 
namespace freetensor {
 
class Derivative : public SymbolTable<Visitor> {
    typedef SymbolTable<Visitor> BaseClass;
 
  public:
    class LazyPartialDerivative {
        SymbolTableData symbolTable_;
        Expr mathExpr_;
        ID rootStmtID_;
 
        // Additional info to replace `y = f(x)`'s derivative to use `y`. The
        // replacement has to be done lazily, because the version of `y` is
        // different from version of `x`
        std::optional<InvertFromStore> invertFromStore_;
        std::optional<bool> usingStore_;
 
        std::optional<std::unordered_set<std::string>> reads_;
 
      public:
        LazyPartialDerivative() {} // Unintialized
 
        LazyPartialDerivative(const SymbolTableData &symbolTable,
                              const Expr &mathExpr, const ID &rootStmtID,
                              const std::optional<InvertFromStore>
                                  &invertFromStore = std::nullopt)
            : symbolTable_(symbolTable), mathExpr_(mathExpr),
              rootStmtID_(rootStmtID), invertFromStore_(invertFromStore) {}
 
        void merge(const LazyPartialDerivative &other) {
            ASSERT(other.rootStmtID_ == rootStmtID_);
            mathExpr_ = makeAdd(mathExpr_, other.mathExpr_);
            usingStore_ = std::nullopt;
        }
 
        const Expr &mathExpr() const { return mathExpr_; }
 
        bool usingStore();
 
        const std::unordered_set<std::string> &reads();
 
        Expr
        replaceExpr(const std::unordered_map<ID, std::string> &intermediatesMap,
                    const std::unordered_map<StmtOrExprID, Expr> &versions,
                    const Expr &expr);
 
        Expr
        genReplaced(const std::unordered_map<ID, std::string> &intermediatesMap,
                    const std::unordered_map<StmtOrExprID, Expr> &versions) {
            return replaceExpr(intermediatesMap, versions, mathExpr_);
        }
    };
 
    class LazyFullDerivative {
        // (x, dy/dx). Instead of using a map, an ordered list is used to ensure
        // that the statements generated by `genGrads` are consistently ordered
        std::vector<std::pair<Load, LazyPartialDerivative>> partials_;
 
        // Union of all partials. These are lazily cached variables
        std::optional<bool> usingStore_;
        std::optional<std::unordered_set<std::string>> reads_;
 
        std::exception_ptr error_;
 
      public:
        void addPartial(const Load &x, const LazyPartialDerivative &partial);
 
        void setError(const std::exception_ptr &error) { error_ = error; }
 
        bool usingStore();
 
        const std::unordered_set<std::string> &reads();
 
        std::vector<Stmt>
        genGrads(const std::unordered_map<ID, std::string> &intermediatesMap,
                 const std::unordered_map<StmtOrExprID, Expr> &versions,
                 const std::unordered_map<std::string, std::string> &gradNames,
                 const Expr &gradY);
    };
 
  private:
    std::unordered_map<StmtOrExprID, LazyFullDerivative> derivatives_;
    std::unordered_map<Expr, LazyPartialDerivative> partials_;
    StmtOrExprID rootExpr_;
    std::optional<InvertFromStore> invertFromStore_;
 
    void setPartial(const Expr &expr, const Expr &partial);
 
  public:
    const auto &derivatives() const { return derivatives_; }
 
  protected:
    using BaseClass::visit;
 
    void visitExpr(const Expr &expr) override;
 
    // If we have `y = f(x)` as a `Store` node, we can use `y` in the
    // derivative. Please note that this does not apply to `ReduceTo` nodes: in
    // `y += exp(x)`, we have no variable equals to `exp(x)`
    void visit(const Store &op) override;
 
    void visit(const Load &op) override;
    void visit(const Add &op) override;
    void visit(const Sub &op) override;
    void visit(const Mul &op) override;
    void visit(const RealDiv &op) override;
    void visit(const Min &op) override;
    void visit(const Max &op) override;
    void visit(const IfExpr &op) override;
    void visit(const Sqrt &op) override;
    void visit(const Exp &op) override;
    void visit(const Ln &op) override;
    void visit(const Square &op) override;
    void visit(const Sigmoid &op) override;
    void visit(const Sin &op) override;
    void visit(const Cos &op) override;
    void visit(const Tan &op) override;
    void visit(const Tanh &op) override;
    void visit(const Abs &op) override;
    void visit(const Cast &op) override;
    void visit(const Intrinsic &op) override;
};
 
} // namespace freetensor
 
#endif // FREE_TENSOR_DERIVATIVE_H