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
View as plain text