compare_sve_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  // SVE mask (predicate) tests, in the same utility-function shape as
     8  // compare_amd64_test.go.
     9  
    10  package simd_test
    11  
    12  import (
    13  	"simd/archsimd"
    14  	"testing"
    15  	"unsafe"
    16  )
    17  
    18  // maskBits returns the bits of an SVE mask. In memory a mask is a uint64 with
    19  // one bit per vector byte, bit 0 first, and zeros beyond the vector length.
    20  func maskBits[M archsimd.Mask8s | archsimd.Mask16s | archsimd.Mask32s | archsimd.Mask64s](m M) uint64 {
    21  	return *(*uint64)(unsafe.Pointer(&m))
    22  }
    23  
    24  // testSVECompare drives a scalable compare that returns a mask. active is the
    25  // runtime number of live lanes; elemBytes is the element width. Lane i is true
    26  // iff bit i*elemBytes of the mask's bits is set.
    27  func testSVECompare[T number, V any, M archsimd.Mask8s | archsimd.Mask16s | archsimd.Mask32s | archsimd.Mask64s](t *testing.T, pool []T, elemBytes, active int,
    28  	load func([]T) V, cmp func(V, V) M, want func(T, T) bool) {
    29  	t.Helper()
    30  	count := sveMaxBytes / elemBytes
    31  	forSlicePair(t, pool, count, func(x, y []T) bool {
    32  		t.Helper()
    33  		bits := maskBits(cmp(load(x), load(y)))
    34  		for i := 0; i < active; i++ {
    35  			got := bits>>(i*elemBytes)&1 == 1
    36  			if got != want(x[i], y[i]) {
    37  				t.Errorf("lane %d: got %v, want %v (x=%v y=%v)", i, got, want(x[i], y[i]), x[i], y[i])
    38  				return false
    39  			}
    40  		}
    41  		return true
    42  	})
    43  }
    44  
    45  func testInt8sCompare(t *testing.T, cmp func(_, _ archsimd.Int8s) archsimd.Mask8s, want func(_, _ int8) bool) {
    46  	var z archsimd.Int8s
    47  	testSVECompare(t, int8s, 1, z.Len(), archsimd.LoadInt8s, cmp, want)
    48  }
    49  
    50  func testInt16sCompare(t *testing.T, cmp func(_, _ archsimd.Int16s) archsimd.Mask16s, want func(_, _ int16) bool) {
    51  	var z archsimd.Int16s
    52  	testSVECompare(t, int16s, 2, z.Len(), archsimd.LoadInt16s, cmp, want)
    53  }
    54  
    55  func testInt32sCompare(t *testing.T, cmp func(_, _ archsimd.Int32s) archsimd.Mask32s, want func(_, _ int32) bool) {
    56  	var z archsimd.Int32s
    57  	testSVECompare(t, int32s, 4, z.Len(), archsimd.LoadInt32s, cmp, want)
    58  }
    59  
    60  func testInt64sCompare(t *testing.T, cmp func(_, _ archsimd.Int64s) archsimd.Mask64s, want func(_, _ int64) bool) {
    61  	var z archsimd.Int64s
    62  	testSVECompare(t, int64s, 8, z.Len(), archsimd.LoadInt64s, cmp, want)
    63  }
    64  
    65  func testUint8sCompare(t *testing.T, cmp func(_, _ archsimd.Uint8s) archsimd.Mask8s, want func(_, _ uint8) bool) {
    66  	var z archsimd.Uint8s
    67  	testSVECompare(t, uint8s, 1, z.Len(), archsimd.LoadUint8s, cmp, want)
    68  }
    69  
    70  func testFloat32sCompare(t *testing.T, cmp func(_, _ archsimd.Float32s) archsimd.Mask32s, want func(_, _ float32) bool) {
    71  	var z archsimd.Float32s
    72  	testSVECompare(t, float32s, 4, z.Len(), archsimd.LoadFloat32s, cmp, want)
    73  }
    74  
    75  func testFloat64sCompare(t *testing.T, cmp func(_, _ archsimd.Float64s) archsimd.Mask64s, want func(_, _ float64) bool) {
    76  	var z archsimd.Float64s
    77  	testSVECompare(t, float64s, 8, z.Len(), archsimd.LoadFloat64s, cmp, want)
    78  }
    79  
    80  func testUint16sCompare(t *testing.T, cmp func(_, _ archsimd.Uint16s) archsimd.Mask16s, want func(_, _ uint16) bool) {
    81  	var z archsimd.Uint16s
    82  	testSVECompare(t, uint16s, 2, z.Len(), archsimd.LoadUint16s, cmp, want)
    83  }
    84  
    85  func testUint32sCompare(t *testing.T, cmp func(_, _ archsimd.Uint32s) archsimd.Mask32s, want func(_, _ uint32) bool) {
    86  	var z archsimd.Uint32s
    87  	testSVECompare(t, uint32s, 4, z.Len(), archsimd.LoadUint32s, cmp, want)
    88  }
    89  
    90  func testUint64sCompare(t *testing.T, cmp func(_, _ archsimd.Uint64s) archsimd.Mask64s, want func(_, _ uint64) bool) {
    91  	var z archsimd.Uint64s
    92  	testSVECompare(t, uint64s, 8, z.Len(), archsimd.LoadUint64s, cmp, want)
    93  }
    94  
    95  func gtWant[T number](a, b T) bool { return a > b }
    96  func geWant[T number](a, b T) bool { return a >= b }
    97  func eqWant[T number](a, b T) bool { return a == b }
    98  func neWant[T number](a, b T) bool { return a != b }
    99  
   100  func TestGreaterSVE(t *testing.T) {
   101  	if !archsimd.ARM64.SVE() {
   102  		t.Skip("no SVE")
   103  	}
   104  	testInt8sCompare(t, archsimd.Int8s.Greater, gtWant[int8])
   105  	testInt16sCompare(t, archsimd.Int16s.Greater, gtWant[int16])
   106  	testInt32sCompare(t, archsimd.Int32s.Greater, gtWant[int32])
   107  	testInt64sCompare(t, archsimd.Int64s.Greater, gtWant[int64])
   108  	testUint8sCompare(t, archsimd.Uint8s.Greater, gtWant[uint8])
   109  	testUint16sCompare(t, archsimd.Uint16s.Greater, gtWant[uint16])
   110  	testUint32sCompare(t, archsimd.Uint32s.Greater, gtWant[uint32])
   111  	testUint64sCompare(t, archsimd.Uint64s.Greater, gtWant[uint64])
   112  	testFloat32sCompare(t, archsimd.Float32s.Greater, gtWant[float32])
   113  	testFloat64sCompare(t, archsimd.Float64s.Greater, gtWant[float64])
   114  }
   115  
   116  // TestMaskBitsSVE checks the memory representation of a mask: a uint64 of
   117  // predicate bits, zero beyond the vector length, readable and writable through
   118  // an unsafe pointer.
   119  func TestMaskBitsSVE(t *testing.T) {
   120  	if !archsimd.ARM64.SVE() {
   121  		t.Skip("no sve")
   122  	}
   123  	vlBytes := archsimd.Int8s{}.Len()
   124  	allBytes := uint64(1)<<vlBytes - 1
   125  
   126  	// Storing a mask overwrites all 64 bits, not just the VL/8 bytes of the
   127  	// predicate.
   128  	u := ^uint64(0)
   129  	*(*archsimd.Mask8s)(unsafe.Pointer(&u)) = archsimd.Mask8sAllTrue()
   130  	if u != allBytes {
   131  		t.Errorf("stored Mask8sAllTrue: bits = %#x, want %#x", u, allBytes)
   132  	}
   133  	u = ^uint64(0)
   134  	*(*archsimd.Mask16s)(unsafe.Pointer(&u)) = archsimd.Mask16sAllTrue()
   135  	if want := allBytes & 0x5555555555555555; u != want {
   136  		t.Errorf("stored Mask16sAllTrue: bits = %#x, want %#x", u, want)
   137  	}
   138  	if got := maskBits(archsimd.Mask64s{}); got != 0 {
   139  		t.Errorf("zero Mask64s: bits = %#x, want 0", got)
   140  	}
   141  
   142  	// A mask read from memory governs the lanes its bits select.
   143  	n := archsimd.Int32s{}.Len()
   144  	xs, ys := make([]int32, n), make([]int32, n)
   145  	for i := range xs {
   146  		xs[i], ys[i] = int32(i+1), int32(-i-1)
   147  	}
   148  	bits := uint64(0)
   149  	for i := 0; i < n; i += 2 {
   150  		bits |= 1 << (4 * i) // even lanes
   151  	}
   152  	m := *(*archsimd.Mask32s)(unsafe.Pointer(&bits))
   153  	got := make([]int32, n)
   154  	archsimd.LoadInt32s(xs).IfElse(m, archsimd.LoadInt32s(ys)).Store(got)
   155  	for i := range got {
   156  		want := ys[i]
   157  		if i%2 == 0 {
   158  			want = xs[i]
   159  		}
   160  		if got[i] != want {
   161  			t.Errorf("IfElse under loaded mask: lane %d = %d, want %d", i, got[i], want)
   162  		}
   163  	}
   164  }
   165  
   166  func TestEqualSVE(t *testing.T) {
   167  	if !archsimd.ARM64.SVE() {
   168  		t.Skip("no SVE")
   169  	}
   170  	testInt8sCompare(t, archsimd.Int8s.Equal, eqWant[int8])
   171  	testInt16sCompare(t, archsimd.Int16s.Equal, eqWant[int16])
   172  	testInt32sCompare(t, archsimd.Int32s.Equal, eqWant[int32])
   173  	testInt64sCompare(t, archsimd.Int64s.Equal, eqWant[int64])
   174  	testUint8sCompare(t, archsimd.Uint8s.Equal, eqWant[uint8])
   175  	testUint16sCompare(t, archsimd.Uint16s.Equal, eqWant[uint16])
   176  	testUint32sCompare(t, archsimd.Uint32s.Equal, eqWant[uint32])
   177  	testUint64sCompare(t, archsimd.Uint64s.Equal, eqWant[uint64])
   178  	testFloat32sCompare(t, archsimd.Float32s.Equal, eqWant[float32])
   179  	testFloat64sCompare(t, archsimd.Float64s.Equal, eqWant[float64])
   180  }
   181  
   182  func TestNotEqualSVE(t *testing.T) {
   183  	if !archsimd.ARM64.SVE() {
   184  		t.Skip("no SVE")
   185  	}
   186  	testInt8sCompare(t, archsimd.Int8s.NotEqual, neWant[int8])
   187  	testInt16sCompare(t, archsimd.Int16s.NotEqual, neWant[int16])
   188  	testInt32sCompare(t, archsimd.Int32s.NotEqual, neWant[int32])
   189  	testInt64sCompare(t, archsimd.Int64s.NotEqual, neWant[int64])
   190  	testUint8sCompare(t, archsimd.Uint8s.NotEqual, neWant[uint8])
   191  	testUint16sCompare(t, archsimd.Uint16s.NotEqual, neWant[uint16])
   192  	testUint32sCompare(t, archsimd.Uint32s.NotEqual, neWant[uint32])
   193  	testUint64sCompare(t, archsimd.Uint64s.NotEqual, neWant[uint64])
   194  	testFloat32sCompare(t, archsimd.Float32s.NotEqual, neWant[float32])
   195  	testFloat64sCompare(t, archsimd.Float64s.NotEqual, neWant[float64])
   196  }
   197  
   198  func TestGreaterEqualSVE(t *testing.T) {
   199  	if !archsimd.ARM64.SVE() {
   200  		t.Skip("no SVE")
   201  	}
   202  	testInt8sCompare(t, archsimd.Int8s.GreaterEqual, geWant[int8])
   203  	testInt16sCompare(t, archsimd.Int16s.GreaterEqual, geWant[int16])
   204  	testInt32sCompare(t, archsimd.Int32s.GreaterEqual, geWant[int32])
   205  	testInt64sCompare(t, archsimd.Int64s.GreaterEqual, geWant[int64])
   206  	testUint8sCompare(t, archsimd.Uint8s.GreaterEqual, geWant[uint8])
   207  	testUint16sCompare(t, archsimd.Uint16s.GreaterEqual, geWant[uint16])
   208  	testUint32sCompare(t, archsimd.Uint32s.GreaterEqual, geWant[uint32])
   209  	testUint64sCompare(t, archsimd.Uint64s.GreaterEqual, geWant[uint64])
   210  	testFloat32sCompare(t, archsimd.Float32s.GreaterEqual, geWant[float32])
   211  	testFloat64sCompare(t, archsimd.Float64s.GreaterEqual, geWant[float64])
   212  }
   213
View as plain text