[shsahiti]

Faster CUDA parallel_for: host-computed counts and an xpu kernel

Summary

Two changes to the CUDA path of parallel_for (and the CUDA parallel_reduce_sum index math):

1. Precompute range counts on the host. itr_index used to recompute ceiling_div(end - begin, step) per dimension in every thread. The new itr_counts computes them once per launch and passes them through parallel_for_function and reduction_contribution. This removes one 64-bit divide per dimension per element, along with the device-side printf/trap from ceiling_div's zero check.
2. Launch with an xpu kernel instead of cub::DeviceFor::Bulk. CUB's for_each agent stores a copy of the functor (agent_t{tile_base, op}). When the functor is indexed with a runtime value, every thread copies the whole functor into local memory. The new parallel_for_kernel takes the operation as a const kernel parameter and calls a const operator(), so nothing is copied.

Results (RTX 4060 Laptop, ratio = xpu time / hand-written raw CUDA, 1.00 = no overhead)

Range counts, parallel_for copy:

	before	after
2D, 2M elements (L2)	1.80x	1.42x
3D, 2M elements (L2)	3.10x	2.23x
3D, 64M elements (DRAM)	1.38x	1.00x

Per-element divides: 1.5 → 1.0 (2D), 3 → 2 (3D).

xpu kernel, functor carrying a lookup table (GPU kernel time, Nsight Systems medians):

functor	before	after
2 KB, 2M elements	118x	0.97x
8 KB, 2M elements	449x	0.96x
2 KB, 4K elements	6.2x	1.01x
8 KB, 4K elements	21.6x	1.07x

Host cost per launch over a raw <<<>>>: +0.33 µs → +0.08 µs.

Testing

• ctest: 11/11 on the CUDA and CPU builds, after rebasing onto main (parallel_reduce_sum: drop caller scratch, keep default pool cached #10).
• A standalone correctness check visiting 60 ranges (1D/2D/3D, begin 0/1/4, step 1/2/3/5 with uneven extents). Every expected index was visited exactly once, and nothing outside the range was visited.
• Nsight Systems and Nsight Compute runs with clocks warmed up and 20 warm-up launches per case.

Known limitations / follow-ups

• Large 2D/3D ranges are ~8% slower than the CUB path (DRAM-bound, 64M elements). The grid is capped at one wave, so each thread walks many elements serially and fewer loads are in flight. Using CUB's layout (2 elements per thread, unrolled, full grid) matched CUB exactly in testing.
• The remaining in-cache 2D/3D gap (~1.4x / ~2.5x) is the one linear / count divide per dimension. cuda::fast_mod_div on host-precomputed divisors, plus a 32-bit index path when total <= INT32_MAX, should remove most of it.
• parallel_for still has no stream parameter. It always launches on the legacy default stream, so it waits on other streams' work and is not captured by CUDA Graphs that record on a user stream.

[karl-kes]

LGTM