wkcn commented on a change in pull request #17486: Update CustomOp doc with 
changes for GPU support
URL: https://github.com/apache/incubator-mxnet/pull/17486#discussion_r375598920
 
 

 ##########
 File path: example/extensions/lib_custom_op/README.md
 ##########
 @@ -146,21 +186,160 @@ Let’s take a closer look at those registry functions:
 
 * **inferShape**: This function is similar to the `inferType` function, except 
it is used for populating the output data shapes. You need to figure out the 
shapes of each output tensors for this computation. For example, if the inputs 
are images with shape (224,224,3) and you write a padding operator to make 10px 
borders for the images, then your output shape will be (234,234,3).
 
-* **forward**: This function executes the main forward computation. It takes 
four arguments. The 1st argument is the attributes. The 2nd argument is the 
input `MXTensors` which stores all data and info of input ndarrays. The 3rd 
argument is the output `MXTensors`. The 4th argument is the `OpResource` object 
for memory allocation and other utilities. Additionally, you can use a 
`dltensor` tensor structure stored in the `MXTensor` as a more standardized 
data structure for computing.
+* **forward**: This function executes the main forward computation. It takes 
four arguments. The 1st argument is the attributes. The 2nd argument is the 
input `MXTensors` which stores all data and info of input ndarrays. The 3rd 
argument is the output `MXTensors`. The 4th argument is the `OpResource` object 
for memory allocation and other utilities. The details of `OpResource` are 
covered in the below section.
+Additionally, you can use a `dltensor` tensor structure stored in the 
`MXTensor` as a more standardized data structure for computing.
 
 * **backward**: This function is doing the backward gradient computation. It 
will be similar to the forward function. And you need to figure out the formula 
of the backward gradient computation.
 
 * **mutateInputs**: This function is for marking mutable inputs. It takes two 
arguments. The 1st argument is the attributes. The 2nd argument is a list of 
input indices that are mutable among all input tensors. It is useful when some 
inputs are auxiliary model parameters and might be altered during 
forward/backward computation. Remember, the index number of `input_indices` 
should not exceed the number of inputs.
 
-### Writing Stateful Custom Operator:
+### Writing Stateful Custom Operator
 
 A stateful custom operator is useful when a forward/backward call needs some 
data or ‘state’ from previous forward/backward calls. Normally we create a 
class, and make instance variables store the states used for computing or 
caching.
 
 Most of the building blocks for making a stateful custom operator is the same 
as regular custom operator, except it’ll register `createOpState` instead of a 
`forward` function for the computation.
 
 * [createOpState](./gemm_lib.cc#L204) - Create stateful operator instance:
     * This function takes two arguments. The 1st argument is attributes. The 
2nd argument is a placeholder for `CustomStatefulOp` object. You must [define a 
class that inherits CustomStatefulOp](./gemm_lib.cc#L178) and override the 
forward function (optionally the backward function). Then you need to create an 
instance of your class and assign it to the placeholder. In this way, all of 
the forward/backward calls will use the same methods in that instance, and the 
instance is able to keep the state of the operator.
+```c++
+    MXReturnValue createOpState(
+        std::map<std::string, std::string> attrs,
+        CustomStatefulOp** op_inst)
+```
+
+* The operator registering function will look like this:
+```c++
+    REGISTER_OP(my_state_op)
+    ...
+    .setCreateOpState(createOpState, "cpu");
+```
+
+## Writing A Custom GPU Operator Library
+
+Most of the building blocks for registering GPU custom operators are the 
exactly same as CPU ones, except you need to specify the `"gpu"` context name 
when registering `forward`, `backward` or `createOpState` function.
+
+### Run A GPU Example
+
+For illustration purposes, we provided a `ReLU` (Rectified Linear Unit) 
activation operator that can run on GPU. Make sure you have installed a CUDA 
compatible MXNet build. Go to `lib_custom_op` directory and follow these steps: 
+
+1. Run `make relu_lib`. The Makefile will invoke `NVCC` compiler to compile 
the CUDA kernel along with regular custom operator functions from `relu_lib.cu` 
to generate `librelu_lib.so` library.
+2. Run `python test_relu.py`. It’ll register the GPU `ReLU` operator in the 
MXNet backend, then invoke the operator by feeding an `NDArray` input with GPU 
context, and output the result tensor with GPU context.
+
+### Writing A Regular GPU Custom Operator
+
+Since most of the building blocks for registering GPU custom operators are the 
exactly same as CPU ones, the registering function for an operator supporting 
both GPU and CPU will look like this:
+
+```c++
+    REGISTER_OP(my_op_gpu)
+    ...
+    .setForward(forwardCPU, "cpu")
+    .setForward(forwardGPU, "gpu")
+    .setBackward(backwardCPU, "cpu")
+    .setBackward(backwardGPU, "gpu");
+```
+
+Note that you can write only the GPU funcitons, as operators don’t have to 
support both CPU and GPU functions (can be GPU only).
+
+After you register forward or backward functions with context name `“gpu”`, 
MXNet will dispatch forward or backward GPU functions you just registered when 
the operator is invoked with GPU context.
+
+In the registered `forwardGPU` function, specify the grid and block size and 
launch the CUDA kernel. In the GPU operators input and output tensors are 
pre-allocated on the GPU memory, just like in the CPU operator tensors are 
pre-allocated on the CPU memory. As a result, you don’t need to call 
`cudaMemcpy` to move the tensor data to GPU device.
+
+```c++
+    MXReturnValue forwardGPU(std::map<std::string, std::string> attrs,
+                             std::vector<MXTensor> inputs,
+                             std::vector<MXTensor> outputs,
+                             OpResource res) {
+        float* in_data = inputs[0].data<float>();
+        float* out_data = outputs[0].data<float>();
+        mx_stream_t cuda_stream = res.get_cuda_stream();
+        int64_t N = inputs[0].size();
+        int block = 256;
+        int grid = (N + (block - 1)) / block;
+        my_op_gpu_forward<<<grid,block,0,cuda_stream>>>(out_data, in_data, N);
+        return MX_SUCCESS;
+    }
+```
+
+Note that the `cuda_stream` object used for launching kernels is passed from 
MXNet backend via `OpResource` object. See below for details of `Operator 
Resource`.
+
+At this point all the attribute functions for each operator (`parseAttrs`, 
`inferShape`, etc.) run on the CPU, including the `forwardGPU` function. The 
only part that will actually run on the GPU is the CUDA kernel function that is 
launched from `forwardGPU` function (ie. my_op_gpu_forward).
+
+```c++
+    __global__ void my_op_gpu_forward(float *out, float *in, int64_t N) {
+        int tid = blockIdx.x * blockDim.x + threadIdx.x;
 
 Review comment:
   It is correct and readable to replace `int tid = blockIdx.x * blockDim.x + 
threadIdx.x;`  with the `CUDA_KERNEL_LOOP` macro.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to