Core: Avoid speculative division by zero in complex_divide_smith

libeigen/eigen!3380

Co-authored-by: Antonio Sanchez <cantonios@google.com>

Eigen/src/Core/MathFunctions.h

diff --git a/Eigen/src/Core/MathFunctions.h b/Eigen/src/Core/MathFunctions.h
index 1c7a286..bbd4e46 100644
--- a/Eigen/src/Core/MathFunctions.h
+++ b/Eigen/src/Core/MathFunctions.h
@@ -2540,19 +2540,28 @@
                                                                            const std::complex<T>& b) {
   const T a_real = numext::real(a);
   const T a_imag = numext::imag(a);
-  const T b_real = numext::real(b);
-  const T b_imag = numext::imag(b);
+  T b_real = numext::real(b);
+  T b_imag = numext::imag(b);
   // Smith's complex division (https://arxiv.org/pdf/1210.4539.pdf),
   // guards against over/under-flow.
   const bool scale_imag = numext::abs(b_imag) <= numext::abs(b_real);
-  const T rscale = scale_imag ? T(1) : b_real / b_imag;
-  // d / c = +-1 for |c| = |d|, which the quotient misses as 0 / 0 or inf / inf when b = 0 or both parts of b are
-  // infinite; a zero divisor then gives an infinity and an infinite one zero, as in pdiv_complex. The sign is not taken
-  // from b_real * b_imag, which can overflow or underflow.
-  const T iscale = !scale_imag ? T(1)
-                   : numext::abs(b_imag) == numext::abs(b_real)
-                       ? numext::copysign(T(1), b_real) * numext::copysign(T(1), b_imag)
-                       : b_imag / b_real;
+  T rscale = T(1);
+  T iscale = T(1);
+  if (numext::abs(b_imag) == numext::abs(b_real)) {
+    // d / c = +-1 for |c| = |d|, which the quotient misses as 0 / 0 or inf / inf when b = 0 or both parts of b are
+    // infinite; a zero divisor then gives an infinity and an infinite one zero, as in pdiv_complex. The sign is not
+    // taken from b_real * b_imag, which can overflow or underflow.
+    iscale = numext::copysign(T(1), b_real) * numext::copysign(T(1), b_imag);
+  } else if (scale_imag) {
+    // Each quotient divides by the larger part of b, which is not zero here. The barrier keeps the quotient on its
+    // branch: an if-converted form evaluates both quotients, and the untaken one divides by zero for a divisor on an
+    // axis, raising FE_DIVBYZERO. The barrier is applied in place, so that the branch costs no register copy.
+    volatile_barrier(b_real);
+    iscale = b_imag / b_real;
+  } else {
+    volatile_barrier(b_imag);
+    rscale = b_real / b_imag;
+  }
   const T denominator = b_real * rscale + b_imag * iscale;
   return std::complex<T>((a_real * rscale + a_imag * iscale) / denominator,
                          (a_imag * rscale - a_real * iscale) / denominator);

Eigen/src/Core/util/Macros.h

diff --git a/Eigen/src/Core/util/Macros.h b/Eigen/src/Core/util/Macros.h
index e568698..cff4832 100644
--- a/Eigen/src/Core/util/Macros.h
+++ b/Eigen/src/Core/util/Macros.h
@@ -1230,6 +1230,46 @@
 #endif
 #endif
 
+namespace Eigen {
+namespace internal {
+// A volatile asm that the optimizer must execute exactly where it is written, so a value passed through it is defined
+// only on its branch, and an operation that consumes it stays there instead of being evaluated speculatively.
+// EIGEN_OPTIMIZATION_BARRIER is not volatile, and its "+g,w" and "+g,x" alternatives put a scalar float or double in
+// memory on clang; the float and double overloads keep the value in a register.
+template <typename T>
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void volatile_barrier(T& x) {
+  EIGEN_UNUSED_VARIABLE(x);
+#if !defined(EIGEN_GPU_COMPILE_PHASE) && !defined(SYCL_DEVICE_ONLY) && \
+    (EIGEN_COMP_GNUC || (defined(__clang__) && defined(_MSC_VER)))
+  __asm__ __volatile__("" : "+m"(x));
+#endif
+}
+#if !defined(EIGEN_GPU_COMPILE_PHASE) && !defined(SYCL_DEVICE_ONLY) && \
+    (EIGEN_COMP_GNUC || (defined(__clang__) && defined(_MSC_VER)))
+#if EIGEN_ARCH_i386_OR_x86_64 && EIGEN_COMP_NVHPC
+// nvc++ rejects a long double asm operand under every constraint, so its bytes go through memory instead.
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void volatile_barrier(long double& x) {
+  __asm__ __volatile__("" : "+m"(*reinterpret_cast<unsigned char(*)[sizeof(long double)]>(&x)));
+}
+#endif
+#if EIGEN_ARCH_ARM_OR_ARM64 && defined(__ARM_FP)
+#define EIGEN_FP_REGISTER_CONSTRAINT "+w"
+#elif EIGEN_ARCH_i386_OR_x86_64 && defined(__SSE2__)
+#define EIGEN_FP_REGISTER_CONSTRAINT "+x"
+#endif
+#ifdef EIGEN_FP_REGISTER_CONSTRAINT
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void volatile_barrier(float& x) {
+  __asm__ __volatile__("" : EIGEN_FP_REGISTER_CONSTRAINT(x));
+}
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE void volatile_barrier(double& x) {
+  __asm__ __volatile__("" : EIGEN_FP_REGISTER_CONSTRAINT(x));
+}
+#undef EIGEN_FP_REGISTER_CONSTRAINT
+#endif
+#endif
+}  // namespace internal
+}  // namespace Eigen
+
 #if EIGEN_COMP_MSVC
 // NOTE MSVC often gives C4127 warnings with compiletime if statements. See bug 1362.
 // This workaround suppresses MSVC C4127 warnings for compile-time conditionals.

test/numext.cpp

diff --git a/test/numext.cpp b/test/numext.cpp
index 134716c..91cecd1 100644
--- a/test/numext.cpp
+++ b/test/numext.cpp
@@ -456,6 +456,9 @@
 #if EIGEN_COMP_MSVC
 #pragma float_control(precise, on, push)
 #pragma fenv_access(on)
+#elif EIGEN_COMP_CLANG && (EIGEN_ARCH_i386_OR_x86_64 || EIGEN_ARCH_ARM_OR_ARM64) && !defined(__GPUCC__) && \
+    !defined(__SYCL_DEVICE_ONLY__)
+#pragma STDC FENV_ACCESS ON
 #endif
 
 // The exceptions numext::divide(a, b) raises. Volatile operands and results keep the division between the flag
@@ -541,6 +544,9 @@
 #if EIGEN_COMP_MSVC
 #pragma fenv_access(off)
 #pragma float_control(pop)
+#elif EIGEN_COMP_CLANG && (EIGEN_ARCH_i386_OR_x86_64 || EIGEN_ARCH_ARM_OR_ARM64) && !defined(__GPUCC__) && \
+    !defined(__SYCL_DEVICE_ONLY__)
+#pragma STDC FENV_ACCESS OFF
 #endif
 
 // log2(2^k) = k on the scalar path, as on the vectorized one. A size-1 array takes the scalar path.