mask_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 accessor tests. A mask is built from a uint64 of predicate bits
     8  // through its memory representation (see maskBits), and each accessor is
     9  // checked against a lane-by-lane reference over a set of lane patterns.
    10  
    11  package simd_test
    12  
    13  import (
    14  	"math"
    15  	"simd/archsimd"
    16  	"testing"
    17  	"unsafe"
    18  )
    19  
    20  type sveMask interface {
    21  	archsimd.Mask8s | archsimd.Mask16s | archsimd.Mask32s | archsimd.Mask64s
    22  }
    23  
    24  // maskFromBits returns the mask whose memory representation is bits.
    25  func maskFromBits[M sveMask](bits uint64) M {
    26  	return *(*M)(unsafe.Pointer(&bits))
    27  }
    28  
    29  // lanesOf returns the n lanes of bits for lanes laneBytes wide.
    30  func lanesOf(bits uint64, laneBytes, n int) []bool {
    31  	lanes := make([]bool, n)
    32  	for i := range lanes {
    33  		lanes[i] = bits>>(i*laneBytes)&1 != 0
    34  	}
    35  	return lanes
    36  }
    37  
    38  // bitsOf is the inverse of lanesOf, with the bits between lanes zero.
    39  func bitsOf(lanes []bool, laneBytes int) uint64 {
    40  	var bits uint64
    41  	for i, l := range lanes {
    42  		if l {
    43  			bits |= 1 << (i * laneBytes)
    44  		}
    45  	}
    46  	return bits
    47  }
    48  
    49  // maskPatterns returns lane patterns to test for n lanes laneBytes wide, as
    50  // predicate bits: none, all, each single lane, each prefix and suffix, and
    51  // alternating lanes. For lanes wider than a byte it adds the same patterns
    52  // with every bit between lanes set, bits a mask ignores.
    53  func maskPatterns(laneBytes, n int) []uint64 {
    54  	var lanes [][]bool
    55  	add := func(f func(i int) bool) {
    56  		l := make([]bool, n)
    57  		for i := range l {
    58  			l[i] = f(i)
    59  		}
    60  		lanes = append(lanes, l)
    61  	}
    62  	add(func(int) bool { return false })
    63  	add(func(int) bool { return true })
    64  	add(func(i int) bool { return i%2 == 0 })
    65  	add(func(i int) bool { return i%2 == 1 })
    66  	for k := 0; k < n; k++ {
    67  		add(func(i int) bool { return i == k })
    68  		add(func(i int) bool { return i < k })
    69  		add(func(i int) bool { return i >= k })
    70  	}
    71  	var pats []uint64
    72  	for _, l := range lanes {
    73  		pats = append(pats, bitsOf(l, laneBytes))
    74  	}
    75  	if laneBytes > 1 {
    76  		var between uint64
    77  		for i := range n * laneBytes {
    78  			if i%laneBytes != 0 {
    79  				between |= 1 << i
    80  			}
    81  		}
    82  		for _, l := range lanes {
    83  			pats = append(pats, bitsOf(l, laneBytes)|between)
    84  		}
    85  	}
    86  	return pats
    87  }
    88  
    89  // testMaskToMask checks a mask-to-mask accessor against want, which maps
    90  // input lanes to output lanes. The result must have only lane bits set.
    91  func testMaskToMask[M sveMask](t *testing.T, name string, laneBytes int, op func(M) M, want func([]bool) []bool) {
    92  	t.Helper()
    93  	n := archsimd.Int8s{}.Len() / laneBytes
    94  	for _, bits := range maskPatterns(laneBytes, n) {
    95  		got := maskBits(op(maskFromBits[M](bits)))
    96  		if w := bitsOf(want(lanesOf(bits, laneBytes, n)), laneBytes); got != w {
    97  			t.Errorf("%s(%#x) = %#x, want %#x", name, bits, got, w)
    98  		}
    99  	}
   100  }
   101  
   102  // testMaskToBool checks a mask-to-bool accessor against want.
   103  func testMaskToBool[M sveMask](t *testing.T, name string, laneBytes int, op func(M) bool, want func([]bool) bool) {
   104  	t.Helper()
   105  	n := archsimd.Int8s{}.Len() / laneBytes
   106  	for _, bits := range maskPatterns(laneBytes, n) {
   107  		if got, w := op(maskFromBits[M](bits)), want(lanesOf(bits, laneBytes, n)); got != w {
   108  			t.Errorf("%s(%#x) = %v, want %v", name, bits, got, w)
   109  		}
   110  	}
   111  }
   112  
   113  func firstWant(lanes []bool) []bool {
   114  	out := make([]bool, len(lanes))
   115  	for i, l := range lanes {
   116  		if l {
   117  			out[i] = true
   118  			break
   119  		}
   120  	}
   121  	return out
   122  }
   123  
   124  func TestMaskFirstSVE(t *testing.T) {
   125  	if !archsimd.ARM64.SVE() {
   126  		t.Skip("no SVE")
   127  	}
   128  	testMaskToMask(t, "Mask8s.First", 1, archsimd.Mask8s.First, firstWant)
   129  	testMaskToMask(t, "Mask16s.First", 2, archsimd.Mask16s.First, firstWant)
   130  	testMaskToMask(t, "Mask32s.First", 4, archsimd.Mask32s.First, firstWant)
   131  	testMaskToMask(t, "Mask64s.First", 8, archsimd.Mask64s.First, firstWant)
   132  }
   133  
   134  func nextWant(lanes []bool) []bool {
   135  	last := -1
   136  	for i, l := range lanes {
   137  		if l {
   138  			last = i
   139  		}
   140  	}
   141  	out := make([]bool, len(lanes))
   142  	if last+1 < len(lanes) {
   143  		out[last+1] = true
   144  	}
   145  	return out
   146  }
   147  
   148  func TestMaskNextSVE(t *testing.T) {
   149  	if !archsimd.ARM64.SVE() {
   150  		t.Skip("no SVE")
   151  	}
   152  	testMaskToMask(t, "Mask8s.Next", 1, archsimd.Mask8s.Next, nextWant)
   153  	testMaskToMask(t, "Mask16s.Next", 2, archsimd.Mask16s.Next, nextWant)
   154  	testMaskToMask(t, "Mask32s.Next", 4, archsimd.Mask32s.Next, nextWant)
   155  	testMaskToMask(t, "Mask64s.Next", 8, archsimd.Mask64s.Next, nextWant)
   156  }
   157  
   158  func allWant(lanes []bool) bool {
   159  	for _, l := range lanes {
   160  		if !l {
   161  			return false
   162  		}
   163  	}
   164  	return true
   165  }
   166  
   167  func TestMaskAllSVE(t *testing.T) {
   168  	if !archsimd.ARM64.SVE() {
   169  		t.Skip("no SVE")
   170  	}
   171  	testMaskToBool(t, "Mask8s.All", 1, archsimd.Mask8s.All, allWant)
   172  	testMaskToBool(t, "Mask16s.All", 2, archsimd.Mask16s.All, allWant)
   173  	testMaskToBool(t, "Mask32s.All", 4, archsimd.Mask32s.All, allWant)
   174  	testMaskToBool(t, "Mask64s.All", 8, archsimd.Mask64s.All, allWant)
   175  }
   176  
   177  func noneWant(lanes []bool) bool {
   178  	for _, l := range lanes {
   179  		if l {
   180  			return false
   181  		}
   182  	}
   183  	return true
   184  }
   185  
   186  func TestMaskNoneSVE(t *testing.T) {
   187  	if !archsimd.ARM64.SVE() {
   188  		t.Skip("no SVE")
   189  	}
   190  	testMaskToBool(t, "Mask8s.None", 1, archsimd.Mask8s.None, noneWant)
   191  	testMaskToBool(t, "Mask16s.None", 2, archsimd.Mask16s.None, noneWant)
   192  	testMaskToBool(t, "Mask32s.None", 4, archsimd.Mask32s.None, noneWant)
   193  	testMaskToBool(t, "Mask64s.None", 8, archsimd.Mask64s.None, noneWant)
   194  }
   195  
   196  func anyWant(lanes []bool) bool { return !noneWant(lanes) }
   197  
   198  func TestMaskAnySVE(t *testing.T) {
   199  	if !archsimd.ARM64.SVE() {
   200  		t.Skip("no SVE")
   201  	}
   202  	testMaskToBool(t, "Mask8s.Any", 1, archsimd.Mask8s.Any, anyWant)
   203  	testMaskToBool(t, "Mask16s.Any", 2, archsimd.Mask16s.Any, anyWant)
   204  	testMaskToBool(t, "Mask32s.Any", 4, archsimd.Mask32s.Any, anyWant)
   205  	testMaskToBool(t, "Mask64s.Any", 8, archsimd.Mask64s.Any, anyWant)
   206  }
   207  
   208  func testMaskToIntAndBack[M sveMask, T signed, V any](t *testing.T, name string, laneBytes int,
   209  	toInt func(M) V, toMask func(V) M, store func(V, []T), load func([]T) V) {
   210  	t.Helper()
   211  	n := archsimd.Int8s{}.Len() / laneBytes
   212  	for _, bits := range maskPatterns(laneBytes, n) {
   213  		m := maskFromBits[M](bits)
   214  		v := toInt(m)
   215  		got := make([]T, n)
   216  		store(v, got)
   217  		lanes := lanesOf(bits, laneBytes, n)
   218  		for i, l := range lanes {
   219  			var want T
   220  			if l {
   221  				want = -1
   222  			}
   223  			if got[i] != want {
   224  				t.Errorf("%s.ToInt(%#x)[%d] = %d, want %d", name, bits, i, got[i], want)
   225  			}
   226  		}
   227  		if gotBits, wantBits := maskBits(toMask(v)), bitsOf(lanes, laneBytes); gotBits != wantBits {
   228  			t.Errorf("%s.ToInt(%#x).ToMask() = %#x, want %#x", name, bits, gotBits, wantBits)
   229  		}
   230  		// Check non-canonical non-zero values for ToMask.
   231  		for i := range got {
   232  			if got[i] != 0 {
   233  				got[i] = T(i + 1)
   234  			}
   235  		}
   236  		if gotBits, wantBits := maskBits(toMask(load(got))), bitsOf(lanes, laneBytes); gotBits != wantBits {
   237  			t.Errorf("%s non-canonical ToMask(%#x) = %#x, want %#x", name, bits, gotBits, wantBits)
   238  		}
   239  	}
   240  }
   241  
   242  func TestMaskToIntAndToMaskSVE(t *testing.T) {
   243  	if !archsimd.ARM64.SVE() {
   244  		t.Skip("no SVE")
   245  	}
   246  	testMaskToIntAndBack(t, "Mask8s", 1, archsimd.Mask8s.ToInt8s, archsimd.Int8s.ToMask, archsimd.Int8s.Store, archsimd.LoadInt8s)
   247  	testMaskToIntAndBack(t, "Mask16s", 2, archsimd.Mask16s.ToInt16s, archsimd.Int16s.ToMask, archsimd.Int16s.Store, archsimd.LoadInt16s)
   248  	testMaskToIntAndBack(t, "Mask32s", 4, archsimd.Mask32s.ToInt32s, archsimd.Int32s.ToMask, archsimd.Int32s.Store, archsimd.LoadInt32s)
   249  	testMaskToIntAndBack(t, "Mask64s", 8, archsimd.Mask64s.ToInt64s, archsimd.Int64s.ToMask, archsimd.Int64s.Store, archsimd.LoadInt64s)
   250  }
   251  
   252  // testMaskFromCount checks a FromCount constructor against the prefix of
   253  // lanes it should activate, through the mask's memory representation.
   254  func testMaskFromCount[M sveMask](t *testing.T, name string, laneBytes, n int, fromCount func(int) M) {
   255  	t.Helper()
   256  	for _, c := range []int{math.MinInt, -1, 0, 1, n / 2, n - 1, n, n + 1, math.MaxInt} {
   257  		m := fromCount(c)
   258  		got := *(*uint64)(unsafe.Pointer(&m))
   259  		lanes := make([]bool, n)
   260  		for i := range lanes {
   261  			lanes[i] = i < c
   262  		}
   263  		if want := bitsOf(lanes, laneBytes); got != want {
   264  			t.Errorf("%s(%d) = %#x, want %#x", name, c, got, want)
   265  		}
   266  	}
   267  }
   268  
   269  func TestMaskFromCountSVE(t *testing.T) {
   270  	if !archsimd.ARM64.SVE() {
   271  		t.Skip("no SVE")
   272  	}
   273  	testMaskFromCount(t, "Mask8sFromCount", 1, archsimd.Int8s{}.Len(), archsimd.Mask8sFromCount)
   274  	testMaskFromCount(t, "Mask16sFromCount", 2, archsimd.Int16s{}.Len(), archsimd.Mask16sFromCount)
   275  	testMaskFromCount(t, "Mask32sFromCount", 4, archsimd.Int32s{}.Len(), archsimd.Mask32sFromCount)
   276  	testMaskFromCount(t, "Mask64sFromCount", 8, archsimd.Int64s{}.Len(), archsimd.Mask64sFromCount)
   277  }
   278
View as plain text