Skip to content

internal/rewrite: emit whole-vector loops over windows the compiler can prove - #4

Open
thesyncim wants to merge 4 commits into
mauri870:masterfrom
thesyncim:tune-full-vector-loop
Open

thesyncim wants to merge 4 commits into
mauri870:masterfrom
thesyncim:tune-full-vector-loop

Conversation

@thesyncim

@thesyncim thesyncim commented Nov 1, 2026 •

Copy link
Copy Markdown

Stacked on #3; this PR's own changes. The code is in internal/rewrite/emit.go; bench/ops_simd.go and the goldens are regenerated.

This changes the shape of the generated element-wise loop. Across bench/ it's about 2.75x faster than what master generates on an M4 Max, and 27-43% faster on the amd64 runners. Most of the gain comes from getting the compiler to drop every check inside the loop.

Right now each iteration is a LoadPart/StorePart, and each call works out how many elements are left. arm64 has no masked loads, so that becomes a cascade of compares and register moves per vector, about 50 instructions per 4 lanes in AccumRowInt32s. On top of that, the bounds checks for slicing each operand at _i stay in the loop on every target.

The new loop:

_w1 := dst[:len(dst):len(dst)]
_count := len(_w1)
_w2, _w3 := a[:_count:_count], b[:_count:_count]
_i, _lanes := 0, simd.BroadcastFloat32s(0).Len()
for ; _i <= _count-4*_lanes; _i += 4 * _lanes {
	_c1, _c2, _c3 := _w1[_i:][:4*_lanes], _w2[_i:][:4*_lanes], _w3[_i:][:4*_lanes]
	simd.LoadFloat32s(_c2).Add(simd.LoadFloat32s(_c3)).Store(_c1)
	// ...the same at [_lanes:], [2*_lanes:] and [3*_lanes:]
}
for ; _i <= _count-_lanes; _i += _lanes {
	// one vector
}
for _i < _count {
	// one partial vector, as before
}

Every operand is cut to its window up front. Each window's length and capacity are _count, and _count is a len(), so the compiler knows it isn't negative. Given _i <= _count-k*_lanes, prove can then show that every slice taken inside the loop is in bounds, and the main loop comes down to the loads, the operation, the stores, the increment and the branch. Each of those details matters. I measured the alternatives on arm64, and each costs 30-50%: writing _i+_lanes <= limit, slicing the original operands at _i+off, taking the count from an expression like len(dst)-1, or using two-index windows.

The unrolled body runs four copies one after another, each with all the statements in order. So the chunks happen in the same order as they would one vector at a time, and a read ahead of a stored slice still happens before that copy's stores. Reductions keep their own loop but use the same condition. Integer reductions now also get the four accumulators that float reductions already had: an integer vector add takes two cycles on arm64, and SumInt32s gets 3x faster. Integer folds wrap, so regrouping them gives exactly the same result.

Since partial loads now only ever see less than a full vector, the emulated Load*Part from golang/go#82001 (which returns len(s)) gives the right count, and GODEBUG=simd=0 builds produce the scalar result. I updated BUGS.md for that.

Numbers are from bench/, best of 6 interleaved runs, in ns/op on an M4 Max:

n=4096 master this PR clang -O3
AddFloat32s 851 326 190
AccumRowInt32s 1535 316 196
DiffInt32s 1740 325 158
ClampInt32s 1056 286 138
SumInt32s 524 168 115
ReLUFloat32s 755 353 503

The geomean against master is 0.48x at n=64, 0.31x at 4096 and 0.32x at 1M. On the GitHub runners, measured together with #5, it's -32% on an EPYC 9V45 (Zen 5, AVX-512) and -43% on an EPYC 7763 (Zen 3, AVX2). The only kernel slower than master there is FillUint8s/64 (+10%): that's a single 64-lane store, where the setup now costs more than the store does.

For context, I also compiled the same loops as C with clang -O3 -mcpu=apple-m4. Master's code took 4.4x as long as clang's; with this it's 1.58x (1.50x at 4096, 1.16x at 1M). Most of what remains is on the compiler side: Go doesn't pair adjacent 128-bit loads and stores into ldp/stp, and each inlined simd call leaves NOPs in the loop.

On tests: the goldens are regenerated. A new loop_shape.txt runs sizes 0-300, and byte loops past 4×64, through every part of the loop for several body shapes, and both builds must print the same checksums. The bench/ fuzzers ran for 15-30s each and TSVC TestGolden passes through toolexec.

It needs #3, which keeps the new names (_count, _w1, ...) from capturing the function's own, and #2: the windows are cut from the loop's bounds, and #2 rejects a constant bound or offset that puts an index below zero instead of rewriting it.

I started on this because the generated code was slower than scalar for a lot of loops on arm64 when I tried loopvec on goav1.

A read at a negative constant offset in a loop with a constant limit, such
as dst[i] = a[i-20] + 1 in a loop that stops before 16, was rewritten into
code that does not compile: the length check of the vector loop is
_ = a[16-21], a negative constant index. -split wrote that file, and the
toolexec wrapper fell back to the original package.

Such a loop panics as soon as it runs, as does one whose first index is
a negative constant (a[i-1] in a loop from 0). Both are now left alone
with their own reason, the way a constant limit of zero or less already
is. A read whose indexes can all be zero or more (a[i-1] from 1) is
still rewritten.

The checks read the start and limit as text, so a named constant or a
constant expression (const s0 = 0; for i := s0; ...) slipped past them.
A constant start or limit is now spelled as its value, which also lets
the short-loop rule see a named constant trip count: goav1's
for x := range restorationExtraHorz, three iterations, is no longer
rewritten.

The vectorized set of std and gonum does not change.
… capture

The rewritten code adds names that begin with an underscore (_i,
_count, _w1, broadcast vectors declared before the loop, ...) and relies
on predeclared names (len, min, max, copy, uint64, the element types) and
on the simd and unsafe packages. Each can mean something else in the
function, and then the vector loop silently computes something else.
All of these already go wrong on master:

  func F(dst, a []int32, _i int) { for i := 0; i < _i; i++ { ... } }

reads loopvec's _i in the loop condition and never runs; with a
package-level func len(s []int32) int the vector loop takes its window
length from that len; and a reduction written as for n := range n, or
for n := n - 1; n >= 0; n--, finishes with for n := _i; n < n; n++,
which never runs.

A loop that would otherwise be rewritten is now left alone, each with
its own reason, when its function uses an identifier that begins with an
underscore, other than _; when one of the names the code relies on,
looked up from the loop, is declared again around it, in the package, or
by an import; or when the loop declares such a name itself, or a name
its range operand or initial values use. The names the overlap helper
uses (uintptr, int, bool, any, false) are checked at package level,
where loopvec declares it. The lookup needs scopes, so the toolexec
wrapper records them. No loop in std, gonum or
github.com/thesyncim/goav1 stops being vectorized.
fixBlankParams works around a compiler bug that rejects _ parameters in
functions with simd code by renaming them _p0, _p1, ..., but it renamed
them in every function of a rewritten file. In a function that reads a
package-level _p0, the parameter then hid it:

  var _p0 = 5
  func Offset(_ int) int { return _p0 }

became func Offset(_p0 int) int { return _p0 } and returned its
argument. Only the functions loopvec rewrote are renamed now. They
mention no name that begins with an underscore, so the new names cannot
capture anything.
…an prove

Every element-wise loop advanced one partial vector at a time: LoadPart
and StorePart on every iteration, each working out how many elements
remain. On arm64, which has no masked loads, that costs more than the
operation itself (about 50 instructions per 4 lanes of AccumRowInt32s),
and the bounds checks of slicing each operand at _i stay in the loop on
every target.

The loop now cuts every slice it touches to its window before it starts:
the elements from the first iteration at that slice's offset, as many as
the loop runs, with a capacity of as many, indexed from zero. _count is
the first window's length, so the compiler knows it is not negative.
Under conditions of the form _i <= _count-k*_lanes it then proves every
slice the loop takes of a window in bounds and drops all bounds and
length checks from the loop.

The main loop runs four copies of the body per iteration over
consecutive vectors, each copy all the statements in order, so the
chunks run in the same order as with one vector per iteration; each
window is cut to the iteration's four vectors, so every copy addresses
its vector at a constant offset. Whole vectors then run one at a time,
and one partial vector takes what is left. Integer reductions get the
four accumulators float reductions have: an integer vector add takes two
cycles on arm64, where they make SumInt32s three times as fast, and
integer folds wrap, so regrouping them gives the same result.

Because a partial load now only ever sees fewer elements than a vector
holds, the emulated Load*Part of golang/go#82001, which returns len(s),
returns the right count, and GODEBUG=simd=0 builds produce the scalar
result.
@thesyncim
thesyncim force-pushed the tune-full-vector-loop branch from ec80e71 to ab91ac1 Compare November 1, 2026 13:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant