Optimize sum_rows_by_key() kernel

cpp/include/raft/linalg/detail/reduce_rows_by_key.cuh

@@ -9,10 +9,17 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/kernel_launch.hpp>
 
+#include <rmm/device_buffer.hpp>
+#include <rmm/device_uvector.hpp>
+
 #include <cub/block/block_reduce.cuh>
+#include <cub/device/device_radix_sort.cuh>
 #include <cub/util_ptx.cuh>
 
+#include <algorithm>
+#include <cstdint>
 #include <limits>
+#include <type_traits>
 
 #define MAX_BLOCKS 65535u
 namespace raft {
@@ -276,18 +283,22 @@ void sum_rows_by_key_large_nkeys_colmajor(const DataIteratorT d_A,
                       d_sums);
 }
 
+/**
+ * Large-nkeys reduction with one thread per element, used for where a warp per row would leave
+ * lanes idle.
+ */
 template <typename DataIteratorT,
           typename KeysIteratorT,
           typename WeightT,
           typename SumsT,
           typename IdxT>
-RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor(const DataIteratorT d_A,
-                                                        IdxT lda,
-                                                        const WeightT* d_weights,
-                                                        KeysIteratorT d_keys,
-                                                        IdxT nrows,
-                                                        IdxT ncols,
-                                                        SumsT* d_sums)
+RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor_per_element(const DataIteratorT d_A,
+                                                                    IdxT lda,
+                                                                    const WeightT* d_weights,
+                                                                    KeysIteratorT d_keys,
+                                                                    IdxT nrows,
+                                                                    IdxT ncols,
+                                                                    SumsT* d_sums)
 {
   IdxT gid = threadIdx.x + (blockDim.x * static_cast<IdxT>(blockIdx.x));
   IdxT j   = gid % ncols;
@@ -299,35 +310,236 @@ RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor(const DataIteratorT d_A,
   raft::myAtomicAdd(&d_sums[j + ncols * l], val);
 }
 
+/**
+ * Large-nkeys reduction with one warp per row: the key and weight are loaded once per row and the
+ * warp's lanes add consecutive columns of the row into d_sums[key].
+ * If d_row_ids is not null, the i-th visited row is d_row_ids[i] and its key is d_keys[i].
+ */
 template <typename DataIteratorT,
           typename KeysIteratorT,
           typename WeightT,
           typename SumsT,
           typename IdxT>
-void sum_rows_by_key_large_nkeys_rowmajor(const DataIteratorT d_A,
-                                          IdxT lda,
-                                          const KeysIteratorT d_keys,
-                                          const WeightT* d_weights,
-                                          IdxT nrows,
-                                          IdxT ncols,
-                                          SumsT* d_sums,
-                                          cudaStream_t st)
+RAFT_KERNEL sum_rows_by_key_large_nkeys_kernel_rowmajor_per_warp(const DataIteratorT d_A,
+                                                                 IdxT lda,
+                                                                 const WeightT* d_weights,
+                                                                 KeysIteratorT d_keys,
+                                                                 const uint32_t* d_row_ids,
+                                                                 IdxT nrows,
+                                                                 IdxT ncols,
+                                                                 SumsT* d_sums)
 {
-  uint32_t block_dim = 128;
-  auto grid_dim      = static_cast<uint32_t>(ceildiv<IdxT>(nrows * ncols, (IdxT)block_dim));
+  // Plain pointers are read through the read-only cache; other iterators (e.g. a transform
+  // iterator converting the input type) are dereferenced as is.
+  constexpr bool kIsPlainPointer =
+    std::is_same_v<DataIteratorT, const SumsT*> || std::is_same_v<DataIteratorT, SumsT*>;
+  const IdxT warp_id = (static_cast<IdxT>(blockIdx.x) * blockDim.x + threadIdx.x) / WarpSize;
+  const IdxT n_warps = static_cast<IdxT>(gridDim.x) * blockDim.x / WarpSize;
+  for (IdxT i = warp_id; i < nrows; i += n_warps) {
+    const IdxT row = d_row_ids != nullptr ? static_cast<IdxT>(d_row_ids[i]) : i;
+    const auto key = static_cast<IdxT>(d_keys[i]);
+    const SumsT w  = d_weights != nullptr ? static_cast<SumsT>(d_weights[row]) : SumsT(1);
+    for (IdxT j = laneId(); j < ncols; j += WarpSize) {
+      SumsT a;
+      if constexpr (kIsPlainPointer) {
+        a = __ldg(d_A + lda * row + j);
+      } else {
+        a = static_cast<SumsT>(d_A[lda * row + j]);
+      }
+      raft::myAtomicAdd(d_sums + ncols * key + j, a * w);
+    }
+  }
+}
+
+template <typename DataIteratorT,
+          typename KeysIteratorT,
+          typename WeightT,
+          typename SumsT,
+          typename IdxT>
+void sum_rows_by_key_large_nkeys_rowmajor_per_warp(const DataIteratorT d_A,
+                                                   IdxT lda,
+                                                   const KeysIteratorT d_keys,
+                                                   const uint32_t* d_row_ids,
+                                                   const WeightT* d_weights,
+                                                   IdxT nrows,
+                                                   IdxT ncols,
+                                                   SumsT* d_sums,
+                                                   cudaStream_t st)
+{
+  constexpr uint32_t block_dim       = 128;
+  constexpr uint32_t warps_per_block = block_dim / WarpSize;
+  const auto grid_dim =
+    static_cast<uint32_t>(std::min<IdxT>(ceildiv<IdxT>(nrows, warps_per_block), IdxT{1} << 30));
   raft::launch_kernel(st,
                       grid_dim,
                       block_dim,
-                      sum_rows_by_key_large_nkeys_kernel_rowmajor,
+                      sum_rows_by_key_large_nkeys_kernel_rowmajor_per_warp,
                       d_A,
                       lda,
                       d_weights,
                       d_keys,
+                      d_row_ids,
                       nrows,
                       ncols,
                       d_sums);
 }
 
+template <typename KeysIteratorT, typename IdxT>
+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<IdxT>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (i < nrows) {
+    keys[i]    = static_cast<uint32_t>(d_keys[i]);
+    row_ids[i] = static_cast<uint32_t>(i);
+  }
+}
+
+/**
+ * Large-nkeys reduction over the rows in key order: the row ids are radix-sorted by key, so the
+ * sums rows being updated at any moment stay in L2 even when all sums do not fit there.
+ * Requires nrows and nkeys to fit in uint32_t; uses 16 bytes per row of temporary memory.
+ */
+template <typename DataIteratorT,
+          typename KeysIteratorT,
+          typename WeightT,
+          typename SumsT,
+          typename IdxT>
+void sum_rows_by_key_large_nkeys_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)
+{
+  rmm::device_uvector<uint32_t> keys(nrows, st), row_ids(nrows, st);
+  rmm::device_uvector<uint32_t> sorted_keys(nrows, st), sorted_row_ids(nrows, st);
+  constexpr uint32_t block_dim = 256;
+  raft::launch_kernel(st,
+                      static_cast<uint32_t>(ceildiv<IdxT>(nrows, IdxT{block_dim})),
+                      block_dim,
+                      sum_rows_by_key_sort_input_kernel,
+                      d_keys,
+                      nrows,
+                      keys.data(),
+                      row_ids.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.data(),
+                                                sorted_keys.data(),
+                                                row_ids.data(),
+                                                sorted_row_ids.data(),
+                                                static_cast<int>(nrows),
+                                                0,
+                                                end_bit,
+                                                st));
+  rmm::device_buffer tmp(tmp_bytes, st);
+  RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(tmp.data(),
+                                                tmp_bytes,
+                                                keys.data(),
+                                                sorted_keys.data(),
+                                                row_ids.data(),
+                                                sorted_row_ids.data(),
+                                                static_cast<int>(nrows),
+                                                0,
+                                                end_bit,
+                                                st));
+  sum_rows_by_key_large_nkeys_rowmajor_per_warp(d_A,
+                                                lda,
+                                                static_cast<const uint32_t*>(sorted_keys.data()),
+                                                sorted_row_ids.data(),
+                                                d_weights,
+                                                nrows,
+                                                ncols,
+                                                d_sums,
+                                                st);
+}
+
+/**
+ * Whether to visit the rows in key order. Sorting pays off when the sums do not fit in L2, so the
+ * atomics of unsorted rows would go to DRAM, when the input is large enough to amortize the sort
+ * and its temporary allocations (16 bytes per row), and when rows are long enough (>= 2 KB) that
+ * the per-row sort cost and the contention of consecutive rows on the same key stay small. The
+ * size thresholds scale with the L2 size; all were measured on A100, H100 and B100.
+ */
+template <typename SumsT, typename IdxT>
+bool sum_rows_by_key_use_sorted(IdxT nrows, IdxT ncols, IdxT nkeys, int l2_bytes, int cc_major)
+{
+  // The sort indexes rows with int and stores keys and row ids as uint32_t.
+  if (static_cast<uint64_t>(nrows) > static_cast<uint64_t>(std::numeric_limits<int>::max()) ||
+      static_cast<uint64_t>(nkeys) > std::numeric_limits<uint32_t>::max()) {
+    return false;
+  }
+  const double row_bytes   = static_cast<double>(ncols) * sizeof(SumsT);
+  const double input_bytes = row_bytes * nrows;
+  const double sums_bytes  = row_bytes * nkeys;
+  // On Ampere, unsorted atomics still mostly hit L2 until the sums exceed the whole L2.
+  const double sums_limit = cc_major >= 9 ? 0.5 * l2_bytes : 1.0 * l2_bytes;
+  return row_bytes >= 2048 && input_bytes >= 40.0 * l2_bytes && sums_bytes > sums_limit;
+}
+
+template <typename DataIteratorT,
+          typename KeysIteratorT,
+          typename WeightT,
+          typename SumsT,
+          typename IdxT>
+void sum_rows_by_key_large_nkeys_rowmajor(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)
+{
+  // Kernel selection is heuristically tunned on A100, H100 and B100.
+  int dev_id, cc_major, l2_bytes;
+  RAFT_CUDA_TRY(cudaGetDevice(&dev_id));
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&cc_major, cudaDevAttrComputeCapabilityMajor, dev_id));
+  RAFT_CUDA_TRY(cudaDeviceGetAttribute(&l2_bytes, cudaDevAttrL2CacheSize, dev_id));
+
+  const IdxT min_cols_per_warp = sizeof(SumsT) > 4 ? 64 : cc_major >= 9 ? 32 : 48;
+  if (ncols < min_cols_per_warp) {
+    // per-element kernel for narrow rows (small ncols), avoiding idle lanes
+    constexpr uint32_t block_dim = 128;
+    const auto grid_dim = static_cast<uint32_t>(ceildiv<IdxT>(nrows * ncols, IdxT{block_dim}));
+    raft::launch_kernel(st,
+                        grid_dim,
+                        block_dim,
+                        sum_rows_by_key_large_nkeys_kernel_rowmajor_per_element,
+                        d_A,
+                        lda,
+                        d_weights,
+                        d_keys,
+                        nrows,
+                        ncols,
+                        d_sums);
+    return;
+  }
+
+  if (sum_rows_by_key_use_sorted<SumsT>(nrows, ncols, nkeys, l2_bytes, cc_major)) {
+    // sorted kernel for highest ncols and nrows by remediating the L2 bottleneck
+    sum_rows_by_key_large_nkeys_rowmajor_sorted(
+      d_A, lda, d_keys, d_weights, nrows, ncols, nkeys, d_sums, st);
+  } else {
+    // per-warp kernel for wider rows until L2 bottleneck hits
+    sum_rows_by_key_large_nkeys_rowmajor_per_warp(
+      d_A, lda, d_keys, static_cast<const uint32_t*>(nullptr), d_weights, nrows, ncols, d_sums, st);
+  }
+}
+
 /**
  * @brief Computes the weighted reduction of matrix rows for each given key
  *
@@ -380,7 +592,8 @@ void reduce_rows_by_key(const DataIteratorT d_A,
     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);
+    sum_rows_by_key_large_nkeys_rowmajor(
+      d_A, lda, d_keys, d_weights, nrows, ncols, nkeys, d_sums, stream);
   }
 }