Page History

...

Lastly, we want custom operators to be first-class operators and have access to all the capabilities that internal MXNet operators do. One example is enabling custom operators to leverage the MXNet resource manager for storage and memory.

Proposed Approach

Compiling Custom Operators

To support compiling custom operators, we need to construct a simple API/file-set that users will compile their custom operators with. The result of this compilation will be a dynamic library (Linux: so, Windows: dll). We will need to provide unit tests that allow users to test their operator registration outside of MXNet to ease debugging.

Just how operators are registered in MXNet with NNVM, we propose a similar lightweight approach that doesnt require compiling custom operators with NNVM.

Dynamically Loading Operator Libraries

After a user compiles custom operator(s) into a library, we need to construct an API to

• load user-specified libraries
• register operators from each library in MXNet so that they can be called/executed

Calling Custom Operators

After a library is loaded, users need to call their operators from their application. Currently the Python CustomOp provides a mechanism that we can re-use for C++/CUDA operators: mx.nd.CustomOp(op_type='my_op', *args). We'll use the same approach here to provide a similar user experience.

Architecture

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 Added

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 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 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 Added

New MXNet APIs

These are new APIs that will be added to MXNet to

C APIs

• MXLoadCustomOpLib - API to load operator libraries
• Load the customOp library
• Go through the operators in the library
• Check that each operator defines required functions
• ParseAttrs, InferType, InferShape, FCompute
• Register each operator found

Python APIs

• load_op_lib - API to load operator libraries
• Takes a path to the operator library
• checks if the path exists and if points to file
• calls C API MXLoadCustomOpLib to perform actual loading
• mx.operator.load_op_lib('/path/to/libtest.so')

New CustomOp Operator

• CustomOp - new operator that executes custom operators loaded from the library
• Takes op_type to identify custom operator name
• Takes any number of kwargs as attributes/parameters
• Takes any number of in-order args as input arrays
• b = mx.nd.CustomOp(a,op_type='sam',myParam='2')

APIs for implementing Custom Operators

These are the new APIs that users will implement for their custom operators

• parseAttrs - takes a set of key/value pairs for attributes and gives users an opportunity to validate the attributes passed to their custom operator.
• int parseAttrs(std::map<std::string,std::string> attrs,
  int* num_in,
  int* num_out);
• Inputs: the map of attributes passed to the operator from the user
• Outputs: num_in, num_out - the number of input/output arrays required for this operator
• returns 1 if success, or zero if failure
• inferType - performs type inference for this operator
• int inferType(std::map<std::string,std::string> attrs,
  std::vector<int> &intypes,
  std::vector<int> &outtypes);
• Inputs: the map of attributes
• Inputs/Outputs: intypes, outtypes - the list of input/output types that should be inferred. Values of of -1 should be defined by this operator as a specific type
• returns 1 if success, or zero if failure
• inferShape - performs shape inference for this operator
• int inferShape(std::map<std::string,std::string> attrs,
  std::vector<std::vector<unsigned int>> &inshapes,
  std::vector<std::vector<unsigned int>> &outshapes);
• Inputs: the map of attributes
• Inputs: inshapes - the shapes of the input arrays
• Outputs: outshapes - the shapes of output arrays
• fcompute - performs computation of this operator
• int myFCompute(std::map<std::string,std::string> attrs,
  std::vector<MXTensor> inputs,
  std::vector<MXTensor> outputs);
• Inputs: the map of attributes
• Input data: inputs, input tensors
• Output data: outputs, output tensors

API for registering operator and its functions:

• REGISTER_OP - registers the operator in the library
• REGISTER_OP(sam)
  .setFCompute_cpu(myFCompute)
  .setParseAttrs(parseAttrs)
  .setInferType(inferType)
  .setInferShape(inferShape);
• REGISTER_OP - macro that defines an custom operator object with given name
• setFCompute_cpu - sets the FCompute function for CPU context
• setFCompute_gpu - sets the FCompute function for GPU context
• setParseAttrs - sets the parse attributes function
• setInferType - sets the infer types function
• setInferShape - sets the infer shapes function

Goals/Usecases

MXNet Java Inference API#Goals

Open Questions

Proposed Approach

Initial PoC in this branch: https://github.com/samskalicky/incubator-mxnet/tree/custom_op/

...