Port C++ psis code to loo? #260

Activity

1. jgabry commented on Jul 3, 2024

   

   jgabry

   on Jul 3, 2024

   Last edited by jgabry

   Member

   More actions

   Thanks @SteveBronder, just saw this (finally catching up on GitHub issues!). Yeah I think the only reason not to do this would be adding the dependency on Rcpp. I guess, like you said, it all depends on whether the added speed would be a game changer. PSIS smoothing is definitely the biggest bottleneck when running loo(), but I've never tried comparing the R and C++ versions to see how much of a speedup there would be. @avehtari what do you think?
2. avehtari commented on Oct 21, 2026

   

   avehtari

   on Oct 21, 2026

   Member

   More actions

   @SteveBronder @jgabry

   I measured this. I compiled Stan's psis.hpp from develop without changes, with a small shim for the three Stan Math functions it uses, and called psis_weights() per column through Rcpp. Its k values and weights agree with loo::psis() to about 1e-12 (with tied draws, the two assign the smoothed values to the tied draws in a different order).

   The C++ core is 3x to 4x faster than the R code for the same work, for example 0.12 s instead of 0.31 s for 4000 draws and 1000 observations on one laptop core. psis() also normalizes the weights and computes n_eff after the smoothing, so with a C++ core psis() itself would be about 1.6x to 2x faster: 0.33 s instead of 0.52 s in that example. loo() often also computes relative_eff(), which takes longer than psis() (0.65 s in the same example), and then the whole call would be about 1.2x faster.

   So just using psis.hpp, the gain is not big enough to justify a dependency on Rcpp and a compiled package.

   This inspired #413 to change do_psis_i() in R to sort only the tail instead of the whole column, with identical results, and psis() is then 1.1x to 1.4x faster than now, making the benefit of C++ even smaller.