Backends of MNN are abstraction of computing devices. MNN currently supports Backend for CPU, Vulkan, OpenCL, Metal, etc. Add new backend only if the computing device is not supported. To add Op, see Add Op document.

Declaration

All new backends need to inherit the Backend class and implement all pure virtual functions.

  1. class XPUBackend final : public Backend {
  2.   XPUBackend(MNNForwardType type, MemoryMode mode);
  3.   virtual ~XPUBackend();
  4.   virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op) override;
  5.   virtual void onExecuteBegin() const override;
  6.   virtual void onExecuteEnd() const override;
  7.   virtual bool onAcquireBuffer(const Tensor* tensor, StorageType storageType) override;
  8.   virtual bool onReleaseBuffer(const Tensor* tensor, StorageType storageType) override;
  9.   virtual bool onClearBuffer() override;
  10.   virtual void onCopyBuffer(const Tensor* srcTensor, const Tensor* dstTensor) const override;
  11. }

Construction and Deconstruction

When backend is constructed, you can specify an additional memory option. In a memory-constrained environment, you should avoid unnecessary memory usage. In the constructor, you can complete the necessary initialization of the access to the computing device, such as preloading shaders running on GPU.

  1. /** backend memory mode */
  2. enum MemoryMode {
  3.     /** use memory without limit. */
  4.     NORMAL = 0,
  5.     /** use memory thriftily. */
  6.     LIMIT = 1
  7. };
  8. /**
  9.  * @brief initializer.
  10.  * @param type  forward type.
  11.  * @param mode  memory mode.
  12.  */
  13. Backend(MNNForwardType type, MemoryMode mode = NORMAL);

Execution Create

Backend creates an exection instance with method onCreate:

  1. virtual Execution* onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op) override;

An execution instance should be created and returned in this method. To separate type of op, providing a registration interface is more recommended than using the switch-case pattern:

  1. class XPUBackend final : public Backend {
  2.     // ...
  3.     class Creator {
  4.     public:
  5.         /**
  6.          * @brief create execution for given input, op on metal backend.
  7.          * @param inputs    given input tensors.
  8.          * @param op        given op.
  9.          * @param backend   metal backend.
  10.          * @return created execution if supported, NULL otherwise.
  11.          */
  12.         virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const MNN::Op *op,
  13.                                     Backend *backend) const = 0;
  14.     };
  15.     /**
  16.      * @brief register creator for given op type.
  17.      * @param type      given op type.
  18.      * @param creator   registering creator.
  19.      */
  20.     static void addCreator(OpType type, Creator *creator);
  21.     // ...
  22. };
  23. template <class T>
  24. class XPUCreatorRegister {
  25. public:
  26.     /**
  27.      * @brief initializer. register T creator for given op type.
  28.      * @param type  given op type.
  29.      */
  30.     XPUCreatorRegister(OpType type) {
  31.         T *test = new T;
  32.         XPUBackend::addCreator(type, test);
  33.     }
  34. };

In this way, you can append the Op type by registration in each Op Execution file:

  1. class XPUPoolingCreator : public XPUBackend::Creator {
  2. public:
  3.     virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const MNN::Op *op, Backend *backend) const {
  4.         return new XPUPooling(backend, op->main_as_Pool());
  5.     }
  6. };
  7. static XPUCreatorRegister<XPUPoolingCreator> __reg(OpType_Pooling);

Memory Management

Backend allocates memory for tensor via onAcquireBuffer and frees memory via onReleaseBuffer. There are three storage modes in memory: STATIC memory is not reused, generally used for op constant storage; DYNAMIC memory can be reused, generally used for variable storage; DYNAMIC_SEPERATE memory can be reused between pipelines, generally used for pipeline Constant storage. There is a better way to allocate/release memory –- record memory usage changes only in onAcquireBuffer and onReleaseBuffer, allocate/release memory until onAllocateBuffer is called, so that all memory needed could be merged into continuous memory and the allocate/release could be completed in one call.

  1. /** backend buffer storage type */
  2. enum StorageType {
  3.     /**
  4.      use NOT reusable memory.
  5.      - allocates memory when `onAcquireBuffer` is called.
  6.      - releases memory when `onReleaseBuffer` is called or when the backend is deleted.
  7.      - do NOTHING when `onClearBuffer` is called.
  8.      */
  9.     STATIC,
  10.     /**
  11.      use reusable memory.
  12.      - allocates or reuses memory when `onAcquireBuffer` is called. prefers reusing.
  13.      - collects memory for reuse when `onReleaseBuffer` is called
  14.      - releases memory when `onClearBuffer` is called or when the backend is deleted.
  15.      */
  16.     DYNAMIC,
  17.     /**
  18.      use NOT reusable memory.
  19.      - allocates memory when `onAcquireBuffer` is called.
  20.      - do NOTHING when `onReleaseBuffer` is called.
  21.      - releases memory when `onClearBuffer` is called or when the backend is deleted.
  22.      */
  23.     DYNAMIC_SEPERATE
  24. };
  25. /**
  26.  * @brief allocate buffer of tensor for given storage type.
  27.  * @param tensor        buffer provider.
  28.  * @param storageType   buffer storage type.
  29.  * @return success or not.
  30.  */
  31. virtual bool onAcquireBuffer(const Tensor* tensor, StorageType storageType) = 0;
  32. /**
  33.  * @brief release buffer of tensor for given storage type.
  34.  * @param tensor        buffer provider.
  35.  * @param storageType   buffer storage type.
  36.  * @return success or not.
  37.  */
  38. virtual bool onReleaseBuffer(const Tensor* tensor, StorageType storageType) = 0;

After all memory is allocated, Backend will receive the onAllocateBuffer callback:

  1. /**
  2.  * @brief callback after all buffers needed by backend ops were allocated.
  3.  * @return success or not. (result not used currently)
  4.  */
  5. virtual bool onAllocateBuffer() {
  6.     return true;
  7. }

Backend needs to release all the memory of DYNAMIC and DYNAMIC_SEPERATE storage mode when onClearBuffer called:

  1. /**
  2.  * @brief clear all dynamic buffers.
  3.  * @return success or not.
  4.  */
  5. virtual bool onClearBuffer() = 0;

In addition, backend is also responsible for copying the tensor data:

  1. /**
  2.  * @brief copy buffer from tensor to tensor.
  3.  * @param srcTensor source buffer provider.
  4.  * @param dstTensor dest buffer provider.
  5.  */
  6. virtual void onCopyBuffer(const Tensor* srcTensor, const Tensor* dstTensor) const = 0;

The copy process may be inside Backend or between Backend and CPU Backend.
The copy process needs to handle the layout conversion between Tensor. With the same layout, you can directly copy the data; different layouts, such as NHWC and NC4HW4, need to do special conversion generally.

Pipeline Callback

Backend receives callbacks in each cycle of the pipeline execution, onResizeBegin and onResizeEnd are called before and after adjusting memory allocation (onResize of op will be called here); onExecuteBegin and onExecuteEnd are called before and after op execution (onExecute of op will be called here); onWaitFinish is relatively special and is called by the user, the asynchronously executed pipeline needs to wait synchronously for completion.

  1. /**
  2.  * @brief callback before resize ops.
  3.  */
  4. virtual void onResizeBegin();
  5. /**
  6.  * @brief callback after resize ops.
  7.  */
  8. virtual void onResizeEnd();
  9. /**
  10.  * @brief callback before executing ops.
  11.  */
  12. virtual void onExecuteBegin() const = 0;
  13. /**
  14.  * @brief callback after executing ops.
  15.  */
  16. virtual void onExecuteEnd() const = 0;
  17. /**
  18.  * @brief wait for all async execution to be finished.
  19.  * @return success or not.
  20.  */
  21. virtual bool onWaitFinish();

Register Backend

Finally, define Backend Creator and call MNNInsertExtraBackendCreator in register method to complete the registration of Backend. The register method should be declared and called in BackendRegister.cpp:

  1. class XPUBackendCreator : public BackendCreator {
  2.     virtual Backend *onCreate(const Backend::Info &info) const {
  3.         return new MetalBackend;
  4.     }
  5. };
  6. void registerCPUBackendCreator() {
  7.     MNNInsertExtraBackendCreator(MNN_FORWARD_CPU, new CPUBackendCreator);
  8. };

When compiling with cmake, after completing code modifications, you also need to modify CMakeLists.txt accordingly.