Transforming User Code To Stan Math Specialized Functions?

[SteveBronder]

Say a user writes

parameters {
vector[N] B;
matrix[N, N] A;
}
transformed parameters {
  real quaded = B' * A * B;
}

In Stan math we have a specialized function quad_form(matrix, vector) to do this

real quaded = quad_form(A, B);

And if we know that the matrix is symmetric, for instance in the case of

parameters {
vector[N] B;
cov_matrix[N] A;
}
transformed parameters {
  real quaded = B' * A * B;
}

quaded would then become

real quaded = quad_form_sym(A, B);

It would be nice for users if they could just write their math then in the compiler we can check if we have specializations for certain types of statements. There's actually a lot of these, like for example a dot product of a vector with itself can use dot_self(B). fma() and others also being use cases here that would give much more optimized code. Mixing this with #877 we could take something like

for (i in 1:N) {
  a[i] = b[i] * c[i] + d[i].
}

we would spit out

a = fma(b, c, d);

I'm not sure how to go about this though. It feels like what I'd like to do is look over the list of statements in log_prob searching for assigns where if the expression on the right side matches some condition, like I see an elementwise multiply + addition then I make a new assign where the right hand side is the optimized stan math signature. There's definitely some weirds here, like how to do this for expressions of subexpressions. Like how to handle

a = b * c + d + e * f + g;

which can become

a = fma(b, c, d + fma(e, f, g));

The mix of #877 and this stuff here is pretty common inside of brms. Like for a random effects model brms prints the following

  // add more terms to the linear predictor
  for (n in 1:N) {
    mu[n] += r_1_1[J_1[n]] * Z_1_1[n] + r_2_1[J_2[n]] * Z_2_1[n] + r_2_2[J_2[n]] * Z_2_2[n];
  }

It would be a bit more advanced but we could turn the above into something like

  mu += fma(r_1_1[J_1], Z_1_1, fma(r_2_1[J_2], Z_2_1, r_2_2[J_2] * Z_2_2));

[nhuurre]

Your first quad_form example almost works already. Compile the following model with stanc --O and you'll see a quad_form call in the C++.

data { int N; }
parameters {
vector[N] B;
matrix[N, N] A;
}
transformed parameters {
  real quaded = transpose(B) * A * B;
}

This is implemented in the optimization suite

stanc3/src/analysis_and_optimization/Partial_evaluator.ml

Lines 595 to 603 in 8626779

	| ( "Times__"
	, [ { pattern=
	FunApp
	( StanLib "Times__"
	, [ {pattern= FunApp (StanLib "transpose", [b]); _}
	; a ] ); _ }
	; c ] )
	when Expr.Typed.equal b c ->
	FunApp (StanLib "quad_form", [a; b])

It doesn't work with your original example because apparently the operator ' is not the same thing as the function transpose and the above match fails to recognize it.

The partial evaluator also knows how to create fma but does not vectorize the loop. That example optimizes to

for (i in 1:N)
  a[i] = fma(b[i], c[i], d[i]);

There are many ways to regroup b * c + d + e * f + g. Currently the optimizer chooses

a = fma(e, f, fma(b, c, d)) + g;

Extending the partial evaluator to produce quad_form_sym is more difficult because the MIR only tracks the type (matrix, vector, etc) and does not propagate constraints (symmetric, ordered, etc).

[rok-cesnovar]

There are a bunch of those done in Optimization as Niko notes. We really need to finish #549 to create the optimization levels which should make these less "dangerous" to use as opposed to all or nothing.

[SteveBronder]

Oh awesome! Yeah I think we should get #549 up and running, we can run some of the benchmarks in cmdstan-performance-tests with the different optimization levels there. I think at the least we should make an O1 with optimizations we know are cool and good

For the transpose miss for B' maybe in the AST or MIR we could convert B' to be wrapped in a tranpose function. Or we could convert all transpose(A) functions for the Transpose__

[SteveBronder]

Also seeing a couple of little places there we could be a bit fancier. Like right now dot_self() is only caught if we are doing a dot_product(B, B). But we could probably look for B' * B when B is a vector and call dot_self() and do like tcrossprod() when it's two vectors.

Would anyone disagree if I went through there and tidied up a bit? It's a bit hard to follow because a lot of the names of things are like a, b, c, l etc. and it would be a bit easier to see what's going on if we just made those into things like lhs, rhs, inner_rhs, inner_lhs, exprs, expr, etc.

[SteveBronder]

Also that partial evaluator may be able to be cleaned up a good bit with something like the query scheme I made in #868