Skip unnecessary assertions and enable non-blocking data transfers

[nistath]

Summary

This backwards-compatible contribution enables higher performance by skipping assertions when they are unnecessary and allowing for truly non-blocking data transfers to the GPU.

Problem

When profiling the performance of a torch-geometric code using nvprof, I found out that setting non_blocking=True when doing data transfers still blocks the current CUDA stream. The underlying reason is that when the to method is called to initiate a data transfer, we construct a new SparseStorage object which repeats some assertions that have already been evaluated on this data. The particular assertions run when the row or col tensors are present, checking that their content conforms with sparse_sizes. The problem with such assertions is that they require a blocking round-trip of communication between the CPU and GPU, which defeats the purpose of non_blocking=True.

Solution

The solution to the assertion problem is to use a trust_data construction argument as an invariant that tells us whether we need to run these assertions. Data transfers or dtype changes do not need these assertions, because the matrix structure was already checked upon construction. Certain other operations such as eye do not need them either, because they are correct by construction.

In addition, I refactored the dtype and data transfer API to align with torch.Tensor and avoid the construction of dummy tensor objects, removing wasted time in the PyTorch allocator.

Code changes

• Uses the trust_data invariant to skip blocking assertions, when unnecessary, during construction of SparseStorage objects.
• Refactors the dtype and device transfer APIs to align with torch.Tensor while maintaining backward compatibility.
• No longer constructs dummy tensors when changing dtype or device.
• Adds WITH_SYMBOLS option to allow for linking without stripping symbols in setup.py

[codecov-commenter]

Codecov Report

Merging #195 (1ca4a90) into master (88c6ceb) will not change coverage.
The diff coverage is 0.00%.

@@           Coverage Diff           @@
##           master     #195   +/-   ##
=======================================
  Coverage   72.32%   72.32%           
=======================================
  Files          28       28           
  Lines        1120     1120           
=======================================
  Hits          810      810           
  Misses        310      310           

Impacted Files	Coverage Δ
torch_sparse/storage.py	20.58% <ø> (ø)
torch_sparse/tensor.py	49.15% <0.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 88c6ceb...1ca4a90. Read the comment docs.

[rusty1s]

This looks absolutely perfect. Thank you for catching this :)

The type_as and device_as functionalities are indeed a relict from supporting an older PyTorch version in which torch.dtype and torch.device type hints were not allowed by the JIT compiler. Glad that this is now supported :)