Where would I put code to vectorize for loops?

[SteveBronder]

I'm trying to parse through transform_mir.ml and write something where, for example

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

we can take that code and transform it to

x = exp(y);

if the following conditions are met:

1. It's just one statement in the loop (for now at least to make things simple)
2. Each statement with indexed types only uses the For type's loopvar: i (so no (i - 1) or anything else.
3. The loops lower and upper make a linear integer sequence starting from 1 whose upper is the same size as the indexed types' size
   • In this case x and y must be of size N
4. The function f has a vectorized equivalent which can be called.

I'm having a hard time figuring out where this would happen? Would this be something when going from the AST to MIR? Or is this a thing that would happen in transform_mir? The only close thing I could find like this is whenever we take a DeclVar and create both an Assignment and Decl.

It feels like what I'd like to do is for a list of the statements in log_prob

1. Scan the list of statements that make up log_prob for any For whose lower and upper is a linear integer sequence. Any statements that are not loops just get put into the returned list as is
2. For the loop, check that it's body matches with the assumptions above.
3. If (2) is true, return back an Assignment type where the lhs and rhs will be the inner indexed type and the vectorized function applied to the rhs.

For (3) I think I could do something like I did in #868 where I could just have a query_body function taking in a check for the above conditions and returning back an Optional type that for Some has a tuple of the stuff I need to build the assign.

Has anyone thought of this more and would have some pointers? This would be a very nice optimization with the new matrix type since now all these simple loops become candidates for it. And also I like the idea of not forcing users to have to write vectorized functions, loops are nice and make sense! But it would be cool to let folks write loops and then for simple cases we can make them into their vectorized equivalents. This is also nice for things like brms where it applies a for loop over the link function for sigma since users can supply their own custom link function. But often times they are doing exactly the loop above or something else that has a vectorized equivalent.

For multi statements I think it would get more complicated than I want to tackle right now. Like we would need to scan each statement of the loop satisfying the checks above, then for the ones that pass the check we pull them out as an assignment and for the ones that don't pass we write a new for loop excluding the ops that passed.

[SteveBronder]

Oh and for log_prob could also be extended to the case of multi-indexing

data {
int N;
int K;
int my_idx[K]
}

parameters {
vector[N] y;
}
// stuff..

model {
vector[N] x;
for (i in my_idx) {
  x[i] = exp(y[i]);
}

Which would just become

  x[my_idx] = exp(y[my_idx]);

And if we preserve knowledge of the lower upper this would probably work with the same scheme but only indexing 1:K and things like that

[rok-cesnovar]

I think this should be a part of one of the optimization levels.

[nhuurre]

I'd split this into two parts. First, a new optimizer pass that transforms

for (i in 1:N) {
  VECTORIZABLE_STMT(i);
}

into

array[N] int i = linspaced_int_array(N,1,N);
VECTORIZABLE_STMT(i);

and second, make the partial evaluator simplify the indexing for expressions like x[linspaced_int_array(N,1,N)].
It can already turn x[:] into x but the tricky part is figuring out whether the N in x[1:N] is equal to size(x) or not.
The list of vectorizable functions goes in Stan_math_signatures.ml.

[SteveBronder]

If anyone has a minute to sort something out for me, I wanted to goof around with this and so far think I have everything setup to just do a no-op optimization pass over the statements, but I can't figure out how to plugin in the new optimization step

https://github.com/stan-dev/stanc3/compare/master...SteveBronder:feature/loop-reducer?expand=1

When I run make I'm getting an error

make
dune build src/stanc/stanc.exe
File "src/analysis_and_optimization/Loop_reduction.ml", line 174, characters 4-13:
Error: The type of this expression,
       ('_weak1 -> '_weak2) ->
       (('_weak3, '_weak4) Stmt.Fixed.t, '_weak1) Program.t ->
       (('_weak3, '_weak4) Stmt.Fixed.t, '_weak2) Program.t,
       contains type variables that cannot be generalized
make: *** [Makefile:2: all] Error 1

Which is from not iterating over the statements in

let eval_prog = Program.map squish_stmt_loops

How do I iterate over those statements?

[nhuurre]

I'm not sure what's causing the problem but a type annotation fixes it

let eval_prog :
       (Expr.Typed.t, Stmt.Located.t) Program.t
    -> (Expr.Typed.t, Stmt.Located.t) Program.t =
  Program.map Fn.id squish_stmt_loops

I think your squish_stmt_loops could be simplified.

let rec squish_stmt_loops Stmt.Fixed.({pattern; meta}) =
  match pattern with
  | Stmt.Fixed.Pattern.For {loopvar; lower; upper; body} ->
      let new_body = zip_for_loop (loopvar, lower, upper, body) in
      Stmt.Fixed.
        { pattern= For {loopvar; lower; upper; body= squish_stmt_loops new_body}
        ; meta }
  | _ -> {pattern= Stmt.Fixed.Pattern.map Fn.id squish_stmt_loops pattern; meta}

[SteveBronder]

Awesome much appreciated! Wonder why it needs that type annotation there?

And yeah the code is a bit of a mess right now. I'm going to get the simple case working and then pretty things up

[SteveBronder]

It can already turn x[:] into x but the tricky part is figuring out whether the N in x[1:N] is equal to size(x) or not.

@seantalts do functions hold sized types or unsized types? It looks like they hold unsized types which is going to make this for loop thing v hard since then we lose size information and I can't check that the upper of the loop is the same size as the inner types of functions

[nhuurre]

I think the types in the MIR are always unsized. But it's probably fine; the vectorized version is always equivalent to a loop

for (i in K:N) y[i] = exp(x[i]);

y[K:N] = exp(x[K:N]);

Looks like even out-of-range access errors are practically identical.
Unless multi-indexing has significant overhead you don't need to worry about size mismatch.

[SteveBronder]

Oh nice! I was doing y = exp(x); for cases of 1:N, but I think in the rvalue() functions we could have a check that the range is the same size