switch to rsqrt

Author: Intron7

docs/release-notes/0.15.1.md

@@ -8,7 +8,7 @@
 ```{rubric} Bug fixes
 ```
 * Fixes `tl.rank_genes_groups` returning NaN/zero `logfoldchanges`/`pvals` with `groups=[subset]` and `reference='rest'` {pr}`651` {smaller}`S Dicks`
-* Fixes float64 precision loss in the sparse Pearson-residual kernels: ``pp.normalize_pearson_residuals`` (and ``pp.highly_variable_genes`` with ``flavor='pearson_residuals'``) on CSR/CSC input divided by the single-precision intrinsic ``sqrtf``, silently capping the float64 path at ~7 significant digits {pr}`658` {smaller}`A Mikaeili`
+* Fixes float64 precision loss in `pp.normalize_pearson_residuals` on CSR/CSC input {pr}`658` {smaller}`A Mikaeili & S Dicks`
 
 ```{rubric} Misc
 ```

src/rapids_singlecell/_cuda/pr/kernels_pr.cuh

@@ -24,7 +24,7 @@ __global__ void sparse_norm_res_csc_kernel(
             ++sparse_idx;
         }
         residuals[res_index] -= mu;
-        residuals[res_index] /= sqrt(mu + mu * mu * inv_theta);
+        residuals[res_index] *= rsqrt(mu + mu * mu * inv_theta);
         // clamp to [-clip, clip]
         if (residuals[res_index] < -clip) residuals[res_index] = -clip;
         if (residuals[res_index] > clip) residuals[res_index] = clip;
@@ -53,7 +53,7 @@ __global__ void sparse_norm_res_csr_kernel(
             ++sparse_idx;
         }
         residuals[res_index] -= mu;
-        residuals[res_index] /= sqrt(mu + mu * mu * inv_theta);
+        residuals[res_index] *= rsqrt(mu + mu * mu * inv_theta);
 
         if (residuals[res_index] < -clip) residuals[res_index] = -clip;
         if (residuals[res_index] > clip) residuals[res_index] = clip;