SoA optimization fail?

[avehtari]

Using --O1 the following code fails to use SoA for beta and mu

data {
  int N;
}
parameters {
  real alpha;
  vector[N] beta;
}
model{
  vector[N] mu = alpha + rep_vector(0.0, N) + beta;
  target += sum(mu);
}

but the following equivalent code does work (showing only the model block)

model{
  vector[N] mu = alpha + (rep_vector(0.0, N) + beta);
  target += sum(mu);
}

but it's not just the order of summing as following works, too

model{
  vector[N] tmp = rep_vector(0.0, N);
  vector[N] mu = alpha + tmp + beta;
  target += sum(mu);
}

@WardBrian commented

I believe the problem is the formulation alpha + rep_vector(0.0, N) + beta leads to the first being a Plus((AutoDiffable UReal) (DataOnly UVector)), which does not contain any eigen autodiff variables so the optimization says oh, SoA can not be used

whereas alpha + (rep_vector(0.0, N) + beta) is a Plus((DataOnly UVector) (AutoDiffable UVector)) which does contain it

tmp is a model-block level variable which means it is always a AutoDiffable

Is there anything to be done to make the first code to use SoA, as figuring out these subtleties can be challenging for users

[WardBrian]

@SteveBronder could you take a look at this? I've tried to nail down where the demotion is happening but have a hard time finding my way around the mem pattern pass

[SteveBronder]

Yes I'll take a look this week!

[SteveBronder]

@avehtari this should be good now if you see any other edge cases plz lmk!!

[avehtari]

Great! Pinging also @paul-buerkner to let him know it's fixed as this came up in brms generated code

[paul-buerkner]

Thank you!