rtx-backend-cuda: native index_select / index_add — HetGAT training 2.5x, inference 4.3x

The Backend trait's index_select and index_add have default bodies that
round-trip through host memory. That is correct everywhere and was the only
implementation CUDA had. Graph message passing is made of these two ops, so
dg-gnn's HetGAT paid a device->host->device copy per layer per pass and the
RTX 5060 Ti sat at ~10% utilisation during training.

Design follows rtx-backend-metal's ops::index: gather is one thread per output
element; scatter-add walks the CSR of the adjoint selection matrix S^T, built
host-side by counting sort, so it needs NO atomics and is deterministic with
duplicate indices — the training loss is bit-identical to the host reference.
Device index buffers are cached per thread keyed by the exact index list, so a
static graph topology uploads once. Two small NVRTC kernels; no cuSPARSE.

Measured on dg-gnn, Harris 42,955 links, v8 recipe, RTX 5060 Ti:

  training batch 8      9,042 -> 3,562 ms/step   (2.5x)
  inference single p50   55.4 ->  12.9 ms        (4.3x)
  inference batch 8       257 ->    20 ms/scen   (13x; batching helps again)
  GPU utilisation       median 10% -> 21%, p90 17% -> 43%

Tests: gather with repeats, scatter-add with duplicates and untouched rows,
the adjoint identity <S x, y> == <x, S^T y> (what autograd relies on), a
hub-heavy pattern against the host reference, and the range-check panic.
rtx-backend-cuda --features cuda: 60 + 16 passed, 0 failed.

Co-Authored-By: Claude Opus 5 <[email protected]>

crates/core/rtx-tensor/src/cuda_kernels/element_wise.cu

@@ -1337,6 +1337,56 @@ __global__ void transpose_2d_kernel(
     }
 }
 
+// =============================================================================
+// Row Index Kernels (index_select / index_add) — graph message passing
+// =============================================================================
+// Both are one thread per OUTPUT element, so writes are coalesced and there
+// are no atomics: index_add walks a CSR row of the adjoint selection matrix
+// (built host-side by counting sort) and sums in a fixed order, so results
+// are deterministic even with duplicate indices. Same design as the Metal
+// backend's one-hot SpMM, without a sparse library.
+
+// out[i, :] = in[indices[i], :]     out is [n_out_rows x row_len]
+__global__ void index_select_rows_kernel(
+    const float* input,
+    float* output,
+    const unsigned int* indices,
+    size_t n_out_rows,
+    size_t row_len
+) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t numel = n_out_rows * row_len;
+    if (idx >= numel) return;
+    size_t i = idx / row_len;
+    size_t j = idx - i * row_len;
+    output[idx] = input[(size_t)indices[i] * row_len + j];
+}
+
+// out[r, :] = sum_{k in row_ptr[r]..row_ptr[r+1]} in[col_idx[k], :]
+// out is [num_rows x row_len]; (row_ptr, col_idx) is the CSR of S^T where
+// S[i, indices[i]] = 1, i.e. col_idx lists the INPUT rows that land on r.
+__global__ void index_add_csr_kernel(
+    const float* input,
+    float* output,
+    const unsigned int* row_ptr,
+    const unsigned int* col_idx,
+    size_t num_rows,
+    size_t row_len
+) {
+    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t numel = num_rows * row_len;
+    if (idx >= numel) return;
+    size_t r = idx / row_len;
+    size_t j = idx - r * row_len;
+    float acc = 0.0f;
+    unsigned int start = row_ptr[r];
+    unsigned int end = row_ptr[r + 1];
+    for (unsigned int k = start; k < end; ++k) {
+        acc += input[(size_t)col_idx[k] * row_len + j];
+    }
+    output[idx] = acc;
+}
+
 // Generic permutation kernel for arbitrary dimension swaps
 // input_strides and output_strides are arrays of size ndim
 __global__ void permute_kernel(