Bit shifting functions

Eigen/src/Core/GenericPacketMath.h

@@ -709,33 +709,21 @@ EIGEN_DEVICE_FUNC inline Packet parg(const Packet& a) {
 }
 
 /** \internal \returns \a a arithmetically shifted by N bits to the right */
-template <int N>
+template <int N, typename T>
-EIGEN_DEVICE_FUNC inline int parithmetic_shift_right(const int& a) {
+EIGEN_DEVICE_FUNC inline T parithmetic_shift_right(const T& a) {
-  return a >> N;
+  return numext::arithmetic_shift_right(a, N);
-}
-template <int N>
-EIGEN_DEVICE_FUNC inline long int parithmetic_shift_right(const long int& a) {
-  return a >> N;
 }
 
 /** \internal \returns \a a logically shifted by N bits to the right */
-template <int N>
+template <int N, typename T>
-EIGEN_DEVICE_FUNC inline int plogical_shift_right(const int& a) {
+EIGEN_DEVICE_FUNC inline T plogical_shift_right(const T& a) {
-  return static_cast<int>(static_cast<unsigned int>(a) >> N);
+  return numext::logical_shift_right(a, N);
-}
-template <int N>
-EIGEN_DEVICE_FUNC inline long int plogical_shift_right(const long int& a) {
-  return static_cast<long>(static_cast<unsigned long>(a) >> N);
 }
 
 /** \internal \returns \a a shifted by N bits to the left */
-template <int N>
+template <int N, typename T>
-EIGEN_DEVICE_FUNC inline int plogical_shift_left(const int& a) {
+EIGEN_DEVICE_FUNC inline T plogical_shift_left(const T& a) {
-  return a << N;
+  return numext::logical_shift_left(a, N);
-}
-template <int N>
-EIGEN_DEVICE_FUNC inline long int plogical_shift_left(const long int& a) {
-  return a << N;
 }
 
 /** \internal \returns the significant and exponent of the underlying floating point numbers

Eigen/src/Core/MathFunctions.h

@@ -1746,6 +1746,23 @@ EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE double fmod(const double& a, const double&
 #undef SYCL_SPECIALIZE_BINARY_FUNC
 #endif
 
+template <typename Scalar, typename Enable = std::enable_if_t<std::is_integral<Scalar>::value>>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar logical_shift_left(const Scalar& a, int n) {
+  return a << n;
+}
+
+template <typename Scalar, typename Enable = std::enable_if_t<std::is_integral<Scalar>::value>>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar logical_shift_right(const Scalar& a, int n) {
+  using UnsignedScalar = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  return bit_cast<Scalar, UnsignedScalar>(bit_cast<UnsignedScalar, Scalar>(a) >> n);
+}
+
+template <typename Scalar, typename Enable = std::enable_if_t<std::is_integral<Scalar>::value>>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar arithmetic_shift_right(const Scalar& a, int n) {
+  using SignedScalar = typename numext::get_integer_by_size<sizeof(Scalar)>::signed_type;
+  return bit_cast<Scalar, SignedScalar>(bit_cast<SignedScalar, Scalar>(a) >> n);
+}
+
 }  // end namespace numext
 
 namespace internal {

Eigen/src/Core/NumTraits.h

@@ -101,10 +101,10 @@ namespace numext {
 template <typename Tgt, typename Src>
 EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Tgt bit_cast(const Src& src) {
   // The behaviour of memcpy is not specified for non-trivially copyable types
-  EIGEN_STATIC_ASSERT(std::is_trivially_copyable<Src>::value, THIS_TYPE_IS_NOT_SUPPORTED);
+  EIGEN_STATIC_ASSERT(std::is_trivially_copyable<Src>::value, THIS_TYPE_IS_NOT_SUPPORTED)
   EIGEN_STATIC_ASSERT(std::is_trivially_copyable<Tgt>::value && std::is_default_constructible<Tgt>::value,
-                      THIS_TYPE_IS_NOT_SUPPORTED);
+                      THIS_TYPE_IS_NOT_SUPPORTED)
-  EIGEN_STATIC_ASSERT(sizeof(Src) == sizeof(Tgt), THIS_TYPE_IS_NOT_SUPPORTED);
+  EIGEN_STATIC_ASSERT(sizeof(Src) == sizeof(Tgt), THIS_TYPE_IS_NOT_SUPPORTED)
 
   Tgt tgt;
   // Load src into registers first. This allows the memcpy to be elided by CUDA.

Eigen/src/Core/functors/UnaryFunctors.h

@@ -219,7 +219,9 @@ struct functor_traits<core_cast_op<SrcType, DstType>> {
  */
 template <typename Scalar, int N>
 struct scalar_shift_right_op {
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(const Scalar& a) const { return a >> N; }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(const Scalar& a) const {
+    return numext::arithmetic_shift_right(a);
+  }
   template <typename Packet>
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(const Packet& a) const {
     return internal::parithmetic_shift_right<N>(a);
@@ -237,7 +239,9 @@ struct functor_traits<scalar_shift_right_op<Scalar, N>> {
  */
 template <typename Scalar, int N>
 struct scalar_shift_left_op {
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(const Scalar& a) const { return a << N; }
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Scalar operator()(const Scalar& a) const {
+    return numext::logical_shift_left(a);
+  }
   template <typename Packet>
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Packet packetOp(const Packet& a) const {
     return internal::plogical_shift_left<N>(a);

test/array_cwise.cpp

@@ -1068,24 +1068,45 @@ void min_max(const ArrayType& m) {
   }
 }
 
-template <int N>
+template <typename Scalar>
-struct shift_left {
+struct shift_imm_traits {
-  template <typename Scalar>
+  enum { Cost = 1, PacketAccess = internal::packet_traits<Scalar>::HasShift };
-  Scalar operator()(const Scalar& v) const {
+};
-    return (v << N);
+
+template <int N, typename Scalar>
+struct logical_left_shift_op {
+  Scalar operator()(const Scalar& v) const { return numext::logical_shift_left(v, N); }
+  template <typename Packet>
+  Packet packetOp(const Packet& v) const {
+    return internal::plogical_shift_left<N>(v);
+  }
+};
+template <int N, typename Scalar>
+struct logical_right_shift_op {
+  Scalar operator()(const Scalar& v) const { return numext::logical_shift_right(v, N); }
+  template <typename Packet>
+  Packet packetOp(const Packet& v) const {
+    return internal::plogical_shift_right<N>(v);
+  }
+};
+template <int N, typename Scalar>
+struct arithmetic_right_shift_op {
+  Scalar operator()(const Scalar& v) const { return numext::arithmetic_shift_right(v, N); }
+  template <typename Packet>
+  Packet packetOp(const Packet& v) const {
+    return internal::parithmetic_shift_right<N>(v);
   }
 };
 
-template <int N>
+template <int N, typename Scalar>
-struct arithmetic_shift_right {
+struct internal::functor_traits<logical_left_shift_op<N, Scalar>> : shift_imm_traits<Scalar> {};
-  template <typename Scalar>
+template <int N, typename Scalar>
-  Scalar operator()(const Scalar& v) const {
+struct internal::functor_traits<logical_right_shift_op<N, Scalar>> : shift_imm_traits<Scalar> {};
-    return (v >> N);
+template <int N, typename Scalar>
-  }
+struct internal::functor_traits<arithmetic_right_shift_op<N, Scalar>> : shift_imm_traits<Scalar> {};
-};
 
 template <typename ArrayType>
-struct signed_shift_test_impl {
+struct shift_test_impl {
   typedef typename ArrayType::Scalar Scalar;
   static constexpr size_t Size = sizeof(Scalar);
   static constexpr size_t MaxShift = (CHAR_BIT * Size) - 1;
@@ -1099,20 +1120,24 @@ struct signed_shift_test_impl {
 
     ArrayType m1 = ArrayType::Random(rows, cols), m2(rows, cols), m3(rows, cols);
 
-    m2 = m1.unaryExpr(internal::scalar_shift_right_op<Scalar, N>());
+    m2 = m1.unaryExpr([](const Scalar& v) { return numext::logical_shift_left(v, N); });
-    m3 = m1.unaryExpr(arithmetic_shift_right<N>());
+    m3 = m1.unaryExpr(logical_left_shift_op<N, Scalar>());
     VERIFY_IS_CWISE_EQUAL(m2, m3);
 
-    m2 = m1.unaryExpr(internal::scalar_shift_left_op<Scalar, N>());
+    m2 = m1.unaryExpr([](const Scalar& v) { return numext::logical_shift_right(v, N); });
-    m3 = m1.unaryExpr(shift_left<N>());
+    m3 = m1.unaryExpr(logical_right_shift_op<N, Scalar>());
+    VERIFY_IS_CWISE_EQUAL(m2, m3);
+
+    m2 = m1.unaryExpr([](const Scalar& v) { return numext::arithmetic_shift_right(v, N); });
+    m3 = m1.unaryExpr(arithmetic_right_shift_op<N, Scalar>());
     VERIFY_IS_CWISE_EQUAL(m2, m3);
 
     run<N + 1>(m);
   }
 };
 template <typename ArrayType>
-void signed_shift_test(const ArrayType& m) {
+void shift_test(const ArrayType& m) {
-  signed_shift_test_impl<ArrayType>::run(m);
+  shift_test_impl<ArrayType>::run(m);
 }
 
 template <typename ArrayType>
@@ -1361,9 +1386,9 @@ EIGEN_DECLARE_TEST(array_cwise) {
         ArrayXXi(internal::random<int>(1, EIGEN_TEST_MAX_SIZE), internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
     CALL_SUBTEST_7(array_generic(Array<Index, Dynamic, Dynamic>(internal::random<int>(1, EIGEN_TEST_MAX_SIZE),
                                                                 internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
-    CALL_SUBTEST_8(signed_shift_test(
+    CALL_SUBTEST_8(shift_test(
         ArrayXXi(internal::random<int>(1, EIGEN_TEST_MAX_SIZE), internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
-    CALL_SUBTEST_9(signed_shift_test(Array<Index, Dynamic, Dynamic>(internal::random<int>(1, EIGEN_TEST_MAX_SIZE),
+    CALL_SUBTEST_9(shift_test(Array<Index, Dynamic, Dynamic>(internal::random<int>(1, EIGEN_TEST_MAX_SIZE),
                                                              internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));
     CALL_SUBTEST_10(array_generic(Array<uint32_t, Dynamic, Dynamic>(internal::random<int>(1, EIGEN_TEST_MAX_SIZE),
                                                                     internal::random<int>(1, EIGEN_TEST_MAX_SIZE))));

test/numext.cpp

@@ -292,6 +292,27 @@ void check_signbit() {
   check_signbit_impl<T>::run();
 }
 
+template <typename T>
+void check_shift() {
+  using SignedT = typename numext::get_integer_by_size<sizeof(T)>::signed_type;
+  using UnsignedT = typename numext::get_integer_by_size<sizeof(T)>::unsigned_type;
+  constexpr int kNumBits = CHAR_BIT * sizeof(T);
+  for (int i = 0; i < 1000; ++i) {
+    const T a = internal::random<T>();
+    for (int s = 1; s < kNumBits; s++) {
+      T a_bsll = numext::logical_shift_left(a, s);
+      T a_bsll_ref = a << s;
+      VERIFY_IS_EQUAL(a_bsll, a_bsll_ref);
+      T a_bsrl = numext::logical_shift_right(a, s);
+      T a_bsrl_ref = numext::bit_cast<T, UnsignedT>(numext::bit_cast<UnsignedT, T>(a) >> s);
+      VERIFY_IS_EQUAL(a_bsrl, a_bsrl_ref);
+      T a_bsra = numext::arithmetic_shift_right(a, s);
+      T a_bsra_ref = numext::bit_cast<T, SignedT>(numext::bit_cast<SignedT, T>(a) >> s);
+      VERIFY_IS_EQUAL(a_bsra, a_bsra_ref);
+    }
+  }
+}
+
 EIGEN_DECLARE_TEST(numext) {
   for (int k = 0; k < g_repeat; ++k) {
     CALL_SUBTEST(check_negate<signed char>());
@@ -354,5 +375,15 @@ EIGEN_DECLARE_TEST(numext) {
     CALL_SUBTEST(check_signbit<int16_t>());
     CALL_SUBTEST(check_signbit<int32_t>());
     CALL_SUBTEST(check_signbit<int64_t>());
+
+    CALL_SUBTEST(check_shift<int8_t>());
+    CALL_SUBTEST(check_shift<int16_t>());
+    CALL_SUBTEST(check_shift<int32_t>());
+    CALL_SUBTEST(check_shift<int64_t>());
+
+    CALL_SUBTEST(check_shift<uint8_t>());
+    CALL_SUBTEST(check_shift<uint16_t>());
+    CALL_SUBTEST(check_shift<uint32_t>());
+    CALL_SUBTEST(check_shift<uint64_t>());
   }
 }