dynamic custom operator support

[samskalicky]

Description

Enhancements to dynamic library loading to support custom operators in libraries.

• added MXTensor/MXDtype structure, and versioning to lib_api.h
• added similar NNVM register op-like capability for custom ops
• operators are found in the library, and re-registered in MXNet during library loading
• operators are re-registered from mx.nd.op to mx.nd and mx.sym.op to mx.sym shortcuts
• Created a new example library in "example/lib_ops" with a GEMM operator

Initially, this project was proposed on the CWiki, however the design has evolved since the initial proposal. The current design is described below.

Design

The goal of this PR to to enable operators to be implemented in separate libraries and loaded at runtime

The main constraint is to maintain a low-level C-types only boundary between MXNet and the library to simplify the building and compiling of external libraries.

Working backwards from the user, users register operators with easy-to-use function prototypes like:

int myForward(std::map<std::string,std::string> attrs,
               std::vector<MXTensor> inputs, 
               std::vector<MXTensor> outputs);

Users' Forward (ie. FCompute) functions are called by a helper function _opCallFCompute that converts the C-types passed across the library boundary to the familiar STL types. This function is implemented in the lib_api.h header file that users compile with their library.

int _opCallFCompute(fcomp_t fcomp, const char* const* keys, const char* const* vals, int num,
                    const int64_t** inshapes, int* indims, void** indata, int* intypes, int num_in,
                    const int64_t** outshapes, int* outdims, void** outdata, int* outtypes, int num_out);

In MXNet's C API, the _opCallFCompute function is found in the library. A lambda function fcomp_conv is created for each operator loaded from the library to convert from MXNet-types to C-types. Then these C-types are passed to _opCallFCompute.

auto fcomp_conv = [=](const nnvm::NodeAttrs& attrs,
                      const OpContext& ctx,
                      const std::vector<TBlob>& inputs,
                      const std::vector<OpReqType>& req,
                      const std::vector<TBlob>& outputs);

The same design is used for: parseAttrs, inferShape, inferType, etc.

Finally, an operator is re-registered in MXNet with the lambda function like:

nnvm::Op &regOp = dmlc::Registry<nnvm::Op>::Get()->__REGISTER_OR_GET__(name);
regOp.set_attr<FCompute>("FCompute<cpu>",fcomp_conv);

Once the C API returns back to Python in the load function in library.py, we regenerate the Python bindings and re-register the operator shortcuts to mx.nd and mx.sym.

After the load function returns back to the user's Python code, they can then use their operators just like any other operate:

mx.nd.myCustomOp(A,B,C)

Current Features

• custom CPU operators
• stateless & stateful operators
• custom subgraph operators
• Memory resource request

Future/Next-steps

(to be done in a separate PR)

• custom GPU operators
• Random number generator resource request
• sparse data types
• migrate lambda functions in MXLoadLib in src/c_api/c_api.cc to classes defined elsewhere
• Documentation, add the "library" python package to the namespace to the doc: https://mxnet.apache.org/api/python/docs/api/ ?

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• updated path to be absolute in example/lib_api/test.py
• update the Makefile/CMakeLists.txt to build the new example in lib_ops instead of lib_api
• moved library import in python/mxnet/init.py to after ndarray/symbol
• added operator reregistration to python/mxnet/library.py
• added operator discovery/registration in MXLoadLib in src/c_api/c_api.cc

Comments

[samskalicky]

@wkcn while this PR is not quite done yet, it would be great to get some early feedback since the design/implementation has changed since our initial discussion. Let me know what you think, thanks!

[rondogency]

@wkcn 1.6 code freeze is tomorrow, so are you ok with this one not going into the 1.6 release? It is because none of us have time to maintain it until late November. After code freeze then we can merge it on Friday, and user can use nightly build to access this feature.

[wkcn]

@rondogency No problem : )

[wkcn]

Hi @samskalicky and @rondogency , is it ready to merge this PR after CI passes?

[samskalicky]

Hi @samskalicky and @rondogency , is it ready to merge this PR after CI passes?

Yes! We're soooooo ready to merge :)

Thanks @zachgk for reruning the unix_cpu job!

[wkcn]

I will merge this PR after the CI passes.
Thank all contributors!

[rondogency]

@wkcn Big thank to Jackie for the merging work!

[TaoLv]

It has a strong assumption that the case will be called from mxnet root folder. Otherwise, the libsample_lib.so will not be found.

$ cd tests/python/unittest/
$ nosetests -v test_extensions:test_custom_op
test_extensions.test_custom_op ... ERROR

======================================================================
ERROR: test_extensions.test_custom_op
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/lvtao/miniconda3/envs/mxnet/lib/python3.6/site-packages/nose/case.py", line 198, in runTest
    self.test(*self.arg)
  File "/home/lvtao/Workspace/mxnet-official/tests/python/unittest/test_extensions.py", line 41, in test_custom_op
    raise MXNetError("library %s not found " % lib)
mxnet.base.MXNetError: library libsample_lib.so not found

----------------------------------------------------------------------
Ran 1 test in 0.005s

FAILED (errors=1)

[wkcn]

Is it possible to use mx.libinfo.find_lib_path to find the library?
https://github.com/apache/incubator-mxnet/blob/93228649340bcacb8056d47d8f6f8a78a9805ae4/python/mxnet/libinfo.py#L26

[rondogency]

Gotcha I will fix it in the next PR