Dynamic custom operator GPU support

[rondogency]

Description

Add GPU support for custom operators.
This is a continuation of custom operators project, initial CPU support is implemented here: #15921

Design

Working backward from the user. Here is an example of custom operator forward function for GPU.

Notice the function interface is the same as CPU operator. The input and output tensors are already in GPU, so you don't need to memcpy them to GPU.

You need the CUDA stream from OpResource in order to launch your CUDA kernel on the correct GPU, and mx_stream_t is defined as cudaStream_t when you compile the code with NVCC.

MXReturnValue forwardGPU(std::map<std::string, std::string> attrs,
                         std::vector<MXTensor> inputs,
                         std::vector<MXTensor> outputs,
                         OpResource res) {
    ... 
    mx_stream_t cuda_stream = res.get_cuda_stream();
    gpu_forward<<<grid, block, 0, cuda_stream>>>(out_data, in_data, N);
    return MX_SUCCESS;
}

Then user will register a single operator with both CPU and GPU computation logic, by specifying the device type in register function. This registration works for any context as user only passes in a string.

REGISTER_OP(my_op)
.setForward(forwardCPU, "cpu")
.setForward(forwardGPU, "gpu")
.setBackward(backwardCPU, "cpu")
.setBackward(backwardGPU, "gpu");

Then user compiles the library similar to compiling the CPU operator library. Python usage is exactly the same as CPU custom operators

import mxnet as mx
mx.library.load(os.path.abspath('libmy_op_lib.so'))
a = mx.nd.array([[1,2,3],[4,5,6]], ctx=mx.gpu())
b = mx.nd.array([[7],[8],[9]], ctx=mx.gpu())
mx.nd.my_op(a,b)

This PR is not

• Supporting CPU function is one library and GPU function in another library. Both CPU and GPU functions have to be in the same library.
• Making inferShape and inferType context aware
• Supporting dynamic loading of custom context

Checklist

Essentials

• 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 https://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Add Fcompute registration, and pass NDArray context to custom library in c_api.cc
• Add context info to MXTensor class in lib_api.h
• Add lib_custom_op/relu.cu example file containing full registration of custom operator "my_relu", and add both CPU and GPU kernel functions in that file
• Modify lib_custom_op/Makefile to compile .cu file using nvcc to custom library

Comments

[eric-haibin-lin]

Some interfaces require more discussion.

[rondogency]

@ptrendx thanks for your comments! I have resolved those comments, and I will appreciate if you could take another quick look and approve this.

[samskalicky]

@rondogency looks like the windows build/test is working now with those cmake changes:

test_operator_gpu.test_custom_op_gpu ... 
MXNet version 10600 supported
[
10:16:08] C:\jenkins_slave\workspace\build-gpu\src\c_api\c_api.cc:286: 
Found 2 operators in library
[
10:16:08] C:\jenkins_slave\workspace\build-gpu\src\c_api\c_api.cc:350: 	Op[0] my_relu
[
10:16:08] C:\jenkins_slave\workspace\build-gpu\src\c_api\c_api.cc:350: 	Op[1] my_state_relu
[
10:16:08] C:\jenkins_slave\workspace\build-gpu\src\c_api\c_api.cc:785: Found 0 partitioners in library
ok (0.6834s)

Now we just need to work through the flaky tests

[wkcn]

Great! LGTM. Thank you!

[rondogency]

@wkcn Thanks! Can you help me to merge it?

[wkcn]

Merged. : )