schedule.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_SCHEDULE_H
#define FREE_TENSOR_SCHEDULE_H
 
#include <functional>
#include <unordered_map>
 
#include <analyze/find_stmt.h>
#include <codegen/native_code.h>
#include <driver/target.h>
#include <func.h>
#include <probability/rand_ctx.h>
#include <random.h>
#include <schedule/as_matmul.h>
#include <schedule/fission.h>
#include <schedule/memoized_schedules.h>
#include <schedule/reorder.h>
#include <schedule/schedule_log.h>
#include <schedule/var_split.h>
#include <stmt.h>
 
namespace freetensor {
 
enum class MoveToSide : int { Before, After };
 
struct AutoScheduleTuneTrial {
    Ref<RandTrace> trace_;
    Func lowered_;
    NativeCode code_;
    double time_, stddev_;
};
 
class Schedule {
    struct Transaction {
        Stmt ast_;
        ScheduleLog logs_;
 
        Transaction(const Stmt &ast, const ScheduleLog &logs)
            : ast_(ast), logs_(logs) {}
    };
 
    Func func_; 
 
    std::vector<Transaction> openTrans_; 
 
    int verbose_ = 0;
 
    Ref<MemoizedSchedules> memoized_;
 
    Ref<OpenMPRandomEngine> rng_;
    Ref<RandCtx<OpenMPRandomEngine>> randCtx_;
 
  private:
    void setAst(const Stmt &ast);
    void setLogs(const ScheduleLog &log);
 
    static Stmt quickOptimizations(const Stmt &ast);
 
    template <class F> auto futureSchedule(const F &sched) {
        return [&](auto &&...args) {
            auto ret = sched(ast(), std::forward<decltype(args)>(args)...);
            if constexpr (std::convertible_to<decltype(ret), Stmt>) {
                return quickOptimizations(ret);
            } else { // pair(Stmt, other info)
                return std::make_pair(quickOptimizations(ret.first),
                                      ret.second);
            }
        };
    }
 
    template <class T> T appendLog(const T &_log) {
        auto log = _log;
        setLogs(memoized_->lookupOrCreate(logs().push(log)));
        ASSERT(logs().top()->type() == log->type());
        log = logs().top().as<typename decltype(log)::Object>();
        log->run();
        return log;
    }
 
    template <class T> auto applyLog(const T &log) {
        auto ret = log->getResult();
        if constexpr (std::convertible_to<decltype(ret), Stmt>) {
            setAst(ret);
            return;
        } else { // pair(Stmt, other info)
            setAst(ret.first);
            return ret.second;
        }
    }
 
  public:
    Schedule() = default;
    Schedule(const Stmt &ast, int verbose = 0);
    Schedule(const Func &func, int verbose = 0)
        : Schedule(func->body_, verbose) {
        func_ = func;
    }
 
    // Copy by default, which means `Ref`s in a `Schedule` object is shared
    Schedule(const Schedule &) = default;
    Schedule &operator=(const Schedule &) = default;
 
    Schedule fork() const { return *this; }
 
    void beginTransaction();
    void commitTransaction();
    void abortTransaction();
    Func func() const {
        ASSERT(func_.isValid());
        return makeFunc(func_->name_, func_->params_, func_->returns_, ast());
    }
 
    const Stmt &ast() const;
 
    const ScheduleLog &logs() const;
 
    int verbose() const { return verbose_; }
 
    template <class T> std::vector<Stmt> findAll(const T &filter) const {
        return findAllStmt(ast(), filter);
    }
 
    template <class T> std::vector<Stmt> findAtLeastOne(const T &filter) const {
        auto ret = findAllStmt(ast(), filter);
        if (ret.empty()) {
            throw InvalidSchedule(ast(), "No statement found by filter");
        }
        return ret;
    }
 
    template <class T> Stmt find(const T &filter) const {
        try {
            return findStmt(ast(), filter);
        } catch (const UnexpectedQueryResult &e) {
            throw InvalidSchedule(ast(), e.what());
        }
    }
 
    std::pair<ID, ID> split(const ID &id, int factor = -1, int nparts = -1,
                            int shift = 0);
 
    void reorder(const std::vector<ID> &order,
                 ReorderMode mode = ReorderMode::PerfectOnly);
 
    ID merge(const ID &loop1, const ID &loop2);
 
    std::vector<ID>
    permute(const std::vector<ID> &loopsId,
            const std::function<std::vector<Expr>(std::vector<Expr>)>
                &transformFunc);
 
    typedef std::unordered_map<ID, ID> IDMap;
    std::pair<IDMap, IDMap> fission(const ID &loop, FissionSide side,
                                    const ID &splitter,
                                    bool allowEnlarge = true,
                                    const std::string &suffix0 = ".0",
                                    const std::string &suffix1 = ".1");
 
    ID fuse(const ID &loop0, const ID &loop1, bool strict = false);
    ID fuse(const ID &loop0, bool strict = false);
    void swap(const std::vector<ID> &order);
 
    void blend(const ID &loop);
 
    std::tuple<ID, ID, std::string, ID>
    cache(const ID &stmt, const std::string &var, MemType mtype);
 
    std::tuple<ID, ID, std::string, ID>
    cacheReduction(const ID &stmt, const std::string &var, MemType mtype);
 
    void setMemType(const ID &def, MemType mtype);
    void setMemType(const ID &def, MemType mtype, bool rejectIndirectAccess);
    void varSplit(const ID &def, int dim, VarSplitMode mode, int factor = -1,
                  int nparts = -1);
 
    void varMerge(const ID &def, int dim);
 
    void varReorder(const ID &def, const std::vector<int> &order);
 
    void varUnsqueeze(const ID &def, int dim);
 
    void varSqueeze(const ID &def, int dim);
 
    std::pair<ID, ID> moveTo(const ID &stmt, MoveToSide side, const ID &dst);
 
    void inlining(const ID &def);
 
    void parallelize(const ID &loop, const ParallelScope &parallel,
                     bool allowReduction = true);
 
    void parallelizeAs(const ID &nest, const ID &reference, const ID &defId);
 
    void unroll(const ID &loop, bool immediate = false);
 
    void vectorize(const ID &loop);
 
    void separateTail(bool noDuplicateVarDefs = false);
 
    void asMatMul(const ID &loop, AsMatMulMode mode, const Ref<Target> &target,
                  MatMulBackend backend);
    void asMatMul(const ID &loop, AsMatMulMode mode, const Ref<Target> &target);
    void asMatMul(const ID &loop,
                  AsMatMulMode mode = AsMatMulMode::KeepMemLayout);
 
    std::pair<ID, int> plutoFuse(const ID &loop0, const ID &loop1,
                                 int nestLevel0 = 0, int nestLevel1 = 0,
                                 int fusableOverlapThreshold = 1,
                                 int fusableNonOverlapTolerance = 4,
                                 bool doSimplify = true);
 
    std::pair<ID, int> plutoPermute(const ID &loop, int nestLevel = 0,
                                    bool doSimplify = true);
 
    void autoSchedule(const Ref<Target> &target,
                      const Ref<RandTrace> &trace = nullptr);
 
    void autoInline(const Ref<Target> &target);
 
    void autoUseLib(const Ref<Target> &target);
 
    void autoReorder(const Ref<Target> &target);
 
    void autoSwap(const Ref<Target> &target);
 
    void autoPluto(const Ref<Target> &target);
 
    void autoFissionFuse(const Ref<Target> &target,
                         const Ref<RandTrace> &trace = nullptr);
 
    void autoMemLayout(const Ref<Target> &target);
 
    void autoParallelize(const Ref<Target> &target);
 
    void autoSetMemType(const Ref<Target> &target);
 
    void autoUnroll(const Ref<Target> &target);
 
    std::vector<AutoScheduleTuneTrial> tuneAutoSchedule(
        int nBatch, int batchSize, const Ref<Device> &device,
        const std::vector<Ref<Array>> &args,
        const std::unordered_map<std::string, Ref<Array>> &kws = {},
        const std::regex &toLearn = std::regex{".*"});
};
 
} // namespace freetensor
 
#endif // FREE_TENSOR_SCHEDULE_H