自定义Backend

    所有新增Backend都需继承类,并实现所有纯虚函数。

    构造与销毁

    Backend构造时,可以额外指定内存环境,在内存受限环境中,应避免非必要的内存使用。可以在构造函数中,完成对计算设备访问的必要初始化,如GPU下预加载shader等。

    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);

    Backend需要通过onCreate为op创建出exection实例:

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

    这样,Op Execution中,就可以通过注册追加Op类型:

    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);

    内存管理

    Backend通过onAcquireBuffer为tensor分配内存,通过onReleaseBuffer为tensor释放内存。内存有三种存储模式:STATIC内存不复用,一般用于op常量存储;DYNAMIC内存可复用,一般用于变量存储;DYNAMIC_SEPERATE内存在pipeline间可复用,一般用于pipeline常量存储。onAcquireBufferonReleaseBuffer中可以不实际分配/释放内存,只记录内存用量变更,在onAllocateBuffer调用时,再根据用量计算出优化方案,一次性完成分配/释放。

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

    在所有内存都分配完成后,backend会收到onAllocateBuffer回调:

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

    此外,backend还需要负责tensor数据的拷贝:

    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;

    拷贝可能在backend内部,也可能在backend与CPU backend之间。 拷贝需要处理Tensor间的布局转换,相同布局时,可以直接拷贝数据;不同布局,如NHWCNC4HW4,则一般需要做特殊转换。

    Backend在pipeline执行的各个周期都会收到回调,onResizeBeginonResizeEnd在调整内存分配前后调用(op的onResize会在此间调用);onExecuteBeginonExecuteEnd在op执行前后调用(op的onExecute会在此间调用);onWaitFinish相对特殊,由用户主动调用,异步执行的pipeline需要同步等待完成。

    注册Backend

    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);