Householder: Preserve reflectors for large components

libeigen/eigen!2869

Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
Co-authored-by: Rasmus Munk Larsen <rlarsen@nvidia.com>

Eigen/src/Householder/Householder.h

@@ -147,6 +147,9 @@ EIGEN_DEVICE_FUNC void MatrixBase<Derived>::makeHouseholder(EssentialPart& essen
   Scalar c0 = coeff(0);
   const RealScalar tol = (std::numeric_limits<RealScalar>::min)();
   RealScalar unscaledNormThreshold = tol;
+  // Whether the direct construction's abs2(c0) + tailSqNorm would exceed the range. Integer scalars keep the direct
+  // path they have always taken; the scaled path divides by a component maximum, which does not apply to them.
+  bool unscaledSqNormOverflows = false;
   EIGEN_IF_CONSTEXPR (!NumTraits<RealScalar>::IsInteger) {
     const RealScalar precision = RealScalar(NumTraits<RealScalar>::epsilon());
     // With flush-to-zero arithmetic, every tail component square below tol can be lost. Account for every component
@@ -155,9 +158,21 @@ EIGEN_DEVICE_FUNC void MatrixBase<Derived>::makeHouseholder(EssentialPart& essen
     // intentional.
     const RealScalar componentCount = RealScalar(size() - 1) * RealScalar(NumTraits<Scalar>::IsComplex ? 2 : 1);
     unscaledNormThreshold = (tol / precision) * componentCount;
+
+    // Both terms overflow well before the reflector stops being representable, so classify the input before the
+    // squares are formed: abs2(c0) is at most twice the square of the larger component of c0, and the tail's own
+    // reduction has already overflowed if tailSqNorm exceeds the bound. Testing the sum with isinf() instead would
+    // not survive -ffinite-math-only, which folds that test away, whereas a comparison against a finite bound is
+    // still evaluated.
+    const RealScalar sqNormBound = NumTraits<RealScalar>::highest() / RealScalar(2);
+    const RealScalar componentBound = numext::sqrt(sqNormBound / RealScalar(2));
+    const RealScalar c0Max = numext::maxi(numext::abs(numext::real(c0)), numext::abs(numext::imag(c0)));
+    unscaledSqNormOverflows = !(c0Max <= componentBound) || !(tailSqNorm <= sqNormBound);
   }
 
-  if (tailSqNorm <= unscaledNormThreshold && !(numext::isnan)(c0)) {
+  // The scaled path forms the reflector from ratios of the largest component and never squares an unscaled
+  // coefficient, so it is also the path for inputs the direct construction cannot square.
+  if ((tailSqNorm <= unscaledNormThreshold || unscaledSqNormOverflows) && !(numext::isnan)(c0)) {
     using Accumulator = typename internal::householder_norm_accumulator<RealScalar>::type;
     const auto tailView = tail.unwind();
     const auto tailComponents = tailView.realView();

test/CMakeLists.txt

@@ -493,6 +493,7 @@ endif()
 ei_add_test(fastmath "${EIGEN_FASTMATH_FLAGS}")
 ei_add_test(bdcsvd_fastmath "${EIGEN_FASTMATH_FLAGS}")
 ei_add_test(stable_norm_fastmath "${EIGEN_FASTMATH_FLAGS}")
+ei_add_test(householder_fastmath "${EIGEN_FASTMATH_FLAGS}")
 
 # # ei_add_test(denseLM)
 

test/householder.cpp

@@ -682,6 +682,107 @@ void householder_small_tail() {
   VERIFY_IS_APPROX(2.0f * tau, tau * tau * (1.0f + essential.squaredNorm()));
 }
 
+// tau and the essential vector are scale invariant and beta is homogeneous, so rescaling by a power of two is exact.
+// It keeps the long double reference below the squaring overflow that this path exists to avoid, which matters where
+// long double is only as wide as double.
+template <typename VectorType, typename EssentialType>
+void verify_large_householder_result(const VectorType& vector, const EssentialType& essential,
+                                     const typename VectorType::Scalar& tau,
+                                     const typename VectorType::RealScalar& beta) {
+  typedef typename VectorType::RealScalar RealScalar;
+  int exponent = 0;
+  (void)std::frexp(vector.cwiseAbs().maxCoeff(), &exponent);
+  const RealScalar scale = std::ldexp(RealScalar(1), exponent);
+  verify_real_householder_result((vector / scale).eval(), essential, tau, RealScalar(beta / scale));
+}
+
+void householder_large_components() {
+  {
+    // The head coefficient squares out of the float range; the tail coefficient does not. The direct
+    // construction forms their sum regardless, so it produced beta = -inf and tau = NaN.
+    Vector2f vector(4e19f, -5.2e18f);
+    Matrix<float, 1, 1> essential;
+    float tau;
+    float beta;
+
+    vector.makeHouseholder(essential, tau, beta);
+
+    VERIFY((numext::isfinite)(tau));
+    VERIFY((numext::isfinite)(beta));
+    VERIFY((numext::isfinite)(essential[0]));
+    verify_large_householder_result(vector, essential, tau, beta);
+  }
+
+  {
+    VectorXf vector(3);
+    vector << 1e20f, 2e20f, 3e20f;
+    VectorXf essential(2);
+    float tau;
+    float beta;
+
+    vector.makeHouseholder(essential, tau, beta);
+
+    verify_large_householder_result(vector, essential, tau, beta);
+  }
+
+  {
+    // Accumulated overflow: every square is representable, their sum is not.
+    VectorXf vector = VectorXf::Constant(65, 1e19f);
+    VectorXf essential(64);
+    float tau;
+    float beta;
+
+    vector.makeHouseholder(essential, tau, beta);
+
+    verify_large_householder_result(vector, essential, tau, beta);
+  }
+
+  {
+    VectorXd vector(3);
+    vector << 1e160, -2e160, 3e160;
+    VectorXd essential(2);
+    double tau;
+    double beta;
+
+    vector.makeHouseholder(essential, tau, beta);
+
+    verify_large_householder_result(vector, essential, tau, beta);
+  }
+
+  {
+    const float largest = (std::numeric_limits<float>::max)();
+    const Vector2cf vector(std::complex<float>(largest, 0.0f), std::complex<float>(largest, 0.0f));
+    Matrix<std::complex<float>, 1, 1> essential;
+    std::complex<float> tau;
+    float beta;
+
+    vector.makeHouseholder(essential, tau, beta);
+
+    // The true norm exceeds the float range, so beta cannot be represented; the reflector itself still must be.
+    VERIFY((numext::isfinite)(numext::real(tau)));
+    VERIFY((numext::isfinite)(numext::imag(tau)));
+    VERIFY((numext::isfinite)(numext::real(essential[0])));
+    VERIFY((numext::isfinite)(numext::imag(essential[0])));
+    Vector2cf householder;
+    householder << std::complex<float>(1.0f, 0.0f), essential;
+    const Matrix2cf transform = Matrix2cf::Identity() - tau * householder * householder.adjoint();
+    VERIFY_IS_APPROX(transform.adjoint() * transform, Matrix2cf::Identity());
+  }
+
+  {
+    // Reflectors this large must still compose into a usable decomposition.
+    MatrixXf matrix(6, 4);
+    for (Index i = 0; i < matrix.rows(); ++i)
+      for (Index j = 0; j < matrix.cols(); ++j) matrix(i, j) = 1e19f * float(internal::random<double>(-1.0, 1.0));
+    const HouseholderQR<MatrixXf> qr(matrix);
+    const MatrixXf q = qr.householderQ() * MatrixXf::Identity(6, 4);
+    const MatrixXf r = qr.matrixQR().topRows(4).template triangularView<Upper>();
+    const MatrixXd scaled = (matrix / 1e19f).cast<double>();
+    const double relative_residual = ((q * r).cast<double>() / 1e19 - scaled).norm() / scaled.norm();
+    VERIFY(relative_residual <= 64 * double(NumTraits<float>::epsilon()));
+  }
+}
+
 EIGEN_DECLARE_TEST(householder) {
   for (int i = 0; i < g_repeat; i++) {
     CALL_SUBTEST_1(householder(Matrix<double, 2, 2>()));
@@ -705,4 +806,5 @@ EIGEN_DECLARE_TEST(householder) {
   CALL_SUBTEST_10(householder_blocked_right_regression<double>());
   CALL_SUBTEST_11(householder_blocked_right_regression<std::complex<double>>());
   CALL_SUBTEST_12(householder_small_tail());
+  CALL_SUBTEST_13(householder_large_components());
 }

test/householder_fastmath.cpp

@@ -0,0 +1,86 @@
+// SPDX-FileCopyrightText: The Eigen Authors
+// SPDX-License-Identifier: MPL-2.0
+
+#include "main.h"
+#include <Eigen/QR>
+
+// -ffinite-math-only lets the compiler assume every result is finite, so isfinite(), isinf(), and isnan() fold to
+// constants in this translation unit and cannot be used to check a result. A comparison against a finite bound is
+// still evaluated, and rejects infinities and NaNs alike.
+template <typename RealScalar>
+bool is_representable(const RealScalar& x) {
+  return numext::abs(x) <= NumTraits<RealScalar>::highest();
+}
+
+// makeHouseholder must reach the scaled path whenever the direct construction would square the input out of range.
+// The dispatch has to be decided from the input magnitudes, because a test on the overflowed sum does not survive
+// here.
+template <typename RealScalar>
+void householder_overflow(const Matrix<RealScalar, Dynamic, 1>& vector) {
+  typedef Matrix<RealScalar, Dynamic, 1> VectorType;
+  typedef Matrix<RealScalar, Dynamic, Dynamic> MatrixType;
+  const Index size = vector.size();
+
+  VectorType essential(size - 1);
+  RealScalar tau;
+  RealScalar beta;
+  vector.makeHouseholder(essential, tau, beta);
+
+  VERIFY(is_representable(tau));
+  VERIFY(is_representable(beta));
+  VERIFY(is_representable(essential.cwiseAbs().maxCoeff()));
+
+  // tau and the essential vector are scale invariant and beta is homogeneous, so applying the reflector to a
+  // power-of-two rescaling of the input is exact and keeps the check itself inside the range.
+  int exponent = 0;
+  (void)std::frexp(vector.cwiseAbs().maxCoeff(), &exponent);
+  const RealScalar scale = std::ldexp(RealScalar(1), exponent);
+  const VectorType scaled = vector / scale;
+
+  VectorType v(size);
+  v[0] = RealScalar(1);
+  v.tail(size - 1) = essential;
+  const MatrixType reflector = MatrixType::Identity(size, size) - tau * v * v.transpose();
+
+  VERIFY_IS_APPROX(reflector.transpose() * reflector, MatrixType::Identity(size, size));
+  const VectorType reflected = reflector * scaled;
+  VERIFY_IS_APPROX(reflected[0], beta / scale);
+  VERIFY_IS_MUCH_SMALLER_THAN(reflected.tail(size - 1).cwiseAbs().maxCoeff(), RealScalar(1));
+}
+
+template <typename RealScalar>
+void householder_overflow_cases() {
+  typedef Matrix<RealScalar, Dynamic, 1> VectorType;
+  const RealScalar huge = numext::sqrt(NumTraits<RealScalar>::highest()) * RealScalar(4);
+
+  // The head alone squares out of range.
+  VectorType head(2);
+  head << huge, huge / RealScalar(8);
+  householder_overflow<RealScalar>(head);
+
+  // A tail component alone squares out of range, with a zero head.
+  VectorType tail(3);
+  tail << RealScalar(0), huge * RealScalar(0.6), huge * RealScalar(0.8);
+  householder_overflow<RealScalar>(tail);
+
+  // Accumulated overflow: every square is representable, their sum is not.
+  householder_overflow<RealScalar>(VectorType::Constant(65, huge / RealScalar(4)));
+
+  // A QR factorization of a matrix at this scale has to compose from those reflectors.
+  const Index rows = 6, cols = 4;
+  Matrix<RealScalar, Dynamic, Dynamic> matrix(rows, cols);
+  for (Index j = 0; j < cols; ++j)
+    for (Index i = 0; i < rows; ++i) matrix(i, j) = huge * RealScalar(internal::random<double>(-1.0, 1.0));
+  const HouseholderQR<Matrix<RealScalar, Dynamic, Dynamic> > qr(matrix);
+  const Matrix<RealScalar, Dynamic, Dynamic> q =
+      qr.householderQ() * Matrix<RealScalar, Dynamic, Dynamic>::Identity(rows, cols);
+  const Matrix<RealScalar, Dynamic, Dynamic> r =
+      qr.matrixQR().topRows(cols).template triangularView<Upper>().toDenseMatrix();
+  const Matrix<RealScalar, Dynamic, Dynamic> scaledMatrix = matrix / huge;
+  VERIFY_IS_APPROX((q * r) / huge, scaledMatrix);
+}
+
+EIGEN_DECLARE_TEST(householder_fastmath) {
+  CALL_SUBTEST_1(householder_overflow_cases<float>());
+  CALL_SUBTEST_2(householder_overflow_cases<double>());
+}