Adding sparse support to MXTensor for custom operators

[guanxinq]

Description

Add support for sparse custom operators. It will support row sparse and CSR formats.
This is a continuation of custom operators project, initial CPU support is implemented here: #15921 and GPU support is implemented here: #17270 .

• Added the MXSparse structure. The data_ptr of MXTensor points to an object of this structure for sparse tensors and point to data for dense tensors.
• Added data storage type enum for MXTensor. It supports dense, row sparse and CSR formats.
• Added alloc_sparse() function within OpResource which fixed the sparse output size issue at run time.
• Added inferSType() function which enables customized storage type inference implementation.
• Created new example libraries in "example/extensions/lib_custom_op" for CSR and row sparse tensor transpose operators.

Design

The function alloc_sparse() in lower level call function CheckAndAlloc(). To call this member function of NDArray, we added lambda functions just as what we did for alloc_cpu().

auto sparse_alloc = [&](int index, int indices_len, int idxptr_len,
                      void** data, int64_t** indices, int64_t** indptr) {
  // Row Sparse
  if(idxptr_len == 0) {
    outputs[index].CheckAndAlloc({mshadow::Shape1(indices_len)});
    *data = outputs[index].data().dptr_;
    *indices = (int64_t*)outputs[index].aux_data(rowsparse::kIdx).dptr_;
  }
  // CSR
  else {
    outputs[index].CheckAndAlloc({mshadow::Shape1(idxptr_len), mshadow::Shape1(indices_len)});
    *data = outputs[index].data().dptr_;
    *indices = (int64_t*)outputs[index].aux_data(csr::kIdx).dptr_;
    *indptr = (int64_t*)outputs[index].aux_data(csr::kIndPtr).dptr_;
  }
};

typedef decltype(sparse_alloc) alloc_type_sparse;
auto sparse_malloc = [](void* _sparse_alloc, int index, int indices_len, int idxptr_len,
                       void** data, int64_t** indices, int64_t** indptr) {
  alloc_type_sparse* sparsealloc = static_cast<alloc_type_sparse*>(_sparse_alloc);
  (*sparsealloc)(index, indices_len, idxptr_len, data, indices, indptr);
};

The lambda function could be called by alloc_sparse() within OpResource.

void alloc_sparse(MXSparse* sparse, int index, int indices_len, int indptr_len = 0) {
  sparse_malloc(sparse_alloc, index, indices_len, indptr_len,
               &(sparse->data), &(sparse->indices), &(sparse->indptr));
}

In the customized implementation, users are able to set output tensor size by

MXReturnValue forward(std::map<std::string, std::string> attrs,
                      std::vector<MXTensor> inputs,
                      std::vector<MXTensor> outputs,
                      OpResource res) {
  // implementation of operators.
  MXSparse*  ptr = outputs[index].data<MXSparse>();
  res.alloc_sparse(ptr, index , indices_len, indptr_len);
  // implementation of operators reply on output size.
  return MX_SUCCESS;
}

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• 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

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[eric-haibin-lin]

Look forward to the complete API with example and documentation :)

[mseth10]

Thanks for the contribution @guanxinq . Left a few comments.

[eric-haibin-lin]

can u resolve conflicts ?

[samskalicky]

can you update MX_LIBRARY_VERSION to 5?

[rondogency]

LGTM! Thanks for the contribution!

[guanxinq]

can you update MX_LIBRARY_VERSION to 5?

Updated.

[rondogency]

@wkcn @eric-haibin-lin this PR is ready to merge!

[wkcn]

LGTM. Thank you for the contribution!

[mseth10]

LGTM