From f7d74a0c0a354a44a64305c1d8e74c8cca49ec9c Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Mon, 31 Aug 2026 21:50:01 +0200 Subject: [PATCH] Keep the sparse Jensen-Shannon post-processing in the input precision The final step of the sparse Jensen-Shannon distance is [=] __device__(value_t input) { return raft::sqrt(0.5 * input); } `0.5` is a `double`, so for `value_t == float` the multiply promotes the expression and `raft::sqrt` resolves to the fp64 overload. Unlike a stray `double` in an add or multiply, an fp64 sqrt is not one instruction: it is a Newton-Raphson refinement sequence. The generated SASS for the float instantiation of the map kernel was DMUL R14, R4, 0.5 MUFU.RSQ64H R13, R15 DMUL R16, R12, R12 DFMA R16, R14, -R16, 1 ... 6 more DMUL/DFMA ten fp64 instructions per element, which run at 1/64 the fp32 rate on consumer parts. Writing the constant as `value_t(0.5)` picks the float overload and reduces this to `FMUL` + `MUFU.RSQ` + four float ops. `0.5` is exactly representable in binary, so the `double` instantiation is bit-for-bit unchanged. The `float` instantiation now rounds once instead of twice, so results move by at most an ulp. Measured on an RTX 5090 (sm_120a, CUDA 13.2), random CSR inputs, median of 5: rows x cols, nnz/row map kernel full pairwise_distance 4096 x 4096, 32 263 -> 25 us 6.44 -> 6.20 ms 1.04x 8192 x 4096, 32 1046 -> 319 us 25.31 -> 24.57 ms 1.03x 16384 x 4096, 8 4167 -> 1395 us 29.21 -> 26.42 ms 1.11x 16384 x 16384, 32 4169 -> 1394 us 130.30 -> 127.53 ms 1.02x The map kernel itself is 3.0-10.4x faster. At 16384 rows it now moves 2.1 GB in 1.39 ms, i.e. it has gone from fp64-throughput-bound to sitting at the memory roofline; the 4096 case is faster still because its output fits in L2. End-to-end gains are smaller because the balanced COO SpMV that produces the input dominates the call. Output checksums are unchanged to six decimal places across all four shapes. --- cpp/src/distance/detail/sparse/lp_distance.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/distance/detail/sparse/lp_distance.cuh b/cpp/src/distance/detail/sparse/lp_distance.cuh index 38025329b9..f396b647ef 100644 --- a/cpp/src/distance/detail/sparse/lp_distance.cuh +++ b/cpp/src/distance/detail/sparse/lp_distance.cuh @@ -266,7 +266,7 @@ class jensen_shannon_unexpanded_distances_t : public distances_t { raft::linalg::map( config_->handle, raft::make_device_vector_view(out_dists, n), - [=] __device__(value_t input) { return raft::sqrt(0.5 * input); }, + [=] __device__(value_t input) { return raft::sqrt(value_t(0.5) * input); }, raft::make_const_mdspan(raft::make_device_vector_view(out_dists, n))); }