Refactoring and improving constraint codegen

[rybern]

There are some issues with how the constraint codegen works. See #783 and discussion under #675. My goals for this refactor are:

• Simplify the code - it's quite challenging to understand current
• Implement either one or both of @SteveBronder's suggestions in Compile error, possibly a problem for stan-dev/stan #783. I might start with just Command to output supported function signatures #1
• Make the code flexible enough to be able to handle tuples

[rybern]

Here's my understanding of how we currently handle constrains:

1. In Ast_to_Mir, for a constrained variable v, we generate:
   • v_free__ declaration (with the same type as v) (here)
   • nested for-loops to assign each element of either v or v_free__ to an internal function (here), e.g.:
     v = fn_constrain__(v[..], "ub", ..)
2. In Transform_Mir, we generate:
   • read in parameters in genquant, log_prob, transform_inits, and prepare_data (here):
     • when a variable's constrained and unconstrained types are different, create a new variable _in__ and read to that; otherwise read to v (here)
   • Search through every statement and find instances where fn_constrain__ is called on a variable that has a _in__ variant, and change it to use the _in__ version on the RHS (here)
   • nested for-loops to call each element of either v or v_free__ with an internal function (here):
     fn_write_param__(v)
3. In Expression_gen, we re-parse the fn_ calls via string matching and then translate them into actual functions (here), for example:
   fn_constrain__(v, "ub", ..) -> FnConstrain -> constrain_ub(v, ...)

[rybern]

I think the main issue with the current setup is the separation of concerns, in two ways:

• The details of how we handle read and constraints with the Stan math backend are split across those three different parts of the compiler (especially between Ast_to_Mir and Transform_Mir), e.g. we manually create free_ variables in Ast_to_Mir and then again in Transform_Mir
• Reading and constraints should actually be tightly coupled tasks, so they should be done together, rather than constraints being partially generated in Ast_to_Mir and reads generated in Transform_Mir

So it seems like the strategy should be either:

• Move the logic for applying constraints from ast_to_mir into transform_mir
• Move the logic for reading (and writing?) parameters from transform_mir to ast_to_mir. From what I've seen of transform_mir, this would probably be harder. Also, it seems more appropriate for that to be in the "backend" code since it's Stan math-specific

@seantalts @nhuurre it would be great to get your thoughts

[seantalts]

Reading and constraints should actually be tightly coupled tasks, so they should be done together, rather than constraints being partially generated in Ast_to_Mir and reads generated in Transform_Mir

Can you say more about this? I think the MIR should represent code that is portable across a variety of backends. In the ones I can think of (other than Stan Math), reading and constraining are different tasks.

[rybern]

So one relevant issue is #786

If I understand correctly @SteveBronder's new deserializer prefers constrains and reads to happen in one call: stan-dev/stan#3013

@nhuurre gave a good example where we'd like reading and constraining int the same for-loop, here: #675 (comment)

They also seem coupled because reads need to know about _free__ variables

Edit: Sorry, I realize now that you weren't asking how they're coupled in Stan math. Maybe ignore the first two points here, but I think they'd still be coupled by the for-loop nesting and _free__ variables in other backends. In any case, my current thinking is that the level of coupling should be determined by the backend, so Ast_to_mir shouldn't have the logic about for-loops or _free__ variables.

[rybern]

I agree that ideally the MIR should be as platform independent as possible prior to Transform_mir, so I'm leaning toward handing this in Transform_mir

[seantalts]

What does it look like moving it to Transform_mir?

A related opinion of mine is that I am fully in favor of vectorizing everything possible in the MIR now - maybe that makes all of the mentioned issues happy? I was against it originally thinking it wasn't the level we wanted for backends, but I discovered that ~all backends have vectorized ~everything :D

I currently lean against adding a MIR node that is a combined Read/Constrain or Read/Unconstrain, but perhaps with vectorized versions of Read and Constrain we can get everything we want pretty easily?

[SteveBronder]

If I understand correctly @SteveBronder's new deserializer prefers constrains and reads to happen in one call:

Yep!

Are there issues with the free variables used in transform_init()? On the backend we could write a context that has similar functions to the deserializer class for call the equivalent free functions. So the call would look something like

double x = context.get_free_lb("x");

A related opinion of mine is that I am fully in favor of vectorizing everything possible in the MIR now - maybe that makes all of the mentioned issues happy? I was against it originally thinking it wasn't the level we wanted for backends, but I discovered that ~all backends have vectorized ~everything :D

Yeah in general I think it's a good idea for the compiler to dish-out to the backends as much as possible. I think that would let the compiler work on higher level optimizations instead of things the C++ compiler / stan math / eigen already does. V unrelated but it would be nice if the stan compiler could see a loop like

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

and know there's a vectorized version of exp() so it can just call x = exp(y). I'd also guess the one liner would be easier for more advanced compiler optimizations like @VMatthijs 's reverse mode compiler stuff

[rybern]

What does it look like moving it to Transform_mir?

I probably won't have time to fully think this out until Monday-ish, but:
What if we don't add FnConstrains to the MIR prior to Transform_mir? I think all of the necessary information is already added in output_vars, because we're using it in Transform_mir to re-derive the _free__ variables. We're already adding e.g. the FnReads and FnWrites to the appropriate blocks in Transform_mir, so it wouldn't be a huge refactor.

[rybern]

@SteveBronder Could you give example code of reading e.g. an array[5,5] simplex using deserializer? Do you need to loop, or is that what the vecsize argument is for?

[SteveBronder]

@SteveBronder Could you give example code of reading e.g. an array[5,5] simplex using deserializer?

Yes if you check out the comment here where I have a stan program with a few different types in using deserializer

For instance here is for an array of correlation matrices.

Do you need to loop, or is that what the vecsize argument is for?

vecsize!

I don't have simplex there but for something like

parameters {
  simplex[N] xx[K];
}

That would become

// just making alias here so line is easier to read
using stdvec_vec_t = std::vector<Eigen::Matrix<local_scalar_t__, -1, 1>> ;
stdvec_vec_t xx = in__.read_simplex<stdvec_vec_t, jacobian__>(lp__, K, N);

All of the constraints have the signature

read_{constrain_type}<{return_type}, {bool_jacobian}>({constrain_params}, lp__, {top-to-bottom dims});

And by top-to-bottom dims there I mean the outer dimensions are passed as the first arg, then inner args recursively. So an

parameters {
  simplex[N] xx[K, J];
}

becomes

// just making alias here so line is easier to read
using stdvec_stdvec_vec_t = std::vector<std::vector<Eigen::Matrix<local_scalar_t__, -1, 1>>>;
stdvec_stdvec_vec_t xx = in__.read_simplex<stdvec_stdvec_vec_t, jacobian__>(lp__, J, K, N);

On the backend it just recursively calls the std::vector specialization of read_simplex until it needs to call the specialization for the eigen vector. Also the function always takes in lp__ but it only ever does stuff with it if jacobian__ is true

[bbbales2]

stdvec_stdvec_vec_t xx = in__.read_simplex<stdvec_stdvec_vec_t, jacobian__>(lp__, J, K, N);

Should be K, J, N? I think with simplex[N] xx[K, J]; k is the first index (the outer-most std::vector)

[SteveBronder]

Oh shoot yes, sorry that should be flipped on the dims so in to out is (lp__, K, J, N);

[rybern]

I think I'm going to leave the refactors for another day/PR, because the new deserializer interface is even simpler than I planned for. New PR at #856