diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 26af9002597..b776e4bddc4 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -1,8 +1,10 @@ #include "argsort.cuh" #ifdef GGML_CUDA_USE_CUB +# ifndef GGML_USE_HIP // CUB comes from hipCUB via common.cuh on HIP # include -# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) +# endif +# if !defined(GGML_USE_HIP) && (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) # define STRIDED_ITERATOR_AVAILABLE # include # endif @@ -83,12 +85,19 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, is_capturing = (capture_status != cudaStreamCaptureStatusNone); #endif // USE_CUDA_GRAPH +#ifdef GGML_USE_HIP + // hipCUB does not provide cub::DeviceSegmentedSort - use the radix segmented sort + const bool use_segmented_radix = true; +#else + const bool use_segmented_radix = is_capturing; +#endif + if (order == GGML_SORT_ORDER_ASC) { if (nrows == 1) { CUDA_CHECK(DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); - } else if (is_capturing) { + } else if (use_segmented_radix) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs( nullptr, temp_storage_bytes, temp_keys, temp_keys, // keys (in-place) temp_indices, dst, // values (indices) @@ -107,7 +116,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, temp_keys, // keys (in-place) temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); - } else if (is_capturing) { + } else if (use_segmented_radix) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( nullptr, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); @@ -127,7 +136,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, temp_keys, // keys (in-place) temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); - } else if (is_capturing) { + } else if (use_segmented_radix) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); @@ -142,7 +151,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, temp_keys, // keys (in-place) temp_indices, dst, // values (indices) ncols, 0, sizeof(float) * 8, stream)); - } else if (is_capturing) { + } else if (use_segmented_radix) { CUDA_CHECK(DeviceSegmentedRadixSort::SortPairsDescending( d_temp_storage, temp_storage_bytes, temp_keys, temp_keys, temp_indices, dst, ncols * nrows, nrows, offset_iterator, offset_iterator + 1, 0, sizeof(float) * 8, stream)); diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 14dd1098c97..0d432470f46 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -107,9 +107,22 @@ #define GGML_CUDA_CC_IS_QY2(cc) (cc >= GGML_CUDA_CC_QY2 && cc < GGML_CUDA_CC_PH1) #define GGML_CUDA_CC_IS_PH1(cc) (cc >= GGML_CUDA_CC_PH1) -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +#if (!defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070) || defined(GGML_USE_HIP) # 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) || defined(GGML_USE_HIP) + +#ifdef GGML_USE_HIP +// On HIP builds CUB comes from hipCUB: the primitives live in namespace hipcub and +// there is no to include, so alias the namespace and map the few CUDA +// stream-capture names the GGML_CUDA_USE_CUB code paths use onto the HIP ones. +# include +namespace cub = hipcub; +using cudaStreamCaptureStatus = hipStreamCaptureStatus; +constexpr auto cudaStreamCaptureStatusNone = hipStreamCaptureStatusNone; +inline cudaError_t cudaStreamIsCapturing(cudaStream_t stream, cudaStreamCaptureStatus * status) { + return hipStreamIsCapturing(stream, status); +} +#endif // GGML_USE_HIP // PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8. // However, this has been bugged in CTK < 12.3 for MSVC builds, see diff --git a/ggml/src/ggml-cuda/cumsum.cu b/ggml/src/ggml-cuda/cumsum.cu index def9c32955f..327f5eeb470 100644 --- a/ggml/src/ggml-cuda/cumsum.cu +++ b/ggml/src/ggml-cuda/cumsum.cu @@ -5,7 +5,9 @@ #include "ggml.h" #ifdef GGML_CUDA_USE_CUB +# ifndef GGML_USE_HIP // CUB comes from hipCUB via common.cuh on HIP # include +# endif #endif // GGML_CUDA_USE_CUB template diff --git a/ggml/src/ggml-cuda/mean.cu b/ggml/src/ggml-cuda/mean.cu index a8f6046e46d..87afc1ab3dd 100644 --- a/ggml/src/ggml-cuda/mean.cu +++ b/ggml/src/ggml-cuda/mean.cu @@ -2,7 +2,9 @@ #include "reduce_rows.cuh" #ifdef GGML_CUDA_USE_CUB +# ifndef GGML_USE_HIP // CUB comes from hipCUB via common.cuh on HIP #include +# endif using namespace cub; #endif // GGML_CUDA_USE_CUB diff --git a/ggml/src/ggml-cuda/sum.cu b/ggml/src/ggml-cuda/sum.cu index c56257b4406..69a9ac554c0 100644 --- a/ggml/src/ggml-cuda/sum.cu +++ b/ggml/src/ggml-cuda/sum.cu @@ -2,7 +2,9 @@ #include "sumrows.cuh" #ifdef GGML_CUDA_USE_CUB +# ifndef GGML_USE_HIP // CUB comes from hipCUB via common.cuh on HIP #include +# endif using namespace cub; #endif // GGML_CUDA_USE_CUB diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd29333..b7ca37e00b2 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -2,8 +2,10 @@ #include "top-k.cuh" #ifdef GGML_CUDA_USE_CUB +# ifndef GGML_USE_HIP // CUB comes from hipCUB via common.cuh on HIP # include -# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) +# endif +# if !defined(GGML_USE_HIP) && (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) # define CUB_TOP_K_AVAILABLE # include using namespace cub;