simd_arm64_test.go

     1  // Copyright 2026 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build goexperiment.simd && arm64
     6  
     7  package simd_test
     8  
     9  import (
    10  	"fmt"
    11  	"simd/archsimd"
    12  	"strings"
    13  	"testing"
    14  )
    15  
    16  func TestLookupOrZero(t *testing.T) {
    17  	// Out-of-range indices produce zero lane value.
    18  	x := []uint8{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
    19  	indices := []uint8{7, 6, 5, 4, 3, 2, 1, 0, 0xff, 8, 16, 9, 128, 10, 20, 11}
    20  	want := []uint8{8, 7, 6, 5, 4, 3, 2, 1, 0, 9, 0, 10, 0, 11, 0, 12}
    21  	got := make([]uint8, len(x))
    22  	archsimd.LoadUint8x16(x).LookupOrZero(archsimd.LoadUint8x16(indices)).StorePart(got)
    23  	checkSlices(t, got, want)
    24  }
    25  
    26  func TestClMul(t *testing.T) {
    27  	if !archsimd.ARM64.PMULL() {
    28  		t.Skip("no carryless multiply")
    29  	}
    30  	var x = archsimd.LoadUint64x2([]uint64{1, 5})
    31  	var y = archsimd.LoadUint64x2([]uint64{3, 9})
    32  
    33  	foo := func(v archsimd.Uint64x2, s []uint64) {
    34  		r := make([]uint64, 2, 2)
    35  		v.StorePart(r)
    36  		checkSlices[uint64](t, r, s)
    37  	}
    38  
    39  	foo(x.CarrylessMultiplyEven(y), []uint64{3, 0})
    40  	foo(x.CarrylessMultiplyEvenOdd(y), []uint64{9, 0})
    41  	foo(x.CarrylessMultiplyOddEven(y), []uint64{15, 0})
    42  	foo(x.CarrylessMultiplyOdd(y), []uint64{45, 0})
    43  	foo(y.CarrylessMultiplyEven(y), []uint64{5, 0})
    44  }
    45  
    46  //go:noinline
    47  func addInt8sNoinline(a, b archsimd.Int8s) archsimd.Int8s { return a.Add(b) }
    48  
    49  //go:noinline
    50  func blackholeSVE() {}
    51  
    52  // TestAddSVEAcrossCall passes scalable vectors across a real (non-inlined) ABI
    53  // boundary, exercising the register/stack passing that size.go's simdify decides
    54  // for SVE types.
    55  func TestAddSVEAcrossCall(t *testing.T) {
    56  	if !archsimd.ARM64.SVE() {
    57  		t.Skip("no sve")
    58  	}
    59  	var a, b, got [32]int8
    60  	for i := range a {
    61  		a[i] = int8(i)
    62  		b[i] = int8(2*i + 1)
    63  	}
    64  	x := archsimd.LoadInt8s(a[:])
    65  	addInt8sNoinline(x, archsimd.LoadInt8s(b[:])).Store(got[:])
    66  	for i := 0; i < x.Len(); i++ {
    67  		if want := a[i] + b[i]; got[i] != want {
    68  			t.Errorf("lane %d: got %d, want %d", i, got[i], want)
    69  		}
    70  	}
    71  }
    72  
    73  //go:noinline
    74  func greaterInt8sNoinline(a, b archsimd.Int8s) archsimd.Mask8s { return a.Greater(b) }
    75  
    76  // TestGreaterSVEMaskRoundTrip returns a mask across a non-inlined call, exercising
    77  // the predicate memory round-trip (PSTR to return it, PLDR to reload it) that the
    78  // mask ABI relies on, and checks it against a > b lane by lane.
    79  func TestGreaterSVEMaskRoundTrip(t *testing.T) {
    80  	if !archsimd.ARM64.SVE() {
    81  		t.Skip("no sve")
    82  	}
    83  	var a, b [32]int8
    84  	for i := range a {
    85  		a[i] = int8(i - 8)
    86  		b[i] = int8(2*i - 20)
    87  	}
    88  	var z archsimd.Int8s
    89  	m := greaterInt8sNoinline(archsimd.LoadInt8s(a[:]), archsimd.LoadInt8s(b[:]))
    90  	bits := maskBits(m)
    91  	for i := 0; i < z.Len(); i++ {
    92  		want := a[i] > b[i]
    93  		got := bits>>i&1 == 1
    94  		if got != want {
    95  			t.Errorf("lane %d: got %v, want %v (a=%d b=%d)", i, got, want, a[i], b[i])
    96  		}
    97  	}
    98  	if bits>>z.Len() != 0 {
    99  		t.Errorf("bits beyond the vector length are set: %#x", bits)
   100  	}
   101  }
   102  
   103  // TestAddSVESpill keeps a scalable vector live across a call, forcing the
   104  // register allocator to spill and reload it (ZSTR/ZLDR).
   105  func TestAddSVESpill(t *testing.T) {
   106  	if !archsimd.ARM64.SVE() {
   107  		t.Skip("no sve")
   108  	}
   109  	var a, b, got [32]int8
   110  	for i := range a {
   111  		a[i] = int8(i)
   112  		b[i] = int8(100 - i)
   113  	}
   114  	sum := archsimd.LoadInt8s(a[:]).Add(archsimd.LoadInt8s(b[:]))
   115  	blackholeSVE() // clobbers caller-saved regs; sum must survive via a spill
   116  	sum.Store(got[:])
   117  	for i := 0; i < sum.Len(); i++ {
   118  		if want := a[i] + b[i]; got[i] != want {
   119  			t.Errorf("lane %d: got %d, want %d", i, got[i], want)
   120  		}
   121  	}
   122  }
   123  
   124  // TestAddSaturatedSVE checks the generated saturating add: an explicit
   125  // boundary case, then every integer shape against the generic emulation.
   126  func TestAddSaturatedSVE(t *testing.T) {
   127  	if !archsimd.ARM64.SVE() {
   128  		t.Skip("no sve")
   129  	}
   130  	var si, gi [32]int8
   131  	for i := range si {
   132  		si[i] = 100 // 100+100 saturates to +127
   133  	}
   134  	vi := archsimd.LoadInt8s(si[:])
   135  	vi.AddSaturated(vi).Store(gi[:])
   136  	for i := 0; i < vi.Len(); i++ {
   137  		if gi[i] != 127 {
   138  			t.Errorf("int8 lane %d: got %d, want 127", i, gi[i])
   139  		}
   140  	}
   141  	testInt8sBinary(t, archsimd.Int8s.AddSaturated, addSaturatedSlice[int8])
   142  	testInt16sBinary(t, archsimd.Int16s.AddSaturated, addSaturatedSlice[int16])
   143  	testInt32sBinary(t, archsimd.Int32s.AddSaturated, addSaturatedSlice[int32])
   144  	testInt64sBinary(t, archsimd.Int64s.AddSaturated, addSaturatedSlice[int64])
   145  	testUint8sBinary(t, archsimd.Uint8s.AddSaturated, addSaturatedSlice[uint8])
   146  	testUint16sBinary(t, archsimd.Uint16s.AddSaturated, addSaturatedSlice[uint16])
   147  	testUint32sBinary(t, archsimd.Uint32s.AddSaturated, addSaturatedSlice[uint32])
   148  	testUint64sBinary(t, archsimd.Uint64s.AddSaturated, addSaturatedSlice[uint64])
   149  }
   150  
   151  // TestSubSaturatedSVE checks the generated saturating subtract: an explicit
   152  // boundary case, then every integer shape against the generic emulation.
   153  func TestSubSaturatedSVE(t *testing.T) {
   154  	if !archsimd.ARM64.SVE() {
   155  		t.Skip("no sve")
   156  	}
   157  	var sx, sy, gi [32]int8
   158  	for i := range sx {
   159  		sx[i] = 100 // 100 - (-100) saturates to +127
   160  		sy[i] = -100
   161  	}
   162  	x, y := archsimd.LoadInt8s(sx[:]), archsimd.LoadInt8s(sy[:])
   163  	x.SubSaturated(y).Store(gi[:])
   164  	for i := 0; i < x.Len(); i++ {
   165  		if gi[i] != 127 {
   166  			t.Errorf("int8 lane %d: got %d, want 127", i, gi[i])
   167  		}
   168  	}
   169  	testInt8sBinary(t, archsimd.Int8s.SubSaturated, subSaturatedSlice[int8])
   170  	testInt16sBinary(t, archsimd.Int16s.SubSaturated, subSaturatedSlice[int16])
   171  	testInt32sBinary(t, archsimd.Int32s.SubSaturated, subSaturatedSlice[int32])
   172  	testInt64sBinary(t, archsimd.Int64s.SubSaturated, subSaturatedSlice[int64])
   173  	testUint8sBinary(t, archsimd.Uint8s.SubSaturated, subSaturatedSlice[uint8])
   174  	testUint16sBinary(t, archsimd.Uint16s.SubSaturated, subSaturatedSlice[uint16])
   175  	testUint32sBinary(t, archsimd.Uint32s.SubSaturated, subSaturatedSlice[uint32])
   176  	testUint64sBinary(t, archsimd.Uint64s.SubSaturated, subSaturatedSlice[uint64])
   177  }
   178  
   179  func TestStringSVE(t *testing.T) {
   180  	if !archsimd.ARM64.SVE() {
   181  		t.Skip("no sve")
   182  	}
   183  	want := func(v any) string {
   184  		return "{" + strings.ReplaceAll(strings.Trim(fmt.Sprint(v), "[]"), " ", ",") + "}"
   185  	}
   186  
   187  	xs := make([]int8, archsimd.Int8s{}.Len())
   188  	ys := make([]int64, archsimd.Int64s{}.Len())
   189  	for i := range xs {
   190  		xs[i] = int8(i % 2)
   191  	}
   192  	for i := range ys {
   193  		ys[i] = int64(i % 2)
   194  	}
   195  	x := archsimd.LoadInt8s(xs)
   196  	y := archsimd.LoadInt64s(ys)
   197  	mx := x.Greater(archsimd.LoadInt8s(make([]int8, len(xs))))
   198  	my := y.Greater(archsimd.LoadInt64s(make([]int64, len(ys))))
   199  
   200  	if x.String() != want(xs) {
   201  		t.Errorf("x=%s wanted %s", x, want(xs))
   202  	}
   203  	if y.String() != want(ys) {
   204  		t.Errorf("y=%s wanted %s", y, want(ys))
   205  	}
   206  	if mx.String() != want(xs) {
   207  		t.Errorf("mx=%s wanted %s", mx, want(xs))
   208  	}
   209  	if my.String() != want(ys) {
   210  		t.Errorf("my=%s wanted %s", my, want(ys))
   211  	}
   212  	t.Logf("x=%s", x)
   213  	t.Logf("y=%s", y)
   214  	t.Logf("mx=%s", mx)
   215  	t.Logf("my=%s", my)
   216  }
   217  
   218  //go:noinline
   219  func keepAliveInt8s(archsimd.Int8s) {}
   220  
   221  // namedMask8s and maxVia pass masks through types other than archsimd.Mask8s
   222  // itself: a defined type, and the GC shape of a generic instantiation. Both
   223  // must use the memory ABI of the mask they are made from, to agree with the
   224  // concrete functions they call and to reach a P register when used.
   225  type namedMask8s archsimd.Mask8s
   226  
   227  //go:noinline
   228  func greaterNamed(x, y archsimd.Int8s) namedMask8s { return namedMask8s(x.Greater(y)) }
   229  
   230  //go:noinline
   231  func maxVia[M archsimd.Mask8s](greater func(x, y archsimd.Int8s) M, x, y archsimd.Int8s) archsimd.Int8s {
   232  	return x.IfElse(archsimd.Mask8s(greater(x, y)), y)
   233  }
   234  
   235  func TestMaskABISVE(t *testing.T) {
   236  	if !archsimd.ARM64.SVE() {
   237  		t.Skip("no sve")
   238  	}
   239  	n := archsimd.Int8s{}.Len()
   240  	xs, ys := make([]int8, n), make([]int8, n)
   241  	for i := range xs {
   242  		xs[i], ys[i] = int8(i%5), int8(i%3)
   243  	}
   244  	x, y := archsimd.LoadInt8s(xs), archsimd.LoadInt8s(ys)
   245  	check := func(name string, v archsimd.Int8s) {
   246  		t.Helper()
   247  		got := make([]int8, n)
   248  		v.Store(got)
   249  		for i := range got {
   250  			if want := max(xs[i], ys[i]); got[i] != want {
   251  				t.Errorf("%s: lane %d = %d, want %d", name, i, got[i], want)
   252  			}
   253  		}
   254  	}
   255  	check("generic", maxVia(archsimd.Int8s.Greater, x, y))
   256  	check("named", x.IfElse(archsimd.Mask8s(greaterNamed(x, y)), y))
   257  }
   258  
   259  // TestIfElseSVE checks IfElse and Masked, and that the merging peephole keeps
   260  // the same semantics whether or not it fires: x.Add(y).IfElse(m, x) folds into a
   261  // predicated add, x.Add(y).IfElse(m, z) does not, and both must agree with a
   262  // lane-by-lane reference.
   263  func TestIfElseSVE(t *testing.T) {
   264  	if !archsimd.ARM64.SVE() {
   265  		t.Skip("no sve")
   266  	}
   267  	n := archsimd.Int8s{}.Len()
   268  	xs, ys, zs := make([]int8, n), make([]int8, n), make([]int8, n)
   269  	for i := range xs {
   270  		xs[i] = int8(i + 1)
   271  		ys[i] = int8(i % 3) // active where xs[i] > ys[i], which alternates early on
   272  		zs[i] = int8(-i - 1)
   273  	}
   274  	x, y, z := archsimd.LoadInt8s(xs), archsimd.LoadInt8s(ys), archsimd.LoadInt8s(zs)
   275  	m := x.Greater(y)
   276  
   277  	got := make([]int8, n)
   278  	check := func(name string, v archsimd.Int8s, want func(i int, active bool) int8) {
   279  		t.Helper()
   280  		v.Store(got)
   281  		for i := 0; i < n; i++ {
   282  			if w := want(i, xs[i] > ys[i]); got[i] != w {
   283  				t.Errorf("%s: lane %d = %d, want %d (x=%d y=%d)", name, i, got[i], w, xs[i], ys[i])
   284  			}
   285  		}
   286  	}
   287  
   288  	check("IfElse", x.IfElse(m, y), func(i int, active bool) int8 {
   289  		if active {
   290  			return xs[i]
   291  		}
   292  		return ys[i]
   293  	})
   294  	check("Masked", x.Masked(m), func(i int, active bool) int8 {
   295  		if active {
   296  			return xs[i]
   297  		}
   298  		return 0
   299  	})
   300  	// Folds into the merging-predicated add.
   301  	check("Add.IfElse(x)", x.Add(y).IfElse(m, x), func(i int, active bool) int8 {
   302  		if active {
   303  			return xs[i] + ys[i]
   304  		}
   305  		return xs[i]
   306  	})
   307  	// Folds via commutativity.
   308  	check("Add.IfElse(y)", x.Add(y).IfElse(m, y), func(i int, active bool) int8 {
   309  		if active {
   310  			return xs[i] + ys[i]
   311  		}
   312  		return ys[i]
   313  	})
   314  	// Folds behind a merging MOVPRFX: the else operand is neither source.
   315  	check("Add.IfElse(z)", x.Add(y).IfElse(m, z), func(i int, active bool) int8 {
   316  		if active {
   317  			return xs[i] + ys[i]
   318  		}
   319  		return zs[i]
   320  	})
   321  	// ADD has no zeroing-predicated form, so Masked folds into the merging one
   322  	// with the zero vector as its else operand.
   323  	check("Add.Masked", x.Add(y).Masked(m), func(i int, active bool) int8 {
   324  		if active {
   325  			return xs[i] + ys[i]
   326  		}
   327  		return 0
   328  	})
   329  
   330  	// The two merges below differ only in which source the inactive lanes
   331  	// keep. The merging machine op pins its result to that source and is not
   332  	// commutative, or CSE — which canonicalizes the args of a commutative op —
   333  	// would conflate the two values and give one of them the wrong else
   334  	// operand.
   335  	keepX := x.Add(y).IfElse(m, x)
   336  	keepY := y.Add(x).IfElse(m, y)
   337  	check("Add.IfElse cse keepX", keepX, func(i int, active bool) int8 {
   338  		if active {
   339  			return xs[i] + ys[i]
   340  		}
   341  		return xs[i]
   342  	})
   343  	check("Add.IfElse cse keepY", keepY, func(i int, active bool) int8 {
   344  		if active {
   345  			return xs[i] + ys[i]
   346  		}
   347  		return ys[i]
   348  	})
   349  
   350  	// SUB is not commutative, so only an "else" operand that is the destructive
   351  	// one — the minuend — folds into the merging-predicated instruction.
   352  	check("Sub.IfElse(x)", x.Sub(y).IfElse(m, x), func(i int, active bool) int8 {
   353  		if active {
   354  			return xs[i] - ys[i]
   355  		}
   356  		return xs[i]
   357  	})
   358  	// Does not fold, and must not silently become y-x.
   359  	check("Sub.IfElse(y)", x.Sub(y).IfElse(m, y), func(i int, active bool) int8 {
   360  		if active {
   361  			return xs[i] - ys[i]
   362  		}
   363  		return ys[i]
   364  	})
   365  	// Does not fold: there is no prefixed form for a non-commutative operation.
   366  	check("Sub.IfElse(z)", x.Sub(y).IfElse(m, z), func(i int, active bool) int8 {
   367  		if active {
   368  			return xs[i] - ys[i]
   369  		}
   370  		return zs[i]
   371  	})
   372  	check("Sub.Masked", x.Sub(y).Masked(m), func(i int, active bool) int8 {
   373  		if active {
   374  			return xs[i] - ys[i]
   375  		}
   376  		return 0
   377  	})
   378  
   379  	// Abs is predicated-only, so its unpredicated API runs under an all-true
   380  	// predicate that a select can simply replace: IfElse becomes the merging
   381  	// form and Masked the zeroing one, each a single instruction.
   382  	absLane := func(v int8) int8 {
   383  		if v < 0 {
   384  			return -v
   385  		}
   386  		return v
   387  	}
   388  	check("Abs.IfElse(z)", x.Abs().ConvertToInt8().IfElse(m, z), func(i int, active bool) int8 {
   389  		if active {
   390  			return absLane(xs[i])
   391  		}
   392  		return zs[i]
   393  	})
   394  	check("Abs.IfElse(x)", x.Abs().ConvertToInt8().IfElse(m, x), func(i int, active bool) int8 {
   395  		if active {
   396  			return absLane(xs[i])
   397  		}
   398  		return xs[i]
   399  	})
   400  	check("Abs.Masked", x.Abs().Masked(m).ConvertToInt8(), func(i int, active bool) int8 {
   401  		if active {
   402  			return absLane(xs[i])
   403  		}
   404  		return 0
   405  	})
   406  
   407  	var scalar int8 = 42
   408  	check("BroadcastConst.IfElse(z)", archsimd.BroadcastInt8s(7).IfElse(m, z), func(i int, active bool) int8 {
   409  		if active {
   410  			return 7
   411  		}
   412  		return zs[i]
   413  	})
   414  	check("BroadcastGp.IfElse(z)", archsimd.BroadcastInt8s(scalar).IfElse(m, z), func(i int, active bool) int8 {
   415  		if active {
   416  			return scalar
   417  		}
   418  		return zs[i]
   419  	})
   420  
   421  	// The prefixed path with every operand still live afterwards, so the
   422  	// destination can be none of them and the merging MOVPRFX has to place the
   423  	// else operand itself.
   424  	rz := x.Add(y).IfElse(m, z)
   425  	keepAliveInt8s(x)
   426  	keepAliveInt8s(y)
   427  	keepAliveInt8s(z)
   428  	check("Add.IfElse(z) with all live", rz, func(i int, active bool) int8 {
   429  		if active {
   430  			return xs[i] + ys[i]
   431  		}
   432  		return zs[i]
   433  	})
   434  
   435  	// The MOVPRFX path: x must survive the destructive predicated add.
   436  	r := x.Add(y).IfElse(m, x)
   437  	keepAliveInt8s(x)
   438  	check("Add.IfElse(x) with x live", r, func(i int, active bool) int8 {
   439  		if active {
   440  			return xs[i] + ys[i]
   441  		}
   442  		return xs[i]
   443  	})
   444  	x.Store(got)
   445  	for i := 0; i < n; i++ {
   446  		if got[i] != xs[i] {
   447  			t.Errorf("x clobbered by MOVPRFX: lane %d = %d, want %d", i, got[i], xs[i])
   448  		}
   449  	}
   450  }
   451  
   452  // TestMaskAllTrueSVE checks the all-true mask constructors: a select under an
   453  // all-true mask returns its first operand in every lane, at each lane width.
   454  func TestMaskAllTrueSVE(t *testing.T) {
   455  	if !archsimd.ARM64.SVE() {
   456  		t.Skip("no sve")
   457  	}
   458  	check := func(name string, n int, sel func(i int) (got, want any)) {
   459  		t.Helper()
   460  		for i := 0; i < n; i++ {
   461  			if got, want := sel(i); got != want {
   462  				t.Errorf("%s: lane %d = %v, want %v", name, i, got, want)
   463  			}
   464  		}
   465  	}
   466  
   467  	xs8, ys8 := make([]int8, archsimd.Int8s{}.Len()), make([]int8, archsimd.Int8s{}.Len())
   468  	for i := range xs8 {
   469  		xs8[i], ys8[i] = int8(i+1), int8(-i-1)
   470  	}
   471  	g8 := make([]int8, len(xs8))
   472  	archsimd.LoadInt8s(xs8).IfElse(archsimd.Mask8sAllTrue(), archsimd.LoadInt8s(ys8)).Store(g8)
   473  	check("Mask8sAllTrue", len(xs8), func(i int) (any, any) { return g8[i], xs8[i] })
   474  
   475  	xs16, ys16 := make([]int16, archsimd.Int16s{}.Len()), make([]int16, archsimd.Int16s{}.Len())
   476  	for i := range xs16 {
   477  		xs16[i], ys16[i] = int16(i+1), int16(-i-1)
   478  	}
   479  	g16 := make([]int16, len(xs16))
   480  	archsimd.LoadInt16s(xs16).IfElse(archsimd.Mask16sAllTrue(), archsimd.LoadInt16s(ys16)).Store(g16)
   481  	check("Mask16sAllTrue", len(xs16), func(i int) (any, any) { return g16[i], xs16[i] })
   482  
   483  	xs32, ys32 := make([]int32, archsimd.Int32s{}.Len()), make([]int32, archsimd.Int32s{}.Len())
   484  	for i := range xs32 {
   485  		xs32[i], ys32[i] = int32(i+1), int32(-i-1)
   486  	}
   487  	g32 := make([]int32, len(xs32))
   488  	archsimd.LoadInt32s(xs32).IfElse(archsimd.Mask32sAllTrue(), archsimd.LoadInt32s(ys32)).Store(g32)
   489  	check("Mask32sAllTrue", len(xs32), func(i int) (any, any) { return g32[i], xs32[i] })
   490  
   491  	xs64, ys64 := make([]int64, archsimd.Int64s{}.Len()), make([]int64, archsimd.Int64s{}.Len())
   492  	for i := range xs64 {
   493  		xs64[i], ys64[i] = int64(i+1), int64(-i-1)
   494  	}
   495  	g64 := make([]int64, len(xs64))
   496  	archsimd.LoadInt64s(xs64).IfElse(archsimd.Mask64sAllTrue(), archsimd.LoadInt64s(ys64)).Store(g64)
   497  	check("Mask64sAllTrue", len(xs64), func(i int) (any, any) { return g64[i], xs64[i] })
   498  }
   499
View as plain text