Core: Avoid signed overflow in narrow unsigned products

libeigen/eigen!3030

Closes #3148

Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
diff --git a/Eigen/src/Core/GenericPacketMath.h b/Eigen/src/Core/GenericPacketMath.h
index 98ab1a2..f4e4e6d 100644
--- a/Eigen/src/Core/GenericPacketMath.h
+++ b/Eigen/src/Core/GenericPacketMath.h
@@ -368,7 +368,7 @@
 /** \internal \returns a * b (coeff-wise) */
 template <typename Packet>
 EIGEN_DEVICE_FUNC inline Packet pmul(const Packet& a, const Packet& b) {
-  return a * b;
+  return internal::mul(a, b);
 }
 template <>
 EIGEN_DEVICE_FUNC inline bool pmul(const bool& a, const bool& b) {
diff --git a/Eigen/src/Core/MathFunctions.h b/Eigen/src/Core/MathFunctions.h
index 9169fa4..df1fb0f 100644
--- a/Eigen/src/Core/MathFunctions.h
+++ b/Eigen/src/Core/MathFunctions.h
@@ -212,13 +212,42 @@
 struct conj_impl : conj_default_impl<Scalar, IsComplex> {};
 
 /****************************************************************************
+ * Implementation of mul                                                  *
+ ****************************************************************************/
+
+// Unsigned operands narrower than int promote to int, where a product whose wrapped value is representable
+// can still overflow: (unsigned short)0xffff squared is 0xfffe0001 > INT_MAX. Multiplying in unsigned int,
+// wider than every such operand, wraps instead. Mixed operand types, signed types, and bool keep operator*.
+template <typename Lhs, typename Rhs,
+          bool NarrowUnsigned = std::is_same<Lhs, Rhs>::value && std::is_unsigned<Lhs>::value &&
+                                !std::is_same<Lhs, bool>::value && (sizeof(Lhs) < sizeof(int))>
+struct mul_impl {
+  EIGEN_DEVICE_FUNC static constexpr EIGEN_ALWAYS_INLINE auto run(const Lhs& a, const Rhs& b) { return a * b; }
+};
+
+template <typename Scalar>
+struct mul_impl<Scalar, Scalar, true> {
+  EIGEN_DEVICE_FUNC static constexpr EIGEN_ALWAYS_INLINE Scalar run(const Scalar& a, const Scalar& b) {
+    return static_cast<Scalar>(static_cast<unsigned int>(a) * static_cast<unsigned int>(b));
+  }
+};
+
+/** \internal \returns a * b, wrapping rather than overflowing the signed type that integral promotion
+ * would give narrow unsigned operands. The result keeps operator*'s type, which for an expression-template
+ * Scalar is a lazy expression: consume it rather than binding it to auto. */
+template <typename Lhs, typename Rhs>
+EIGEN_DEVICE_FUNC constexpr EIGEN_ALWAYS_INLINE auto mul(const Lhs& a, const Rhs& b) {
+  return mul_impl<Lhs, Rhs>::run(a, b);
+}
+
+/****************************************************************************
  * Implementation of abs2                                                 *
  ****************************************************************************/
 
 template <typename Scalar, bool IsComplex>
 struct abs2_impl_default {
   using RealScalar = typename NumTraits<Scalar>::Real;
-  EIGEN_DEVICE_FUNC static inline RealScalar run(const Scalar& x) { return x * x; }
+  EIGEN_DEVICE_FUNC static inline RealScalar run(const Scalar& x) { return internal::mul(x, x); }
 };
 
 template <typename Scalar>
@@ -516,11 +545,11 @@
   static EIGEN_DEVICE_FUNC inline ScalarX run(ScalarX x, ScalarY y) {
     ScalarX res(1);
     eigen_assert(!NumTraits<ScalarY>::IsSigned || y >= 0);
-    if (y & 1) res *= x;
+    if (y & 1) res = internal::mul(res, x);
     y >>= 1;
     while (y) {
-      x *= x;
-      if (y & 1) res *= x;
+      x = internal::mul(x, x);
+      if (y & 1) res = internal::mul(res, x);
       y >>= 1;
     }
     return res;
diff --git a/Eigen/src/Core/functors/BinaryFunctors.h b/Eigen/src/Core/functors/BinaryFunctors.h
index 5c7a581..8672440 100644
--- a/Eigen/src/Core/functors/BinaryFunctors.h
+++ b/Eigen/src/Core/functors/BinaryFunctors.h
@@ -85,7 +85,7 @@
 #endif
   EIGEN_DEVICE_FUNC constexpr EIGEN_STRONG_INLINE result_type
   operator()(const LhsScalar& a, const RhsScalar& b) const {
-    return a * b;
+    return internal::mul(a, b);
   }
   template <typename Packet>
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet packetOp(const Packet& a, const Packet& b) const {
diff --git a/Eigen/src/Core/functors/UnaryFunctors.h b/Eigen/src/Core/functors/UnaryFunctors.h
index 74ab18c..247eea5 100644
--- a/Eigen/src/Core/functors/UnaryFunctors.h
+++ b/Eigen/src/Core/functors/UnaryFunctors.h
@@ -903,7 +903,7 @@
  */
 template <typename Scalar>
 struct scalar_square_op {
-  EIGEN_DEVICE_FUNC constexpr inline Scalar operator()(const Scalar& a) const { return a * a; }
+  EIGEN_DEVICE_FUNC constexpr inline Scalar operator()(const Scalar& a) const { return internal::mul(a, a); }
   template <typename Packet>
   EIGEN_DEVICE_FUNC inline Packet packetOp(const Packet& a) const {
     return internal::pmul(a, a);
@@ -934,7 +934,9 @@
  */
 template <typename Scalar>
 struct scalar_cube_op {
-  EIGEN_DEVICE_FUNC constexpr inline Scalar operator()(const Scalar& a) const { return a * a * a; }
+  EIGEN_DEVICE_FUNC constexpr inline Scalar operator()(const Scalar& a) const {
+    return internal::mul(a, internal::mul(a, a));
+  }
   template <typename Packet>
   EIGEN_DEVICE_FUNC inline Packet packetOp(const Packet& a) const {
     return internal::pmul(a, pmul(a, a));
diff --git a/test/integer_types.cpp b/test/integer_types.cpp
index a19cfb1..b65d7a8 100644
--- a/test/integer_types.cpp
+++ b/test/integer_types.cpp
@@ -117,6 +117,67 @@
   }
 }
 
+// Unsigned operands narrower than int are promoted to int before multiplication, so a product above
+// INT_MAX overflows the signed intermediate although the wrapped result, the exact product modulo
+// 2^(8*sizeof(Scalar)), is representable. A multiply that keeps the promoted signed type fails this
+// under UBSan by report rather than by comparison.
+template <typename Scalar>
+void unsigned_wraparound_tests() {
+  using WideScalar = unsigned long long;
+  using ArrayType = Array<Scalar, Dynamic, 1>;
+  using MatrixType = Matrix<Scalar, Dynamic, Dynamic>;
+  EIGEN_STATIC_ASSERT(sizeof(Scalar) < sizeof(WideScalar), THIS_TYPE_IS_NOT_SUPPORTED)
+  const Scalar highest = NumTraits<Scalar>::highest();
+  const WideScalar modulus = WideScalar(highest) + 1;
+  const Scalar values[] = {Scalar(0), Scalar(1), Scalar(highest / 2), Scalar(highest - 1), highest};
+  const Index numValues = Index(sizeof(values) / sizeof(values[0]));
+
+  // Long enough to exercise whole packets as well as the scalar tail.
+  const Index size = 8 * internal::packet_traits<Scalar>::size + 3;
+  ArrayType a(size), b(size);
+  for (Index i = 0; i < size; ++i) {
+    a(i) = values[i % numValues];
+    b(i) = values[(size - i) % numValues];
+  }
+
+  const ArrayType product = a * b;
+  const ArrayType square = a.square();
+  const ArrayType cube = a.cube();
+  const ArrayType abs2 = a.abs2();
+  const ArrayType pow2 = a.pow(Scalar(2));
+  const ArrayType pow3 = a.pow(Scalar(3));
+  for (Index i = 0; i < size; ++i) {
+    const WideScalar x = a(i), y = b(i);
+    const WideScalar x2 = x * x % modulus;
+    const WideScalar x3 = x2 * x % modulus;
+    VERIFY_IS_EQUAL(product(i), Scalar(x * y % modulus));
+    VERIFY_IS_EQUAL(square(i), Scalar(x2));
+    VERIFY_IS_EQUAL(cube(i), Scalar(x3));
+    VERIFY_IS_EQUAL(abs2(i), Scalar(x2));
+    VERIFY_IS_EQUAL(pow2(i), Scalar(x2));
+    VERIFY_IS_EQUAL(pow3(i), Scalar(x3));
+    VERIFY_IS_EQUAL(numext::abs2(a(i)), Scalar(x2));
+    VERIFY_IS_EQUAL(numext::pow(a(i), Scalar(2)), Scalar(x2));
+    VERIFY_IS_EQUAL(internal::pmul(a(i), b(i)), Scalar(x * y % modulus));
+  }
+
+  // predux_mul reduces through the same multiply; keep every factor away from zero.
+  const Index reduxSize = 3 * internal::packet_traits<Scalar>::size + 1;
+  ArrayType c(reduxSize);
+  WideScalar expectedProd = 1;
+  for (Index i = 0; i < reduxSize; ++i) {
+    c(i) = Scalar(highest - Scalar(i % 3));
+    expectedProd = expectedProd * WideScalar(c(i)) % modulus;
+  }
+  VERIFY_IS_EQUAL(c.prod(), Scalar(expectedProd));
+
+  // The matrix product multiplies through pmul, at magnitudes integer_type_tests never reaches.
+  const MatrixType m = MatrixType::Constant(4, 4, highest);
+  const MatrixType mm = m * m;
+  const WideScalar h = highest;
+  VERIFY_IS_EQUAL(mm(0, 0), Scalar(4 * (h * h % modulus) % modulus));
+}
+
 template <int>
 void integer_types_extra() {
   VERIFY_IS_EQUAL(int(internal::scalar_div_cost<int>::value), 8);
@@ -131,6 +192,7 @@
   for (int i = 0; i < g_repeat; i++) {
     CALL_SUBTEST_1(integer_type_tests(Matrix<unsigned int, 1, 1>()));
     CALL_SUBTEST_1(integer_type_tests(Matrix<unsigned long, 3, 4>()));
+    CALL_SUBTEST_1(unsigned_wraparound_tests<unsigned int>());
 
     CALL_SUBTEST_2(integer_type_tests(Matrix<long, 2, 2>()));
     CALL_SUBTEST_2(signed_integer_type_tests(Matrix<long, 2, 2>()));
@@ -140,11 +202,13 @@
 
     CALL_SUBTEST_4(integer_type_tests(Matrix<unsigned char, 3, 3>()));
     CALL_SUBTEST_4(integer_type_tests(Matrix<unsigned char, Dynamic, Dynamic>(20, 20)));
+    CALL_SUBTEST_4(unsigned_wraparound_tests<unsigned char>());
 
     CALL_SUBTEST_5(integer_type_tests(Matrix<short, Dynamic, 4>(7, 4)));
     CALL_SUBTEST_5(signed_integer_type_tests(Matrix<short, Dynamic, 4>(7, 4)));
 
     CALL_SUBTEST_6(integer_type_tests(Matrix<unsigned short, 4, 4>()));
+    CALL_SUBTEST_6(unsigned_wraparound_tests<unsigned short>());
 
     CALL_SUBTEST_7(integer_type_tests(Matrix<long long, 11, 13>()));
     CALL_SUBTEST_7(signed_integer_type_tests(Matrix<long long, 11, 13>()));
diff --git a/test/packetmath.cpp b/test/packetmath.cpp
index f445f01..b160ec9 100644
--- a/test/packetmath.cpp
+++ b/test/packetmath.cpp
@@ -2287,6 +2287,23 @@
   }
 }
 
+// At saturated unsigned short operands the scalar pmul family must wrap, not overflow the int that
+// integral promotion would otherwise multiply in.
+void packetmath_unsigned_short() {
+  // Volatile inputs keep sanitizer builds from folding away the promotions.
+  volatile unsigned short values[] = {0, 1, 32768, 40232, 58075, 65535};
+  const unsigned short c = 65535;
+  for (unsigned short a : values) {
+    for (unsigned short b : values) {
+      VERIFY_IS_EQUAL(internal::pmul(a, b), REF_MUL(a, b));
+      VERIFY_IS_EQUAL(internal::pmadd(a, b, c), REF_MADD(a, b, c));
+      VERIFY_IS_EQUAL(internal::pmsub(a, b, c), REF_MSUB(a, b, c));
+      VERIFY_IS_EQUAL(internal::pnmadd(a, b, c), REF_NMADD(a, b, c));
+      VERIFY_IS_EQUAL(internal::pnmsub(a, b, c), REF_NMSUB(a, b, c));
+    }
+  }
+}
+
 namespace Eigen {
 namespace test {
 
@@ -2330,6 +2347,7 @@
     CALL_SUBTEST_4(test::runner<uint8_t>::run());
     CALL_SUBTEST_5(test::runner<int16_t>::run());
     CALL_SUBTEST_6(test::runner<uint16_t>::run());
+    CALL_SUBTEST_6(packetmath_unsigned_short());
     CALL_SUBTEST_7(test::runner<int32_t>::run());
     CALL_SUBTEST_8(test::runner<uint32_t>::run());
     CALL_SUBTEST_9(test::runner<int64_t>::run());