internal/loopir, internal/rewrite: vectorize stores at a loop-invariant offset

AGENTS.md

@@ -94,18 +94,21 @@ safety argument stays structural:
   the loop runs at least once (`x` may be a nil pointer the scalar loop never reads), and
   the overlap check and the vector loop use the locals; the fallback is the original loop.
   Arrays, pointers to slices, maps, and the result of a call stay rejected.
-- A store is `ident[i]` where `i` is *exactly* the loop variable (`index`). A load
-  may also be `ident[i+e]`, `ident[e+i]` or `ident[i-e]` for a loop-invariant integer
-  `e` (`indexAt`, `Load.Off`). A slice that is also stored may be read only at a
-  constant positive offset in a forward loop: iteration i then reads what a later
-  iteration writes, so the scalar loop sees the old value and the emitter loads it
-  before any store of the chunk (`statements` hoists those loads). A negative
-  offset is a recurrence (`ReasonCarriedDependence`); a variable distance or a
-  loop that counts down is `ReasonOffsetOfStored`. The length check is
-  against the last index read (`lastIndex`), and offset loads are not capped at the
-  limit. Every other store and load is `ident[i]` (`index`) — `dst[i+1]`, `dst[2*i]` are rejected outright, regardless of
-  whether they're actually dependent. That is what guarantees no two
-  iterations touch the same element.
+- A store or a load is `ident[i]`, `ident[i+e]`, `ident[e+i]` or `ident[i-e]` where
+  `i` is *exactly* the loop variable and `e` a loop-invariant integer (`indexAt`,
+  `Store.Off`, `Load.Off`). Every store to one slice is at the same offset
+  (`ReasonStoreOffsets` otherwise), so each iteration writes its own element. A slice
+  that is also stored may be read at that offset, or at a constant positive distance
+  past it in a forward loop: iteration i then reads what a later iteration writes, so
+  the scalar loop sees the old value and the emitter loads it before any store of the
+  chunk (`statements` hoists those loads). A negative distance is a recurrence
+  (`ReasonCarriedDependence`); a variable distance or a loop that counts down is
+  `ReasonOffsetOfStored`. The length check is against the last index touched
+  (`lastIndex`); offset loads are not capped at the limit, and a store at an offset
+  is capped at the limit plus the offset (`destination`), so it writes nothing the
+  scalar loop does not. `dst[2*i]` and `dst[i][j]` are rejected outright, regardless
+  of whether they're actually dependent. That is what guarantees no two iterations
+  touch the same element.
 - A non-slice operand becomes an `Invariant` only if it is a variable, a
   constant, a conversion or negation of one, or a sum, difference, product or
   bitwise combination of two (`invariant`). It is evaluated
@@ -159,8 +162,14 @@ reversal), no new dependence checking is needed, just a wider clause match.
 
 `internal/rewrite/emit.go` never looks at the original statement's clauses —
 it only uses `plan.Loop.Node.Pos()/.End()` to know what text span to replace,
-and generates a **fixed forward** `for _i := 0; _i < bound; { ... }` template
-driven entirely by the `Plan` (`Root`, `Dst`, `Bound`, ...).
+and generates a **fixed forward** template driven entirely by the `Plan`
+(`Root`, `Dst`, `Bound`, ...): every operand is cut to its window
+(`x[start+off:][:_count]`) before the loop, then whole vectors while
+`_i <= _count-_lanes`, then one partial vector for what is left. Keep both:
+with every window exactly `_count` long and that condition, the compiler
+proves each `window[_i:]` holds a whole vector and drops all bounds and length
+checks from the main loop; slicing the original operands at `_i+off`, or
+writing `_i+_lanes <= _count`, leaves them in (each worth 30-50% on arm64).
 `Induction.Step` records the direction but the emitter ignores it. This means:
 extending lowering to accept a new *loop clause shape* that still produces
 the same kind of `Loop` (same body semantics) needs **zero** emitter changes.

BUGS.md

@@ -12,7 +12,7 @@ gotip, "1.27.1" is the release toolchain loopvec is tested against.
 | [#81845](https://github.com/golang/go/issues/81845) calls left unresolved when inlined into a package initializer | closed as a duplicate of [#80689](https://github.com/golang/go/issues/80689), which is open | loops in package-level `var x = func() {...}()` | not rewritten (the generated initializer fails to link), see `package_initializer_no_rewrite.txt` and `toolexec_package_initializer.txt` |
 | [#81846](https://github.com/golang/go/issues/81846) `Float32s.Neg` rebuilds its sign constant every iteration | open | speed of negation loops | none; the loop is correct, only slower than it should be |
 | [#81847](https://github.com/golang/go/issues/81847) `GOEXPERIMENT=simd` alone slows async preemption on AVX-512 | open | any benchmark comparing a scalar and a simd build | benchmark with `GODEBUG=asyncpreemptoff=1` |
-| [#82001](https://github.com/golang/go/issues/82001) emulated `Load{Float64,Int64,Uint64}sPart` return `len(s)` | open | generated loops under `GODEBUG=simd=0` | none; the loops use the count `Load*Part` returns, as the package documents |
+| [#82001](https://github.com/golang/go/issues/82001) emulated `Load{Float64,Int64,Uint64}sPart` return `len(s)` | open | nothing loopvec generates | the loops use `Load*Part` only for what is left after the last whole vector, where `len(s)` is the count loaded |
 | [#82002](https://github.com/golang/go/issues/82002) `Min` and `Max` on floats differ from Go's for NaN and -0 | open | float `min`/`max` | not rewritten (`float min and max differ from Go's ...`) |
 
 ## #80657: simd in a method
@@ -76,10 +76,12 @@ regressions.
 ## #82001: emulated partial loads
 
 With `GODEBUG=simd=0` the emulated `Load{Float64,Int64,Uint64}sPart` return `len(s)`
-instead of the number of elements loaded. loopvec's loops advance by that count (the loop
-shape the `simd` package documents), so in that mode a copy of ten `float64`s produces
-`[1 2 0 0 0 0 0 0 0 0]` instead of the input. The hardware-backed mode gives the right
-answer.
+instead of the number of elements loaded. A loop that advances by that count on every
+iteration (the loop shape the `simd` package documents) skips elements in that mode: a
+copy of ten `float64`s produces `[1 2 0 0 0 0 0 0 0 0]` instead of the input. loopvec's
+loops load whole vectors while one fits and use `Load*Part` only for what is left, fewer
+elements than a vector holds, where `len(s)` is the number loaded, so they are not
+affected. The hardware-backed mode gives the right answer either way.
 
 ## #82002: float `Min` and `Max`
 

README.md

@@ -45,6 +45,7 @@ operations that lower to **AVX-512/AVX2/NEON** depending on the target CPU:
 | `for i := range x { if x[i] < 0 { x[i] = 0 } }`, `if c { dst[i] = u } else if d { dst[i] = v } else { dst[i] = w }` | compare, then select (`IfElse`): ReLU, leaky ReLU, ReLU6, hard tanh, clamp, sign, step; also when the branches assign a temporary (`t := x[i]+3; if t < 0 { t = 0 } else if t > 6 { t = 6 }`: hard sigmoid, hard swish, cubic SiLU) |
 | `for i := range v.Values { v.Values[i] *= a }`, `dst.f[i] += src.f[i]` | a slice reached through a field (`x.f`, `x.f.g`); the root must not be assigned in the loop, and the field is copied into a local once |
 | `for i := range dst { dst[i] = a[i+1] - a[i] }`, `work[j] += a[row*lda+j]` | read at an offset from the index (stencil, row of a 2-D slice); a slice the loop also writes is read only at a constant positive offset (`a[i] = a[i+1] + b[i]`) |
+| `b[i*ldb+j] = a[i*lda+j]`, `c[row+j] += alpha * x[j]`, `line[off+x] = v` | store at an offset from the index (row of a 2-D slice); every store to a slice is at the same offset, and the slice is read only there or at a constant positive distance past it |
 
 Supported element types: `int8`, `int16`, `int32`, `int64`, `uint8`, `uint16`,
 `uint32`, `uint64`, `float32`, `float64`.

bench/bench_test.go

@@ -451,6 +451,21 @@ func BenchmarkAccumRowInt32s(b *testing.B) {
 	}
 }
 
+// StoreRowFloat32s adds alpha*x into one row of a two-dimensional slice.
+func BenchmarkStoreRowFloat32s(b *testing.B) {
+	for _, n := range benchSizes {
+		b.Run(fmt.Sprintf("%d", n), func(b *testing.B) {
+			x := make([]float32, n)
+			c := make([]float32, 2*n)
+			setBytes(b, n, 4, 3)
+			b.ResetTimer()
+			for b.Loop() {
+				StoreRowFloat32s(c, x, 2, 1, n)
+			}
+		})
+	}
+}
+
 func BenchmarkMinUint16s(b *testing.B) {
 	for _, n := range benchSizes {
 		b.Run(fmt.Sprintf("%d", n), func(b *testing.B) {

bench/ops.go

@@ -332,3 +332,9 @@ func HardSigmoidFloat32s(dst, x []float32) {
 		dst[i] = t / 6
 	}
 }
+
+func StoreRowFloat32s(c, x []float32, alpha float32, row, ldc int) {
+	for j := range x {
+		c[row*ldc+j] += alpha * x[j]
+	}
+}