Versions Compared

Key

  • This line was added.
  • This line was removed.
  • Formatting was changed.

...

The figure below shows the high-level architecture proposed. The user will call the load_op_lib API to load their custom operator library. This will result in the operators being discovered/registered from the so/dll into MXNet's CustomOp registry (similar to NNVM op registry, but only for customOps). Then the user will call the CustomOp operator and specify the op_type attribute to be the name of the customOp from their library. This customOp will indirectly call their custom operator in the library at runtime.

Image AddedImage Removed

When building a customOp Library, users will write 4 functions for each operator: FCompute, InferShape, InferType, and ParseAttrs. These are similar to the standard functions required for current Backend C/C++/CUDA operators in MXNet. Similarly, they will register their op (ie. the 4 functions) in the library. As shown above, this “local-registration” will be parsed by MXNet when loading the customOp library at runtime.
Image Removed


Image Added

Runtime Behavior

Heres the overall runtime behavior for CustomOps. Its it is broken down into 2 parts: initial library load, and operator execution.

First, the user writes their custom op functions: FCompute, InferShape, InferType, and ParseAttrs. Then they statically register the functions in their library with REGISTER_OP. Next they compile and produce a shared library (so/dll). Then they run their MXNet model, and load their library. During the initial setup, the user calls load_op_lib in their code to load their shared library. During the loading process, MXNet parses all of the operators that have been registered by getting the number of ops registered with the _opRegSize function. Then it iteratively gets each op by calling the _opRegGet and analyzes it before re-registering it inside MXNet's customOp registry.
Image Removed
Image Added


Then, later when a CustomOp operator is bound/executed the functions from the shared library are executed. During the bind step, the Operator is looked-up by op name and the attributes for the operator are analyzed by the customOp's parseAttrs function in the shared library. Then for type and shape inference, the respective functions are also called through the inferType and inferShape APIs. Lastly, when executing the forward pass, the FCompute function is called for the operator from the shared library.


Image RemovedImage Added

New MXNet APIs

...