improve performance of unique with inverse indices

[ngimel]

Partial fix for #15804, only w/o dim.
For @jcjohnson benchmarking script I'm getting the following results on V100:
Before:

unning with N = 10000, M = 10000
cuda (no inverse): 0.98 ms
cpu (no inverse): 0.96 ms
cuda (with inverse): 1.07 ms
cpu (with inverse): 1.76 ms

Running with N = 10000, M = 100000
cuda (no inverse): 0.76 ms
cpu (no inverse): 1.53 ms
cuda (with inverse): 1.23 ms
cpu (with inverse): 3.02 ms

Running with N = 100000, M = 100000
cuda (no inverse): 1.28 ms
cpu (no inverse): 11.22 ms
cuda (with inverse): 69.76 ms
cpu (with inverse): 20.28 ms

Running with N = 100000, M = 1000000
cuda (no inverse): 0.78 ms
cpu (no inverse): 18.78 ms
cuda (with inverse): 133.45 ms
cpu (with inverse): 34.09 ms

Running with N = 500000, M = 500000
cuda (no inverse): 1.43 ms
cpu (no inverse): 61.13 ms
cuda (with inverse): 3315.18 ms
cpu (with inverse): 104.57 ms

Running with N = 500000, M = 5000000
cuda (no inverse): 0.86 ms
cpu (no inverse): 96.44 ms
cuda (with inverse): 5209.93 ms
cpu (with inverse): 176.10 ms

After

Running with N = 10000, M = 10000
cuda (no inverse): 1.04 ms
cpu (no inverse): 0.94 ms
cuda (with inverse): 0.64 ms
cpu (with inverse): 1.76 ms

Running with N = 10000, M = 100000
cuda (no inverse): 0.77 ms
cpu (no inverse): 1.55 ms
cuda (with inverse): 0.58 ms
cpu (with inverse): 2.79 ms

Running with N = 100000, M = 100000
cuda (no inverse): 1.30 ms
cpu (no inverse): 14.15 ms
cuda (with inverse): 1.63 ms
cpu (with inverse): 20.90 ms

Running with N = 100000, M = 1000000
cuda (no inverse): 0.82 ms
cpu (no inverse): 18.63 ms
cuda (with inverse): 0.61 ms
cpu (with inverse): 33.52 ms

Running with N = 500000, M = 500000
cuda (no inverse): 1.51 ms
cpu (no inverse): 59.81 ms
cuda (with inverse): 1.23 ms
cpu (with inverse): 110.69 ms

Running with N = 500000, M = 5000000
cuda (no inverse): 0.92 ms
cpu (no inverse): 104.26 ms
cuda (with inverse): 0.84 ms
cpu (with inverse): 187.12 ms

[soumith]

before I land this, could you please locally check correctness as well, against CPU.

The unit-tests only cover small_3d test cases.

[jcjohnson]

This is amazing -- thank you so much @ngimel for the quick fix!

[ngimel]

@soumith tested for Justin's benchmark's cases, all clean. Do you want some tests added to the test suite?

import torch
for N in [10000, 100000, 500000]:
   for M in [N, 10 * N]:
      print('Running with N = %d, M = %d' % (N, M))
      a = torch.randint(M, size=(N,))
      unique, indices = a.unique(return_inverse=True)
      a=a.cuda()
      unique_gpu, indices_gpu = a.unique(return_inverse=True)
      unique = unique.cuda()
      indices = indices.cuda()
      print((unique-unique_gpu).abs().max(), (indices - indices_gpu).abs().max())

[soumith]

@ngimel your testing sounds good. and there are smoke tests with small_3d matrices in test_torch / test_cuda which are hopefully sufficient. If you see anything worth adding as a unit-test, do send a follow-up PR, I'll land this now

[facebook-github-bot]

@soumith is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[gkioxari]

@ngimel @jcjohnson Is it possible to have the unique_index be returned as well, when the appropriate flag is turned on? This would be a similar functionality to np.unique()

[ngimel]

Yes, it's possible, my original commit was returning unique_index instead of inverse_index, so just restore thrust::unique_by_key call with appropriate handling of options? 883edac