optimizer.h

Go to the documentation of this file.

 
#ifndef _NNDEPLOY_NET_OPTIMIZER_H_
#define _NNDEPLOY_NET_OPTIMIZER_H_
 
#include <unordered_map>
#include <unordered_set>
 
#include "nndeploy/ir/ir.h"
#include "nndeploy/net/util.h"
#include "nndeploy/op/op.h"
 
namespace nndeploy {
namespace net {
// help types
using OpSet = std::unordered_set<ir::OpType>;
 
enum OptPassType : int {
  // Operator fusion
  kOptPassTypeFuseConvBias,
  kOptPassTypeFuseConvBatchNorm,
  kOptPassTypeFuseConvRelu,
  kOptPassTypeFuseConvAct,
 
  // Eliminate useless op
  kOptPassTypeEliminateCommonSubexpression,
  kOptPassTypeEliminateDeadOp,
 
  // Constant Folding
  kOptPassTypeFoldConstant,
 
  // QDQ fuse
  kOptPassTypeFuseQdq,
};
 
 
class Net;
 
class OptPass {
 public:
  OptPass(std::string name);
  virtual ~OptPass();
 
  virtual int seqPatternMatch(std::vector<TensorWrapper*>& tensor_repository,
                              std::vector<OpWrapper*>& op_repository,
                              const std::vector<ir::OpType>& types,
                              int begin_op_index);
 
  virtual int seqPatternMatch(std::vector<TensorWrapper*>& tensor_repository,
                              std::vector<OpWrapper*>& op_repository,
                              const std::vector<OpSet>& types,
                              std::vector<ir::OpType>& matched_types,
                              int begin_op_index);
 
  virtual base::Status seqPatternMatchUpateTensorRepository(
      std::vector<TensorWrapper*>& tensor_repository,
      std::vector<OpWrapper*>& op_repository,
      const std::vector<ir::OpType>& types, int begin_op_index);
 
  virtual base::Status seqPatternMatchUpateOpRepository(
      std::vector<TensorWrapper*>& tensor_repository,
      std::vector<OpWrapper*>& op_repository,
      const std::vector<ir::OpType>& types, int begin_op_index);
 
  virtual base::Status rmOpFromPredecessor(OpWrapper* op_wrapper);
 
  virtual base::Status rmOpFromSuccessors(OpWrapper* op_wrapper);
 
  virtual base::Status rmInputTensorAndMaybeDelete(
      OpWrapper* op_wrapper, std::vector<TensorWrapper*>& tensor_repository);
 
  virtual base::Status rmOutputTensorAndMaybeDelete(
      OpWrapper* op_wrapper, std::vector<TensorWrapper*>& tensor_repository);
 
  std::string getName();
  base::Status setNet(Net* net);
 
  virtual base::Status optimize(std::vector<TensorWrapper*>& tensor_repository,
                                std::vector<OpWrapper*>& op_repository,
                                int begin_op_index) = 0;
 
 protected:
  std::string name_;  // pass名称
 
  Net* net_ =
      nullptr;  // 该pass所属的Net，可能要修改Net内部的数据，例如释放某些tensor
};
 
class OptPassCreator {
 public:
  virtual ~OptPassCreator() {};
 
  virtual std::shared_ptr<OptPass> createOptPass() = 0;
};
 
template <typename T>
class TypeOptPassCreator : public OptPassCreator {
  virtual std::shared_ptr<OptPass> createOptPass() {
    return std::shared_ptr<T>(new T());
  }
};
 
std::map<base::DeviceTypeCode,
         std::map<int, std::map<OptPassType, std::shared_ptr<OptPassCreator>>>>&
getGlobalOptPassCreatorMap();
 
template <typename T>
class TypeOptPassRegister {
 public:
  explicit TypeOptPassRegister(base::DeviceTypeCode device_type_code,
                               OptPassType type, int level) {
    auto& creator_map = getGlobalOptPassCreatorMap();
    auto device_map = creator_map.find(device_type_code);
    if (device_map == creator_map.end()) {
      creator_map[device_type_code] =
          std::map<int,
                   std::map<OptPassType, std::shared_ptr<OptPassCreator>>>();
    }
    auto level_map = creator_map[device_type_code].find(level);
    if (level_map == creator_map[device_type_code].end()) {
      creator_map[device_type_code][level] =
          std::map<OptPassType, std::shared_ptr<OptPassCreator>>();
    }
 
    auto creator = creator_map[device_type_code][level].find(type);
    if (creator == creator_map[device_type_code][level].end()) {
      creator_map[device_type_code][level][type] = std::shared_ptr<T>(new T());
    }
  }
};
 
std::shared_ptr<OptPass> createOptPass(base::DeviceType device_type, int level,
                                       OptPassType type);
 
class NNDEPLOY_CC_API Optimizer {
 public:
  Optimizer();
  ~Optimizer();
 
  base::Status init(base::DeviceType device_type,
                    std::set<OptPassType> enable_pass,
                    std::set<OptPassType> disable_pass);
  base::Status deinit();
 
  base::Status addPass(OptPassType type, int level);
  base::Status removePass(OptPassType type);
 
  base::Status optimize(std::vector<TensorWrapper*>& tensor_repository,
                        std::vector<OpWrapper*>& op_repository, Net* net);
 
 protected:
  base::DeviceType device_type_;
  std::map<int, std::map<OptPassType, std::shared_ptr<OptPass>>>
      opt_passes_;  // 第一个key是优先级，数字越小， 优先级越高，
                    // 在图优化时首先执行这个pass
};
 
extern NNDEPLOY_CC_API std::string optPassTypeToString(OptPassType type);
extern NNDEPLOY_CC_API OptPassType stringToOptPassType(const std::string &src);
 
 
}  // namespace net
}  // namespace nndeploy
 
#endif /* _NNDEPLOY_NET_OPTIMIZER_H_ */