[ROCM] Add Thrust support

[masahi]

This adds rocThrust support for our rocm backend, which can be used for sorting and scan.
https://github.com/ROCmSoftwarePlatform/rocThrust

Use "rocm -libs=thrust" to enable thrust based sort and scan.

contrib/rocthrust/thrust.cc code is identical to contrib/thrust/thrust.cu used by CUDA backend. The same thrust code works on both CUDA and AMD.

@mbrookhart @csullivan @antinucleon

[masahi]

@tkonolige @mbrookhart Thanks for review, it should be ready to go. This can be used for benchmarking the new TIR sort @mbrookhart is cooking.

I'll follow up with another PR to enforce -libs=thrust option for other thrust usages.

[mbrookhart]

Devil's advocate:

If I reformat the thrust.cu file with clang-format-10, like you said, they are identical except for the datatype errors you added and #include <thrust/sequence.h> in the amd version.

Is there any reason to keep two copies around?

(otherwise LGTM, thank you!)

[csullivan]

Thanks @masahi for adding rocthrust support!

[masahi]

@mbrookhart Yes, I wanted to reuse thrust.cu as is, by adding the following to ROCM.cmake:

file(GLOB CONTRIB_THRUST_SRC src/runtime/contrib/thrust/*.cu)

But unlike for CUDA, cmake doesn't recognize that this .cu file needs to be compiled by the same CXX compiler (hipcc) together with other .cc files, so it will not be compiled (it is ignored during compilation).

I think CMake has a special support for NVCC and .cu file, to correctly compile .cu with NVCC and others with gcc/clang. On rocm, everything needs to be compiled by the same compiler, hipcc (a wrapper around clang).

[masahi]

Maybe pytorch people have a solution to the duplication problem. But my cmake-fu is too low and I have no interest in learning about cmake :)

[mbrookhart]

@tkonolige @csullivan Any other change requests or are you happy?

[tkonolige]

@masahi You can use set_source_files_properties(src/runtime/contrib/thrust/thrust.cu PROPERTIES LANGUAGE CXX).

[masahi]

@masahi You can use set_source_files_properties(src/runtime/contrib/thrust/thrust.cu PROPERTIES LANGUAGE CXX).

Thanks I'll try this ASAP

[masahi]

Great set_source_files_properties worked! I removed rocthrust/thrust.cc, the same thrust.cu is used by CUDA and rocm now cc @mbrookhart

[masahi]

Thanks @mbrookhart @tkonolige @csullivan