diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 53eccdd65..31f5aeeac 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5273,6 +5273,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_SUM: return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: +#if defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB) + return true; +#else + return op->src[0]->ne[0] <= 1024; +#endif // defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB) case GGML_OP_ARGSORT: #ifndef GGML_CUDA_USE_CUB return op->src[0]->ne[0] <= 1024; diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index 9681cd293..c7a0c8317 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -48,6 +48,168 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE +#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) + +static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { + const uint32_t bits = __float_as_uint(value); + const uint32_t mask = (uint32_t) (-(int32_t) (bits >> 31)) | 0x80000000U; + return bits ^ mask; +} + +struct top_k_radix_state { + uint32_t prefix; + uint32_t prefix_mask; + int rank; + int greater_count; + int equal_count; +}; + +static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < nrows) { + states[row] = {0, 0, k, 0, 0}; + } +} + +template +static __global__ void top_k_radix_histogram( + const float * __restrict__ src, + const top_k_radix_state * __restrict__ states, + int * __restrict__ block_histograms, + int ncols, + int blocks_per_row, + int shift) { + constexpr int NBINS = 1 << RADIX_BITS; + + const int row = blockIdx.x / blocks_per_row; + const int row_block = blockIdx.x % blocks_per_row; + const int tid = threadIdx.x; + const float * row_src = src + (size_t) row * ncols; + __shared__ int histogram[NBINS]; + + histogram[tid] = 0; + __syncthreads(); + + const top_k_radix_state state = states[row]; + for (int col = row_block * BLOCK_SIZE + tid; + col < ncols; + col += blocks_per_row * BLOCK_SIZE) { + const uint32_t key = top_k_float_to_ordered(row_src[col]); + if ((key & state.prefix_mask) == state.prefix) { + atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1); + } + } + __syncthreads(); + + const size_t histogram_offset = + ((size_t) row * blocks_per_row + row_block) * NBINS; + block_histograms[histogram_offset + tid] = histogram[tid]; +} + +template +static __global__ void top_k_radix_select( + const int * __restrict__ block_histograms, + top_k_radix_state * __restrict__ states, + int blocks_per_row, + int shift) { + constexpr int NBINS = 1 << RADIX_BITS; + + const int row = blockIdx.x; + const int tid = threadIdx.x; + __shared__ int histogram[NBINS]; + + int count = 0; + for (int row_block = 0; row_block < blocks_per_row; ++row_block) { + const size_t offset = ((size_t) row * blocks_per_row + row_block) * NBINS; + count += block_histograms[offset + tid]; + } + histogram[tid] = count; + __syncthreads(); + + if (tid == 0) { + top_k_radix_state state = states[row]; + int bin = NBINS - 1; + while (bin > 0 && histogram[bin] < state.rank) { + state.rank -= histogram[bin--]; + } + state.prefix |= (uint32_t) bin << shift; + state.prefix_mask |= (uint32_t) (NBINS - 1) << shift; + states[row] = state; + } +} + +static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < nrows) { + states[row].greater_count = 0; + states[row].equal_count = 0; + } +} + +template +static __global__ void top_k_radix_gather( + const float * __restrict__ src, + int * __restrict__ dst, + top_k_radix_state * __restrict__ states, + int ncols, + int k, + int blocks_per_row) { + const int row = blockIdx.x / blocks_per_row; + const int row_block = blockIdx.x % blocks_per_row; + const int tid = threadIdx.x; + const float * row_src = src + (size_t) row * ncols; + int * row_dst = dst + (size_t) row * k; + top_k_radix_state * state = &states[row]; + + for (int col = row_block * BLOCK_SIZE + tid; + col < ncols; + col += blocks_per_row * BLOCK_SIZE) { + const uint32_t key = top_k_float_to_ordered(row_src[col]); + if (key > state->prefix) { + const int pos = atomicAdd(&state->greater_count, 1); + row_dst[pos] = col; + } else if (key == state->prefix) { + const int pos = atomicAdd(&state->equal_count, 1); + if (pos < state->rank) { + row_dst[k - state->rank + pos] = col; + } + } + } +} + +static void top_k_radix_cuda( + ggml_cuda_pool & pool, + const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) { + constexpr int BLOCK_SIZE = 256; + constexpr int RADIX_BITS = 8; + constexpr int NBINS = 1 << RADIX_BITS; + const int blocks_per_row = std::min((ncols + 1023) / 1024, 64); + + ggml_cuda_pool_alloc states_alloc(pool, nrows); + ggml_cuda_pool_alloc histograms_alloc(pool, (size_t) nrows * blocks_per_row * NBINS); + top_k_radix_state * states = states_alloc.get(); + int * histograms = histograms_alloc.get(); + + top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k); + + const dim3 row_grid(blocks_per_row * nrows); + for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) { + top_k_radix_histogram + <<>>( + src, states, histograms, ncols, blocks_per_row, shift); + top_k_radix_select + <<>>(histograms, states, blocks_per_row, shift); + } + + top_k_radix_reset_counters + <<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows); + top_k_radix_gather + <<>>( + src, dst, states, ncols, k, blocks_per_row); +} + +#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -96,10 +258,18 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { dst_d += k * iter_nrows; } #else // GGML_CUDA_USE_CUB - ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); - int * tmp_dst = temp_dst_alloc.get(); - argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); - CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, - cudaMemcpyDeviceToDevice, stream)); +#if defined(GGML_USE_HIP) + if (ncols > 1024) { + top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); + } else { +#endif // defined(GGML_USE_HIP) + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + int * tmp_dst = temp_dst_alloc.get(); + argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); + CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, + cudaMemcpyDeviceToDevice, stream)); +#if defined(GGML_USE_HIP) + } +#endif // defined(GGML_USE_HIP) #endif }