Matrix comprehensions (`build_matrix` function)

[WardBrian]

This would be a variadic function with signature:
matrix build_matrix(int I, int J, Fn f, T1 arg1,...TN argN) (name up for debate)

where Fn is a function (int, int, T1, ..., TN) => real

It returns a matrix M such that M[i,j] = f(i,j, arg1, ... argN)

The real advantage to this function is that it can do local autodiff, which will improve memory usage/locality and allow us to have the returned matrix be a SoA matrix (cf #2344 (comment)).

Once implemented, build_[row]_vector is the obvious extension.

[syclik]

Sounds awesome. My first thought was isn't this close to "apply" in R?

[WardBrian]

I don't think the function in apply gets access to the indices, but broadly speaking it's a similar idea. It's also close to python's list comprehensions -- in fact, semantically, we basically want build_vector to be equivalent to

def build_vector(f, N, *args):
   return [f(i, *args) for i in range(N)]

[SteveBronder]

So I started toying with this today. I was looking at how integrate_1d andlaplace handle precalculating the gradients. One issue is that both of those functions always return a scalar where we want back a matrix. Since laplace and integrate_1d only have one output we only need one pass of reverse mode to pre calculate the jacobian adjoint product info we need. But this build_matrix needs to return a matrix based on several several inputs. So I tried doing an adjoint method like cvodes_integrator_adjoint.hpp, but running benchmarks I'm finding that to be much slower.

I have my version here that uses a reverse pass inside of the reverse pass to calculate the adjoints we need.

https://github.com/stan-dev/math/compare/feat/build_matrix?expand=1

The graph below shows the slowdown relative to a simple loop. So the simple loop is almost always faster.

[WardBrian]

I’m not necessarily surprised it’s slower than a simple loop in a microbenchmark, it would need to also be compared with the benefit of then using an SoA result in a further computation vs the loop being stuck in AoS

[SteveBronder]

icic, if the goal is just to allow more promotion to SoA then I think we should get rid of all of the adjoint tricks I have and just have a simple for loop version. Then at the end I'll promote it to SoA. That will always be good.

tbh I think a thing more beneficial than this would be a for loop collapser optimization pass in the compiler. We can use the monotonic framework over for loops to detect if we can collapse simple for loops into vectorized instructions

vector[N] x;
// ...
vector[N] y;
for (i in 1:N) {
  y[i] = exp(x[i]);
}

into

vector[N] y = exp(x);

If we want to force promotion to SoA we could also do it through stan-dev/stanc3#1439 with an optimization pass that checks if a vector/matrix after a loop should be promoted to SoA (which would just make a new {name}_soa_tmp__ used in the rest of the program.

[WardBrian]

@SteveBronder I talked to @avehtari about this in the language meeting this morning. He said he was unsurprised that there would not be any speedup if the number of parameters was the same as the size of the matrix.

The use case he is interested in is for something like a gp covariance matrix, which has a smaller number of free parameters going in than the NxM matrix coming out (e.g., the entire matrix could be computed from data, except for the diagonal elements which depend on some parameters)

It would be worth benchmarking some examples that align with this use case more, if possible. I've asked @avehtari to supply some