make_heap_alloc.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_MAKE_HEAP_ALLOC_H
#define FREE_TENSOR_MAKE_HEAP_ALLOC_H
 
#include <unordered_set>
 
#include <func.h>
#include <mutator.h>
#include <pass/const_fold.h>
 
namespace freetensor {
 
class InsertAlloc : public Mutator {
    std::string var_;
    bool inserted_ = false, delayed_ = false;
 
  public:
    InsertAlloc(const std::string &var) : var_(var) {}
 
    bool delayed() const { return delayed_; }
 
  protected:
    Stmt visit(const StmtSeq &op) override;
    Stmt visit(const For &op) override { return op; }
    Stmt visit(const If &op) override { return op; }
};
 
class InsertFree : public Mutator {
    std::string var_;
    bool inserted_ = false, madeEarly_ = false;
 
  public:
    InsertFree(const std::string &var) : var_(var) {}
 
    bool madeEarly() const { return madeEarly_; }
 
  protected:
    Stmt visit(const StmtSeq &op) override;
    Stmt visit(const For &op) override { return op; }
    Stmt visit(const If &op) override { return op; }
};
 
class MakeHeapAlloc : public ConstFold {
    // Inherit ConstFold for determine dynamic sizes
    typedef ConstFold BaseClass;
 
    bool inKernel_ = false;
 
  private:
    bool inKernel() const { return inKernel_; }
    bool isDynamicSized(const VarDef &op) const;
 
  protected:
    using BaseClass::visit;
    Stmt visit(const VarDef &op) override;
    Stmt visit(const For &op) override;
    Stmt visit(const MatMul &op) override { return op; }
};
 
Stmt makeHeapAlloc(const Stmt &op);
 
DEFINE_PASS_FOR_FUNC(makeHeapAlloc);
 
} // namespace freetensor
 
#endif // FREE_TENSOR_MAKE_HEAP_ALLOC_H