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