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

AGENTS.md

@@ -104,22 +104,25 @@ 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`. A constant offset that puts an
-  index below zero, the lowest one read with a constant start or every one with a
-  constant limit, is `ReasonNegativeIndex`: the scalar loop panics and the emitted
-  code would index at a negative constant. A positive one that pushes a constant
-  start or limit past the largest int is rejected as that bound, since the emitted
-  code adds them. 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`. A constant offset that puts an index below zero, the lowest
+  one touched with a constant start or every one with a constant limit, is
+  `ReasonNegativeIndex`: the scalar loop panics and the emitted code would index at a
+  negative constant. A positive one that pushes a constant start or limit past the
+  largest int is rejected as that bound, since the emitted code adds them. 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

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]
+	}
+}

bench/ops_simd.go

@@ -1994,6 +1994,52 @@ func HardSigmoidFloat32s(dst, x []float32) {
 	}
 }
 
+func StoreRowFloat32s(c, x []float32, alpha float32, row, ldc int) {
+	_vcAFloat32s46 := simd.BroadcastFloat32s(alpha)
+	if _loopvecOverlap(c, x, 0+(min(0, row*ldc)), len(x)+(max(0, row*ldc)), 0, len(x)) {
+		for j := range x {
+			c[row*ldc+j] += alpha * x[j]
+		}
+	} else {
+		if len(x) > 0 {
+			_ = c[len(x)+(row*ldc)-1]
+			_ = x[len(x)-1]
+			{
+				_w1 := c[row*ldc:][:len(x):len(x)]
+				_count := len(_w1)
+				_w2 := x[:_count:_count]
+				_i, _lanes := 0, simd.BroadcastFloat32s(0).Len()
+				for ; _i <= _count-4*_lanes; _i += 4 * _lanes {
+					_c1, _c2 := _w1[_i:][:4*_lanes], _w2[_i:][:4*_lanes]
+					_v1 := simd.LoadFloat32s(_c2)
+					_v2 := simd.LoadFloat32s(_c1)
+					_v1.MulAdd(_vcAFloat32s46, _v2).Store(_c1)
+					_v3 := simd.LoadFloat32s(_c2[_lanes:])
+					_v4 := simd.LoadFloat32s(_c1[_lanes:])
+					_v3.MulAdd(_vcAFloat32s46, _v4).Store(_c1[_lanes:])
+					_v5 := simd.LoadFloat32s(_c2[2*_lanes:])
+					_v6 := simd.LoadFloat32s(_c1[2*_lanes:])
+					_v5.MulAdd(_vcAFloat32s46, _v6).Store(_c1[2*_lanes:])
+					_v7 := simd.LoadFloat32s(_c2[3*_lanes:])
+					_v8 := simd.LoadFloat32s(_c1[3*_lanes:])
+					_v7.MulAdd(_vcAFloat32s46, _v8).Store(_c1[3*_lanes:])
+				}
+				for ; _i <= _count-_lanes; _i += _lanes {
+					_v9 := simd.LoadFloat32s(_w2[_i:])
+					_v10 := simd.LoadFloat32s(_w1[_i:])
+					_v9.MulAdd(_vcAFloat32s46, _v10).Store(_w1[_i:])
+				}
+				for _i < _count {
+					_v11, _n := simd.LoadFloat32sPart(_w2[_i:])
+					_v12, _ := simd.LoadFloat32sPart(_w1[_i:])
+					_v11.MulAdd(_vcAFloat32s46, _v12).StorePart(_w1[_i:])
+					_i += _n
+				}
+			}
+		}
+	}
+}
+
 func _loopvecOverlap[T any](a, b []T, aLo, aHi, bLo, bHi int) bool {
 	if aLo >= aHi || bLo >= bHi {
 		return false

internal/loopir/ir.go

@@ -172,9 +172,12 @@ type Loop struct {
 // Stmt is a statement of a loop body: a Store or a Let.
 type Stmt interface{ isStmt() }
 
-// Store writes Val to Dst[i].
+// Store writes Val to Dst[i+Off]. Off is source text for an integer expression
+// that is the same on every iteration, "" for zero, as in Load. Every store to
+// one slice is at the same offset, so no two iterations write the same element.
 type Store struct {
 	Dst *Ref
+	Off string
 	Val Value
 }
 
@@ -213,9 +216,10 @@ type Value interface{ isValue() }
 
 // Load reads Ref[i+Off]. Off is source text for an integer expression that is the
 // same on every iteration, "" for zero. A slice that is also stored in the loop is
-// read only at a constant positive offset: iteration i then reads the element a
-// later iteration writes, and the vector loop, which loads the chunk before any
-// store of the body, sees the same old value.
+// read only at the offset it is stored at, or a constant positive distance past
+// it: iteration i then reads the element a later iteration writes, and the vector
+// loop, which loads the chunk before any store of the body, sees the same old
+// value.
 // FromRange is set when the source is the range value variable, which holds the
 // element as it was when the iteration began: a store to the same slice earlier in
 // the body does not change it.

internal/loopir/lower.go

@@ -616,7 +616,7 @@ func (l *lowerer) body(node ast.Stmt, block *ast.BlockStmt, trip Trip, step int6
 				// dst[i] = src[i] is a copy, emitted as the copy builtin. Copying a
 				// slice onto itself does nothing; at an offset it shifts the elements,
 				// which copy does correctly where the ranges overlap.
-				if v.Ref.Obj == stmt.Dst.Obj && v.Off == "" {
+				if v.Ref.Obj == stmt.Dst.Obj && v.Off == stmt.Off {
 					return nil, ReasonUnsupportedOperand
 				}
 			case *Invariant:
@@ -652,58 +652,75 @@ func (l *lowerer) body(node ast.Stmt, block *ast.BlockStmt, trip Trip, step int6
 		}
 	}
 
-	// A slice that is both stored and read at an offset is a dependence between
-	// iterations: iteration i reads the element iteration i+d writes. The vector
-	// loop runs a whole chunk of iterations statement by statement, so it keeps the
-	// scalar result only when the read is of a later iteration's element (d > 0, in a
-	// forward loop): that iteration has not run yet, so the scalar loop reads the
-	// value from before the loop, and the emitted loop loads it before any store of
-	// the chunk. A read of an earlier iteration's element (d < 0) is a recurrence,
-	// and a distance that is not a constant, or a loop that counts down, has no
-	// order loopvec can check.
+	// Every store to a slice is at one offset, so each iteration writes its own
+	// element. Two offsets would have different iterations write the same element,
+	// in an order the vector loop, which runs one statement over a whole chunk
+	// before the next, does not keep.
 	//
-	// A constant offset that puts an index below zero makes the scalar loop panic
-	// when it runs. Start is the lowest index the loop visits, whichever way it
-	// counts, so a constant start puts the lowest index read at Start+off, and a
-	// constant limit puts the highest at Limit-1+off. Such a loop is left alone,
-	// since the vector loop would index or slice at that negative constant, which
-	// does not compile. Start is never negative and Limit never below one (see
-	// rangeReason), so only a negative offset can do this, and it is at least
-	// -maxOffset, so negating it cannot overflow. A positive offset that pushes a
-	// constant start or limit past the largest int would not compile either, as the
-	// emitted code adds them.
-	stored := map[types.Object]bool{}
+	// A slice that is both stored and read at another offset is a dependence
+	// between iterations: iteration i reads the element iteration i+d writes, d
+	// the read offset less the store offset. The vector loop runs a whole chunk of
+	// iterations statement by statement, so it keeps the scalar result only when
+	// the read is of a later iteration's element (d > 0, in a forward loop): that
+	// iteration has not run yet, so the scalar loop reads the value from before the
+	// loop, and the emitted loop loads it before any store of the chunk. A read of
+	// an earlier iteration's element (d < 0) is a recurrence, and a distance that
+	// is not a constant, or a loop that counts down, has no order loopvec can check.
+	stored := map[types.Object]string{}
 	for _, stmt := range body {
 		if store, ok := stmt.(Store); ok {
-			stored[store.Dst.Obj] = true
+			if off, ok := stored[store.Dst.Obj]; ok && off != store.Off {
+				return nil, ReasonStoreOffsets
+			}
+			stored[store.Dst.Obj] = store.Off
 		}
 	}
+	// A constant offset that puts an index below zero makes the scalar loop panic
+	// when it reaches that index. Start is the lowest index the loop visits,
+	// whichever way it counts, so a constant start puts the lowest index touched at
+	// Start+off, and a constant limit puts the highest at Limit-1+off. Such a loop
+	// is left alone, since the vector loop would index or slice at that negative
+	// constant, which does not compile. Start is never negative and Limit never
+	// below one (see rangeReason), so only a negative offset can do this, and it is
+	// at least -maxOffset, so negating it cannot overflow. A positive offset that
+	// pushes a constant start or limit past the largest int would not compile
+	// either, as the emitted code adds them.
+	offsetReason := func(offset string) Reason {
+		off, err := strconv.ParseInt(offset, 10, 64)
+		switch {
+		case err != nil:
+			return ""
+		case off < 0 && (trip.StartConst && trip.StartValue < -off || trip.LimitConst && trip.LimitValue <= -off):
+			return ReasonNegativeIndex
+		case off > 0 && trip.StartConst && trip.StartValue > math.MaxInt64-off:
+			return ReasonUnsupportedStart
+		case off > 0 && trip.LimitConst && trip.LimitValue > math.MaxInt64-off:
+			return ReasonUnsupportedBound
+		}
+		return ""
+	}
 	for _, stmt := range body {
+		if store, ok := stmt.(Store); ok {
+			if reason := offsetReason(store.Off); reason != "" {
+				return nil, reason
+			}
+		}
 		for _, leaf := range Leaves(stmtValue(stmt)) {
 			load, ok := leaf.(*Load)
-			if !ok || load.Off == "" {
+			if !ok {
 				continue
 			}
-			distance, err := strconv.ParseInt(load.Off, 10, 64)
-			if stored[load.Ref.Obj] {
+			if off, isStored := stored[load.Ref.Obj]; isStored && load.Off != off {
+				d, ok := distance(load.Off, off)
 				switch {
-				case err != nil || step != 1:
+				case !ok || step != 1:
 					return nil, ReasonOffsetOfStored
-				case distance < 0:
+				case d < 0:
 					return nil, ReasonCarriedDependence
 				}
 			}
-			if err == nil && distance < 0 &&
-				(trip.StartConst && trip.StartValue < -distance || trip.LimitConst && trip.LimitValue <= -distance) {
-				return nil, ReasonNegativeIndex
-			}
-			if err == nil && distance > 0 {
-				switch {
-				case trip.StartConst && trip.StartValue > math.MaxInt64-distance:
-					return nil, ReasonUnsupportedStart
-				case trip.LimitConst && trip.LimitValue > math.MaxInt64-distance:
-					return nil, ReasonUnsupportedBound
-				}
+			if reason := offsetReason(load.Off); reason != "" {
+				return nil, reason
 			}
 		}
 	}
@@ -879,7 +896,7 @@ func (l *lowerer) deadOutsideLoop(v *types.Var) bool {
 	return true
 }
 
-// store lowers one assignment to dst[i].
+// store lowers one assignment to dst[i], or to dst[i+e] for a loop-invariant e.
 func (l *lowerer) store(assign *ast.AssignStmt) (Store, Reason) {
 	var op Op
 	switch assign.Tok {
@@ -892,7 +909,7 @@ func (l *lowerer) store(assign *ast.AssignStmt) (Store, Reason) {
 			return Store{}, ReasonUnsupportedOperand
 		}
 	}
-	dst, ok := l.index(assign.Lhs[0])
+	dst, off, ok := l.indexAt(assign.Lhs[0])
 	if !ok {
 		return Store{}, ReasonUnsupportedDestination
 	}
@@ -908,18 +925,18 @@ func (l *lowerer) store(assign *ast.AssignStmt) (Store, Reason) {
 		if assign.Tok == token.SHR_ASSIGN {
 			shift = token.SHR
 		}
-		val, reason = l.shiftOf(shift, &Load{Ref: dst}, assign.Rhs[0])
+		val, reason = l.shiftOf(shift, &Load{Ref: dst, Off: off}, assign.Rhs[0])
 	default:
 		// dst[i] op= rhs is dst[i] = dst[i] op rhs.
 		var rhs Value
 		if rhs, reason = l.value(assign.Rhs[0]); reason == "" {
-			val = &Binary{Op: op, X: &Load{Ref: dst}, Y: rhs}
+			val = &Binary{Op: op, X: &Load{Ref: dst, Off: off}, Y: rhs}
 		}
 	}
 	if reason != "" {
 		return Store{}, reason
 	}
-	return Store{Dst: dst, Val: val}, ""
+	return Store{Dst: dst, Off: off, Val: val}, ""
 }
 
 // loadKey identifies the element a load reads: the range value variable reads the
@@ -1060,7 +1077,7 @@ func (l *lowerer) ifStore(s *ast.IfStmt) (Store, Reason) {
 	}
 	sel := store.Val.(*Select)
 	known := loadKeys(sel.Cond)
-	known[loadKey{store.Dst.Obj, ""}] = true
+	known[loadKey{store.Dst.Obj, store.Off}] = true
 	for _, leaf := range Leaves(sel) {
 		if load, ok := leaf.(*Load); ok && !known[loadKey{load.Ref.Obj, load.Off}] {
 			return Store{}, ReasonConditionalLoad
@@ -1087,16 +1104,16 @@ func (l *lowerer) ifChain(s *ast.IfStmt) (Store, Reason) {
 	case nil:
 		// No else: the element keeps its value. The vector loop stores it back, which
 		// is only unobservable when the condition already reads the element.
-		if !loadKeys(cond)[loadKey{then.Dst.Obj, ""}] {
+		if !loadKeys(cond)[loadKey{then.Dst.Obj, then.Off}] {
 			return Store{}, ReasonConditionalStore
 		}
-		other = &Load{Ref: then.Dst}
+		other = &Load{Ref: then.Dst, Off: then.Off}
 	case *ast.BlockStmt:
 		store, reason := l.branch(e)
 		if reason != "" {
 			return Store{}, reason
 		}
-		if store.Dst.Obj != then.Dst.Obj {
+		if store.Dst.Obj != then.Dst.Obj || store.Off != then.Off {
 			return Store{}, ReasonUnsupportedIf
 		}
 		other = store.Val
@@ -1105,12 +1122,12 @@ func (l *lowerer) ifChain(s *ast.IfStmt) (Store, Reason) {
 		if reason != "" {
 			return Store{}, reason
 		}
-		if store.Dst.Obj != then.Dst.Obj {
+		if store.Dst.Obj != then.Dst.Obj || store.Off != then.Off {
 			return Store{}, ReasonUnsupportedIf
 		}
 		other = store.Val
 	}
-	return Store{Dst: then.Dst, Val: &Select{Cond: cond, Then: then.Val, Else: other}}, ""
+	return Store{Dst: then.Dst, Off: then.Off, Val: &Select{Cond: cond, Then: then.Val, Else: other}}, ""
 }
 
 // branch lowers the body of an if: one assignment to an element.
@@ -1396,10 +1413,22 @@ func (l *lowerer) load(expr ast.Expr) (*Load, bool) {
 	return &Load{Ref: ref, Off: off}, true
 }
 
-// index matches slice[i] where slice is an identifier and i the loop index.
-func (l *lowerer) index(expr ast.Expr) (*Ref, bool) {
-	ref, off, ok := l.indexAt(expr)
-	return ref, ok && off == ""
+// distance is the constant read offset less the constant store offset, each ""
+// for zero; ok is false when either is not a constant.
+func distance(read, store string) (d int64, ok bool) {
+	r, s := int64(0), int64(0)
+	var err error
+	if read != "" {
+		if r, err = strconv.ParseInt(read, 10, 64); err != nil {
+			return 0, false
+		}
+	}
+	if store != "" {
+		if s, err = strconv.ParseInt(store, 10, 64); err != nil {
+			return 0, false
+		}
+	}
+	return r - s, true
 }
 
 // maxOffset bounds a constant offset, so that adding it to a limit cannot