scalar_prop_const.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_SCALAR_PROP_CONST_H
#define FREE_TENSOR_SCALAR_PROP_CONST_H
 
#include <analyze/symbol_table.h>
#include <func.h>
#include <hash.h>
#include <mutator.h>
#include <pass/const_fold.h>
 
#include <map>
 
namespace freetensor {
 
class ScalarPropConst : public SymbolTable<ConstFold> {
  protected:
    typedef SymbolTable<ConstFold> BaseClass;
 
    struct ScalarIndex : public std::variant<int64_t, std::string> {
        template <class... Args>
        ScalarIndex(Args &&...args)
            : std::variant<int64_t, std::string>(std::forward<Args>(args)...) {}
 
        std::strong_ordering operator<=>(const ScalarIndex &other) const {
            switch (index() * 2 + other.index()) {
            case 0: // int int
                return std::get<int64_t>(*this) <=> std::get<int64_t>(other);
            case 1: // int str
                return std::strong_ordering::less;
            case 2: // str int
                return std::strong_ordering::greater;
            case 3: // str str
                return std::get<std::string>(*this) <=>
                       std::get<std::string>(other);
            default:
                ASSERT(false);
            }
        }
    };
 
    struct ScalarIndices {
        std::vector<ScalarIndex> offset;
 
        std::strong_ordering operator<=>(const ScalarIndices &other) const {
            ASSERT(offset.size() == other.offset.size() &&
                   "Index count should be identical for same tensor");
            return offset <=> other.offset;
        }
 
        bool operator==(const ScalarIndices &) const = default;
    };
 
    std::optional<ScalarIndices> tryToScalar(const std::vector<Expr> &exprs);
 
    std::unordered_map<std::string, std::map<ScalarIndices, Expr>> constants_;
 
    std::unordered_multimap<std::string, std::pair<std::string, ScalarIndices>>
        iter_dep_constants_;
 
    void gen_constant(const std::string &name,
                      const std::optional<ScalarIndices> &indices,
                      const Expr &value);
    void kill_iter_dep_entry(const std::string &name,
                             const ScalarIndices &indices);
    void kill_constant(const std::string &name,
                       const std::optional<ScalarIndices> &indices);
    void kill_iter(const std::string &it_var);
 
    bool intersect_constants_with(
        std::unordered_map<std::string, std::map<ScalarIndices, Expr>> other);
 
  protected:
    using BaseClass::visit;
    Stmt visit(const Store &store_orig) override;
    Stmt visit(const ReduceTo &op) override;
    Expr visit(const Load &load_orig) override;
    Stmt visit(const If &op) override;
    Stmt visit(const VarDef &vd) override;
    Stmt visit(const For &op) override;
};
 
Stmt scalarPropConst(const Stmt &op);
 
DEFINE_PASS_FOR_FUNC(scalarPropConst)
 
} // namespace freetensor
 
#endif // FREE_TENSOR_SCALAR_PROP_CONST_H