eric-haibin-lin commented on a change in pull request #10081: [MXNET-82] [WIP] Sparse op tutorial for developers URL: https://github.com/apache/incubator-mxnet/pull/10081#discussion_r175205030
########## File path: docs/how_to/add_sparse_op_in_backend.md ########## @@ -0,0 +1,429 @@ +# A Guide to Implementing Sparse Operators in MXNet Backend + +## Prerequisites +- Basic knowledge of [how to implement a dense operator in MXNet backend](https://mxnet.incubator.apache.org/versions/master/how_to/add_op_in_backend.html) +- Basic knowledge of [CSRNDArray](http://mxnet.incubator.apache.org/tutorials/sparse/csr.html) and [RowSparseNDArray](http://mxnet.incubator.apache.org/tutorials/sparse/row_sparse.html) in MXNet + +## Introduction +In the [previous tutorial](https://mxnet.incubator.apache.org/versions/master/how_to/add_op_in_backend.html), +we went through the steps to implementing an operator using C++ in the MXNet backend. +In this tutorial, we will cover how sparse operators are implemented +in the backend. Specifically, we will practice adding CSRNDArray support to the forward function of the `quadratic` operator. + +## Implementation +### A Sparse Operator Example + +Let's consider the quadratic function `f(x) = ax^2+bx+c` when x is a CSRNDArray. +Notice that if the input x is sparse and c is 0.0, the output is also sparse. +If c is non-zero, the output is dense. In MXNet frontend, the operator works like this: + +```python +>>> x = mx.nd.array([[0,1],[2,0]).tostype('csr') +>>> x +<CSRNDArray 2x2 @cpu(0)> +>>> y = mx.nd.sparse.quadratic(x, a=1, b=2, c=0) +>>> y +<CSRNDArray 2x2 @cpu(0)> +>>> z = mx.nd.quadratic(x, a=1, b=2, c=3) +>>> z +[[ 3. 6.] + [ 11. 3.]] +<NDArray 2x2 @cpu(0)> +``` + +The statement `z = mx.nd.quadratic(x, a=1, b=2, c=3)` generates a warning message which says +the sparse input is converted to dense storage, and the dense operator is used to compute the dense output. +This is the "storage fallback" mechanism in MXNet, where a dense operator is automatically used for +inputs that a sparse operator doesn't have special kernels for. + +In this tutorial, we will implement the forward function of the sparse quadratic operator. +The storage type of the output depends on the inputs: +- quadratic('csr', a, b, 0.0) outputs 'csr' +- otherwise, outputs 'default' + +To implement this, we first register the storage type inference property of the operator, from which the operator +infers the output storage type based on operator arguments and inputs types. Then we implement the forward +function for the case where c is 0.0 and x is a CSRNDArray. + +Next, we are going to + +- Understand the FComputeEx and relevant NDArray interfaces in backend. +- Define storage type inference functions in quadratic_op-inl.h. +- Define the forward function in quadratic_op-inl.h. +- Register the sparse operator using nnvm in quadratic_op.cc and quadratic_op.cu for CPU and GPU computing, respectively. + +Now let's walk through the process step by step. + +### The FComputeEx and Relevant NDArray Interfaces in Backend + +Before we dive into the details of relevant interfaces, here are two differences between +dense and sparse operators: +- Dense operators only handle dense inputs and outputs. Sparse operators support various combinations of +storage types. +- Memories of inputs and outputs are pre-allocated based their shapes for dense operators. However, with sparse representations, memories for sparse inputs and outputs depend on the number of non-zero elements they have, +which is only known at runtime. + +With these differences in mind, let's review the `FCompute` interface introduced in the previous operator tutorial: +```cpp +void (const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector<TBlob>& inputs, + const std::vector<OpReqType>& req, + const std::vector<TBlob>& outputs); +``` +Notice the `FCompute` interface doesn't include data structures that could be used to query storage +types of inputs, nor manipulate auxiliary arrays like `indices` and `indptr`. +Therefore, instead of the `FCompute` interface, sparse operators are registered with the following `FComputeEx` interface: +```cpp +void (const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector<NDArray>& inputs, + const std::vector<OpReqType>& req, + const std::vector<NDArray>& outputs); +``` +where the vectors of TBlobs are replaced with vectors of NDArrays. Now, let's go through a few important methods in the NDArray class. + +In the python frontend, there are three types of NDArrays, namely `mx.nd.NDArray`, `mx.nd.sparse.RowSparseNDArray` and `mx.nd.sparse.CSRNDArray`. In the C++ backend, however, all of them are represented by the `mxnet::NDArray` class. +The `storage_type()` method indicates the storage type of the NDArray: +```cpp +enum NDArrayStorageType { + kUndefinedStorage = -1, // undefined storage + kDefaultStorage, // dense + kRowSparseStorage, // row sparse + kCSRStorage, // csr +}; + +// return the type of storage format +inline NDArrayStorageType storage_type() const; +``` + +On the other hand, from python one could inspect the auxiliary array of a sparse ndarray via +`RowSparseNDArray.indices`, `CSRNDArray.indices` and `CSRNDArray.indptr`, and the actual data array +via `RowSparseNDArray.data` and `CSRNDArray.data`. + +In the backend, auxliary arrays such as `indices` and `indptr` are retrieved by +the `aux_data` method, while the actual data array is retrived by the +`data` method. + +```cpp +namespace csr { +enum CSRAuxType {kIndPtr, kIdx}; +} + +namespace rowsparse { +enum RowSparseAuxType {kIdx}; +} + +// return the i-th aux data TBlob +inline TBlob aux_data(size_t i) const; + +// return the data TBlob +inline const TBlob& data() const; +``` + +Finally, the `CheckAndAlloc` method comes in handy when memory allocations for +the data and auxiliary arrays are needed for sparse NDArrays at run time. +```cpp +// allocate memory for non-default storage ndarrays based on auxliary array shapes +inline void CheckAndAlloc(const std::vector<TShape> &aux_shapes) +``` + +### Storage Type Inference +Storage type inference is the process of deducing storage types of `NDArray`s Review comment: Ok ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on 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
