op.h

Go to the documentation of this file.

 
#ifndef _NNDEPLOY_OP_OP_H_
#define _NNDEPLOY_OP_OP_H_
 
#include "nndeploy/base/any.h"
#include "nndeploy/base/common.h"
#include "nndeploy/base/glic_stl_include.h"
#include "nndeploy/base/log.h"
#include "nndeploy/base/macro.h"
#include "nndeploy/base/object.h"
#include "nndeploy/base/param.h"
#include "nndeploy/base/shape.h"
#include "nndeploy/base/status.h"
#include "nndeploy/base/string.h"
#include "nndeploy/base/time_profiler.h"
#include "nndeploy/device/buffer.h"
#include "nndeploy/device/device.h"
#include "nndeploy/device/memory_pool.h"
#include "nndeploy/device/tensor.h"
#include "nndeploy/ir/ir.h"
 
namespace nndeploy {
namespace op {
 
class NNDEPLOY_CC_API Op {
 public:
  Op();
 
  virtual ~Op();
 
  base::Status setName(std::string name);
  std::string getName();
 
  base::Status setOpType(ir::OpType op_type);
  ir::OpType getOpType();
 
  virtual base::Status setParam(std::shared_ptr<base::Param> param);
  virtual std::shared_ptr<base::Param> getParam();
 
  base::Status setDeviceType(base::DeviceType device_type);
  base::DeviceType getDeviceType();
 
  void setStream(device::Stream *stream);
  device::Stream *getStream();
 
  virtual base::Status setPrecisionType(base::PrecisionType precision_type);
  base::PrecisionType getPrecisionType();
 
  std::string getInputName(int index = 0);
  std::string getOutputName(int index = 0);
 
  device::Tensor *getInput(int index = 0);
  device::Tensor *getOutput(int index = 0);
 
  virtual base::Status setInput(device::Tensor *input);
  virtual base::Status setOutput(device::Tensor *output);
 
  base::Status setInput(device::Tensor *input, int index);
  base::Status setOutput(device::Tensor *output, int index);
 
  base::Status setAllInputName(std::initializer_list<std::string>);
  base::Status setAllOutputName(std::initializer_list<std::string>);
 
  base::Status setAllInputName(std::vector<std::string> &);
  base::Status setAllOutputName(std::vector<std::string> &);
 
  std::vector<std::string> getAllInputName();
  std::vector<std::string> getAllOutputName();
 
  std::vector<device::Tensor *> getAllInput();
  std::vector<device::Tensor *> getAllOutput();
 
  device::Tensor *getInputTensor(const std::string &name);
  device::Tensor *getOutputTensor(const std::string &name);
  base::Status replaceInputTensor(const std::string &name,
                                  device::Tensor *tensor);
  base::Status replaceOutputTensor(const std::string &name,
                                   device::Tensor *tensor);
 
  base::Status rmInput(device::Tensor *tensor);
 
  base::Status setAllInput(std::vector<device::Tensor *> inputs);
  base::Status setAllOutput(std::vector<device::Tensor *> outputs);
 
  bool getConstructed();
 
  base::Status setParallelType(const base::ParallelType &paralle_type);
  base::ParallelType getParallelType();
 
  void setInnerFlag(bool flag);
 
  void setInitializedFlag(bool flag);
  bool getInitialized();
 
  void setTimeProfileFlag(bool flag);
  bool getTimeProfileFlag();
 
  void setDebugFlag(bool flag);
  bool getDebugFlag();
 
  void setRunningFlag(bool flag);
  bool isRunning();
 
  virtual base::Status inferDataType();
  virtual base::Status inferShape();
  virtual base::Status inferDataFormat();
 
  virtual base::Status init();
  virtual base::Status deinit();
 
  virtual base::Status reshape(base::ShapeMap &shape_map);
 
  virtual base::Status preRun();
  virtual uint64_t getWorkspaceSize();
  virtual void setWorkspace(void *workspace);
  virtual base::Status allocateWorkspace();
  virtual uint64_t getFlops();
 
  virtual base::Status checkOrAllocOutput();
 
  virtual base::Status run() = 0;
  virtual base::Status postRun();
 
 protected:
  ir::OpDesc op_desc_;
 
  base::DeviceType device_type_;
  bool is_external_stream_ = false;
  device::Stream *stream_ = nullptr;
 
  base::PrecisionType precision_type_ = base::kPrecisionTypeFp32;
 
  std::vector<device::Tensor *> inputs_;
  std::vector<device::Tensor *> outputs_;
 
  bool workspace_is_external_ = false;  // workspace是否是外部传入
  uint64_t workspace_size_ = 0;         // workspace大小
  void *workspace_ = nullptr;           // op的workspace
  uint64_t flops_ = 0;                  // op的flops
 
  // 是否是图中内部节点
  bool is_inner_ = false;
  // 并行类型
  base::ParallelType parallel_type_ = base::kParallelTypeNone;
  // 是否 可以是inplace op
  bool is_inplace_ = false;
  // 参数&输入是否发生变化
  bool is_changed_ = false;
 
  bool constructed_ = false;
  bool initialized_ = false;
  bool is_running_ = false;
  bool is_time_profile_ = false;
  bool is_debug_ = false;
};
 
class NNDEPLOY_CC_API OpCreator {
 public:
  virtual ~OpCreator(){};
 
  virtual Op *createOp(base::DeviceType device_type, const std::string &name,
                       ir::OpType op_type, std::vector<std::string> &inputs,
                       std::vector<std::string> &outputs) = 0;
 
  virtual std::shared_ptr<Op> createOpSharedPtr(
      base::DeviceType device_type, const std::string &name, ir::OpType op_type,
      std::vector<std::string> &inputs, std::vector<std::string> &outputs) = 0;
};
 
template <typename T>
class TypeOpCreator : public OpCreator {
  virtual Op *createOp(base::DeviceType device_type, const std::string &name,
                       ir::OpType op_type, std::vector<std::string> &inputs,
                       std::vector<std::string> &outputs) {
    auto op = new T();
    op->setDeviceType(device_type);
    op->setName(name);
    op->setOpType(op_type);
    op->setAllInputName(inputs);
    op->setAllOutputName(outputs);
    return op;
  }
 
  virtual std::shared_ptr<Op> createOpSharedPtr(
      base::DeviceType device_type, const std::string &name, ir::OpType op_type,
      std::vector<std::string> &inputs, std::vector<std::string> &outputs) {
    auto op = std::make_shared<T>();
    op->setDeviceType(device_type);
    op->setName(name);
    op->setOpType(op_type);
    op->setAllInputName(inputs);
    op->setAllOutputName(outputs);
    return op;
  }
};
 
extern NNDEPLOY_CC_API std::map<base::DeviceTypeCode, std::map<ir::OpType, std::shared_ptr<OpCreator>>>
    &getGlobalOpCreatorMap();
 
template <typename T>
class TypeOpRegister {
 public:
  explicit TypeOpRegister(base::DeviceTypeCode device_type_code,
                          ir::OpType op_type) {
    getGlobalOpCreatorMap()[device_type_code][op_type] =
        std::shared_ptr<T>(new T());
  }
};
 
extern NNDEPLOY_CC_API Op *createOp(base::DeviceType device_type, const std::string &name,
             ir::OpType op_type);
 
extern NNDEPLOY_CC_API Op *createOp(base::DeviceType device_type, const std::string &name,
             ir::OpType op_type, std::initializer_list<std::string> inputs,
             std::initializer_list<std::string> outputs);
 
extern NNDEPLOY_CC_API Op *createOp(base::DeviceType device_type, const std::string &name,
             ir::OpType op_type, std::vector<std::string> &inputs,
             std::vector<std::string> &outputs);
 
extern NNDEPLOY_CC_API Op *createOp(base::DeviceType device_type, const std::string &name,
             ir::OpType op_type, std::vector<std::string> &inputs,
             std::vector<std::string> &outputs,
             std::shared_ptr<base::Param> param);
 
extern NNDEPLOY_CC_API Op *createOp(base::DeviceType device_type, std::shared_ptr<ir::OpDesc> op_desc);
 
extern NNDEPLOY_CC_API std::shared_ptr<Op> createOpSharedPtr(base::DeviceType device_type,
                                      const std::string &name,
                                      ir::OpType op_type);
 
extern NNDEPLOY_CC_API std::shared_ptr<Op> createOpSharedPtr(
    base::DeviceType device_type, const std::string &name, ir::OpType op_type,
    std::initializer_list<std::string> inputs,
    std::initializer_list<std::string> outputs);
 
extern NNDEPLOY_CC_API std::shared_ptr<Op> createOpSharedPtr(base::DeviceType device_type,
                                      const std::string &name,
                                      ir::OpType op_type,
                                      std::vector<std::string> &inputs,
                                      std::vector<std::string> &outputs);
 
extern NNDEPLOY_CC_API std::shared_ptr<Op> createOpSharedPtr(base::DeviceType device_type,
                                      const std::string &name,
                                      ir::OpType op_type,
                                      std::vector<std::string> &inputs,
                                      std::vector<std::string> &outputs,
                                      std::shared_ptr<base::Param> param);
 
extern NNDEPLOY_CC_API std::shared_ptr<Op> createOpSharedPtr(base::DeviceType device_type,
                                      std::shared_ptr<ir::OpDesc> op_desc);
 
using SISOOpFunc =
    std::function<base::Status(device::Tensor *input, device::Tensor *output,
                               std::shared_ptr<base::Param> op_param)>;
 
using SIMOOpFunc = std::function<base::Status(
    device::Tensor *input, std::initializer_list<device::Tensor *> outputs,
    std::shared_ptr<base::Param> op_param)>;
 
using MISOOpFunc = std::function<base::Status(
    std::initializer_list<device::Tensor *> inputs, device::Tensor *output,
    std::shared_ptr<base::Param> op_param)>;
 
using MIMOOpFunc =
    std::function<base::Status(std::initializer_list<device::Tensor *> inputs,
                               std::initializer_list<device::Tensor *> outputs,
                               std::shared_ptr<base::Param> op_param)>;
 
using namespace base;
using namespace ir;
#define REGISTER_OP_IMPLEMENTION(device_type_code, op_type, op_class) \
  TypeOpRegister<TypeOpCreator<op_class>>                             \
      g_##device_type_code##op_class##_register(device_type_code, op_type);
 
}  // namespace op
}  // namespace nndeploy
 
#endif