invert_stmts.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_INVERT_STMTS_H
#define FREE_TENSOR_INVERT_STMTS_H
 
#include <unordered_map>
#include <unordered_set>
 
#include <analyze/symbol_table.h>
#include <autograd/derivative.h>
#include <visitor.h>
 
namespace freetensor {
 
struct InversionInfo {
    Stmt inv_;
    Expr cond_; // null for always invert
};
 
class FindInvertibles : public SymbolTable<Visitor> {
    typedef SymbolTable<Visitor> BaseClass;
 
    // Original ID -> Inverse Statement
    std::unordered_map<ID, Stmt> invertibles_;
 
  public:
    const auto &invertibles() const { return invertibles_; }
 
  protected:
    using BaseClass::visit;
    // TODO 1: Support invertibles in Store.
    // TODO 2: We now always invert first, and then compute gradient. Things
    // will go wrong after we support Store, where we may use y for gradient.
    void visit(const ReduceTo &op) override;
};
 
std::tuple<Stmt, std::unordered_map<ID, InversionInfo>>
invertStmts(const Stmt &op,
            std::unordered_map<ID, std::unordered_set<ID>> *idsNeeded,
            std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
                *derivatives);
 
} // namespace freetensor
 
#endif // FREE_TENSOR_INVERT_STMTS_H