Formalize the loop vectorization code

[SteveBronder]

Note: I am posting this now as I think it is close to done, but I still need to review about 20% of the code (the hardest bits). I wanted to post it now to see if there is interest in this and I think this will allow for a lot more loop vectorization optimizations. Marking as a draft until I've done the final review of the nasty bits and would like another human to read it over.

This PR modifies the loop vectorization optimization to use the dependency analysis graph we have inside of the analysis and optimization directory. Dependency analysis is how loop vectorization is often handled in compilers since most of the analysis depends on whether code motion will mess up a read after write, write after write, etc. dependency.

I had claude read through how llvm and gcc handle loop vectorization and it turns out they use a dependency analysis that is very similar to what we do in dependency_analysis.ml. So I had claude scan the test folders for gcc and llvm and backport a lot of the examples that each library uses so we could check that we catch the cases that gcc and llvm test for. You can see these in test/integration/good/compiler-optimizations/loop-vectorization/models.

The motivating stan model for this branch was radon_hierarchical_intercept_centered which has two assignments that need lifted out of the loop so that everything can be moved.

Here is a summary of the code

for (n in 1:N) {
  muj[n] = ...; 
  mu[n] = muj[n] + ...; 
  target += normal_lpdf(y[n] | mu[n], sigma),
}

This can be compressed into the following

muj = ...; 
mu = muj + ...; 
target += normal_lpdf(y | mu, sigma),

Before this branch, the vectorize_loops pass in Optimize.ml decided what to hoist out of a for loop with private, name-level bookkeeping: it collected the variables a loop body read and wrote, refused to widen any statement that mentioned a written variable or assigned a read one, and always emitted the residual loop after all statements could be lifted out of the loop. That guard cannot express the read/write assignment that is used in the radon example. In order to handle more complex dependencies like with the radon model we need to use an actual dependency analysis procedure.

This branch replaces the previous heuristic vectorize_loops with the classical framework for vectorization by loop distribution: data dependence analysis followed by pi-block code generation[1]. Loop_dependence.ml classifies each subscript as Invariant, Affine in the loop variable, or Varying, tests pairs of accesses with the Zero Index Variable (ZIV) and strong Single Index Variable (SIV) tests [2] extended with symbolic offsets. It then builds a per-loop dependence graph whose strongly connected components are emitted in topological order and allow for the loop vectorization. The main flow breaks down into the 3 following pieces

1. Acyclic components become whole-container statements
2. Cyclic components stay in a residual loop
3. Typed fusion [3] merges adjacent residual loops

Loop_vectorize.ml is the client that widens statements. The pass is now on at O1, ships a --debug-loop-vectorization report explaining every per-statement decision, and is checked against a corpus of 89 loops translated from the GCC vectorizer tests (gcc.dg/vect/no-vfa-vect-depend-.c, gcc.dg/tree-ssa/ldist-.c), LLVM's LoopAccessAnalysis, LoopDistribute and LoopVectorize/memdep.ll tests, and the TSVC benchmark loops [4], each with pinned MIR and C++.

[1] (Allen and Kennedy, "Automatic Translation of FORTRAN Programs to Vector Form", ACM TOPLAS 9(4), 1987; Kennedy and Allen, Optimizing Compilers for Modern Architectures, ch. 2, 3 and 6)
[2] Goff, Kennedy and Tseng ("Practical Dependence Testing", PLDI 1991)
[3] (Kennedy and Allen §6.2.5)
[4] (Callahan, Dongarra and Levine 1988; Maleki et al. 2011, via UoB-HPC/TSVC_2)

---

ZIV, zero index variable. Neither subscript mentions the loop variable, as in v[k] versus v[3] or v[k + 1] versus v[k + 2]. If the two constant expressions are provably different the accesses are independent. If they are the same expression, both accesses hit the same element every iteration, which the code reports as the confused direction set {<, =, >}. In Loop_dependence.ml this is the Invariant o1, Invariant o2 arm of subscript_dependence: symbolic terms are subtracted and only the constant remainder is compared, so v[k + 1] versus v[k + 2] is independent, while v[k] versus v[m] stays confused because nothing relates k to m.

---

Strong SIV, single index variable with equal coefficients. Both subscripts have the form c * n + offset with the same coefficient c, as in a[n + 1] versus a[n]. Then the two accesses coincide exactly when the iterations differ by d = (offset1 - offset2) / c, so the test yields a single dependence distance. If the division is not exact the accesses are independent; d = 0 is a loop-independent dependence with direction {=}; d > 0 or d < 0 gives {<} or {>} with the distance attached. This branch extends the test with symbolic offsets: x[n + k] versus x[n + k] has the symbol k cancel to distance 0, whereas a[n] versus a[n + k] leaves a symbolic remainder and is reported confused, which is why TSVC s431 stays sequential while design example 12 vectorizes.

---

Note: In a departure from how we normally write integration tests, test/integration/good/compiler-optimizations/loop-vectorization has expected and models folders that hold C++/mir files per stan file that we test over. imo I think this is easier to read than one giant MIR/C++ file. Though it requires some weird dune-fu.

Submission Checklist

• Run unit tests
• Documentation
  • OR, no user-facing changes were made

Release notes

Formalize loop vectorization

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)

[WardBrian]

I think this is a massive code burden for an optimization that works pretty well with the simpler analysis and is entirely the kind of thing a user could/should just write themselves (in contrast to say, SoA deduction).

I'm not saying we definitely shouldn't do it, but it does feel a bit like benchmark chasing on poorly written models

[SteveBronder]

Is the code itself is that massive? 99% of the code here is in the tests.

I think writing code with for loops is way way easier to reason about. So if a user can write a bunch of loops that we then optimize out later that feels like a big W to me.

[WardBrian]

It's still a 2000 line diff if you ignore the tests, but I assume you also want the reviewer to actually review the test output (or else why have it)

[codecov]

Codecov Report

❌ Patch coverage is 80.40816% with 144 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.59%. Comparing base (de7dfd7) to head (7606dd6).
⚠️ Report is 14 commits behind head on master.

Files with missing lines	Patch %	Lines
src/analysis_and_optimization/Loop_vectorize.ml	78.19%	58 Missing ⚠️
src/analysis_and_optimization/Loop_dependence.ml	83.77%	55 Missing ⚠️
src/analysis_and_optimization/Dataflow_types.ml	9.09%	20 Missing ⚠️
...c/analysis_and_optimization/Dependence_analysis.ml	86.27%	7 Missing ⚠️
src/middle/Operator.ml	92.00%	2 Missing ⚠️
src/middle/Stmt.ml	81.81%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1718      +/-   ##
==========================================
- Coverage   92.53%   91.59%   -0.94%     
==========================================
  Files          69       71       +2     
  Lines       10369    10963     +594     
==========================================
+ Hits         9595    10042     +447     
- Misses        774      921     +147     

Files with missing lines	Coverage Δ
src/analysis_and_optimization/Optimize.ml	94.56% <ø> (+0.25%)	⬆️
src/driver/Entry.ml	96.55% <100.00%> (+0.21%)	⬆️
src/driver/Flags.ml	100.00% <100.00%> (ø)
src/stanc/CLI.ml	100.00% <100.00%> (ø)
src/stanc/stanc.ml	87.09% <100.00%> (ø)
src/middle/Operator.ml	91.78% <92.00%> (+0.87%)	⬆️
src/middle/Stmt.ml	80.80% <81.81%> (-9.81%)	⬇️
...c/analysis_and_optimization/Dependence_analysis.ml	95.13% <86.27%> (-4.87%)	⬇️
src/analysis_and_optimization/Dataflow_types.ml	12.00% <9.09%> (-13.00%)	⬇️
src/analysis_and_optimization/Loop_dependence.ml	83.77% <83.77%> (ø)
... and 1 more

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[SteveBronder]

I'm fine and happy to break this up into multiple pull requests. I think part of it also got larger because I included two loop optimizations we currently do not do. I have a list of about 10 other loop optimizations that I think would be useful in Stan but also need the dependency analysis

[nhuurre]

I have a list of about 10 other loop optimizations that I think would be useful in Stan

This PR contains a couple mentions of unpublished(?) design-docs/active/vectorize-loop-fission.md.
I think you should write down your overarching plan for these optimizations somewhere others can read it.

---

This PR modifies the loop vectorization optimization to use the dependency analysis graph we have inside of the analysis and optimization directory.
...
llvm and gcc handle loop vectorization and it turns out they use a dependency analysis that is very similar to what we do in dependency_analysis.ml.

It appears Loop_vectorize.ml does not use Dependence_analysis.ml but the completely disjoint Loop_dependence.ml. Does that mean the existing code was not similar enough to re-use/adapt? And why did you change Dependence_analysis.ml anyway?

---

Dataflow_types.subscript.Affine.coeff is always 1 for vectorizable terms. Tracking non-one values just adds useless complexity. I also do not see why .terms needs factoring instead of just being an expr.

---

--debug-loop-vectorization report explaining every per-statement decision

To what end? Do you expect a situation where a loop fails to vectorize for some trivial reason and the user is encouraged to modify their code? Or is this just for developers to debug the optimization?

[SteveBronder]

I have a list of about 10 other loop optimizations that I think would be useful in Stan

This PR contains a couple mentions of unpublished(?) design-docs/active/vectorize-loop-fission.md.
I think you should write down your overarching plan for these optimizations somewhere others can read it.

Yes sorry I have a little doc here

---

This PR modifies the loop vectorization optimization to use the dependency analysis graph we have inside of the analysis and optimization directory.
...
llvm and gcc handle loop vectorization and it turns out they use a dependency analysis that is very similar to what we do in dependency_analysis.ml.

It appears Loop_vectorize.ml does not use Dependence_analysis.ml but the completely disjoint Loop_dependence.ml. Does that mean the existing code was not similar enough to re-use/adapt? And why did you change Dependence_analysis.ml anyway?

We definitely can reuse Dependence_analysis.ml. The agent was just being stupid. Right now I'm breaking this PR into several smaller PRs where this will be fixed.

---

Dataflow_types.subscript.Affine.coeff is always 1 for vectorizable terms. Tracking non-one values just adds useless complexity. I also do not see why .terms needs factoring instead of just being an expr.

I am certainly a novice at this. I asked claude about it and it hoo'd and haw'd in a way that was distrustful anyway so I'll remove these.

---

--debug-loop-vectorization report explaining every per-statement decision

To what end? Do you expect a situation where a loop fails to vectorize for some trivial reason and the user is encouraged to modify their code? Or is this just for developers to debug the optimization?

This is just for developers to read. Though now I think for the optimization steps we should have a --debug-optimize with an optional value after for only printing debugging for a particular section.