Feature/opencl deserializer

[SteveBronder]

Submission Checklist

• Run unit tests: ./runTests.py src/test/unit
• Run cpplint: make cpplint
• Declare copyright holder and open-source license: see below

Summary

This PR adds a deserializer that works with matrix_cl and var_value<matrix_cl> types. This will allow Stan models to have log_prob methods which accept a gpu matrix type as the serialized input and then return back the gradient on the GPU.

The main tricky part of this pull request is the gpu's min alignment of subbuffers for OpenCL. In OpenCL, if you have a buffer on the gpu then you are allowed to make sub buffers out of that buffer as long as the sub buffers are aligned on the correct boundary. This is usually 128 bytes, but can vary from device to device.

We call log_prob with a vector that is a serialized representation of the parameters of the stan model. When we move that vector over to the GPU we need add padding to the GPU buffer such that when we make the sub buffers they have the correct alignment. So this PR is a mix of the deserializer class we need in the stan model the stan compiler generates and utilities for making sure the memory on the GPU is padded correctly.

Another alternative design would be to just have a vector of matrix_cl buffers and have the log_prob signature take in some sort of context type. Since the matrix_cl types will need their own log_prob we can totally rethink how log_prob works. For instance, the signature for the gpu code could be something like the below where we just pass a map with the buffers we need stored in it

virtual void log_prob(context<var_value<matrix_cl<double>>> ctx);

Then internally we would not use a deserializer at all but have code like

auto&& alpha = ctx["alpha"];

I'm very open to this design. It would complicate things in the compiler but simplify the data passing which I think would be nice.

How to Verify

./runTests.py 

Copyright and Licensing

Please list the copyright holder for the work you are submitting (this will be you or your assignee, such as a university or company): Simons Foundation

By submitting this pull request, the copyright holder is agreeing to license the submitted work under the following licenses:

• Code: BSD 3-clause (https://opensource.org/licenses/BSD-3-Clause)
• Documentation: CC-BY 4.0 (https://creativecommons.org/licenses/by/4.0/)

[WardBrian]

auto&& alpha = ctx["alpha"];

In this alternative, would the context type basically be something like our current var_context class? What would the return type of operator[] be here?

I'm also curious how the usages of log_prob would need to be adjusted, e.g. to take a leapfrog step would now be something more than just some vector math, right? It almost sounds like Jax's PyTrees, which are super powerful in general but I imagine a bear to implement.

[SteveBronder]

In this alternative, would the context type basically be something like our current var_context class?

The context class would mostly just be a wrapper around an std::unordered_map.

So in the algorithm we start with a vector that we need to break up into chunks on the GPU. We can either break it up into chunks by the scheme here where we have one giant buffer and then correctly pad each parameter such that we can make sub buffers representing the parameters. That requires knowing the dimensions of each parameter such that we can do the sub buffer padding correctly. The sub buffer padding condition is kind of a bummer because it means that, even if the algorithm we used was totally on the gpu, we would still need to do this chunked copy over to the buffer with the correct padding in it so that we could then make sub buffers.

So the flow is like

(parameter dim info, vector) -> padded buffer -> (subbuffer, subbuffer, subbuffer)

Alternatively, we could have some sort of context which holds a std::unordered_map<std::string, matrix_cl>. We would query the model to give us the parameter names and dimensions, then create the unordered map holding the buffers of the parameters. As I'm typing this out I'm realizing that for std::vector<matrix_cl> parameter types you would still need subbuffers and the whole subbuffer padding scheme. Unless you stored everything as std::unordered_map<std::string, std::vector<matrix_cl>>. Now that I have written out the context scheme I'm not sure I like it as much

What would the return type of operator[] be here?

Now that I have typed it out I think we would need something like context.get<{PARAMETER_TYPE}>("param_name");. I'm not loving that

I'm also curious how the usages of log_prob would need to be adjusted, e.g. to take a leapfrog step would now be something more than just some vector math, right? It almost sounds like Jax's PyTrees, which are super powerful in general but I imagine a bear to implement.

So I think the actual impl of log_prob would not change much. The transfer over to the GPU would require some model querying for dimension information (and also parameter names in the context version). The only thing that would change in the compiler code for log_prob would be the way we "deserialize". The function signature would just be the context instead of the vector.

Either way, I think this would require the model making a new log_prob that is specialized for matrix_cl types

[WardBrian]

Now that I have typed it out I think we would need something like context.get<{PARAMETER_TYPE}>("param_name");. I'm not loving that

I don't hate that, but it would mean that the new log prob would be for a specific concrete context class, since you can't do inheritance on templated members like that. Maybe that's fine (GPU code is going to be pretty specific, anyway), it's just that our other interfaces tend to be more generic (e.g. var_context can be a whole bunch of different impls)

[SteveBronder]

Yeah but now that I'm golfing the context pattern I just don't really see what the pros would be relative to just the big buffer approach we have here. For instance, even if the algorithm was on the gpu and the input vector was already an opencl buffer I don't think the context here provides a lot of benefit over the buffer version.