stan-dev / stan-dev/stanc3

[FR] Allow SoA for UDFs

Open
#1,237 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature optimization
Dominant language
OCaml
Stars
160
Forks
59
Avg merge
21h 45m
Merged PRs (30d)
26

Description

Is your feature request related to a problem? Please describe.

It would be nice to extend the Struct of Arrays (SoA) framework to support UDFs so that they can be used in reduce_sum and other higher order functions.

Right now if a user calls a higher order function we have to demote every matrix / vector passed to that function to Array of Structs (AoS). This is unfortunate since reduce_sum is very powerful for large independent blocks of data and parameters.

Describe the solution you'd like

I think we can do this by the following

  1. During the SoA optimization pass, when the optimization hits a UDF or a higher order function it starts a sub-call of the SoA optimization for the UDF. It will just return the list of inputs that cannot be SoA and then continue the rest of the larger optimization pass.
  2. At the end of the SoA optimization pass the program runs another pass over the program collecting which matrices are SoA. Then when it comes to a UDF in the program it looks at that call of the UDFs argument memory type (Either SoA or AoS) and appends that set of argument memory types to a list in the UDFs meta record. So now each UDF defined in the functions block knows what combinations of AoS and SoA arguments it needs to generate.
  3. When the program starts printing out the C++, it will go through each UDF's list of memory patterns and generate a signature and body for each

I think the above will work? It sounds like it's only 3 steps but there's a lot of little things to do in all of those.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the SoA optimization pass, UDF and higher-order-function handling, and the C++ printing stage described in the issue. Determine how argument memory patterns are represented in UDF metadata and how generated signatures and bodies are currently selected. Done means UDFs used by reduce_sum and other higher-order functions can retain SoA arguments without demotion to AoS.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, ocaml
Domain
compilers, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.