Keep the sparse Jensen-Shannon post-processing in the input precision

[maxwbuckley]

Problem

The final step of the sparse Jensen-Shannon distance is a raft::linalg::map over the whole a_nrows x b_nrows output:

[=] __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 a multiply, an fp64 sqrt is not one instruction — it is a Newton-Raphson refinement sequence. The SASS emitted for the float instantiation of the map kernel was:

DMUL        R14, R4, 0.5      <-- the literal
MUFU.RSQ64H R13, R15
DMUL        R16, R12, R12
DFMA        R16, R14, -R16, 1
DFMA        R18, R16, R18, 0.5
DMUL        R16, R12, R16
DFMA        R16, R18, R16, R12
DMUL        R18, R14, R16
DFMA        R22, R18, -R18, R14
DFMA        R4,  R22, R20, R18

Ten fp64 instructions per output element, each running at 1/64 the fp32 rate on consumer parts.

Fix

Write the constant as value_t(0.5) so overload resolution picks the float path. The same region becomes:

FMUL     R0, R4, 0.5
MUFU.RSQ R7, R0
FMUL.FTZ R5, R0, R7
FMUL.FTZ R7, R7, 0.5
FFMA     R0, -R5, R5, R0
FFMA     R5, R0, R7, R5

Numerics

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 individual results move by at most an ulp.

Measurements

RTX 5090 (sm_120a), CUDA 13.2, random CSR inputs, median of 5 runs, libcuvs.so swapped between the two builds and interleaved:

rows x cols, nnz/row	map kernel	speedup	full pairwise_distance	speedup
4096 x 4096, 32	263 -> 25 us	10.4x	6.40 -> 6.16 ms	1.04x
8192 x 4096, 32	1046 -> 319 us	3.28x	25.25 -> 24.54 ms	1.03x
16384 x 4096, 8	4168 -> 1396 us	2.99x	29.16 -> 26.38 ms	1.11x
16384 x 16384, 32	4168 -> 1394 us	2.99x	130.08 -> 127.32 ms	1.02x

Output checksums (mean over the full distance matrix) are identical to six decimal places in all four shapes.

At 16384 rows the map kernel now moves 2.1 GB in 1.39 ms, i.e. it sits at the memory roofline rather than being fp64-throughput-bound — which is why it stops at 3x. The 4096 case reaches 10.4x because its output fits in L2.

End-to-end gains are smaller because the balanced COO SpMV that produces the map's input is 85-95% of the call.

Same class of issue as #2531.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxwbuckley]

@lowener thank you :)