native_code.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_NATIVE_CODE_H
#define FREE_TENSOR_NATIVE_CODE_H
 
#include <optional>
#include <string>
 
#include <driver/array.h>
#include <driver/target.h>
#include <func.h>
#include <type/access_type.h>
#include <type/data_type.h>
#include <type/mem_type.h>
 
namespace freetensor {
 
struct NativeCodeParam {
    std::string name_;
    std::optional<DataType> dtype_; 
    AccessType atype_;
    std::optional<MemType> mtype_; 
    Ref<Ref<Array>> closure_;      
    bool updateClosure_; 
 
    bool isInClosure() const { return closure_.isValid(); }
 
    NativeCodeParam(const std::string &name,
                    const std::optional<DataType> &dtype,
                    const AccessType &atype,
                    const std::optional<MemType> &mtype,
                    const Ref<Ref<Array>> &closure, bool updateClosure)
        : name_(name), dtype_(dtype), atype_(atype), mtype_(mtype),
          closure_(closure), updateClosure_(updateClosure) {}
};
 
std::ostream &operator<<(std::ostream &os, const NativeCodeParam &p);
 
struct NativeCodeRet {
    std::string name_;
    DataType dtype_;
    Ref<Ref<Array>> closure_; 
    bool returnClosure_;      
 
    bool isInClosure() const { return closure_.isValid(); }
 
    NativeCodeRet(const std::string &name, const DataType &dtype,
                  const Ref<Ref<Array>> &closure, bool returnClosure)
        : name_(name), dtype_(dtype), closure_(closure),
          returnClosure_(returnClosure) {}
};
 
std::ostream &operator<<(std::ostream &os, const NativeCodeRet &r);
 
struct StaticInfo {
    uint64_t gpuGlobalStaticPoolSize_ = 0;
};
 
class NativeCode {
    std::string name_;
    std::vector<NativeCodeParam> params_;
    std::vector<NativeCodeRet>
        returns_; // NOTE: multiple items in `returns_` may share the same name.
                  // In this case, one variable should be returned to multiple
                  // positions
    std::string code_;
    std::string entry_; 
    Ref<Target> target_;
    StaticInfo staticInfo_;
 
  public:
    NativeCode() {} // Uninitialized
    NativeCode(const std::string &name,
               const std::vector<NativeCodeParam> &params,
               const std::vector<NativeCodeRet> &returns,
               const std::string &code, const std::string &entry,
               const Ref<Target> &target,
               const StaticInfo &staticInfo = StaticInfo{})
        : name_(name), params_(params), returns_(returns), code_(code),
          entry_(entry), target_(target), staticInfo_(staticInfo) {}
 
    static NativeCode fromFunc(const Func &func, const std::string &code,
                               const std::string &entry,
                               const Ref<Target> &target,
                               const StaticInfo &staticInfo = StaticInfo{});
 
    const auto &name() const { return name_; }
    const auto &params() const { return params_; }
    const auto &returns() const { return returns_; }
    const auto &code() const { return code_; }
    const auto &entry() const { return entry_; }
    const auto &target() const { return target_; }
    const auto &staticInfo() const { return staticInfo_; }
};
 
} // namespace freetensor
 
#endif // FREE_TENSOR_NATIVE_CODE_H