GPUCommonAlgorithm: Use CUB for soring on device instead of Thrust

Author: davidrohr

GPU/Common/GPUCommonAlgorithmThrust.h

@@ -23,11 +23,16 @@
 #pragma GCC diagnostic pop
 
 #include "GPUCommonDef.h"
+#include "GPUCommonHelpers.h"
 
-#ifdef __CUDACC__
+#ifndef __HIPCC__ // CUDA
 #define GPUCA_THRUST_NAMESPACE thrust::cuda
-#else
+#define GPUCA_CUB_NAMESPACE cub
+#include <cub/cub.cuh>
+#else // HIP
 #define GPUCA_THRUST_NAMESPACE thrust::hip
+#define GPUCA_CUB_NAMESPACE hipcub
+#include <hipcub/hipcub.hpp>
 #endif
 
 namespace o2::gpu
@@ -89,11 +94,20 @@ template <class T, class S>
 GPUhi() void GPUCommonAlgorithm::sortOnDevice(auto* rec, int32_t stream, T* begin, size_t N, const S& comp)
 {
   thrust::device_ptr<T> p(begin);
+#if 0 // Use Thrust
   auto alloc = rec->getThrustVolatileDeviceAllocator();
   thrust::sort(GPUCA_THRUST_NAMESPACE::par(alloc).on(rec->mInternals->Streams[stream]), p, p + N, comp);
+#else // Use CUB
+  size_t tempSize = 0;
+  void* tempMem = nullptr;
+  GPUChkErrS(GPUCA_CUB_NAMESPACE::DeviceMergeSort::SortKeys(tempMem, tempSize, begin, N, comp, rec->mInternals->Streams[stream]));
+  tempMem = rec->AllocateVolatileDeviceMemory(tempSize);
+  GPUChkErrS(GPUCA_CUB_NAMESPACE::DeviceMergeSort::SortKeys(tempMem, tempSize, begin, N, comp, rec->mInternals->Streams[stream]));
+#endif
 }
 } // namespace o2::gpu
 
 #undef GPUCA_THRUST_NAMESPACE
+#undef GPUCA_CUB_NAMESPACE
 
 #endif

GPU/GPUTracking/Base/GPUGeneralKernels.h

@@ -27,9 +27,9 @@
 #endif
 
 #if defined(__HIPCC__)
-#define GPUCA_CUB hipcub
+#define GPUCA_CUB_NAMESPACE hipcub
 #else
-#define GPUCA_CUB cub
+#define GPUCA_CUB_NAMESPACE cub
 #endif
 
 namespace o2::gpu
@@ -54,7 +54,7 @@ class GPUKernelTemplate
   struct GPUSharedMemoryWarpScan64 {
     // Provides the shared memory resources for warp wide CUB collectives
 #if (defined(__CUDACC__) || defined(__HIPCC__)) && defined(GPUCA_GPUCODE) && !defined(GPUCA_GPUCODE_HOSTONLY)
-    typedef GPUCA_CUB::WarpScan<T> WarpScan;
+    typedef GPUCA_CUB_NAMESPACE::WarpScan<T> WarpScan;
     union {
       typename WarpScan::TempStorage cubWarpTmpMem;
     };
@@ -65,9 +65,9 @@ class GPUKernelTemplate
   struct GPUSharedMemoryScan64 {
     // Provides the shared memory resources for CUB collectives
 #if (defined(__CUDACC__) || defined(__HIPCC__)) && defined(GPUCA_GPUCODE) && !defined(GPUCA_GPUCODE_HOSTONLY)
-    typedef GPUCA_CUB::BlockScan<T, I> BlockScan;
-    typedef GPUCA_CUB::BlockReduce<T, I> BlockReduce;
-    typedef GPUCA_CUB::WarpScan<T> WarpScan;
+    typedef GPUCA_CUB_NAMESPACE::BlockScan<T, I> BlockScan;
+    typedef GPUCA_CUB_NAMESPACE::BlockReduce<T, I> BlockReduce;
+    typedef GPUCA_CUB_NAMESPACE::WarpScan<T> WarpScan;
     union {
       typename BlockScan::TempStorage cubTmpMem;
       typename BlockReduce::TempStorage cubReduceTmpMem;
@@ -110,6 +110,6 @@ class GPUitoa : public GPUKernelTemplate
 
 } // namespace o2::gpu
 
-#undef GPUCA_CUB
+#undef GPUCA_CUB_NAMESPACE
 
 #endif