Operator

Add new operator

Operator definition and declaration

nndeploy operator is registered and used in a class-based manner. The implementation and backend binding of the operator are strongly tied, but operators of the same type from different backends can share some common functions, such as shape inference. In nndeploy, operators are divided into two levels, the first parent class level is in nndeploy/framework/include/nndeploy/op, defined in nndeploy/framework/source/nndeploy/op. For example, RMSNorm:

class OpRMSNorm : public Op {
 public:
  OpRMSNorm() : Op() { is_inplace_ = true; }
  virtual ~OpRMSNorm() {}

  virtual base::Status inferShape();

  virtual base::Status run();
};

In the parent class level, two key functions need to be implemented:

• inferShape: Infer the shape of the output Tensor, this process is based on the network’s forward propagation; when the network’s input size is determined, the shape of all Tensors along the network’s continuous propagation path is inferred. When the shape inference is completed, the subsequent internal memory request is made based on the shape and life cycle of the Tensor.

• run: The run function of the parent class implements a non-dependent on any specific instruction set, which can be executed on most CPUs as an unoptimized version for fast implementation.

The implementation of specific backend’s operator subclass is under nndeploy/framework/source/nndeploy/op/xxx, xxx is the backend file folder. For example, if there are no special requirements to rewrite the parent class’s other virtual functions, only the implementation of run is used. In this layer, it is assumed that all inputs, outputs, shapes, value types are correct, only the implementation needs attention.

All Ops must be registered into the factory model:

REGISTER_OP_IMPLEMENTION(base::DeviceTypeCode::kDeviceTypeCodeCpu,
                         ir::kOpTypeSoftmax, OpSoftmax)

Function form operator

The above implementations are all in the form of class, sometimes it is expected to directly perform calculations on Tensor in the form of a function, hence the operator in the form of function is directly applied. The declaration and definition of this function interface are in the same file as the parent class Op. Each operator is basically in a fixed form, creating Op objects, setting inputs and outputs, initialization, and execution. For example:

base::Status rmsNorm(device::Tensor *input1, device::Tensor *input2,
                     device::Tensor *input3, device::Tensor *output) {
  base::Status status = base::kStatusCodeOk;

  Op *op = createOp(input1->getDeviceType(), "", ir::kOpTypeRMSNorm);
  if (op == nullptr) {
    NNDEPLOY_LOGE("createOp failed");
    return base::kStatusCodeErrorNotImplement;
  }
  status = op->setInput(input1, 0);
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "setInput failed");
  status = op->setInput(input2, 1);
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "setInput failed");
  status = op->setInput(input3, 2);
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "setInput failed");
  status = op->setOutput(output, 0);
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "setOutput failed");
  status = op->init();
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "init failed");
  status = op->preRun();
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "preRun failed");
  status = op->run();
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "run failed");
  status = op->postRun();
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "postRun failed");
  status = op->deinit();
  NNDEPLOY_RETURN_ON_NEQ(status, base::kStatusCodeOk, "deinit failed");

  delete op;

  return status;
}

Python interface export

Exporting the operator in the form of function at the Python end, can be directly used by nndeploy.op.xxop(input). The export interface is located in nndeploy/python/src/op/op.cc, for example:

NNDEPLOY_API_PYBIND11_MODULE("op", m) { m.def("rms_norm", rmsNormFunc); }

Indicates exporting to the python op module, nndeploy.op.rms_norm will be directed to execute in the rmsNormFunc function. rmsNormFunc is a middleware between Python and Cpp, its implementation is in nndeploy/python/src/op/op_func.cc. In the Func layer, input legality checks are performed, if it is a non-inplace operator, it also needs to request output Tensor. But do not perform internal memory allocation for the output Tensor, because internal memory allocation depends on shape inference, which is performed inside the function operator.

Test

The operator’s test is located in nndeploy/python/nndeploy/tests/op, using the unittest framework for testing. Data construction can use numpy, and then converted to nndeploy’s Tensor. Standard result comparison can use numpy or PyTorch, converting nndeploy’s Tensor to numpy for comparison.