Fix: Logic for demotion rules around scalar ad types and data matrices

[SteveBronder]

Summary

Fixes #1232 where some mixes of scalars, data matrices, and autodiff matrices would not be promoted to SoA matrices appropriately.

For the following code we would not promote to SoA when we should have

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

but for the following model block it does make SoA matrices correctly

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

This is because of a bug in the logic for detecting whether we should demote expressions when we have a mix of scalars and data matrices.

The difference between the two model blocks above is that, in pseudo mir code, the first has a scalar and data matrix operation inside of the expression while the second, because of the parenthesis, has a scalar and autodiff matrix operation

assign(mu, add(add(alpha, rep_vector(0.0, N)), beta)

assign(mu, add(alpha, add(rep_vector(0.0, N), beta))

The first fails because the optimizer sees a scalar and data matrix operation and fails the whole line. But the second just has a scalar and autodiff matrix operation so it does not fail.

Though we do want the following to fail for trying to promote p. In this code, p is derived from an autodiff scalar and a data matrix, which we can't promote to SoA.

data {
  int J;
  array[J] int n;
  vector[J] x;
  array[J] int y;
  real r;
  real R;
}
transformed data {
  vector[J] threshold_angle = asin((R - r) ./ x);
}
parameters {
  real<lower=0> sigma;
}
model {
  vector[J] p = 2 * Phi(threshold_angle / sigma) - 1;
  y ~ binomial(n, p);
}

I found a bug in the logic where if the optimizer saw any op(autodiff scalar, data matrix) operations before it saw a autodiff matrix it would fail the expression and the left hand side would need to be demoted to a AoS.

This PR fixes that by, for each expression with a autodiff matrix type, it goes through checking if it is derived only by autodiff scalars and data matrices. If so then we do a second pass to check if there are any autodiff matrices in the expression graph for that statement that are derived outside of the expression. If we find an autodiff matrix that is not derived in the expression then we do not force demotion of the left hand side.

Submission Checklist

• [x] Run unit tests
• Documentation
  • [ ] If a user-facing facing change was made, the documentation PR is here:
  • [x] OR, no user-facing changes were made

Release notes

Fix logic for deducing whether mixes of scalars, data matrices, and autodiff matrices can be promoted to SoA

Copyright and Licensing

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the BSD 3-clause license (https://opensource.org/licenses/BSD-3-Clause)

[codecov[bot]]

Codecov Report

Merging #1347 (51d26d1) into master (0f36ee7) will increase coverage by 0.05%. Report is 6 commits behind head on master. The diff coverage is 87.03%.

:exclamation: Current head 51d26d1 differs from pull request most recent head 870a5f4. Consider uploading reports for the commit 870a5f4 to get more accurate results

@@            Coverage Diff             @@
##           master    #1347      +/-   ##
==========================================
+ Coverage   89.32%   89.38%   +0.05%     
==========================================
  Files          64       64              
  Lines       10588    10571      -17     
==========================================
- Hits         9458     9449       -9     
+ Misses       1130     1122       -8     

Files Changed	Coverage Δ
src/analysis_and_optimization/Optimize.ml	93.30% <ø> (ø)
src/analysis_and_optimization/Memory_patterns.ml	89.85% <86.53%> (+1.80%)	:arrow_up:
src/middle/Stan_math_signatures.ml	98.13% <100.00%> (+0.08%)	:arrow_up:

[SteveBronder]

Think this is ready for review!