vectorize_loops: gathers, assignment loops, and widened arguments - #1668
Closed
seantalts wants to merge 9 commits into
Closed
vectorize_loops: gathers, assignment loops, and widened arguments#1668seantalts wants to merge 9 commits into
seantalts wants to merge 9 commits into
Conversation
…zed densities
Loops whose body is a single scalar density statement are rewritten to the
vectorized density call, e.g.
for (n in 1:N) target += normal_lpdf(y[n] | mu[n], sigma);
becomes
target += normal_lpdf(y | mu, sigma);
which shares subcomputations across elements and allocates O(1) autodiff
nodes instead of O(N). Tilde statements are the same MIR shape, and their
proportionality flag is preserved.
Arguments must be scalar invariants or exactly x[n]; x[n] becomes the
slice x[lower:upper], or x alone when the loop range provably spans the
declared size (data and parameters only, since nothing else is immutable).
At least one argument must vary with the loop, invariant containers bail
(the loop broadcasts them, the vectorized call would zip them), and the
rewritten call must typecheck against the Stan Math signatures or the
loop is left untouched. User-defined densities and truncated tildes are
never rewritten.
Enabled at --Oexperimental only for now. Refs stan-dev#356, stan-dev#702.
Declared outer sizes are runtime invariants: whole-variable assignments are size-checked by stan::model::assign (verified empirically for vectors and arrays), and nothing else can resize. So transformed data and all output variables join data and parameters in the trusted-size map, and a transformed data vector spanning the loop range now appears bare. Invariant arguments must also pass cannot_remove_expr: vectorizing drops their evaluation count from N to one, which a side-effecting argument (an _lp call incrementing the target) would observe. Also folds the slice construction into the argument classifier as one pattern match, reuses Expr.Helpers.add_int_index instead of building the indexed expression by hand, and drops the RELEASE-NOTES.txt entry since the changelog is updated at release time.
A matrix-row lane typechecks as a scalar statement but the sliced argument is a matrix and normal_lpdf has no matrix signature, so this exercises the final re-typecheck rejection. The function-body test runs the pass over a loop inside a user-defined function, and the test printer now shows function bodies so the tests can assert they are unchanged.
seantalts
force-pushed
the
feat/vectorize-loops-2
branch
from
August 16, 2026 16:47
ca06812 to
3c04aaf
Compare
seantalts
force-pushed
the
feat/vectorize-loops-2
branch
from
August 16, 2026 16:47
3c04aaf to
08d9842
Compare
seantalts
force-pushed
the
feat/vectorize-loops-2
branch
from
August 16, 2026 16:53
08d9842 to
103aa75
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1668 +/- ##
==========================================
+ Coverage 92.31% 92.35% +0.03%
==========================================
Files 67 67
Lines 9972 10096 +124
==========================================
+ Hits 9206 9324 +118
- Misses 766 772 +6
🚀 New features to boost your workflow:
|
A generated quantity is not in scope in the model block, so a model local may legally reuse its name at another size. Trusting the generated quantity's declared size could then turn a partial-range loop over the local into a whole-variable rewrite. Only data, parameters, and transformed parameters keep trusted sizes. Their names cannot be reused. Regression test from review.
Widening a scalar expression in the loop variable to its whole-range
container form is now recursive. Leaves are x[n], becoming the slice
x[lower:upper] (or x alone when the range provably spans the
declaration), and x[idx[n]], the hierarchical idiom, becoming the
multi-index gather x[idx[lower:upper]]. Interior nodes are plain StanLib
applications of widened children: the call must re-typecheck at a
container return type, since the scalar signature being the one the loop
used makes the container overload its elementwise map by Stan Math
convention, while functions that consume the lane dimension
(dot_product and friends) typecheck at a scalar and are rejected.
Operators retry as their elementwise variant when the plain form does
not fit, and suffixed calls never widen: a vectorized _lpdf value is the
summed lp, not the lanes'.
This unlocks two statement forms at once. Density arguments built from
arithmetic widen, and elementwise assignment loops become sliced
assignments:
for (n in 1:N) mu[n] = alpha + beta * x[n];
becomes
mu[1:N] = alpha + beta * x;
with the written variable barred from the right-hand side and the loop
bounds (which excludes recurrences), and the widened right-hand side
required to match the target's declared type, tracked in an environment
of every declaration (an array target with a vector-widened right-hand
side must not typecheck). Gather nodes are typed as their base
container: Expr.Helpers.infer_type_of_indexed calls a multi-indexed
vector real, and a lying scalar type makes Memory_patterns leave the
gathered base out of the SoA demotion set, generating C++ that does not
compile.
Across posteriordb this claims 11 of 120 models with the other 109
byte-identical; measured per-gradient changes range from 1.58x
(election88_full) to 0.92x (radon_county, where the sliced assignment
costs one extra indexed pass over its scalar loop), with log probability
and gradients agreeing in all of them. Refs stan-dev#356, stan-dev#702.
seantalts
force-pushed
the
feat/vectorize-loops-2
branch
from
August 16, 2026 19:23
103aa75 to
fdd8e9b
Compare
Member
Author
|
I'm just going to close this, the earlier PR was reverted. |
Collaborator
|
It was not reverted, it just required a small fix: #1672 |
Member
|
Apologies for the confusion @seantalts -- it wasn't reverted (and if it had been, it was only going to be for the weekend until I could fix the build issue). This just needs a rebase on master I believe |
Member
Author
|
Ah I see, thanks for fixing! I'm not very familiar with github anymore and used the new automerge button. I find I kinda don't care enough to make the comments pithy and poignant, but I will instruct Fable in that way and see how it goes :P |
seantalts
changed the title
Member
Author
|
Continued in #1678 after a rebase onto master. GitHub can't reopen a PR whose branch was force-pushed after closing. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rebased onto master after #1666 merged. Single commit.
Widening is now recursive.
x[n]becomes the slicex[lower:upper], orxwhen the range provably spans the declaration.x[idx[n]]becomes the multi-index gatherx[idx[lower:upper]]. A StanLib call widens its arguments and is re-typechecked. It has to come back as a container: the loop was using the scalar signature, so the container overload is the elementwise map, and functions like dot_product that typecheck at a scalar are rejected. Operators fall back to their elementwise variant, sox[n] * w[n]becomesx .* w. Suffixed calls are excluded since a vectorized_lpdfvalue is the summed lp.This covers density arguments with arithmetic in them:
and assignment loops:
For assignments the written variable can't appear in the right-hand side or the loop bounds, and the widened right-hand side has to match the target's declared type. Gathers are typed as their base container directly instead of through
infer_type_of_indexed, see #1667.Measurements, same protocol as #1666: 11 of 120 posteriordb models change and the other 109 are byte-identical. election88_full goes 716us to 453us per gradient (1.58x). surgical_model, losscurve_sislob, and pilots are 1.05x to 1.08x. log10earn_height is 1.01x. radon_county is 0.92x, the sliced assignment costs one extra pass over the data compared to its scalar loop, and the fix for that belongs in codegen. Log probability and gradients agree with the pass off in every model. stanc runs a little faster on the changed models and clang compile times don't move.
Submission Checklist
--Oexperimentalbehavior changes)Release notes
The
vectorize_loopsoptimization (enabled at--Oexperimental) now handles indirectly indexed arguments such asa[county[n]], density arguments built from elementwise arithmetic, and elementwise assignment loops.Copyright and Licensing
By submitting this pull request, the copyright holder is agreeing to
license the submitted work under the BSD 3-clause license (https://opensource.org/licenses/BSD-3-Clause)
I used AI.