[TOPI][OP] Use Thrust sort for argsort and topk

[kazum]

The current GPU sort implementation (odd-even transposition sort) is too slow when the number of elements is large. This PR introduces Thrust implementation of sort which is much faster.

Note that this change requires CMake 3.8 3.13 or later since we have to use nvcc to compile a thrust code.

• benchmark script

  import tvm
  from tvm import relay
  from tvm.contrib import graph_runtime
  import numpy as np
  
  target = 'cuda'
  ctx = tvm.gpu(0)
  n = 100000
  
  x = relay.var("x", shape=(n,))
  out = relay.topk(x)
  func = relay.Function([x], out[0])
  
  with relay.build_config(opt_level=3):
      graph, lib, params = relay.build(func, target)
  
  module = graph_runtime.create(graph, lib, ctx)
  
  print("Evaluate inference time cost...")
  ftimer = module.module.time_evaluator("run", ctx, number=1, repeat=3)
  prof_res = np.array(ftimer().results) * 1000
  print("Mean inference time (std dev): %.2f ms (%.2f ms)" %
        (np.mean(prof_res), np.std(prof_res)))

• result (NVIDIA P100, without thrust)

  Evaluate inference time cost...
  Mean inference time (std dev): 2058.89 ms (0.07 ms)
  

• result (NVIDIA P100, with thrust)

  Evaluate inference time cost...
  Mean inference time (std dev): 1.11 ms (0.03 ms)
  

@icemelon9 @vinx13 @masahi could you help to review?

[Laurawly]

Is this cuda 10.1 compatible? I tried to build it on my machine with cuda 10.1 and it gives me the following errors:

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(42): error: namespace "thrust" has no member "device_ptr"

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(42): error: type name is not allowed

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(43): error: namespace "thrust" has no member "device_ptr"

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(43): error: type name is not allowed

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(44): error: namespace "thrust" has no member "device_ptr"

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(44): error: type name is not allowed

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(52): error: identifier "data_ptr" is undefined

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(52): error: identifier "values_ptr" is undefined

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(55): error: identifier "indices_ptr" is undefined

/home/ubuntu/workplace/tvm-1/src/runtime/contrib/thrust/thrust.cu(42): error: identifier "data_ptr" is undefined

[kazum]

@Laurawly Thanks, I've update the code and confirmed that it can be compiled on Ubuntu 18.04 and CUDA 10.1.

[masahi]

Looks great!

[icemelon]

lgtm

[Laurawly]

Could you also add USE_THRUST option in config.cmake file?

[Laurawly]

Thanks @kazum @masahi @icemelon9