...
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.
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.
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.
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.
New MXNet APIs
...