Enable CUB

Author: hjc4869

.github/workflows/build.yml

@@ -558,7 +558,7 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev libcurl4-openssl-dev rocwmma-dev
+          sudo apt-get install -y build-essential git cmake rocblas-dev hipblas-dev hipcub-dev libcurl4-openssl-dev rocwmma-dev
 
       - name: ccache
         uses: ggml-org/[email protected]

ggml/src/ggml-cuda/argsort.cu

@@ -1,8 +1,16 @@
 #include "argsort.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
+#if defined(GGML_USE_HIP)
+
+
+#include <hipcub/hipcub.hpp>
+using namespace hipcub;
+#else
 #    include <cub/cub.cuh>
 using namespace cub;
+#endif // GGML_USE_HIP
+
 #endif  // GGML_CUDA_USE_CUB
 
 static __global__ void init_indices(int * indices, const int ncols, const int nrows) {

ggml/src/ggml-cuda/common.cuh

@@ -91,9 +91,9 @@
 #define GGML_CUDA_CC_IS_QY2(cc)      (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_NG)
 #define GGML_CUDA_CC_IS_NG(cc)       (cc >= GGML_CUDA_CC_NG)
 
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#if defined(GGML_USE_HIP) || (!defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070)
 #    define GGML_CUDA_USE_CUB
-#endif  // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#endif  // defined(GGML_USE_HIP) || (!defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070)
 
 #ifdef __CUDA_ARCH_LIST__
 constexpr bool ggml_cuda_has_arch_impl(int) {

ggml/src/ggml-cuda/mean.cu

@@ -2,8 +2,15 @@
 #include "reduce_rows.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
+
+#if defined(GGML_USE_HIP)
+#include <hipcub/hipcub.hpp>
+using namespace hipcub;
+#else
 #include <cub/cub.cuh>
 using namespace cub;
+#endif // GGML_USE_HIP
+
 #endif  // GGML_CUDA_USE_CUB
 
 template <typename T> __global__ void divide_by_count(T * result, size_t count) {

ggml/src/ggml-cuda/ssm-scan.cu

@@ -1,10 +1,17 @@
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#if defined(GGML_USE_HIP) || (!defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070)
 #define USE_CUB
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
+#endif // defined(GGML_USE_HIP) || (!defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070)
 
 #ifdef USE_CUB
+
+#if defined(GGML_USE_HIP)
+#include <hipcub/hipcub.hpp>
+using namespace hipcub;
+#else
 #include <cub/cub.cuh>
 using namespace cub;
+#endif // GGML_USE_HIP
+
 #endif // USE_CUB
 
 #include "ssm-scan.cuh"
@@ -48,8 +55,8 @@ __global__ void __launch_bounds__(splitD, 1)
     __shared__ float smemC[N];
 
 #ifdef USE_CUB
-    using BlockLoad = cub::BlockLoad<float, splitD, N, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
-    using BlockStore = cub::BlockStore<float, splitD, N, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockLoad = BlockLoad<float, splitD, N, BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockStore = BlockStore<float, splitD, N, BLOCK_STORE_WARP_TRANSPOSE>;
 
     union CubTempStorage {
         typename BlockLoad::TempStorage load_temp;

ggml/src/ggml-cuda/sum.cu

@@ -2,8 +2,16 @@
 #include "sumrows.cuh"
 
 #ifdef GGML_CUDA_USE_CUB
+
+#if defined(GGML_USE_HIP)
+#include <hipcub/hipcub.hpp>
+using namespace hipcub;
+#else
 #include <cub/cub.cuh>
 using namespace cub;
+
+#endif // GGML_USE_HIP
+
 #endif  // GGML_CUDA_USE_CUB
 
 #include <cstdint>
@@ -16,7 +24,6 @@ void sum_f32_cuda(ggml_cuda_pool & pool, const float * x, float * dst, const int
     DeviceReduce::Sum(tmp_alloc.ptr, tmp_size, x, dst, ne, stream);
 #else
     // Use (inefficient) sum_rows implementation as a fallback.
-    // For AMD there is rocPRIM which could be used as a drop-in replacement via hipcub but this would require C++11 -> C++14.
     sum_rows_f32_cuda(x, dst, ne, 1, stream);
     GGML_UNUSED(pool);
 #endif // GGML_CUDA_USE_CUB

ggml/src/ggml-cuda/vendors/hip.h

@@ -106,6 +106,9 @@
 #define cudaStreamPerThread hipStreamPerThread
 #define cudaStreamSynchronize hipStreamSynchronize
 #define cudaStreamWaitEvent(stream, event, flags) hipStreamWaitEvent(stream, event, flags)
+#define cudaStreamIsCapturing hipStreamIsCapturing
+#define cudaStreamCaptureStatus hipStreamCaptureStatus
+#define cudaStreamCaptureStatusNone hipStreamCaptureStatusNone
 #define cudaGraphExec_t hipGraphExec_t
 #define cudaGraphNode_t hipGraphNode_t
 #define cudaKernelNodeParams hipKernelNodeParams