shift_amd64_test.go

     1  // Copyright 2025 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 && amd64
     6  
     7  package simd_test
     8  
     9  import (
    10  	"simd/archsimd"
    11  	"testing"
    12  )
    13  
    14  func TestRotateAllLeftAMD64(t *testing.T) {
    15  	x := uint8(0x81)
    16  	if y := rotl(x, 1); y != 3 {
    17  		t.Errorf("Expected 3, got 0x%x", y)
    18  	}
    19  	if y := rotl(x, 7); y != 0xc0 {
    20  		t.Errorf("Expected 0xc0, got 0x%x", y)
    21  	}
    22  	if y := rotr(x, 4); y != 0x18 {
    23  		t.Errorf("Expected 0x18, got 0x%x", y)
    24  	}
    25  
    26  	for i := uint64(0); i < 65; i++ {
    27  		testUint64x4Unary(t, curry2(archsimd.Uint64x4.RotateAllLeft, i), rotlOfSlice[uint64](i))
    28  		testUint32x8Unary(t, curry2(archsimd.Uint32x8.RotateAllLeft, i), rotlOfSlice[uint32](i))
    29  		//		testUint16x16Unary(t, curry2(archsimd.Uint16x16.RotateAllLeft, i), rotlOfSlice[uint16](i))
    30  		//		testUint8x32Unary(t, curry2(archsimd.Uint8x32.RotateAllLeft, i), rotlOfSlice[uint8](i))
    31  	}
    32  
    33  }
    34  
    35  func TestRotateAllRightAMD64(t *testing.T) {
    36  	x := uint8(0x81)
    37  	if y := rotr(x, 1); y != 0xc0 {
    38  		t.Errorf("Expected 0xc0, got 0x%x", y)
    39  	}
    40  	if y := rotr(x, 7); y != 3 {
    41  		t.Errorf("Expected 3, got 0x%x", y)
    42  	}
    43  	if y := rotr(x, 4); y != 0x18 {
    44  		t.Errorf("Expected 0x18, got 0x%x", y)
    45  	}
    46  
    47  	for i := uint64(0); i < 65; i++ {
    48  		testUint64x4Unary(t, curry2(archsimd.Uint64x4.RotateAllRight, i), rotrOfSlice[uint64](i))
    49  		testUint32x8Unary(t, curry2(archsimd.Uint32x8.RotateAllRight, i), rotrOfSlice[uint32](i))
    50  		//		testUint16x16Unary(t, curry2(archsimd.Uint16x16.RotateAllLeft, i), rotlOfSlice[uint16](i))
    51  		//		testUint8x32Unary(t, curry2(archsimd.Uint8x32.RotateAllLeft, i), rotlOfSlice[uint8](i))
    52  	}
    53  }
    54  
    55  func TestShift(t *testing.T) {
    56  	if !archsimd.X86.AVX2() {
    57  		t.Skip("requires AVX2")
    58  	}
    59  
    60  	testInt32x4Binary(t,
    61  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftLeft(y.AsUint32x4()) },
    62  		map2(func(x, y int32) int32 { return x << uint32(y) }))
    63  	testInt32x4Binary(t,
    64  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftRight(y.AsUint32x4()) },
    65  		map2(func(x, y int32) int32 { return x >> uint32(y) }))
    66  	testUint32x4Binary(t,
    67  		func(x, y archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftRight(y) },
    68  		map2(func(x, y uint32) uint32 { return x >> y }))
    69  }
    70  
    71  func concatInt32s(x, y int32) int64 {
    72  	return (int64(x) << 32) | int64(uint32(y))
    73  }
    74  
    75  func concatUint32s(x, y uint32) uint64 {
    76  	return (uint64(x) << 32) | uint64(y)
    77  }
    78  
    79  func TestShiftAllConcat(t *testing.T) {
    80  	if !archsimd.X86.AVX512VBMI2() {
    81  		t.Skip("requires AVX512-VBMI2")
    82  	}
    83  
    84  	// Note that unlike their non-Concat counterparts, these wrap the shift count.
    85  
    86  	hide := hideConst[uint64]
    87  
    88  	// ShiftAllLeftConcat
    89  	salc := func(shift uint64) func(x, y int32) int32 {
    90  		return func(x, y int32) int32 {
    91  			return int32(concatInt32s(x, y) >> (32 - shift%32))
    92  		}
    93  	}
    94  
    95  	testInt32x4Binary(t,
    96  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllLeftConcatMod32(y, 2) },
    97  		map2(salc(2)))
    98  	testInt32x4Binary(t,
    99  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllLeftConcatMod32(y, hide(2)) },
   100  		map2(salc(hide(2))))
   101  
   102  	testInt32x4Binary(t,
   103  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllLeftConcatMod32(y, 128) },
   104  		map2(salc(128)))
   105  	testInt32x4Binary(t,
   106  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllLeftConcatMod32(y, hide(128)) },
   107  		map2(salc(hide(128))))
   108  
   109  	// Signed ShiftAllRightConcat
   110  	sarc := func(shift uint64) func(x, y int32) int32 {
   111  		return func(x, y int32) int32 {
   112  			return int32(concatInt32s(y, x) >> (shift % 32))
   113  		}
   114  	}
   115  
   116  	testInt32x4Binary(t,
   117  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, 2) },
   118  		map2(sarc(2)))
   119  	testInt32x4Binary(t,
   120  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, hide(2)) },
   121  		map2(sarc(hide(2))))
   122  
   123  	testInt32x4Binary(t,
   124  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, 128) },
   125  		map2(sarc(128)))
   126  	testInt32x4Binary(t,
   127  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, hide(128)) },
   128  		map2(sarc(hide(128))))
   129  
   130  	testInt32x4Binary(t,
   131  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, 256) },
   132  		map2(sarc(256)))
   133  	testInt32x4Binary(t,
   134  		func(x, y archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftAllRightConcatMod32(y, hide(256)) },
   135  		map2(sarc(hide(256))))
   136  
   137  	// Unsigned ShiftAllRightConcat
   138  	usarc := func(shift uint64) func(x, y uint32) uint32 {
   139  		return func(x, y uint32) uint32 {
   140  			return uint32(concatUint32s(y, x) >> (shift % 32))
   141  		}
   142  	}
   143  
   144  	testUint32x4Binary(t,
   145  		func(x, y archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftAllRightConcatMod32(y, 2) },
   146  		map2(usarc(2)))
   147  	testUint32x4Binary(t,
   148  		func(x, y archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftAllRightConcatMod32(y, hide(2)) },
   149  		map2(usarc(hide(2))))
   150  
   151  	testUint32x4Binary(t,
   152  		func(x, y archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftAllRightConcatMod32(y, 128) },
   153  		map2(usarc(128)))
   154  	testUint32x4Binary(t,
   155  		func(x, y archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftAllRightConcatMod32(y, hide(128)) },
   156  		map2(usarc(hide(128))))
   157  }
   158  
   159  func TestShiftConcat(t *testing.T) {
   160  	if !archsimd.X86.AVX512VBMI2() {
   161  		t.Skip("requires AVX512-VBMI2")
   162  	}
   163  
   164  	// Note that unlike their non-Concat counterparts, these wrap the shift count.
   165  
   166  	testInt32x4Ternary(t,
   167  		func(x, y, z archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftLeftConcatMod32(y, z.AsUint32x4()) },
   168  		map3(func(x, y, z int32) int32 {
   169  			return int32(concatInt32s(x, y) >> (32 - uint32(z)%32))
   170  		}))
   171  
   172  	testInt32x4Ternary(t,
   173  		func(x, y, z archsimd.Int32x4) archsimd.Int32x4 { return x.ShiftRightConcatMod32(y, z.AsUint32x4()) },
   174  		map3(func(x, y, z int32) int32 {
   175  			return int32(concatInt32s(y, x) >> (uint32(z) % 32))
   176  		}))
   177  
   178  	testUint32x4Ternary(t,
   179  		func(x, y, z archsimd.Uint32x4) archsimd.Uint32x4 { return x.ShiftRightConcatMod32(y, z) },
   180  		map3(func(x, y, z uint32) uint32 {
   181  			return uint32(concatUint32s(y, x) >> (z % 32))
   182  		}))
   183  }
   184  
   185  func TestConcatShiftBytesRight(t *testing.T) {
   186  	hide := hideConst[uint64]
   187  
   188  	csbr := func(shift uint64) func(x, y []uint8) []uint8 {
   189  		return func(x, y []uint8) []uint8 {
   190  			z := make([]uint8, len(x))
   191  			for i := range z {
   192  				target := i + int(shift)
   193  				if target < 16 {
   194  					z[i] = y[target]
   195  				} else if target < 32 {
   196  					z[i] = x[(target - 16)]
   197  				}
   198  			}
   199  			return z
   200  		}
   201  	}
   202  
   203  	t.Run("Uint8x16", func(t *testing.T) {
   204  		if !archsimd.X86.AVX() {
   205  			t.Skip("requires AVX")
   206  		}
   207  		for _, shift := range []uint64{0, 2, 16, 20, 32, 128} {
   208  			t.Log("shift", shift)
   209  			testUint8x16Binary(t,
   210  				func(x, y archsimd.Uint8x16) archsimd.Uint8x16 { return x.ConcatShiftBytesRight(y, shift) },
   211  				csbr(shift))
   212  			testUint8x16Binary(t,
   213  				func(x, y archsimd.Uint8x16) archsimd.Uint8x16 { return x.ConcatShiftBytesRight(y, hide(shift)) },
   214  				csbr(hide(shift)))
   215  		}
   216  	})
   217  
   218  	t.Run("Uint8x32", func(t *testing.T) {
   219  		if !archsimd.X86.AVX2() {
   220  			t.Skip("requires AVX2")
   221  		}
   222  		for _, shift := range []uint64{0, 2, 16, 20, 32, 128} {
   223  			t.Log("shift", shift)
   224  			testUint8x32Binary(t,
   225  				func(x, y archsimd.Uint8x32) archsimd.Uint8x32 { return x.ConcatShiftBytesRightGrouped(y, shift) },
   226  				grouped2(csbr(shift)))
   227  			testUint8x32Binary(t,
   228  				func(x, y archsimd.Uint8x32) archsimd.Uint8x32 { return x.ConcatShiftBytesRightGrouped(y, hide(shift)) },
   229  				grouped2(csbr(hide(shift))))
   230  		}
   231  	})
   232  
   233  	t.Run("Uint8x64", func(t *testing.T) {
   234  		if !archsimd.X86.AVX512() {
   235  			t.Skip("requires AVX512")
   236  		}
   237  		for _, shift := range []uint64{0, 2, 16, 20, 32, 128} {
   238  			t.Log("shift", shift)
   239  			testUint8x64Binary(t,
   240  				func(x, y archsimd.Uint8x64) archsimd.Uint8x64 { return x.ConcatShiftBytesRightGrouped(y, shift) },
   241  				grouped2(csbr(shift)))
   242  			testUint8x64Binary(t,
   243  				func(x, y archsimd.Uint8x64) archsimd.Uint8x64 { return x.ConcatShiftBytesRightGrouped(y, hide(shift)) },
   244  				grouped2(csbr(hide(shift))))
   245  		}
   246  	})
   247  }
   248  
   249  func TestShiftAllAMD64(t *testing.T) {
   250  	if archsimd.X86.AVX2() {
   251  		// ShiftAllLeft
   252  		testInt16x16ShiftAll(t, archsimd.Int16x16.ShiftAllLeft, shiftAllLeftSlice[int16])
   253  		testInt32x8ShiftAll(t, archsimd.Int32x8.ShiftAllLeft, shiftAllLeftSlice[int32])
   254  		testInt64x4ShiftAll(t, archsimd.Int64x4.ShiftAllLeft, shiftAllLeftSlice[int64])
   255  		testUint16x16ShiftAll(t, archsimd.Uint16x16.ShiftAllLeft, shiftAllLeftSlice[uint16])
   256  		testUint32x8ShiftAll(t, archsimd.Uint32x8.ShiftAllLeft, shiftAllLeftSlice[uint32])
   257  		testUint64x4ShiftAll(t, archsimd.Uint64x4.ShiftAllLeft, shiftAllLeftSlice[uint64])
   258  
   259  		// ShiftAllRight signed
   260  		testInt16x16ShiftAll(t, archsimd.Int16x16.ShiftAllRight, shiftAllRightSlice[int16])
   261  		testInt32x8ShiftAll(t, archsimd.Int32x8.ShiftAllRight, shiftAllRightSlice[int32])
   262  		// Int64x4 ShiftAllRight requires AVX-512
   263  
   264  		// ShiftAllRight unsigned
   265  		testUint16x16ShiftAll(t, archsimd.Uint16x16.ShiftAllRight, shiftAllRightSlice[uint16])
   266  		testUint32x8ShiftAll(t, archsimd.Uint32x8.ShiftAllRight, shiftAllRightSlice[uint32])
   267  		testUint64x4ShiftAll(t, archsimd.Uint64x4.ShiftAllRight, shiftAllRightSlice[uint64])
   268  	}
   269  
   270  	if archsimd.X86.AVX512() {
   271  		// 512-bit vectors (AVX512)
   272  		// ShiftAllLeft
   273  		testInt16x32ShiftAll(t, archsimd.Int16x32.ShiftAllLeft, shiftAllLeftSlice[int16])
   274  		testInt32x16ShiftAll(t, archsimd.Int32x16.ShiftAllLeft, shiftAllLeftSlice[int32])
   275  		testInt64x8ShiftAll(t, archsimd.Int64x8.ShiftAllLeft, shiftAllLeftSlice[int64])
   276  		testUint16x32ShiftAll(t, archsimd.Uint16x32.ShiftAllLeft, shiftAllLeftSlice[uint16])
   277  		testUint32x16ShiftAll(t, archsimd.Uint32x16.ShiftAllLeft, shiftAllLeftSlice[uint32])
   278  		testUint64x8ShiftAll(t, archsimd.Uint64x8.ShiftAllLeft, shiftAllLeftSlice[uint64])
   279  
   280  		// ShiftAllRight signed
   281  		testInt16x32ShiftAll(t, archsimd.Int16x32.ShiftAllRight, shiftAllRightSlice[int16])
   282  		testInt32x16ShiftAll(t, archsimd.Int32x16.ShiftAllRight, shiftAllRightSlice[int32])
   283  		testInt64x8ShiftAll(t, archsimd.Int64x8.ShiftAllRight, shiftAllRightSlice[int64])
   284  		// 256-bit Int64x4 ShiftAllRight (requires AVX-512)
   285  		testInt64x4ShiftAll(t, archsimd.Int64x4.ShiftAllRight, shiftAllRightSlice[int64])
   286  
   287  		// ShiftAllRight unsigned
   288  		testUint16x32ShiftAll(t, archsimd.Uint16x32.ShiftAllRight, shiftAllRightSlice[uint16])
   289  		testUint32x16ShiftAll(t, archsimd.Uint32x16.ShiftAllRight, shiftAllRightSlice[uint32])
   290  		testUint64x8ShiftAll(t, archsimd.Uint64x8.ShiftAllRight, shiftAllRightSlice[uint64])
   291  	}
   292  }
   293
View as plain text