How to add static matrix?

[SteveBronder]

Description

@bob-carpenter has a nice discourse thread discussing how a static_matrix type would work in the stan language. This idea has come up again now for efficient autodiff with Eigen #1785 and I think is the path forward to use the OpenCL backend for autodiff #1639

My main question is how we would go about implementing this. Should we make a stan Matrix type that for doubles just inherits from Eigen and for var is two Eigen matrices? Something like

// pseudocode
class Matrix<double> : public Eigen::Matrix<double, -1, -1> {
// stuff
}

class Matrix<var> {
  Eigen::Matrix<double, -1, -1> val_;
  Eigen::Matrix<double, -1, -1> adj_;
// stuff
}

One other fun idea with this is to add an ExecutionPolicy template argument sort of like in std parallel. This would be an enum with values like sequential, parallel_opencl, parallel_tbb etc. and could allow for pretty simple swaps of execution backends.

I'm open to any suggestions!

Current Version:

v3.1.0

[t4c1]

I think the first suggestion could be simplified a bit by writing a specialization for Eigen::Matrix<var>. Also we can make non-static version with second member Eigen::Matrix<vari*, -1, -1> adj_;. This is not as efficient as what you suggested, but can be modified.

The second suggestion (ExecutionPolicy) seems very complex.

[bob-carpenter]

Why not just something like the following, which would work for matrices, vectors, and row vectors?

template <typename T, int R, int C>
class static_matrix {
  Eigen::Matrix<T, R, C> val_;
  Eigen::Matrix<T, R, C> adj_;
};

For the dynamic case, both members need to be Matrix<vari*, ...> because both would change under an assignment.

[SteveBronder]

I think the first suggestion could be simplified a bit by writing a specialization for Eigen::Matrix. Also we can make non-static version with second member Eigen::Matrix<vari*, -1, -1> adj_;. This is not as efficient as what you suggested, but can be moddified.

I'm not sure what you have in mind here can you elaborate?

The second suggestion (ExecutionPolicy) seems very complex.

Maybe I need to think about it

Why not just something like the following, which would work for matrices, vectors, and row vectors?

Yeah mine was just pseudocode

For the dynamic case, both members need to be Matrix<vari*, ...> because both would change under an assignment.

I thought for static matrices we wanted to not allow assignment after the first declaration? In my mind this would never be used by users, but the compiler would do an optimization sweep and if a matrix was never assigned after first initialization it would change it from a normal matrix to a static matrix

[bob-carpenter]

I thought for static matrices we wanted to not allow assignment after the first declaration?

That's for the dynamic case, not the static case. We do.

I'm in favor of enabling users to write efficient code rather than hoping the compiler can sort it out. But this might be a case where we can guarantee the compiler finds all the relevant optimizations.

Also, I was wrong about need two vari* --- just one of them has the value and the adjoint in it. That's essentially what we do with var, which is just a pointer to implementation (PIMPL) of the vari.

[SteveBronder]

Was going to make a separate issue, but in a side email Bob brought up a point that if we templated var and vari ala

template <typename T>
struct var {
  vari<T>* vi_;
}

template <typename T>
struct vari {
  T val_;
  T adj_;
}

This would work nicely for static matrices via var<Eigen::Matrix<double, -1, -1>>. One of the reasons I've harked and barked about C++17 is that making this change now would break essentially everything, but with c++17 templates can be deduced from the constructor so having a template on var would be fully backwards compatible

The example below is a little more far out / groovy, but we can also pass through our vari classes as a template. I think that would let us do some pretty fun things

https://godbolt.org/z/rZJJld

[SteveBronder]

Hmm,

var<matrix_cl<double>>

lol, that seems pretty cool to me!

[t4c1]

I guess there will be some rough edgess to smooth. But this doees look like the simplest way forward.

[SteveBronder]

@t4c1 do you think our first step here is just to make the var and vari templates and generically just do var<double> / vari<double> everywhere? Then once we do that we can slowly unwind everything to actually take template types

[SteveBronder]

Also I'm going to bring this up in the stan meeting tmrw, maybe we could also host an impromptu stan math meeting about it

[bob-carpenter]

make the var and vari templates and generically just do var / vari everywhere

I'm not sure what you mean. What would the var and vari classes look like?

I don't think we want any backward-compatibility breaking moves from this.

[SteveBronder]

Ack markdown got rid of my brackets, just added them

[SteveBronder]

I don't think we want any backward-compatibility breaking moves from this.

tbh I'm not really sure how we do this without breaking backwards compatibility in some form. We either have to move up to C++17 for the template constructor deduction or have at least var<> / var<double> everywhere

[bob-carpenter]

Then I think we should wait for C++17.

[SteveBronder]

I'm mixed about it, we can discuss it in the meeting tmrw. This is a really nice solution to have pretty much everything use the OpenCL backend as well

[t4c1]

We could maintain backwards compatibility by doing this:

template <typename T>
struct var_template {
  vari<T>* vi_;
}

template <typename T>
struct vari_template {
  T val_;
  T adj_;
}

using var = var_template<double>;
using vari = vari_template<double>;

[SteveBronder]

! Excellent!

[SteveBronder]

Took me about 5 minutes to throw up a basic implementation using the aliases!

https://github.com/stan-dev/math/compare/feature/var-template

So that's cool, we can use the aliases as a shim till we generalize stuff to take in generic var_types

This feels like a thing that should have an RFC associated with it to talk about how the static matrix stuff will work. Looking over our ops I think this solution satisfies a lot of what we want!

[bob-carpenter]

Yes, a design doc would be good. But your port didn't use that shim everywhere and I didn't understand where the template <> is coming from for functions where it's not a specialization (like operator*=).

The fvar could also use unfolding for matrices, but I'd take it one direction at a time.

[SteveBronder]

Lol is this supposed to be legal code?

  fvar<Eigen::Matrix<var, -1, 1>> A(Eigen::Matrix<var, -1, 1>::Zero(10));
  fvar<Eigen::Matrix<var, -1, 1>> B(Eigen::Matrix<var, -1, 1>::Zero(10));
  auto xx = multiply(A, B);

Because that compiles

[bob-carpenter]

We only want to construct fvar out of fvar or double, I think.