vectorize_loops refactor - #1683
Merged
Merged
Conversation
Widening 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]], becoming the multi-index gather x[idx[lower:upper]].
Interior nodes are StanLib applications of widened children, re-typechecked
at a container return type. The scalar signature is the one the loop used,
so the container overload is its elementwise map by Stan Math convention,
and requiring a container result rejects functions like dot_product that
consume the lane dimension and typecheck at a scalar. Operators retry as
their elementwise variant. Suffixed calls never widen: a vectorized _lpdf
value is the summed lp, and _rng would collapse N draws.
This covers density arguments built from arithmetic and elementwise
assignment loops:
for (n in 1:N) mu[n] = alpha + beta * x[n];
becomes
mu[1:N] = alpha + beta * x;
Soundness notes, in the order the code checks them:
- Invariant subtrees pass through but must be side-effect-free, since the
rewrite evaluates them once instead of N times.
- The written variable may appear in neither the right-hand side (a
recurrence) nor the bounds (the generated for statement re-evaluates its
condition every iteration).
- The widened right-hand side must equal the target's declared unsized
type: the loop assigned real into real, and an array target with a
vector-widened right-hand side must not typecheck. Declared types and
outer sizes are tracked in an environment threaded through each block,
seeded from data, parameters, and transformed parameters. Generated
quantity names are excluded since they are not in scope in the model
block and a model local may reuse the name at another size.
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). Log probability and
gradients agree in all of them. Refs stan-dev#356, stan-dev#702.
Tests for row_vector lanes and a side-effecting call inside an assignment right-hand side. The EltPow retry could never fire: Pow and EltPow resolve to the same pow signature list, so any application that fails as ^ also fails as .^. The elementwise power test shows plain ^ widening directly.
v[1:N] = rhs becomes v[:] = rhs when the range provably spans the declaration. The colon form reaches the plain three-argument assign in the generated code, the same path hand-written v[:] = rhs takes, so it is exercised for every type mix including SoA right-hand sides into AoS targets. The sliced form cost one extra indexed pass over the data: radon_county measured 0.92x against its own scalar loop before this and 1.08x after, and election88_full moves from 1.56x to 1.62x. An earlier attempt emitted no index at all and was reverted: that statement shape is one the frontend never produces, and it walked untested codegen paths that do not compile for some type mixes. The colon index keeps the rewrite inside shapes reachable from source.
A promotion round on a Mac picked up last-ulp float differences from this machine's libm. These values belong to the CI platform.
Track outer dimensions through lexical scopes, invalidate facts after assignments, and use the assignment's target type directly. Vectorize multiple independent loop statements while retaining loops whose statements interfere, and preserve RNG call multiplicity.
Keep the EltPow case in elt_variant even though ^ already typechecks on containers, so the map stays complete if that changes. Check the scalar lane once with UnsizedType.is_scalar_type instead of guarding every branch of the match. Vectorize statements inside profile blocks. The interference analysis now walks the whole subtree of each body statement, so a write hidden inside a profile is still seen. Drop the singleton SList special case and let list_collapsing handle it. Vectorized statements now report their own source location rather than the whole loop's, which is what multi-statement bodies already did. Fold the option guards into the patterns they belong to and into the returned value, and use Option.Syntax. More edge cases in the codegen fixture: multi-statement bodies, a profile block, a recurrence, an interfering write, an invariant container argument, and a reassigned transformed-data size.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #1683 +/- ##
==========================================
+ Coverage 92.37% 92.50% +0.13%
==========================================
Files 69 69
Lines 10188 10359 +171
==========================================
+ Hits 9411 9583 +172
+ Misses 777 776 -1
🚀 New features to boost your workflow:
|
Member
|
I didn't exactly follow all the review discussion in #1678, but did #1681 get folded in at some point? Just curious where this sits in terms of the overall plan On a related note, I don't think I clocked that this was only at I guess I don't understand the risk behind putting it in O1 directly, either. We could do that immediately following the next release (2wks) if we wanted more time for developer testing |
Collaborator
Author
|
It's As for #1681 no, it is not included here, but offers three changes:
|
nhuurre
commented
Aug 31, 2026
WardBrian
reviewed
Aug 31, 2026
Contributor
|
I think once all of it is in we move it to |
WardBrian
reviewed
Sep 2, 2026
WardBrian
left a comment
WardBrian
left a comment
Member
There was a problem hiding this comment.
Mostly minor comments, the one about the polymorphic variants is the only one preventing me from saying I fully understand this code. The behavior looks good based on the tests however
Comment on lines
+4194
to
+4197
| let%expect_test "vectorize bail: array target with a vector right-hand side" = | ||
| (* The loop assigned real into real, but the widened rhs is a vector and the | ||
| target an array, so the sliced assignment would not typecheck. | ||
| GLM_Binomial_model's generated quantities have this shape. *) |
Member
There was a problem hiding this comment.
future: this special case might be common enough to justify inserting a to_array call?
| target += std_normal_lupdf(y[i, x[i]]); | ||
| w2[:, i] = y[i]; | ||
| w3[i, :] = z[i]; | ||
| target += ((N - 0) * std_normal_lupdf(y[i, i])); |
Member
There was a problem hiding this comment.
As an aside, I'm surprised that x - 0 isn't partial evaluated to x
3 tasks
WardBrian
reviewed
Sep 2, 2026
Comment on lines
+884
to
+888
| An expression widens when every loop-varying node is [x[n]], [x[idx[n]]], or | ||
| a StanLib application of widened children that re-typechecks at a container | ||
| return type, with operators retrying as their elementwise variant. [x[n]] | ||
| becomes [x[lower:upper]], or [x] alone when the range provably spans the | ||
| declaration. Loops that do not match are left unchanged. *) |
Member
There was a problem hiding this comment.
Suggested change
| An expression widens when every loop-varying node is [x[n]], [x[idx[n]]], or | |
| a StanLib application of widened children that re-typechecks at a container | |
| return type, with operators retrying as their elementwise variant. [x[n]] | |
| becomes [x[lower:upper]], or [x] alone when the range provably spans the | |
| declaration. Loops that do not match are left unchanged. *) | |
| *) |
This comment still feels kind of slop-y, and I think the 'provably spans' part isn't true anymore anyway
This was
linked to
issues
Sep 2, 2026
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
4 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.
This is a continuation of / alternatve to #1678 and keeps all tests from there but rewrites
vectorize_loopslogic in a way I think is more maintainable.Although loop vectorization is
--Oexperimentaloptimization, note that this PR also changes--O1by adding index simplifications there.Submission Checklist
Release notes
The vectorize_loops optimization (enabled at --Oexperimental) now handles indirectly indexed arguments such as a[county[n]], density arguments built from elementwise arithmetic, independent statements within one loop, 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)