Skip to content

internal/rewrite: make the overlap check small enough to inline #18

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
thesyncim wants to merge 2 commits into from
Open

internal/rewrite: make the overlap check small enough to inline #18

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6be9764
internal/rewrite: make the overlap check small enough to inline
thesyncim Nov 1, 2026
23ff9ab
testdata: update the goldens for the new overlap call
thesyncim Nov 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 25 additions & 31 deletions bench/ops_simd.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
)

func AddFloat32s(dst, a, b []float32) {
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i] + b[i]
}
Expand Down Expand Up @@ -54,7 +54,7 @@ func AddFloat32s(dst, a, b []float32) {
}

func MulFloat32s(dst, a, b []float32) {
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i] * b[i]
}
Expand Down Expand Up @@ -129,7 +129,7 @@ func ScalFloat32s(x []float32, c float32) {

func AxpyFloat32s(dst, a []float32, alpha float32, b []float32) {
_vcAFloat32s3 := simd.BroadcastFloat32s(alpha)
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i]*alpha + b[i]
}
Expand Down Expand Up @@ -176,7 +176,7 @@ func AxpyFloat32s(dst, a []float32, alpha float32, b []float32) {
func MixFloat32s(dst, a, b []float32, alpha, beta float32) {
_vcAFloat32s4 := simd.BroadcastFloat32s(alpha)
_vcBFloat32s4 := simd.BroadcastFloat32s(beta)
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i]*alpha + b[i]*beta
}
Expand Down Expand Up @@ -221,7 +221,7 @@ func MixFloat32s(dst, a, b []float32, alpha, beta float32) {
}

func NegFloat32s(dst, src []float32) {
if _loopvecOverlap(dst, src) {
if _loopvecOverlap(dst, src, 0, len(dst), 0, len(src)) {
for i := range dst {
dst[i] = -src[i]
}
Expand Down Expand Up @@ -259,7 +259,7 @@ func NegFloat32s(dst, src []float32) {
}

func DivFloat32s(dst, a, b []float32) {
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i] / b[i]
}
Expand Down Expand Up @@ -305,7 +305,7 @@ func DivFloat32s(dst, a, b []float32) {

func DaxpyFloat32s(dst, a []float32, alpha float32) {
_vcAFloat32s7 := simd.BroadcastFloat32s(alpha)
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] += a[i] * alpha
}
Expand Down Expand Up @@ -396,7 +396,7 @@ func FillUint8s(dst []byte) {

func ReverseIncFloat32s(dst, src []float32) {
_vcFloat32s10 := simd.BroadcastFloat32s(1)
if _loopvecOverlap(dst, src) {
if _loopvecOverlap(dst, src, 0, len(dst), 0, len(src)) {
for i := len(dst) - 1; i >= 0; i-- {
dst[i] = src[i] + 1
}
Expand Down Expand Up @@ -434,7 +434,7 @@ func ReverseIncFloat32s(dst, src []float32) {
}

func CopyFloat32s(dst, src []float32) {
if _loopvecOverlap(dst, src) {
if _loopvecOverlap(dst, src, 0, len(dst), 0, len(src)) {
for i := range dst {
dst[i] = src[i]
}
Expand All @@ -447,7 +447,7 @@ func CopyFloat32s(dst, src []float32) {
}

func AndNotUint64s(dst, a, b []uint64) {
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = a[i] &^ b[i]
}
Expand Down Expand Up @@ -492,7 +492,7 @@ func AndNotUint64s(dst, a, b []uint64) {
}

func AbsFloat32s(dst, a []float32) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = float32(math.Abs(float64(a[i])))
}
Expand Down Expand Up @@ -530,7 +530,7 @@ func AbsFloat32s(dst, a []float32) {
}

func SqrtFloat32s(dst, a []float32) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = float32(math.Sqrt(float64(a[i])))
}
Expand Down Expand Up @@ -570,7 +570,7 @@ func SqrtFloat32s(dst, a []float32) {
func ClampInt32s(dst, a []int32, lo, hi int32) {
_vcAInt32s15 := simd.BroadcastInt32s(lo)
_vcBInt32s15 := simd.BroadcastInt32s(hi)
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = min(max(a[i], lo), hi)
}
Expand Down Expand Up @@ -608,7 +608,7 @@ func ClampInt32s(dst, a []int32, lo, hi int32) {
}

func ShrInt32s(dst, a []int32, n uint) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = a[i] >> n
}
Expand Down Expand Up @@ -646,7 +646,7 @@ func ShrInt32s(dst, a []int32, n uint) {
}

func ShlUint32s(dst, a []uint32, n uint) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = a[i] << n
}
Expand Down Expand Up @@ -818,7 +818,7 @@ func ChainInt32s(a, b, c, e []int32) {
func WideInt32s(dst, a, b, c []int32, k, m int32) {
_vcInt32s20_1 := simd.BroadcastInt32s(k)
_vcInt32s20_2 := simd.BroadcastInt32s(k * m)
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) || _loopvecOverlap(dst, c) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) || _loopvecOverlap(dst, c, 0, len(dst), 0, len(c)) {
for i := range dst {
dst[i] = a[i]*b[i] + c[i]*k + (k*m - a[i])
}
Expand Down Expand Up @@ -1405,7 +1405,7 @@ func SmoothInt32s(dst, a []int32) {

func ShiftInt32s(dst, a []int32, off int) {
_vcInt32s33 := simd.BroadcastInt32s(3)
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = a[i+off] * 3
}
Expand Down Expand Up @@ -1443,7 +1443,7 @@ func ShiftInt32s(dst, a []int32, off int) {
}

func AccumRowInt32s(work, a []int32, row, lda int) {
if _loopvecOverlap(work, a) {
if _loopvecOverlap(work, a, 0, len(work), 0, len(a)) {
for j := range work {
work[j] += a[row*lda+j]
}
Expand Down Expand Up @@ -1644,7 +1644,7 @@ func ReLUFloat32s(x []float32) {
func LeakyReLUFloat32s(dst, x []float32, alpha float32) {
_vcAFloat32s38 := simd.BroadcastFloat32s(0)
_vcBFloat32s38 := simd.BroadcastFloat32s(alpha)
if _loopvecOverlap(dst, x) {
if _loopvecOverlap(dst, x, 0, len(dst), 0, len(x)) {
for i := range dst {
if x[i] > 0 {
dst[i] = x[i]
Expand Down Expand Up @@ -1719,7 +1719,7 @@ func SignFloat32s(dst, x []float32) {
_vcFloat32s40_1 := simd.BroadcastFloat32s(0)
_vcFloat32s40_2 := simd.BroadcastFloat32s(1)
_vcFloat32s40_4 := simd.BroadcastFloat32s(-1)
if _loopvecOverlap(dst, x) {
if _loopvecOverlap(dst, x, 0, len(dst), 0, len(x)) {
for i := range dst {
if x[i] > 0 {
dst[i] = 1
Expand Down Expand Up @@ -1825,7 +1825,7 @@ func ScaleFieldFloat32s(v *Vec, a float32) {
func AddFieldFloat32s(dst, a *Vec) {
if len(dst.Values) > 0 {
_sdst_Values, _sa_Values := dst.Values, a.Values
if _loopvecOverlap(_sdst_Values, _sa_Values) {
if _loopvecOverlap(_sdst_Values, _sa_Values, 0, len(_sdst_Values), 0, len(_sa_Values)) {
for i := range dst.Values {
dst.Values[i] += a.Values[i]
}
Expand Down Expand Up @@ -1873,7 +1873,7 @@ func HardSwishFloat32s(dst, x []float32) {
_vcFloat32s44_1 := simd.BroadcastFloat32s(3)
_vcFloat32s44_2 := simd.BroadcastFloat32s(0)
_vcFloat32s44_4 := simd.BroadcastFloat32s(6)
if _loopvecOverlap(x, dst) {
if _loopvecOverlap(x, dst, 0, len(x), 0, len(dst)) {
for i := range dst {
t := x[i] + 3
if t < 0 {
Expand Down Expand Up @@ -1932,7 +1932,7 @@ func HardSigmoidFloat32s(dst, x []float32) {
_vcFloat32s45_1 := simd.BroadcastFloat32s(3)
_vcFloat32s45_2 := simd.BroadcastFloat32s(0)
_vcFloat32s45_4 := simd.BroadcastFloat32s(6)
if _loopvecOverlap(x, dst) {
if _loopvecOverlap(x, dst, 0, len(x), 0, len(dst)) {
for i := range dst {
t := x[i] + 3
if t < 0 {
Expand Down Expand Up @@ -1994,18 +1994,12 @@ func HardSigmoidFloat32s(dst, x []float32) {
}
}

func _loopvecOverlap[T any](a, b []T, window ...int) bool {
aLo, aHi, bLo, bHi := 0, len(a), 0, len(b)
if len(window) == 4 {
aLo, aHi, bLo, bHi = window[0], window[1], window[2], window[3]
}
func _loopvecOverlap[T any](a, b []T, aLo, aHi, bLo, bHi int) bool {
if aLo >= aHi || bLo >= bHi {
return false
}
size := unsafe.Sizeof(a[0])
aBase := uintptr(unsafe.Pointer(unsafe.SliceData(a)))
bBase := uintptr(unsafe.Pointer(unsafe.SliceData(b)))
aStart, aEnd := aBase+uintptr(aLo)*size, aBase+uintptr(aHi)*size
bStart, bEnd := bBase+uintptr(bLo)*size, bBase+uintptr(bHi)*size
return aStart < bEnd && bStart < aEnd
return aBase+uintptr(aLo)*size < bBase+uintptr(bHi)*size && bBase+uintptr(bLo)*size < aBase+uintptr(aHi)*size
}
18 changes: 8 additions & 10 deletions internal/rewrite/rewriter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ func wrapWithOverlapCheck(plan *loopir.Plan, origText, simdText string) string {
if plan.Bound == "" && plan.Start == "0" {
// The loop runs over whole slices, so comparing whole slices covers
// every element it touches.
fmt.Fprintf(&cond, "_loopvecOverlap(%s, %s)", pair[0].Name, pair[1].Name)
fmt.Fprintf(&cond, "_loopvecOverlap(%[1]s, %[2]s, 0, len(%[1]s), 0, len(%[2]s))", pair[0].Name, pair[1].Name)
continue
}
aLo, aHi := window(plan, pair[0])
Expand Down Expand Up @@ -285,22 +285,20 @@ func offsetText(v int) string {
}

// overlapHelperSrc is _loopvecOverlap's source, appended once per rewritten
// file that needs it.
// file that needs it. It reports whether a[aLo:aHi] and b[bLo:bHi] share
// memory. It takes the windows as plain arguments, a whole slice as 0 and its
// length, so that it is small enough to inline: every rewritten loop with two
// slices calls it before it runs, and a call costs a short loop as much as
// the loop itself.
const overlapHelperSrc = `
func _loopvecOverlap[T any](a, b []T, window ...int) bool {
aLo, aHi, bLo, bHi := 0, len(a), 0, len(b)
if len(window) == 4 {
aLo, aHi, bLo, bHi = window[0], window[1], window[2], window[3]
}
func _loopvecOverlap[T any](a, b []T, aLo, aHi, bLo, bHi int) bool {
if aLo >= aHi || bLo >= bHi {
return false
}
size := unsafe.Sizeof(a[0])
aBase := uintptr(unsafe.Pointer(unsafe.SliceData(a)))
bBase := uintptr(unsafe.Pointer(unsafe.SliceData(b)))
aStart, aEnd := aBase+uintptr(aLo)*size, aBase+uintptr(aHi)*size
bStart, bEnd := bBase+uintptr(bLo)*size, bBase+uintptr(bHi)*size
return aStart < bEnd && bStart < aEnd
return aBase+uintptr(aLo)*size < bBase+uintptr(bHi)*size && bBase+uintptr(bLo)*size < aBase+uintptr(aHi)*size
}
`

Expand Down
18 changes: 6 additions & 12 deletions testdata/abs_sqrt.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ import (
)

func Abs64(dst, a []float64) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = math.Abs(a[i])
}
Expand Down Expand Up @@ -114,7 +114,7 @@ func Abs64(dst, a []float64) {
}

func Abs32(dst, a []float32) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = float32(math.Abs(float64(a[i])))
}
Expand Down Expand Up @@ -152,7 +152,7 @@ func Abs32(dst, a []float32) {
}

func Sqrt32(dst, a []float32) {
if _loopvecOverlap(dst, a) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) {
for i := range dst {
dst[i] = float32(math.Sqrt(float64(a[i])))
}
Expand Down Expand Up @@ -247,7 +247,7 @@ func InvSqrt(a []float64) {
}

func AbsDiff(dst, a, b []float64) {
if _loopvecOverlap(dst, a) || _loopvecOverlap(dst, b) {
if _loopvecOverlap(dst, a, 0, len(dst), 0, len(a)) || _loopvecOverlap(dst, b, 0, len(dst), 0, len(b)) {
for i := range dst {
dst[i] = math.Abs(a[i] - b[i])
}
Expand Down Expand Up @@ -291,20 +291,14 @@ func AbsDiff(dst, a, b []float64) {
}
}

func _loopvecOverlap[T any](a, b []T, window ...int) bool {
aLo, aHi, bLo, bHi := 0, len(a), 0, len(b)
if len(window) == 4 {
aLo, aHi, bLo, bHi = window[0], window[1], window[2], window[3]
}
func _loopvecOverlap[T any](a, b []T, aLo, aHi, bLo, bHi int) bool {
if aLo >= aHi || bLo >= bHi {
return false
}
size := unsafe.Sizeof(a[0])
aBase := uintptr(unsafe.Pointer(unsafe.SliceData(a)))
bBase := uintptr(unsafe.Pointer(unsafe.SliceData(b)))
aStart, aEnd := aBase+uintptr(aLo)*size, aBase+uintptr(aHi)*size
bStart, bEnd := bBase+uintptr(bLo)*size, bBase+uintptr(bHi)*size
return aStart < bEnd && bStart < aEnd
return aBase+uintptr(aLo)*size < bBase+uintptr(bHi)*size && bBase+uintptr(bLo)*size < aBase+uintptr(aHi)*size
}
-- testmain/main.go --
package main
Expand Down
16 changes: 5 additions & 11 deletions testdata/activation_functions.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,7 @@ func ReLU(x []float32) {
func LeakyReLU(dst, x []float32, alpha float32) {
_vcAFloat32s1 := simd.BroadcastFloat32s(0)
_vcBFloat32s1 := simd.BroadcastFloat32s(alpha)
if _loopvecOverlap(dst, x) {
if _loopvecOverlap(dst, x, 0, len(dst), 0, len(x)) {
for i := range dst {
if x[i] > 0 {
dst[i] = x[i]
Expand Down Expand Up @@ -312,7 +312,7 @@ func Sign(dst, x []float32) {
_vcFloat32s5_1 := simd.BroadcastFloat32s(0)
_vcFloat32s5_2 := simd.BroadcastFloat32s(1)
_vcFloat32s5_4 := simd.BroadcastFloat32s(-1)
if _loopvecOverlap(dst, x) {
if _loopvecOverlap(dst, x, 0, len(dst), 0, len(x)) {
for i := range dst {
if x[i] > 0 {
dst[i] = 1
Expand Down Expand Up @@ -359,7 +359,7 @@ func Step(dst, x []float32, threshold float32) {
_vcFloat32s6_1 := simd.BroadcastFloat32s(threshold)
_vcFloat32s6_2 := simd.BroadcastFloat32s(1)
_vcFloat32s6_3 := simd.BroadcastFloat32s(0)
if _loopvecOverlap(dst, x) {
if _loopvecOverlap(dst, x, 0, len(dst), 0, len(x)) {
for i := range dst {
if x[i] >= threshold {
dst[i] = 1
Expand Down Expand Up @@ -549,20 +549,14 @@ func NaNToZero(x []float32) {
}
}

func _loopvecOverlap[T any](a, b []T, window ...int) bool {
aLo, aHi, bLo, bHi := 0, len(a), 0, len(b)
if len(window) == 4 {
aLo, aHi, bLo, bHi = window[0], window[1], window[2], window[3]
}
func _loopvecOverlap[T any](a, b []T, aLo, aHi, bLo, bHi int) bool {
if aLo >= aHi || bLo >= bHi {
return false
}
size := unsafe.Sizeof(a[0])
aBase := uintptr(unsafe.Pointer(unsafe.SliceData(a)))
bBase := uintptr(unsafe.Pointer(unsafe.SliceData(b)))
aStart, aEnd := aBase+uintptr(aLo)*size, aBase+uintptr(aHi)*size
bStart, bEnd := bBase+uintptr(bLo)*size, bBase+uintptr(bHi)*size
return aStart < bEnd && bStart < aEnd
return aBase+uintptr(aLo)*size < bBase+uintptr(bHi)*size && bBase+uintptr(bLo)*size < aBase+uintptr(aHi)*size
}
-- ops.go.guarded --
//go:build !goexperiment.simd
Expand Down
Loading