[Performance][CUDA] Faster CSRToCOO

[mfbalin]

Description

CUB recently implemented a DeviceCopy::Batched algorithm, which can be used to implement the DeviceRunLength::Decode algorithm, which in turn can be used to implement the CSRToCOO for the GPU in a more efficient way for the 64-bit code path. The current implementation suffers from O(E log N) complexity, which this PR reduces to O(E).

My motivation was that CSRToCOO is called inside the labor sampler algorithm, so improving the performance will also improve Labor Sampler and make its complexity truly linear in the number of edges.

My experimental evaluations show that the Labor Sampler runtime consistently improves with this change.

Checklist

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
• I've leverage the tools to beautify the python and c++ code.
• The PR is complete and small, read the Google eng practice (CL equals to PR) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
• All changes have test coverage
• Code is well-documented
• To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

[dgl-bot]

Not authorized to trigger CI. Please ask core developer to help trigger via issuing comment:

• @dgl-bot

[mfbalin]

@BarclayII @yaox12, could you take a look? This PR requires using the latest thrust and cub versions.

[yaox12]

@mfbalin

1. Has this feature been included in any CUB release version?
2. CUB 2.x introduced the libcu++ dependency. Do you know if it's compatible with CUDA 10.2, since DGL still supports it?

[yaox12]

@dgl-bot

[mfbalin]

@yaox12

1. Not yet, I just recently contributed DeviceCopy to CUB in this PR, I guess it will be included in the next release. I don't know the timeline of the next release though.
2. It looks like thrust includes cub and libcu++ as submodule dependencies however I don't know if the libcu++ version that comes with thrust is compatible with CUDA 10.2. What I was able to find is that libcu++ was first released for CUDA 10.2. I don't know if they dropped CUDA 10.2 support in later releases.

[dgl-bot]

Commit ID: 92b0b17

Build ID: 3

Status: ❌ CI test failed in Stage [GPU Build].

Report path: link

Full logs path: link

[yaox12]

@dgl-bot

[dgl-bot]

Commit ID: 855f011

Build ID: 5

Status: ❌ CI test failed in Stage [GPU Build].

Report path: link

Full logs path: link

[mfbalin]

It looks like the libcudacxx headers can not be found even though I added it to the cuda include directories, what else needs to be done so that those headers can be found?

Error msg: /root/jenkins/workspace/dgl_PR-5648@2/third_party/thrust/thrust/detail/type_traits.h:27:10: fatal error: cuda/std/type_traits: No such file or directory

Edit: added include to the include path for libcudacxx, I hope it will work now.

[yaox12]

@dgl-bot

[dgl-bot]

Commit ID: 078be43

Build ID: 8

Status: ❌ CI test failed in Stage [Authentication].

Report path: link

Full logs path: link

[yaox12]

@dgl-bot

[dgl-bot]

Commit ID: 078be43

Build ID: 9

Status: ❌ CI test failed in Stage [C++ GPU].

Report path: link

Full logs path: link

[mfbalin]

I am not sure if this new code path or the existing int32_t specialization is faster. One would need to do a benchmark to see if the cusparse csr2coo or cub is better. However, for the int64_t specialization, CUB should outperform the existing implementation in all cases.
Cusparse code path:

dgl/src/array/cuda/csr2coo.cu

Line 44 in 078be43

CUSPARSE_CALL(cusparseXcsr2coo(

[dgl-bot]

Commit ID: 02bdcb5

Build ID: 10

Status: ❌ CI test failed in Stage [C++ GPU].

Report path: link

Full logs path: link

[frozenbugs]

@mfbalin can you help us take a look at the failing test?

[dgl-bot]

Commit ID: befcf8a

Build ID: 11

Status: ❌ CI test failed in Stage [C++ GPU].

Report path: link

Full logs path: link

[mfbalin]

@frozenbugs Partition had a minor bug due to the 1 off specification of histogramEven call parameters. Tests should all pass now.

[dgl-bot]

Commit ID: 8915f8c868838cf41b01cae6097ba88c51302698

Build ID: 12

Status: ⚪️ CI test cancelled due to overrun.

Report path: link

Full logs path: link

[mfbalin]

Actually, we might want to fix this bug in a separate PR so that if there is any problem with the new thrust version (on the off change), we can take this PR back but the bug fix is still stays. see: #6001

[frozenbugs]

Actually, we might want to fix this bug in a separate PR so that if there is any problem with the new thrust version (on the off change), we can take this PR back but the bug fix is still stays.

Solid judgement :)

[frozenbugs]

Partition had a minor bug due to the 1 off specification of histogramEven call parameters. Tests should all pass now.

Thanks for fixing, any reason this was not exposed without your optimization?

[mfbalin]

The histogram has bins, let's say we are computing a histogram for bins with values 0 to 6 (7 bins in total). Without the fix, the bin boundaries were [0, 8) (8 exclusive) vs [0, 7). Documentation says that CUB evenly distributed the values in the range to bins. But 8 doesn't divide 7. So without the fix, with the new update, 0 and 1 was in the first bin. Without the update 6 and 7 were in the last bin (Which didn't expose the bug). I am speculating that this caused a rounding error and a change in the implementation in CUB might have exposed this bug.

[dgl-bot]

Commit ID: 406ec17

Build ID: 13

Status: ✅ CI test succeeded.

Report path: link

Full logs path: link