simd/archsimd: add tests for IsNaN

Change-Id: I374ce84fd21c41a04e2d5964d8aa872545c6a8a7
Reviewed-on: https://go-review.googlesource.com/c/go/+/733661
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: David Chase <drchase@google.com>

5 files changed, 135 insertions, 3 deletions

diff --git a/src/simd/archsimd/_gen/tmplgen/main.go b/src/simd/archsimd/_gen/tmplgen/main.go
index e0d607af12..8db185e1e0 100644
--- a/src/simd/archsimd/_gen/tmplgen/main.go
+++ b/src/simd/archsimd/_gen/tmplgen/main.go
@@ -104,6 +104,11 @@ var uintShapes = &shapes{
 	uints: []int{8, 16, 32, 64},
 }
 
+var floatShapes = &shapes{
+	vecs:   []int{128, 256, 512},
+	floats: []int{32, 64},
+}
+
 var integerShapes = &shapes{
 	vecs:  []int{128, 256, 512},
 	ints:  []int{8, 16, 32, 64},
@@ -507,6 +512,22 @@ func test{{.VType}}Compare(t *testing.T, f func(_, _ archsimd.{{.VType}}) archsi
 }
 `)
 
+var compareUnaryTemplate = shapedTemplateOf(floatShapes, "compare_unary_helpers", `
+// test{{.VType}}UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func test{{.VType}}UnaryCompare(t *testing.T, f func(x archsimd.{{.VType}}) archsimd.Mask{{.WxC}}, want func(x []{{.Etype}}) []int64) {
+	n := {{.Count}}
+	t.Helper()
+	forSlice(t, {{.Etype}}s, n, func(x []{{.Etype}}) bool {
+	 	t.Helper()
+		a := archsimd.Load{{.VType}}Slice(x)
+		g := make([]int{{.EWidth}}, n)
+		f(a).ToInt{{.WxC}}().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() {t.Helper(); t.Logf("x=%v", x)})
+	})
+}
+`)
+
 // TODO this has not been tested yet.
 var compareMaskedTemplate = templateOf("comparemasked_helpers", `
 // test{{.VType}}CompareMasked tests the simd masked comparison method f against the expected behavior generated by want
@@ -941,7 +962,7 @@ func main() {
 		one(*th, curryTestPrologue("ternary simd methods"), ternaryTemplate, ternaryFlakyTemplate)
 	}
 	if *ch != "" {
-		one(*ch, curryTestPrologue("simd methods that compare two operands"), compareTemplate)
+		one(*ch, curryTestPrologue("simd methods that compare two operands"), compareTemplate, compareUnaryTemplate)
 	}
 	if *cmh != "" {
 		one(*cmh, curryTestPrologue("simd methods that compare two operands under a mask"), compareMaskedTemplate)
diff --git a/src/simd/archsimd/internal/simd_test/compare_helpers_test.go b/src/simd/archsimd/internal/simd_test/compare_helpers_test.go
index c520aa7f17..7a33f0ffa4 100644
--- a/src/simd/archsimd/internal/simd_test/compare_helpers_test.go
+++ b/src/simd/archsimd/internal/simd_test/compare_helpers_test.go
@@ -462,3 +462,87 @@ func testFloat64x8Compare(t *testing.T, f func(_, _ archsimd.Float64x8) archsimd
 		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x); t.Logf("y=%v", y) })
 	})
 }
+
+// testFloat32x4UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat32x4UnaryCompare(t *testing.T, f func(x archsimd.Float32x4) archsimd.Mask32x4, want func(x []float32) []int64) {
+	n := 4
+	t.Helper()
+	forSlice(t, float32s, n, func(x []float32) bool {
+		t.Helper()
+		a := archsimd.LoadFloat32x4Slice(x)
+		g := make([]int32, n)
+		f(a).ToInt32x4().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
+
+// testFloat64x2UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat64x2UnaryCompare(t *testing.T, f func(x archsimd.Float64x2) archsimd.Mask64x2, want func(x []float64) []int64) {
+	n := 2
+	t.Helper()
+	forSlice(t, float64s, n, func(x []float64) bool {
+		t.Helper()
+		a := archsimd.LoadFloat64x2Slice(x)
+		g := make([]int64, n)
+		f(a).ToInt64x2().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
+
+// testFloat32x8UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat32x8UnaryCompare(t *testing.T, f func(x archsimd.Float32x8) archsimd.Mask32x8, want func(x []float32) []int64) {
+	n := 8
+	t.Helper()
+	forSlice(t, float32s, n, func(x []float32) bool {
+		t.Helper()
+		a := archsimd.LoadFloat32x8Slice(x)
+		g := make([]int32, n)
+		f(a).ToInt32x8().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
+
+// testFloat64x4UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat64x4UnaryCompare(t *testing.T, f func(x archsimd.Float64x4) archsimd.Mask64x4, want func(x []float64) []int64) {
+	n := 4
+	t.Helper()
+	forSlice(t, float64s, n, func(x []float64) bool {
+		t.Helper()
+		a := archsimd.LoadFloat64x4Slice(x)
+		g := make([]int64, n)
+		f(a).ToInt64x4().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
+
+// testFloat32x16UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat32x16UnaryCompare(t *testing.T, f func(x archsimd.Float32x16) archsimd.Mask32x16, want func(x []float32) []int64) {
+	n := 16
+	t.Helper()
+	forSlice(t, float32s, n, func(x []float32) bool {
+		t.Helper()
+		a := archsimd.LoadFloat32x16Slice(x)
+		g := make([]int32, n)
+		f(a).ToInt32x16().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
+
+// testFloat64x8UnaryCompare tests the simd unary comparison method f against the expected behavior generated by want
+func testFloat64x8UnaryCompare(t *testing.T, f func(x archsimd.Float64x8) archsimd.Mask64x8, want func(x []float64) []int64) {
+	n := 8
+	t.Helper()
+	forSlice(t, float64s, n, func(x []float64) bool {
+		t.Helper()
+		a := archsimd.LoadFloat64x8Slice(x)
+		g := make([]int64, n)
+		f(a).ToInt64x8().StoreSlice(g)
+		w := want(x)
+		return checkSlicesLogInput(t, s64(g), w, 0.0, func() { t.Helper(); t.Logf("x=%v", x) })
+	})
+}
diff --git a/src/simd/archsimd/internal/simd_test/compare_test.go b/src/simd/archsimd/internal/simd_test/compare_test.go
index efafbfdc35..e678676be0 100644
--- a/src/simd/archsimd/internal/simd_test/compare_test.go
+++ b/src/simd/archsimd/internal/simd_test/compare_test.go
@@ -298,3 +298,15 @@ func TestNotEqual(t *testing.T) {
 		testUint64x8Compare(t, archsimd.Uint64x8.NotEqual, notEqualSlice[uint64])
 	}
 }
+
+func TestIsNaN(t *testing.T) {
+	testFloat32x4UnaryCompare(t, archsimd.Float32x4.IsNaN, isNaNSlice[float32])
+	testFloat32x8UnaryCompare(t, archsimd.Float32x8.IsNaN, isNaNSlice[float32])
+	testFloat64x2UnaryCompare(t, archsimd.Float64x2.IsNaN, isNaNSlice[float64])
+	testFloat64x4UnaryCompare(t, archsimd.Float64x4.IsNaN, isNaNSlice[float64])
+
+	if archsimd.X86.AVX512() {
+		testFloat32x16UnaryCompare(t, archsimd.Float32x16.IsNaN, isNaNSlice[float32])
+		testFloat64x8UnaryCompare(t, archsimd.Float64x8.IsNaN, isNaNSlice[float64])
+	}
+}
diff --git a/src/simd/archsimd/internal/simd_test/helpers_test.go b/src/simd/archsimd/internal/simd_test/helpers_test.go
index 37cc98194a..d6963586c0 100644
--- a/src/simd/archsimd/internal/simd_test/helpers_test.go
+++ b/src/simd/archsimd/internal/simd_test/helpers_test.go
@@ -182,12 +182,14 @@ var nzero = -zero
 var inf = 1 / zero
 var ninf = -1 / zero
 var nan = math.NaN()
+var snan32 = math.Float32frombits(0x7f800001)
+var snan64 = math.Float64frombits(0x7ff0000000000001)
 
 // N controls how large the test vectors are
 const N = 144
 
-var float32s = nOf(N, []float32{float32(inf), float32(ninf), 1, float32(nan), float32(zero), 2, float32(nan), float32(zero), 3, float32(-zero), float32(1.0 / zero), float32(-1.0 / zero), 1.0 / 2, 1.0 / 4, 1.0 / 8, 1.0 / 1000, 1.0 / 1000000, 1, -1, 0, 2, -2, 3, -3, math.MaxFloat32, 1 / math.MaxFloat32, 10, -10, 100, 20, -20, 300, -300, -4000, -80, -160, -3200, -64, -4, -8, -16, -32, -64})
-var float64s = nOf(N, []float64{inf, ninf, nan, zero, -zero, 1 / zero, -1 / zero, 0.0001, 0.0000001, 1, -1, 0, 2, -2, 3, -3, math.MaxFloat64, 1.0 / math.MaxFloat64, 10, -10, 100, 20, -20, 300, -300, -4000, -80, -16, -32, -64})
+var float32s = nOf(N, []float32{float32(inf), float32(ninf), 1, float32(nan), snan32, -float32(nan), -snan32, float32(zero), 2, float32(nan), float32(zero), 3, float32(-zero), float32(1.0 / zero), float32(-1.0 / zero), 1.0 / 2, 1.0 / 4, 1.0 / 8, 1.0 / 1000, 1.0 / 1000000, 1, -1, 0, 2, -2, 3, -3, math.MaxFloat32, 1 / math.MaxFloat32, 10, -10, 100, 20, -20, 300, -300, -4000, -80, -160, -3200, -64, -4, -8, -16, -32, -64})
+var float64s = nOf(N, []float64{inf, ninf, nan, snan64, -nan, -snan64, zero, -zero, 1 / zero, -1 / zero, 0.0001, 0.0000001, 1, -1, 0, 2, -2, 3, -3, math.MaxFloat64, 1.0 / math.MaxFloat64, 10, -10, 100, 20, -20, 300, -300, -4000, -80, -16, -32, -64})
 
 var int32s = nOf(N, []int32{1, -1, 0, 2, 4, 8, 1024, 0xffffff, -0xffffff, 0x55555, 0x77777, 0xccccc, -0x55555, -0x77777, -0xccccc, -4, -8, -16, -32, -64})
 var uint32s = nOf(N, []uint32{1, 0, 2, 4, 8, 1024, 0xffffff, ^uint32(0xffffff), 0x55555, 0x77777, 0xccccc, ^uint32(0x55555), ^uint32(0x77777), ^uint32(0xccccc)})
diff --git a/src/simd/archsimd/internal/simd_test/simulation_helpers_test.go b/src/simd/archsimd/internal/simd_test/simulation_helpers_test.go
index da107d8061..ac60b6d377 100644
--- a/src/simd/archsimd/internal/simd_test/simulation_helpers_test.go
+++ b/src/simd/archsimd/internal/simd_test/simulation_helpers_test.go
@@ -29,6 +29,10 @@ func notEqual[T number](x, y T) bool {
 	return x != y
 }
 
+func isNaN[T float](x T) bool {
+	return x != x
+}
+
 func abs[T number](x T) T {
 	// TODO this will need a non-standard FP-equality test.
 	if x == 0 { // true if x is -0.
@@ -299,6 +303,15 @@ func notEqualSlice[T number](x, y []T) []int64 {
 	return mapCompare[T](notEqual)(x, y)
 }
 
+func isNaNSlice[T float](x []T) []int64 {
+	return map1[T](func(x T) int64 {
+		if isNaN(x) {
+			return -1
+		}
+		return 0
+	})(x)
+}
+
 func ceilSlice[T float](x []T) []T {
 	return map1[T](ceil)(x)
 }