I made some modifications to make sparse accumulation into dense gradients on GPU a lot faster (~10x) by avoiding coalescing and calling indexAdd. I also sped up coalesce by using something closer to what LookupTable does. I still need to go back and unify everything under indexAdd, but that will be trickier.

Here are my previous notes from #1147 :

I dug into speeding up sparse for nn.Embedding based on my benchmark on GPU (https://gist.github.com/adamlerer/865c1a09000c7dc8208e1456255209c2). I think my findings apply more broadly though.

• Most importantly, the contiguous (or reorder) operation, which merges rows with the same index, is inevitably slow on GPU because it demands a host synchronization. The reason is that the number of unique rows (calculated on GPU) becomes the new nnz, i.e. indices.size[0] / values.size[0] (which is a CPU field). So, it ends up being much faster to compute operations like spcadd (the main op for Embedding update) directly on the non-'contiguous' tensor, using indexAdd or LookupTable_accGradParameters.
  The potential problem with never compacting sparse tensors is that they can grow to unbounded size if you repeatedly call cadd without ever calling contiguous. Maybe make it the user's responsibility to call contiguous when necessary? This can be problematic when the cadd is buried in a library somewhere e.g. in autograd backward. I don't have a good answer here... We could have heuristics for when to compact a tensor. Something like this would be kinda clever (too clever?):

void THCSTensor_(cadd)(r, a, b) {
...
r->minUnique = max(a->minUnique + b->minUnique);
...
if (r->nnz / r->minUnique > COMPACTION_THRESHOLD) {
  THCSTensor_(contiguous)(r);
  r->minUnique = r->nnz;
}

• LookupTable_accGradParameters is just computing indexAdd, yet it's implemented completely differently. The different implementation was due to a desire a couple years ago for LookupTable to be deterministic therefore not use atomicAdd. But then why does indexAdd still use atomicAdd? Either we care about determinism or we don't... my guess is that we don't any more since even cudnn isn't deterministic.
  I benchmarked then and the non-deterministic one is several times faster, regardless of how many index collision there are. We should pick if we care about determinism, use that to pick an indexAdd kernel in THC, and then delegate LookupTable, spcadd, etc. all to that.

• Autograd backwards on nn.Embedding spends about 1ms in Python autograd code, which after speeding things up takes >90% of the time for batch sizes below 1e4 (it was taking ~50% of the time before my kernel changes). So batch=1e3 isn't interesting, we should look at batch=1e4.

Btw I'm still unsure about whether we should make coalesce an in-place op or not. It's an important decision and has some trade-offs

Pro in-place:

• A coalesced tensor is always faster to operate on than an un-coalesced tensor.
• Suppose you have an uncoalesced sparse tensor S and you repeatedly call coalescing ops on it:

S = S + S # now uncoalesced
for i in range(N):
  T = S * S # coalesces each time

This is going to be extremely inefficient.

Pro out-of-place:

• const operands should not be mutable; therefore, if a const operand (e.g. input to cadd) can get coalesced in place, it must be invisible to the Python user. That means that t.indexes(), t.values(), and t.__str__ should all call coalesce. That hides opportunities for optimizations by munging indices and values directly from Python. Then if there are C bugs, you can get nasty behavior like a Heisenbug where when you call print on the tensor, the bug goes away.

Always coalesce

There's a third option which is to always keep tensors in coalesced form. This is probably hard to make efficient. But it's possible that if we move nnz to be stored on the device, we can avoid host syncs and make it reasonably fast.

Heisenbugs suck. I'll switch to out-of-place coalesce. It will be up to the user to coalesce tensors explicitly if they're going to be operating on them multiple times.

I'm not too qualified to review this diff, but some thoughts did occur to me:

1. Does it make sense to still offer an in-place coalesce, even though operators won't call this directly? The main use case would be a user is performing a bunch of sparse ops in place, and then wants to finally coalesce the tensor before passing it off.
2. It would be great if we had some docs, because there is quite a bit of diversity in how duplicate indexes are handled. SciPy, for example, sums duplicate indexes together (like us, if I understand these diffs correctly), but TensorFlow drops duplicates (Sparse tensor construction given repeated indices tensorflow/tensorflow#371). I know the API is in flux, but it would still be great to have a paragraph or two stating what the basic assumptions in the implementation are.

Thanks @ezyang

1. I don't see what value an in-place op provides... the in-place coalesce is no more efficient than out-of place, because a sparse tensor is just a tiny wrapper for {THLongTensor indices, THTensor values}, so all you're creating is a new wrapper. The difference for the in-place one is that if anyone else was holding a reference to the tensor, it would be changed as well. Is there a reason we want that?

2. Yeah, we definitely need docs. I will at least put a comment in coalesce but we do need to write the full docs for this library soon. I won't be able to get to it until at least late-May, so if you'd like to pick this up, it would be a huge help.

Re: duplicates, the current implementation must sum duplicates because we take advantage of this to do some ops efficiently. cadd just "adds" two sparse tensors by concatenating their indices and values. This way you can add together a bunch of tensors and then accumulate them into a dense tensor, without ever doing a (slow) coalesce.
P.S. We may be able to make coalesce much faster and then wouldn't need this, but that would require some substantial and tricky rewriting (esp. holding nnz on the GPU to avoid a host sync after coalesce).

Note to self: in the torch tests, self.assertEqual calls x.coalesce() on sparse tensors before comparison. But if coalesce is broken (e.g. to always return an empty tensor) these tests will pass.
If it's fast enough, I want to write a simple Python coalesce, and then replace x.coalesce() with checkedCoalesce(x), which runs the Python and torch versions and asserts they're equal.

This is ready for review and merge. @apaszke you'll note that you were right in #1147 to suggest making coalesce out-of-place.

Re: Embedding perf, here's the results of the benchmark script at https://gist.github.com/adamlerer/4159ddb4bd19e87de1b4bb3215a1db6f

$ python benchmark_sparse.py  1
cuda=True
Old done in 0.493923 s
New done in 0.617863 s
New done in 0.382167 s (no optimizer)

cf. 3s for new before this PR.

👍