[Performance] Improve COO to CSR, and sort columns of CSR only when necessary.

[nv-dlasalle]

Description

Avoid unnecessarily sort CSR columns, reduces the amount of work done when rows are sorted, and avoids synchronizing to double check if a COO is not sorted. It also parallelizes the conversion of an unsorted COO to CSR on the CPU, which can significantly reduce startup time (when the formats are generated for the global graph).

On examples/pytorch/rgcn/entity_classify_mp.py with the ogb-mag dataset, using the master branch with the PyTorch allocator, sorting on the GPU looks like:

And here it is with this PR:

With this PR, you can see, the check for if the matrix is sorted is removed, and the number of radix sort iterations is significantly reduced.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented
• To the my best knowledge, examples are either not affected by this change,
  or have been fixed to be compatible with this change

Changes

Changes UnitGraph to return unsorted CSRs for GetInCSR() and GetOutCSR()

Changes COOToCSR to not check if the matrix is sorted, if it claims to be unsorted. This eliminates a D2H copy, which forces synchronization between the CPU and GPU.

Changes COOSort() to respect the sort_columns flag, and merges the separate int32_t and int64_t version into a single one which uses CUB and takes advantage of the small number of rows relative to the width of IdType. That is, if we have 100,000 nodes in our block, we will only perform a radix sort of width 17, rather than the full 32 or 64 bits of the IdType.

Changes the CPU version of COOToCSR to generate the CSR in parallel, by doing two passes over the graph, one to count the number of edges per vertex, and the second to insert the edges in the correct positions.

This also expands the unit tests in tests/cpp/test_spmat_csr.cc to ensure unsorted CSRs are not problematic.

[jermainewang]

The PR is mostly in good shape. Some comments.

[jermainewang]

LGTM