adds optimizations so that (real, data {vector/matrix}) expressions can use SoA

[SteveBronder]

Note this also moves the for loop compaction from --Oexperimental to --O1

For cases like vector fma(var, data vector, var) the compiler cannot use SoA because none of the inputs are var matrix types. So the internals of fma just see var, data vector, var and then use Eigen::Matrix<var>. This pull request adds an additional forward and reverse pass in the monotone framework used by the SoA optimization routine to detect functions where a scalar could be broadcast to a var_value<Matrix> such that the return is the same value, but with a return type that is a var_value<Matrix> instead of a Matrix<var>. So we can take code like the below which is just var scalars and data vectors and call rep_vector on the first argument to get back a var<Matrix> type

vector[N] test = fma(test1, x, test2);   // test1, test2: parameters; x: data

becomes:

vector[N] test = fma(rep_vector(test1, rows(x)), x, test2);

Funnily enough, I have actually caught claude doing this when I ask it to write the program to make it easier to use SoA.

The functions we can do this with are a small but regularly used:

• fma, beta, lmultiply, add, subtract, elt_multiply, elt_divide, and multiply.

let scalar_broadcast_fns : string String.Map.t =
  String.Map.of_list
    [ ("fma", "fma"); ("beta", "beta"); ("lmultiply", "lmultiply")
    ; ("add", "add"); ("subtract", "subtract"); ("elt_multiply", "elt_multiply")
    ; ("elt_divide", "elt_divide"); ("multiply", "elt_multiply") ]

These are the target functions since all of these functions are elementwise. So whether we do fma(real, data vector, real) or fma(rep_vector(real), data vector, real) we get the same result.

The optimization pass does a speculative rewrite where we look for expressions with above functions and data / scalar patterns when we think we can wrap one var scalar into a rep_vector and get the promotion trick to work. We then do a forward and backward pass to check if promoting the result to var<Matrix> would work. If it will work then we keep the promotion and move on, else we skip promoting that expression.

Performance

Here is the gradient of sum(fma(a, x, b)) and sum(fma(rep_vector(a, N), x, b)) for autodiff scalars a and b and a
data vector x of length N.
The AoS row is the code we emitted before this branch and the SoA row is what we emit now.

N	AoS time	AoS varis	SoA time	SoA varis	Speedup
100	2.3 µs	104	0.6 µs	6	4.0×
1 000	20.4 µs	1 004	3.1 µs	6	6.5×
10 000	81.1 µs	10 004	16.4 µs	6	4.9×
100 000	1 038 µs	100 004	430 µs	6	2.4×

Submission Checklist

• Run unit tests
• Documentation
  • no user-facing changes were made

Release notes

Adds optimization to SoA path for allowing certain expressions with (real, data vector) patterns to use 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)

[WardBrian]

did you mean for this folder to get committed?