[FR] reduce_sum does not block SoA

[wds15]

Current Behavior:

Stan models using reduce_sum cannot benefit from the SoA optimisation whenever reduce_sum is used with a data only slicing argument.

Expected Behavior:

If the argument to be sliced for reduce_sum is data only, then SoA should work just fine, given that the shared arguments are simply copied as a whole. Moreover, whenever the slicing argument is an array of parameters, then SoA should also work given that an array ( std::vector) is being chopped into pieces.

Model Code

Model which exhibits the issue

  // likelihood including constants
  if (!prior_only) {
    // reduce_sum is blocking SoA, but it should not since the slice
    // argument is data in this case
    //target += reduce_sum(partial_log_lik_lpmf, seq, grainsize, Y, Xc, b, Intercept, sigma, J_1, Z_1_1, r_1_1);
    // this call instead gives me a full SoA program
    target += partial_log_lik_lpmf(seq| 1, N, Y, Xc, b, Intercept, sigma, J_1, Z_1_1, r_1_1);
  }

Environment:

Stan 2.30.0

Anything else:

Here is the full Stan model, which is a minor modification of a brms generated one. So the goal is to get brms to create reduce_sum programs which can benefit from SoA in addition to within-chain parallelisation.

soa-example-brms-soa_stan.txt

[SteveBronder]

I discussed this in the link below. TL;DR anytime a matrix/vector goes into a UDF we need to run an analysis to see what signatures are valid with SoA/AoS variables. This is very doable but will require a bit of work. I made #1237 to describe what I think needs to happen in the compiler

https://discourse.mc-stan.org/t/30-40-drop-in-sampling-time-using-stanc-o1-optimizations/28347/18?u=stevebronder

[SteveBronder]

@wds15 fyi I update the description of the issue to reflect the problem. I think this is more of a feature request and not a bug as not supporting reduce sum atm is intended behavior

[wds15]

Yeah…looked as a bug to me from a user perspective…but it’s a new feature given where we stand.