diff --git a/cpp/include/raft/linalg/detail/reduce_rows_by_key.cuh b/cpp/include/raft/linalg/detail/reduce_rows_by_key.cuh index dbe3f59b69..b19d40471e 100644 --- a/cpp/include/raft/linalg/detail/reduce_rows_by_key.cuh +++ b/cpp/include/raft/linalg/detail/reduce_rows_by_key.cuh @@ -7,12 +7,21 @@ #include #include +#include #include +#include + +#include +#include #include +#include #include +#include +#include #include +#include #define MAX_BLOCKS 65535u namespace raft { @@ -49,138 +58,16 @@ void convert_array(IteratorT1 dst, IteratorT2 src, int n, cudaStream_t st) raft::launch_kernel(st, grid, block, convert_array_kernel, dst, src, n); } -template -struct quad { - T x, y, z, w; -}; - -// -// Functor for reduce by key, small k -// -template -struct quadSum { - __host__ __device__ __forceinline__ quad operator()(const quad& a, const quad& b) const - { - // wasting a double4.. - quad c; - c.x = a.x + b.x; - c.y = a.y + b.y; - c.z = a.z + b.z; - c.w = a.w + b.w; - - return c; - } -}; - // // Reduce by keys // We need to sum each dimension by labels // The labels are not adjacent // -// -// Reduce by keys - for keys <= 4 -// - -#define SUM_ROWS_SMALL_K_DIMX 256 -#define SUM_ROWS_BY_KEY_SMALL_K_MAX_K 4 -template -__launch_bounds__(SUM_ROWS_SMALL_K_DIMX, 4) - - RAFT_KERNEL sum_rows_by_key_small_nkeys_kernel(const DataIteratorT d_A, - IdxT lda, - const char* d_keys, - const WeightT* d_weights, - IdxT nrows, - IdxT ncols, - IdxT nkeys, - SumsT* d_sums) -{ - typedef typename std::iterator_traits::value_type DataType; - typedef cub::BlockReduce, SUM_ROWS_SMALL_K_DIMX> BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - for (IdxT idim = static_cast(blockIdx.y); idim < ncols; idim += gridDim.y) { - if (idim != static_cast(blockIdx.y)) __syncthreads(); // we're reusing temp_storage - - // threadIdx.x stores partial sum for current dim and key=threadIdx.x in this reg - quad thread_sums; - thread_sums.x = 0.0; - thread_sums.y = 0.0; - thread_sums.z = 0.0; - thread_sums.w = 0.0; - - // May use vectorized load - not necessary for doubles - for (IdxT block_offset_irow = blockIdx.x * blockDim.x; - block_offset_irow < nrows; // we will syncthreads() inside the loop, no CTA divergence - block_offset_irow += blockDim.x * gridDim.x) { - IdxT irow = block_offset_irow + threadIdx.x; - DataType val = (irow < nrows) ? d_A[irow * lda + idim] : DataType(0); - if (d_weights && irow < nrows) { val = val * d_weights[irow]; } - // we are not reusing the keys - after profiling - // d_keys is mainly loaded from L2, and this kernel is DRAM BW bounded - // (experimentation gave a 10% speed up - not worth the many code lines added) - IdxT row_key = (irow < nrows) ? d_keys[irow] : std::numeric_limits::max(); - - thread_sums.x += (row_key == 0) ? static_cast(val) : SumsT(0); - thread_sums.y += (row_key == 1) ? static_cast(val) : SumsT(0); - thread_sums.z += (row_key == 2) ? static_cast(val) : SumsT(0); - thread_sums.w += (row_key == 3) ? static_cast(val) : SumsT(0); - } - - // End of column - // Saving local sums back to global mem - - // Strided access - - // Reducing by key - thread_sums = BlockReduce(temp_storage).Reduce(thread_sums, quadSum()); - - if (threadIdx.x < 32) { - // We only need 4 - thread_sums = cub::ShuffleIndex<32>(thread_sums, 0, 0xffffffff); - if (static_cast(threadIdx.x) < nkeys) { - if (threadIdx.x == 0) raft::myAtomicAdd(&d_sums[threadIdx.x * ncols + idim], thread_sums.x); - if (threadIdx.x == 1) raft::myAtomicAdd(&d_sums[threadIdx.x * ncols + idim], thread_sums.y); - if (threadIdx.x == 2) raft::myAtomicAdd(&d_sums[threadIdx.x * ncols + idim], thread_sums.z); - if (threadIdx.x == 3) raft::myAtomicAdd(&d_sums[threadIdx.x * ncols + idim], thread_sums.w); - } - } - } -} - -template -void sum_rows_by_key_small_nkeys(const DataIteratorT d_A, - IdxT lda, - const char* d_keys, - const WeightT* d_weights, - IdxT nrows, - IdxT ncols, - IdxT nkeys, - SumsT* d_sums, - cudaStream_t st) -{ - dim3 grid, block; - block.x = SUM_ROWS_SMALL_K_DIMX; - block.y = 1; // Necessary - - grid.x = raft::ceildiv(nrows, (IdxT)block.x); - grid.x = std::min(grid.x, 32u); - grid.y = ncols; - grid.y = std::min(grid.y, MAX_BLOCKS); - raft::launch_kernel(st, - grid, - block, - sum_rows_by_key_small_nkeys_kernel, - d_A, - lda, - d_keys, - d_weights, - nrows, - ncols, - nkeys, - d_sums); -} +/** Floating point including the extended types, which std::is_floating_point does not cover. */ +template +inline constexpr bool is_floating_point_v = + std::is_floating_point_v || std::is_same_v || std::is_same_v; // // Reduce by keys - large number of keys @@ -255,7 +142,7 @@ void sum_rows_by_key_large_nkeys_colmajor(const DataIteratorT d_A, cudaStream_t st) { dim3 grid, block; - block.x = SUM_ROWS_SMALL_K_DIMX; + block.x = 256; block.y = 1; // Necessary grid.x = raft::ceildiv(nrows, (IdxT)block.x); @@ -276,12 +163,48 @@ void sum_rows_by_key_large_nkeys_colmajor(const DataIteratorT d_A, d_sums); } -template +DI void sum_rows_by_key_load(const DataIteratorT d_A, IdxT idx, bool valid, SumsT (&v)[VecLen]) +{ + if constexpr (VecLen > 1) { + TxN_t x; + x.fill(SumsT(0)); + if (valid) { x.load(static_cast(d_A), idx); } +#pragma unroll + for (int c = 0; c < VecLen; ++c) { + v[c] = x.val.data[c]; + } + } else { + v[0] = valid ? static_cast(d_A[idx]) : SumsT(0); + } +} + +/** (Vectorized) atomic add */ +template +DI void sum_rows_by_key_atomic_add(SumsT* dst, const SumsT (&v)[VecLen]) +{ + if constexpr (std::is_same_v && VecLen == 4) { + raft::myAtomicAdd(reinterpret_cast(dst), make_float4(v[0], v[1], v[2], v[3])); + } else { +#pragma unroll + for (int c = 0; c < VecLen; ++c) { + raft::myAtomicAdd(dst + c, v[c]); + } + } +} + +/** + * One thread owns one element of a row, multiplies it by the row's weight and adds it to the sums + * of the row's key. When vectorized, one threads owns a vector of contiguous elements. + */ +template -RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor(const DataIteratorT d_A, +RAFT_KERNEL sum_rows_by_key_kernel_rowmajor_elementwise(const DataIteratorT d_A, IdxT lda, const WeightT* d_weights, KeysIteratorT d_keys, @@ -289,45 +212,557 @@ RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor(const DataIteratorT d_A, IdxT ncols, SumsT* d_sums) { - IdxT gid = threadIdx.x + (blockDim.x * static_cast(blockIdx.x)); - IdxT j = gid % ncols; - IdxT i = gid / ncols; + const IdxT q = threadIdx.x + blockDim.x * static_cast(blockIdx.x); + const IdxT vecs_per_row = ncols / VecLen; + const IdxT i = q / vecs_per_row; if (i >= nrows) return; - IdxT l = static_cast(d_keys[i]); - SumsT val = d_A[j + lda * i]; - if (d_weights != nullptr) val *= d_weights[i]; - raft::myAtomicAdd(&d_sums[j + ncols * l], val); + const IdxT j = (q - i * vecs_per_row) * VecLen; + const SumsT w = d_weights != nullptr ? static_cast(d_weights[i]) : SumsT(1); + SumsT a[VecLen]; + sum_rows_by_key_load(d_A, lda * i + j, true, a); +#pragma unroll + for (int c = 0; c < VecLen; ++c) { + a[c] *= w; + } + sum_rows_by_key_atomic_add(d_sums + ncols * static_cast(d_keys[i]) + j, a); } -template -void sum_rows_by_key_large_nkeys_rowmajor(const DataIteratorT d_A, +void sum_rows_by_key_rowmajor_elementwise(const DataIteratorT d_A, IdxT lda, const KeysIteratorT d_keys, const WeightT* d_weights, IdxT nrows, IdxT ncols, SumsT* d_sums, + uint32_t block_dim, cudaStream_t st) { - uint32_t block_dim = 128; - auto grid_dim = static_cast(ceildiv(nrows * ncols, (IdxT)block_dim)); + const IdxT n_items = nrows * (ncols / VecLen); + const auto grid_dim = static_cast(ceildiv(n_items, static_cast(block_dim))); + raft::launch_kernel(st, + grid_dim, + block_dim, + sum_rows_by_key_kernel_rowmajor_elementwise, + d_A, + lda, + d_weights, + d_keys, + nrows, + ncols, + d_sums); +} + +constexpr int kFewKeysMaxThreads = 512; // Max threads per block +constexpr size_t kFewKeysSmemBytes = 48 * 1024; // 48 KiB - Each GPU from Ampere+ should have it. + +/** + * Kernel specialized for few keys (high atomic contention) and narrow rows. + * + * - One copy has nkeys * ncols sums (same as the global sums). Each blocks holds multiple copies in + * its shared memory. + * - A block owns a range of whole rows. + * - Each thread owns one column in one copy. It sums the rows in its column of its copy, using the + * row's key. Therefore, each thread owns different addresses in shared memory, so no atomics are + * needed. + * - At the end, the copies are summed into the global sums. + */ +template +RAFT_KERNEL __launch_bounds__(kFewKeysMaxThreads) + sum_rows_by_key_kernel_rowmajor_small_keys_narrow_rows(const DataIteratorT d_A, + IdxT lda, + const WeightT* d_weights, + KeysIteratorT d_keys, + IdxT nrows, + IdxT ncols, + IdxT nkeys, + IdxT rows_per_block, + SumsT* d_sums) +{ + // Shared memory which holds multiple copies. Chosen layout: smem[n_copies][nkeys][ncols] + /// Layout choice: + /// - Each threads owns one column in one copy. So having ncols in the last dimension + /// makes neighboring threads access neighboring shared memory addresses - no bank conflicts. + /// - nkeys X ncols matches the global sum layout. Writing the sums to global memory is + /// coalesced. + extern __shared__ __align__(16) unsigned char few_keys_smem[]; + SumsT* private_sums = reinterpret_cast(few_keys_smem); + const int n = static_cast(ncols); + const int n_copies = blockDim.x / n; + /// c: column index owned by this thread + const int c = threadIdx.x % n; + /// p: copy index owned by this thread + const int p = threadIdx.x / n; + /// [r0, r1): range of rows owned by this block. + const IdxT r0 = blockIdx.x * rows_per_block; + const IdxT r1 = std::min(r0 + rows_per_block, nrows); + + // Zero-initialize private sums + for (IdxT i = threadIdx.x; i < n_copies * nkeys * ncols; i += blockDim.x) { + private_sums[i] = SumsT(0); + } + __syncthreads(); + + // Load rows and add them to the sums of their keys. + /// private_sums[copies][keys][cols] -> mine[keys][:] + SumsT* mine = private_sums + static_cast(p) * nkeys * ncols + c; + for (IdxT r = r0 + p; r < r1; r += static_cast(kRowsAhead) * n_copies) { + SumsT a[kRowsAhead]; + SumsT w[kRowsAhead]; + IdxT key[kRowsAhead]; + /// Load all kRowsAhead rows beforehand to have enough bytes in flight +#pragma unroll + for (int u = 0; u < kRowsAhead; ++u) { + const IdxT row = r + static_cast(u) * n_copies; + const bool valid = row < r1; + a[u] = valid ? static_cast(d_A[lda * row + c]) : SumsT(0); + w[u] = !valid ? SumsT(0) + : d_weights != nullptr ? static_cast(d_weights[row]) + : SumsT(1); + key[u] = valid ? static_cast(d_keys[row]) : IdxT(0); + } + /// Add each row to the sum of its key +#pragma unroll + for (int u = 0; u < kRowsAhead; ++u) { + mine[key[u] * ncols] += a[u] * w[u]; // mine[key][0] - each key is ncols apart + } + } + __syncthreads(); + + // Sum the copies into the global sums + for (IdxT i = threadIdx.x; i < nkeys * ncols; i += blockDim.x) { + SumsT t = SumsT(0); + for (int q = 0; q < n_copies; ++q) { + t += private_sums[static_cast(q) * nkeys * ncols + i]; + } + raft::myAtomicAdd(d_sums + i, t); + } +} + +/** Whether one private copy of the sums fits the few-keys kernel's shared memory. */ +template +bool sum_rows_by_key_few_keys_fits(IdxT nkeys, IdxT ncols) +{ + const size_t one_copy_bytes = + sizeof(SumsT) * static_cast(nkeys) * static_cast(ncols); + return one_copy_bytes <= kFewKeysSmemBytes; +} + +template +void sum_rows_by_key_rowmajor_few_keys(const DataIteratorT d_A, + IdxT lda, + const KeysIteratorT d_keys, + const WeightT* d_weights, + IdxT nrows, + IdxT ncols, + IdxT nkeys, + SumsT* d_sums, + cudaStream_t st) +{ + RAFT_EXPECTS( + ncols > 0 && ncols <= kFewKeysMaxThreads && sum_rows_by_key_few_keys_fits(nkeys, ncols), + "The sums do not fit in shared memory"); + constexpr int kRowsAhead = 8; + const size_t copy_bytes = std::max(1, sizeof(SumsT) * nkeys * ncols); + const int n_copies = static_cast(std::max( + 1, std::min(kFewKeysMaxThreads / ncols, kFewKeysSmemBytes / copy_bytes))); + const auto block_dim = static_cast(n_copies * ncols); + const int blocks_per_sm = block_dim >= 256 ? 2 : 4; + const IdxT rows_per_block = + ceildiv(nrows, static_cast(blocks_per_sm) * raft::getMultiProcessorCount()); + const auto grid_dim = static_cast(ceildiv(nrows, rows_per_block)); + raft::launch_kernel({st, copy_bytes * n_copies}, + grid_dim, + block_dim, + sum_rows_by_key_kernel_rowmajor_small_keys_narrow_rows, + d_A, + lda, + d_weights, + d_keys, + nrows, + ncols, + nkeys, + rows_per_block, + d_sums); +} + +/** Prepare keys and row ids for radix-sort. */ +template +RAFT_KERNEL sum_rows_by_key_sort_input_kernel(KeysIteratorT d_keys, + IdxT nrows, + uint32_t* keys, + uint32_t* row_ids) +{ + const IdxT i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < nrows) { + keys[i] = static_cast(d_keys[i]); + row_ids[i] = static_cast(i); + } +} + +/** + * Sum sorted-by-key rows in registers and flush to output when the key changes or the chunk ends. + * + * Work split: + * tile 0 tile 1 ... + * chunk 0 | warp 0 | warp 1 | ... + * chunk 1 | warp n_tiles | ... + * + * To have enough bytes in flight, each warp loads kRowsAhead rows before adding any of them to the + * sums. On each row, each lane owns kVecs vectors of VecLen contiguous columns. + */ +template +RAFT_KERNEL sum_rows_by_key_kernel_rowmajor_sorted(const DataIteratorT d_A, + IdxT lda, + const WeightT* d_weights, + const uint32_t* d_keys, + const uint32_t* d_row_ids, + IdxT nrows, + IdxT ncols, + IdxT chunk, + IdxT n_tiles, + SumsT* d_sums) +{ + // Each warp owns: rows [r0, r1) and columns [col0, col0 + kTileCols) + constexpr IdxT kTileCols = kVecs * VecLen * WarpSize; + const int64_t warp_id = (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / WarpSize; + const IdxT r0 = static_cast(warp_id / n_tiles) * chunk; + if (r0 >= nrows) return; // spare warps of the last block + const IdxT r1 = std::min(r0 + chunk, nrows); + const IdxT col0 = static_cast(warp_id % n_tiles) * kTileCols + laneId() * VecLen; + + // Each lane tracks in registers the sums of each vector and the current key of the processed row. + SumsT acc[kVecs][VecLen] = {}; + uint32_t cur = d_keys[r0]; + + // Flush to output + auto flush = [&]() { +#pragma unroll + for (int v = 0; v < kVecs; ++v) { + const IdxT col = col0 + v * VecLen * WarpSize; + if (col < ncols) { + sum_rows_by_key_atomic_add(d_sums + ncols * static_cast(cur) + col, acc[v]); + } +#pragma unroll + for (int c = 0; c < VecLen; ++c) { + acc[v][c] = SumsT(0); // restart at zero for the next run + } + } + }; + + for (IdxT r = r0; r < r1; r += kRowsAhead) { + // Phase 1: issue all loads of the next kRowsAhead sorted rows (keys, weights, data) before + // using any of them, so they are in flight together. + SumsT a[kRowsAhead][kVecs][VecLen]; + uint32_t key[kRowsAhead]; + SumsT w[kRowsAhead]; +#pragma unroll + for (int u = 0; u < kRowsAhead; ++u) { + const IdxT i = r + u; + // Invalid rows: Apply zero weight and unchanged key + const bool valid = i < r1; + const IdxT row = valid ? static_cast(d_row_ids[i]) : IdxT(0); + key[u] = valid ? d_keys[i] : cur; + w[u] = !valid ? SumsT(0) + : d_weights != nullptr ? static_cast(d_weights[row]) + : SumsT(1); +#pragma unroll + for (int v = 0; v < kVecs; ++v) { + const IdxT col = col0 + v * VecLen * WarpSize; + sum_rows_by_key_load(d_A, lda * row + col, valid && col < ncols, a[u][v]); + } + } + // Phase 2: do FMA and flush on key changes +#pragma unroll + for (int u = 0; u < kRowsAhead; ++u) { + if (key[u] != cur) { + flush(); + cur = key[u]; + } +#pragma unroll + for (int v = 0; v < kVecs; ++v) { +#pragma unroll + for (int c = 0; c < VecLen; ++c) { + acc[v][c] += a[u][v][c] * w[u]; + } + } + } + } + flush(); // necessary +} + +template +void sum_rows_by_key_rowmajor_sorted_launch(const DataIteratorT d_A, + IdxT lda, + const WeightT* d_weights, + const uint32_t* d_keys, + const uint32_t* d_row_ids, + IdxT nrows, + IdxT ncols, + SumsT* d_sums, + cudaStream_t st) +{ + constexpr int kVecs = kLaneBytes / (VecLen * sizeof(SumsT)); + constexpr int kRowsAhead = kLaneBytes == 16 ? 4 : 8; + constexpr IdxT kWarpsPerSM = 64; + // n_tiles: how many tiles of kVecs*VecLen columns fit in a row. Each warp owns one tile. + const IdxT n_tiles = ceildiv(ncols, IdxT{kVecs * VecLen * WarpSize}); + // n_chunks: how many chunks of rows to split the work into. Each warp owns one chunk. + const IdxT n_chunks = + std::max(1, IdxT{kWarpsPerSM} * raft::getMultiProcessorCount() / n_tiles); + // chunk: how many rows each warp owns. The last chunk may be smaller. + const IdxT chunk = ceildiv(nrows, n_chunks); + // n_warps: how many warps to cover everything. + const IdxT n_warps = ceildiv(nrows, chunk) * n_tiles; + constexpr uint32_t block_dim = 128; + const auto grid_dim = static_cast(ceildiv(n_warps, IdxT{block_dim / WarpSize})); raft::launch_kernel(st, grid_dim, block_dim, - sum_rows_by_key_large_nkeys_kernel_rowmajor, + sum_rows_by_key_kernel_rowmajor_sorted, d_A, lda, d_weights, d_keys, + d_row_ids, nrows, ncols, + chunk, + n_tiles, d_sums); } +/** + * Radix-sort the row ids by key and then sum runs of equal keys in registers. + * Batching defines upper limit for peak memory allocation. + */ +template +void sum_rows_by_key_rowmajor_sorted(const DataIteratorT d_A, + IdxT lda, + const KeysIteratorT d_keys, + const WeightT* d_weights, + IdxT nrows, + IdxT ncols, + IdxT nkeys, + SumsT* d_sums, + cudaStream_t st, + bool dry_run) +{ + // Limit batch size to 4M rows + const IdxT batch_rows = std::min(nrows, 4000000); + rmm::device_uvector keys(batch_rows, st), keys_alt(batch_rows, st); + rmm::device_uvector row_ids(batch_rows, st), row_ids_alt(batch_rows, st); + // Double buffering avoids additional scratch space when radix sorts undergoes > 1 pass. + // 1 pass corresponds to 8-bit (selected by the CUB library) + cub::DoubleBuffer keys_db(keys.data(), keys_alt.data()); + cub::DoubleBuffer row_ids_db(row_ids.data(), row_ids_alt.data()); + // Only the bits that can be set in a key in [0, nkeys) take part in the sort. + int end_bit = 1; + while (end_bit < 32 && (IdxT{1} << end_bit) < nkeys) { + ++end_bit; + } + size_t tmp_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + nullptr, tmp_bytes, keys_db, row_ids_db, static_cast(batch_rows), 0, end_bit, st)); + rmm::device_buffer tmp(tmp_bytes, st); + if (dry_run) { return; } + + // Pick 16-byte vector when the input is a plain pointer and rows and sums are aligned. + constexpr int kVecLen = 16 / sizeof(SumsT); + constexpr bool kIsPlainPointer = + std::is_same_v || std::is_same_v; + bool vectorized = false; + if constexpr (kIsPlainPointer) { + vectorized = ncols % kVecLen == 0 && lda % kVecLen == 0 && raft::is_aligned(d_A, 16) && + raft::is_aligned(d_sums, 16); + } + + // A warp covers kLaneBytes * 32 bytes of a row: 32 bytes per lane (a 1 KiB tile), or 16 bytes + // per lane when the row is narrower than that, so fewer lanes idle. + const bool narrow_rows = ncols * sizeof(SumsT) < 32 * WarpSize; + + for (IdxT batch_start = 0; batch_start < nrows; batch_start += batch_rows) { + // Rows [batch_start, batch_start + n_batch_rows); row ids are local to the batch. + const IdxT n_batch_rows = std::min(batch_rows, nrows - batch_start); + keys_db = cub::DoubleBuffer(keys.data(), keys_alt.data()); + row_ids_db = cub::DoubleBuffer(row_ids.data(), row_ids_alt.data()); + constexpr uint32_t block_dim = 256; + const auto grid_dim = static_cast(ceildiv(n_batch_rows, IdxT{block_dim})); + raft::launch_kernel(st, + grid_dim, + block_dim, + sum_rows_by_key_sort_input_kernel, + d_keys + batch_start, + n_batch_rows, + keys_db.Current(), + row_ids_db.Current()); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + tmp.data(), tmp_bytes, keys_db, row_ids_db, static_cast(n_batch_rows), 0, end_bit, st)); + const auto batch_A = d_A + batch_start * lda; + const WeightT* batch_w = d_weights != nullptr ? d_weights + batch_start : nullptr; + const uint32_t* s_keys = keys_db.Current(); + const uint32_t* s_rowids = row_ids_db.Current(); + auto launch = [&](auto vec_len) { + constexpr int VecLen = decltype(vec_len)::value; + if (narrow_rows) { + sum_rows_by_key_rowmajor_sorted_launch( + batch_A, lda, batch_w, s_keys, s_rowids, n_batch_rows, ncols, d_sums, st); + } else { + sum_rows_by_key_rowmajor_sorted_launch( + batch_A, lda, batch_w, s_keys, s_rowids, n_batch_rows, ncols, d_sums, st); + } + }; + if constexpr (kIsPlainPointer) { + if (vectorized) { + launch(std::integral_constant{}); + continue; + } + } + launch(std::integral_constant{}); + } +} + +enum class sum_rows_by_key_kernel { elementwise, few_keys, sorted }; + +struct sum_rows_by_key_plan { + sum_rows_by_key_kernel kernel; + int vec_len; + uint32_t block_dim; +}; + +/** + * Kernel selection based on the problem size + */ +template +sum_rows_by_key_plan sum_rows_by_key_select_kernel( + const DataIteratorT d_A, IdxT lda, IdxT nrows, IdxT ncols, IdxT nkeys, const SumsT* d_sums) +{ + constexpr double kSortedMinRowsPerKey = 5000; + constexpr double kSortedMinRowBytes = 512; + constexpr double kSortedMinInputBytes = 128.0 * 1024 * 1024; + // Few keys and narrow rows: per-block sums in shared memory. + constexpr IdxT kFewKeysMaxCols = 128; + constexpr IdxT kFewKeysMaxKeys = 100; + constexpr size_t kFewKeysMaxCopyBytes = kFewKeysSmemBytes / 2; // at least 2 copies + + int dev_id, cc_major; + RAFT_CUDA_TRY(cudaGetDevice(&dev_id)); + RAFT_CUDA_TRY(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, dev_id)); + + // Select vectorized loads and vectorized atomic writes + bool whole_vectors = false; + /// plain float pointer + if constexpr (std::is_same_v && (std::is_same_v || + std::is_same_v)) { + /// 16-byte alignments + whole_vectors = + ncols % 4 == 0 && lda % 4 == 0 && raft::is_aligned(d_A, 16) && raft::is_aligned(d_sums, 16); + } + /// Hopper+ (sm_9x) supports float4 atomics + const int vec_len = whole_vectors && cc_major >= 9 ? 4 : 1; + + // Select kernel + if (ncols > 0 && ncols <= kFewKeysMaxCols && nkeys <= kFewKeysMaxKeys && + sizeof(SumsT) * static_cast(nkeys) * static_cast(ncols) <= + kFewKeysMaxCopyBytes) { + return {sum_rows_by_key_kernel::few_keys, 1, 0}; + } + const double row_bytes = static_cast(ncols) * sizeof(SumsT); + const bool contended = static_cast(nrows) >= kSortedMinRowsPerKey * nkeys; + const bool wide_large = + row_bytes >= kSortedMinRowBytes && row_bytes * nrows >= kSortedMinInputBytes; + if (static_cast(nkeys) <= std::numeric_limits::max() && + (contended || wide_large)) { + return {sum_rows_by_key_kernel::sorted, vec_len, 0}; + } + return {sum_rows_by_key_kernel::elementwise, vec_len, 256u}; +} + +template +void sum_rows_by_keyrowmajor(const DataIteratorT d_A, + IdxT lda, + const KeysIteratorT d_keys, + const WeightT* d_weights, + IdxT nrows, + IdxT ncols, + IdxT nkeys, + SumsT* d_sums, + cudaStream_t st, + bool dry_run) +{ + const auto plan = sum_rows_by_key_select_kernel(d_A, lda, nrows, ncols, nkeys, d_sums); + if (dry_run && plan.kernel != sum_rows_by_key_kernel::sorted) { return; } + + constexpr bool kCanVectorize = + std::is_same_v && + (std::is_same_v || std::is_same_v); + + switch (plan.kernel) { + case sum_rows_by_key_kernel::sorted: + sum_rows_by_key_rowmajor_sorted( + d_A, lda, d_keys, d_weights, nrows, ncols, nkeys, d_sums, st, dry_run); + break; + case sum_rows_by_key_kernel::few_keys: + sum_rows_by_key_rowmajor_few_keys( + d_A, lda, d_keys, d_weights, nrows, ncols, nkeys, d_sums, st); + break; + case sum_rows_by_key_kernel::elementwise: + if constexpr (kCanVectorize) { + if (plan.vec_len == 4) { + sum_rows_by_key_rowmajor_elementwise<4>( + d_A, lda, d_keys, d_weights, nrows, ncols, d_sums, plan.block_dim, st); + break; + } + } + sum_rows_by_key_rowmajor_elementwise<1>( + d_A, lda, d_keys, d_weights, nrows, ncols, d_sums, plan.block_dim, st); + break; + } +} + /** * @brief Computes the weighted reduction of matrix rows for each given key * @@ -342,7 +777,7 @@ void sum_rows_by_key_large_nkeys_rowmajor(const DataIteratorT d_A, * @param[in] lda Real row size for input data, d_A * @param[in] d_keys Keys for each row (1 x nrows) * @param[in] d_weights Weights for each observation in d_A (1 x nrows) - * @param[out] d_keys_char Scratch memory for conversion of keys to char + * @param[out] d_keys_char UNUSED, kept for compatibility * @param[in] nrows Number of rows in d_A and d_keys * @param[in] ncols Number of data columns in d_A * @param[in] nkeys Number of unique keys in d_keys @@ -365,23 +800,22 @@ void reduce_rows_by_key(const DataIteratorT d_A, IdxT nkeys, SumsT* d_sums, cudaStream_t stream, - bool reset_sums) + bool reset_sums, + bool dry_run = false) { typedef typename std::iterator_traits::value_type KeyType; + static_assert( + !is_floating_point_v || + (is_floating_point_v && sizeof(WeightT) <= sizeof(SumsT)), + "Floating weight type is broader than the sum type. Instead of silently casting down the " + "weight precision, we throw an error here."); // Following kernel needs memset - if (reset_sums) { cudaMemsetAsync(d_sums, 0, ncols * nkeys * sizeof(SumsT), stream); } - - if (d_keys_char != nullptr && nkeys <= SUM_ROWS_BY_KEY_SMALL_K_MAX_K) { - // sum_rows_by_key_small_k is BW bounded. d_keys is loaded ncols time - avoiding wasting BW - // with doubles we have ~20% speed up - with floats we can hope something around 2x - // Converting d_keys to char - convert_array(d_keys_char, d_keys, nrows, stream); - sum_rows_by_key_small_nkeys( - d_A, lda, d_keys_char, d_weights, nrows, ncols, nkeys, d_sums, stream); - } else { - sum_rows_by_key_large_nkeys_rowmajor(d_A, lda, d_keys, d_weights, nrows, ncols, d_sums, stream); - } + if (reset_sums && !dry_run) { cudaMemsetAsync(d_sums, 0, ncols * nkeys * sizeof(SumsT), stream); } + + (void)d_keys_char; + sum_rows_by_keyrowmajor( + d_A, lda, d_keys, d_weights, nrows, ncols, nkeys, d_sums, stream, dry_run); } /** @@ -395,12 +829,13 @@ void reduce_rows_by_key(const DataIteratorT d_A, * @param[in] d_A Input data array (lda x nrows) * @param[in] lda Real row size for input data, d_A * @param[in] d_keys Keys for each row (1 x nrows) - * @param d_keys_char Scratch memory for conversion of keys to char + * @param d_keys_char UNUSED, kept for compatibility * @param[in] nrows Number of rows in d_A and d_keys * @param[in] ncols Number of data columns in d_A * @param[in] nkeys Number of unique keys in d_keys * @param[out] d_sums Row sums by key (ncols x d_keys) * @param[in] stream CUDA stream + * @param[in] reset_sums Whether to reset the output sums to zero before reducing */ template void reduce_rows_by_key(const DataIteratorT d_A, @@ -412,7 +847,8 @@ void reduce_rows_by_key(const DataIteratorT d_A, IdxT nkeys, SumsT* d_sums, cudaStream_t stream, - bool reset_sums) + bool reset_sums, + bool dry_run = false) { typedef typename std::iterator_traits::value_type DataType; reduce_rows_by_key(d_A, @@ -425,7 +861,8 @@ void reduce_rows_by_key(const DataIteratorT d_A, nkeys, d_sums, stream, - reset_sums); + reset_sums, + dry_run); } }; // end namespace detail diff --git a/cpp/include/raft/linalg/reduce_rows_by_key.cuh b/cpp/include/raft/linalg/reduce_rows_by_key.cuh index 32005b14d1..9023d087dd 100644 --- a/cpp/include/raft/linalg/reduce_rows_by_key.cuh +++ b/cpp/include/raft/linalg/reduce_rows_by_key.cuh @@ -42,7 +42,7 @@ void convert_array(IteratorT1 dst, IteratorT2 src, int n, cudaStream_t st) * @param[in] lda Real row size for input data, d_A * @param[in] d_keys Keys for each row (1 x nrows) * @param[in] d_weights Weights for each observation in d_A (1 x nrows) - * @param[out] d_keys_char Scratch memory for conversion of keys to char + * @param[out] d_keys_char Unused, kept for compatibility (formerly scratch for keys as char) * @param[in] nrows Number of rows in d_A and d_keys * @param[in] ncols Number of data columns in d_A * @param[in] nkeys Number of unique keys in d_keys @@ -82,7 +82,7 @@ void reduce_rows_by_key(const DataIteratorT d_A, * @param[in] d_A Input data array (lda x nrows) * @param[in] lda Real row size for input data, d_A * @param[in] d_keys Keys for each row (1 x nrows) - * @param d_keys_char Scratch memory for conversion of keys to char + * @param d_keys_char Unused, kept for compatibility (formerly scratch for keys as char) * @param[in] nrows Number of rows in d_A and d_keys * @param[in] ncols Number of data columns in d_A * @param[in] nkeys Number of unique keys in d_keys @@ -133,7 +133,8 @@ void reduce_rows_by_key(const DataIteratorT d_A, * @param[in] d_keys Keys for each row raft::device_vector_view (1 x nrows) * @param[out] d_sums Row sums by key raft::device_matrix_view (ncols x d_keys) * @param[in] n_unique_keys Number of unique keys in d_keys - * @param[out] d_keys_char Scratch memory for conversion of keys to char, raft::device_vector_view + * @param[out] d_keys_char Unused, kept for compatibility (formerly scratch for keys as char), + * raft::device_vector_view * @param[in] d_weights Weights for each observation in d_A raft::device_vector_view optional (1 * x nrows) * @param[in] reset_sums Whether to reset the output sums to zero before reducing @@ -149,36 +150,38 @@ void reduce_rows_by_key( std::optional> d_weights = std::nullopt, bool reset_sums = true) { - if (resource::get_dry_run_flag(handle)) { return; } RAFT_EXPECTS(d_A.extent(0) == d_A.extent(0) && d_sums.extent(1) == n_unique_keys, "Output is not of size ncols * n_unique_keys"); RAFT_EXPECTS(d_keys.extent(0) == d_A.extent(1), "Keys is not of size nrows"); + const bool dry_run = resource::get_dry_run_flag(handle); if (d_weights) { RAFT_EXPECTS(d_weights.value().extent(0) == d_A.extent(1), "Weights is not of size nrows"); - reduce_rows_by_key(d_A.data_handle(), - d_A.extent(0), - d_keys.data_handle(), - d_weights.value().data_handle(), - d_keys_char.data_handle(), - d_A.extent(1), - d_A.extent(0), - n_unique_keys, - d_sums.data_handle(), - resource::get_cuda_stream(handle).get(), - reset_sums); + detail::reduce_rows_by_key(d_A.data_handle(), + d_A.extent(0), + d_keys.data_handle(), + d_weights.value().data_handle(), + d_keys_char.data_handle(), + d_A.extent(1), + d_A.extent(0), + n_unique_keys, + d_sums.data_handle(), + resource::get_cuda_stream(handle).get(), + reset_sums, + dry_run); } else { - reduce_rows_by_key(d_A.data_handle(), - d_A.extent(0), - d_keys.data_handle(), - d_keys_char.data_handle(), - d_A.extent(1), - d_A.extent(0), - n_unique_keys, - d_sums.data_handle(), - resource::get_cuda_stream(handle).get(), - reset_sums); + detail::reduce_rows_by_key(d_A.data_handle(), + d_A.extent(0), + d_keys.data_handle(), + d_keys_char.data_handle(), + d_A.extent(1), + d_A.extent(0), + n_unique_keys, + d_sums.data_handle(), + resource::get_cuda_stream(handle).get(), + reset_sums, + dry_run); } } diff --git a/cpp/include/raft/util/cuda_utils.cuh b/cpp/include/raft/util/cuda_utils.cuh index da56be7bf5..0a95fd9dff 100644 --- a/cpp/include/raft/util/cuda_utils.cuh +++ b/cpp/include/raft/util/cuda_utils.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -43,6 +43,22 @@ DI void myAtomicAdd(int64_t* address, int64_t val) } while (assumed != old); } +// Specialization for float4 +template <> +DI void myAtomicAdd(float4* address, float4 val) +{ +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + // Vectorized atomic write is only available from Hopper. + atomicAdd(address, val); +#else + float* elements = reinterpret_cast(address); + myAtomicAdd(elements, val.x); + myAtomicAdd(elements + 1, val.y); + myAtomicAdd(elements + 2, val.z); + myAtomicAdd(elements + 3, val.w); +#endif +} + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 600) // Ref: // http://on-demand.gputechconf.com/gtc/2013/presentations/S3101-Atomic-Memory-Operations.pdf diff --git a/cpp/tests/linalg/reduce_rows_by_key.cu b/cpp/tests/linalg/reduce_rows_by_key.cu index 43af82279e..005326685e 100644 --- a/cpp/tests/linalg/reduce_rows_by_key.cu +++ b/cpp/tests/linalg/reduce_rows_by_key.cu @@ -159,7 +159,7 @@ class ReduceRowTest : public ::testing::TestWithParam> { rmm::device_uvector scratch_buf; }; -// ReduceRowTestF +// ReduceRowTestFloat // 128 Obs, 32 cols, 6 clusters const std::vector> inputsf2 = {{0.000001f, 128, 32, 6, 1234ULL, false}, {0.000001f, 128, 32, 6, 1234ULL, true, 1.0}, @@ -174,7 +174,7 @@ TEST_P(ReduceRowTestF, Result) } INSTANTIATE_TEST_CASE_P(ReduceRowTests, ReduceRowTestF, ::testing::ValuesIn(inputsf2)); -// ReduceRowTestD +// ReduceRowTestDouble // 128 Obs, 32 cols, 6 clusters, double precision const std::vector> inputsd2 = { {0.00000001, 128, 32, 6, 1234ULL, false}, @@ -190,14 +190,15 @@ TEST_P(ReduceRowTestD, Result) } INSTANTIATE_TEST_CASE_P(ReduceRowTests, ReduceRowTestD, ::testing::ValuesIn(inputsd2)); -// ReduceRowTestSmallnKey -// 128 Obs, 32 cols, 3 clusters +// ReduceRowTestFewKeys +// 128 Obs, 32-128 cols, 3-10 clusters const std::vector> inputsf_small_nkey = { {0.000001f, 128, 32, 3, 1234ULL, false}, {0.000001f, 128, 32, 3, 1234ULL, true, 5.0}, - {0.000001f, 128, 32, 3, 1234ULL, true, 8.0}}; -typedef ReduceRowTest ReduceRowTestSmallnKey; -TEST_P(ReduceRowTestSmallnKey, Result) + {0.000001f, 128, 32, 3, 1234ULL, true, 8.0}, + {0.00001f, 10000, 128, 10, 1234ULL, true, 4.0}}; // widest few-keys rows +typedef ReduceRowTest ReduceRowTestFewKeys; +TEST_P(ReduceRowTestFewKeys, Result) { ASSERT_TRUE(raft::devArrMatch(out_ref.data(), out.data(), @@ -205,7 +206,7 @@ TEST_P(ReduceRowTestSmallnKey, Result) raft::CompareApprox(params.tolerance))); } INSTANTIATE_TEST_CASE_P(ReduceRowTests, - ReduceRowTestSmallnKey, + ReduceRowTestFewKeys, ::testing::ValuesIn(inputsf_small_nkey)); // ReduceRowTestBigSpace @@ -262,5 +263,61 @@ INSTANTIATE_TEST_CASE_P(ReduceRowTests, ReduceRowTestManyClusters, ::testing::ValuesIn(inputsf_many_cluster)); +// ReduceRowTestLargeClustersAndCols +// 100000 Obs, 5000 cols, 3072 clusters: Should trigger the vectorized kernel on GPU newer than 9.0. +const std::vector> inputsf_vectorized = { + {0.00001f, 100000, 5000, 3072, 1234ULL, false}, + {0.00001f, 100000, 5000, 3072, 1234ULL, true, 32.0}, + {0.00001f, 100000, 5000, 3072, 1234ULL, true, 16.0}}; +typedef ReduceRowTest ReduceRowTestLargeClustersAndCols; +TEST_P(ReduceRowTestLargeClustersAndCols, Result) +{ + ASSERT_TRUE(raft::devArrMatch(out_ref.data(), + out.data(), + params.cols * params.nkeys, + raft::CompareApprox(params.tolerance))); +} +INSTANTIATE_TEST_CASE_P(ReduceRowTests, + ReduceRowTestLargeClustersAndCols, + ::testing::ValuesIn(inputsf_vectorized)); + +TEST(ReduceRowsByKeyDryRun, SortedPathAllocations) +{ + raft::resources handle; + auto stream = resource::get_cuda_stream(handle).get(); + constexpr int nrows = 100000; + constexpr int ncols = 64; + constexpr int nkeys = 1024; + constexpr int end_bit = 10; // > 8 bits so that >1 radix passes + rmm::device_uvector in(nrows * ncols, stream); + rmm::device_uvector sums(nkeys * ncols, stream); + rmm::device_uvector keys(nrows, stream); + raft::random::RngState r(1234ULL); + uniformInt(handle, r, keys.data(), nrows, 0u, uint32_t{nkeys}); + RAFT_CUDA_TRY(cudaMemsetAsync(in.data(), 0, in.size() * sizeof(float), stream)); + + size_t sort_bytes = 0; + cub::DoubleBuffer sort_keys, sort_values; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + nullptr, sort_bytes, sort_keys, sort_values, nrows, 0, end_bit, stream)); + const size_t expected = 4 * nrows * sizeof(uint32_t) + sort_bytes; + + auto action = [&](raft::resources const& res) { + detail::sum_rows_by_key_rowmajor_sorted(in.data(), + ncols, + keys.data(), + static_cast(nullptr), + nrows, + ncols, + nkeys, + sums.data(), + resource::get_cuda_stream(res).get(), + resource::get_dry_run_flag(res)); + }; + const auto dry = raft::util::dry_run_execute(handle, action); + EXPECT_EQ(dry.device_global, expected); + execute_with_dry_run_check(handle, action, alloc_behavior::ARGUMENT_DRIVEN, expected); +} + } // end namespace linalg } // end namespace raft