ENH: Add SIMD operation signbit

howjmay · howjmay · commit 08fa7e6cc4e1 · 2021-08-26T16:30:50.000+08:00
diff --git a/numpy/core/src/_simd/_simd.dispatch.c.src b/numpy/core/src/_simd/_simd.dispatch.c.src
@@ -316,6 +316,10 @@ SIMD_IMPL_INTRIN_2(@intrin@_@sfx@, v@sfx@, v@sfx@, v@sfx@)
 
 SIMD_IMPL_INTRIN_1(not_@sfx@, v@sfx@, v@sfx@)
 
+#if @fp_only@
+SIMD_IMPL_INTRIN_1(signbit_@sfx@, v@bsfx@, v@sfx@)
+#endif
+
 /**begin repeat1
  * #intrin = cmpeq, cmpneq, cmpgt, cmpge, cmplt, cmple#
  */
@@ -551,6 +555,10 @@ SIMD_INTRIN_DEF(@intrin@_@sfx@)
 /**end repeat1**/
 #endif // shl_imm
 
+#if @fp_only@
+SIMD_INTRIN_DEF(signbit_@sfx@)
+#endif
+
 /**begin repeat1
  * #intrin = and, or, xor, not, cmpeq, cmpneq, cmpgt, cmpge, cmplt, cmple#
  */
diff --git a/numpy/core/src/common/simd/avx2/operators.h b/numpy/core/src/common/simd/avx2/operators.h
@@ -219,4 +219,12 @@ NPY_FINLINE npyv_b32 npyv_notnan_f32(npyv_f32 a)
 NPY_FINLINE npyv_b64 npyv_notnan_f64(npyv_f64 a)
 { return _mm256_castpd_si256(_mm256_cmp_pd(a, a, _CMP_ORD_Q)); }
 
+// signbit
+NPY_FINLINE npyv_b32 npyv_signbit_f32(npyv_f32 a) {
+    return (npyv_shri_s32(_mm256_castps_si256(a), 31));
+}
+NPY_FINLINE npyv_b64 npyv_signbit_f64(npyv_f64 a) {
+    return (npyv_shri_s64(_mm256_castpd_si256(a), 63));
+}
+
 #endif // _NPY_SIMD_AVX2_OPERATORS_H
diff --git a/numpy/core/src/common/simd/avx512/operators.h b/numpy/core/src/common/simd/avx512/operators.h
@@ -5,6 +5,8 @@
 #ifndef _NPY_SIMD_AVX512_OPERATORS_H
 #define _NPY_SIMD_AVX512_OPERATORS_H
 
+#include "conversion.h"
+
 /***************************
  * Shifting
  ***************************/
@@ -321,4 +323,12 @@ NPY_FINLINE npyv_b32 npyv_notnan_f32(npyv_f32 a)
 NPY_FINLINE npyv_b64 npyv_notnan_f64(npyv_f64 a)
 { return _mm512_cmp_pd_mask(a, a, _CMP_ORD_Q); }
 
+// signbit
+NPY_FINLINE npyv_b32 npyv_signbit_f32(npyv_f32 a) {
+    return npyv_cvt_b32_s32(npyv_shri_s32(_mm512_castps_si512(a), 31));
+}
+NPY_FINLINE npyv_b64 npyv_signbit_f64(npyv_f64 a) {
+    return npyv_cvt_b64_s64(npyv_shri_s64(_mm512_castpd_si512(a), 63));
+}
+
 #endif // _NPY_SIMD_AVX512_OPERATORS_H
diff --git a/numpy/core/src/common/simd/neon/operators.h b/numpy/core/src/common/simd/neon/operators.h
@@ -5,6 +5,8 @@
 #ifndef _NPY_SIMD_NEON_OPERATORS_H
 #define _NPY_SIMD_NEON_OPERATORS_H
 
+#include "conversion.h"
+
 /***************************
  * Shifting
  ***************************/
@@ -240,4 +242,14 @@ NPY_FINLINE npyv_b32 npyv_notnan_f32(npyv_f32 a)
     { return vceqq_f64(a, a); }
 #endif
 
+// signbit
+NPY_FINLINE npyv_b32 npyv_signbit_f32(npyv_f32 a) {
+    return (vrshrq_n_u32(vreinterpretq_u32_f32(a), 31));
+}
+#if NPY_SIMD_F64
+NPY_FINLINE npyv_b64 npyv_signbit_f64(npyv_f64 a) {
+    return (vrshrq_n_u64(vreinterpretq_u64_f64(a), 63));
+}
+#endif
+
 #endif // _NPY_SIMD_NEON_OPERATORS_H
diff --git a/numpy/core/src/common/simd/sse/operators.h b/numpy/core/src/common/simd/sse/operators.h
@@ -277,4 +277,12 @@ NPY_FINLINE npyv_b32 npyv_notnan_f32(npyv_f32 a)
 NPY_FINLINE npyv_b64 npyv_notnan_f64(npyv_f64 a)
 { return _mm_castpd_si128(_mm_cmpord_pd(a, a)); }
 
+// signbit
+NPY_FINLINE npyv_b32 npyv_signbit_f32(npyv_f32 a) {
+    return (npyv_shri_s32(_mm_castps_si128(a), 31));
+}
+NPY_FINLINE npyv_b64 npyv_signbit_f64(npyv_f64 a) {
+    return (npyv_shri_s64(_mm_castpd_si128(a), 63));
+}
+
 #endif // _NPY_SIMD_SSE_OPERATORS_H
diff --git a/numpy/core/src/common/simd/vsx/operators.h b/numpy/core/src/common/simd/vsx/operators.h
@@ -5,6 +5,8 @@
 #ifndef _NPY_SIMD_VSX_OPERATORS_H
 #define _NPY_SIMD_VSX_OPERATORS_H
 
+#include "conversion.h"
+
 /***************************
  * Shifting
  ***************************/
@@ -241,4 +243,12 @@ NPY_FINLINE npyv_b32 npyv_notnan_f32(npyv_f32 a)
 NPY_FINLINE npyv_b64 npyv_notnan_f64(npyv_f64 a)
 { return vec_cmpeq(a, a); }
 
+// signbit
+NPY_FINLINE npyv_b32 npyv_signbit_f32(npyv_f32 a) {
+    return npyv_cvt_b32_f32(npyv_and_f32(a, npyv_cvt_f32_b32(npyv_setall_u32(0x80000000))));
+}
+NPY_FINLINE npyv_b64 npyv_signbit_f64(npyv_f64 a) {
+    return npyv_cvt_b64_f64(npyv_and_f64(a, npyv_cvt_f64_b64(npyv_setall_u64(0x8000000000000000))));
+}
+
 #endif // _NPY_SIMD_VSX_OPERATORS_H
diff --git a/numpy/core/tests/test_simd.py b/numpy/core/tests/test_simd.py
@@ -361,6 +361,27 @@ def test_max(self):
             _max = self.max(vdata_a, vdata_b)
             assert _max == data_max
 
+    def test_signbit(self):
+        pinf, ninf, nan = self._pinfinity(), self._ninfinity(), self._nan()
+        data = self._data()
+        vdata = self.load(data)
+
+        mask_true = self._true_mask()
+
+        def to_bool(vector):
+            return [lane == mask_true for lane in vector]
+        
+        signbit_cases = ((nan, False), (-nan, True), (pinf, False),
+         (ninf, True))
+        for case, desired in signbit_cases:
+            _data_signbit = [desired]*self.nlanes
+            vsignbit = to_bool(self.signbit(self.setall(case)))
+            assert vsignbit == pytest.approx(_data_signbit, nan_ok=True)
+
+        data_signbit = [a < 0 for a in data]
+        vadata_signbit = to_bool(self.signbit(vdata))
+        assert vadata_signbit == data_signbit
+
     def test_min(self):
         """
         Test intrinsics:
