perf: fuse banded TriMultiply neighbour-correction loops

dal-cpp/dal/math/matrix/banded.cpp

@@ -17,13 +17,32 @@ namespace Dal {
             const Vector_<>& x, const Vector_<>& diag, const Vector_<>& above, const Vector_<>& below, Vector_<>* r) {
             REQUIRE(x.size() == diag.size(), "");
             r->Resize(x.size());
+            // Pass 1 (vectorizable elementwise, identical to the unfused Transform) seeds r = diag * x.
             Transform(x, diag, std::multiplies<>(), r);
-            auto pr = r->begin();
-            for (auto px = x.begin() + 1, pa = above.begin(); pa != above.end(); ++px, ++pa, ++pr)
-                *pr += *px * *pa;
-            pr = r->begin() + 1;
-            for (auto px = x.begin(), pb = below.begin(); pb != below.end(); ++px, ++pb, ++pr)
-                *pr += *px * *pb;
+            // Pass 2 fuses the two strided neighbour corrections into one sweep:
+            //   r[i] += above[i] * x[i+1]  (superdiagonal, i = 0..n-2)
+            //   r[i] += below[i-1] * x[i-1] (subdiagonal, i = 1..n-1)
+            // The accumulator order (above term first, then below) matches the unfused code so
+            // -ffp-contract=fast forms the same FMAs.
+            const int n = static_cast<int>(x.size());
+            if (n < 2)
+                return;
+            auto xNext = x.begin() + 1; // x[i+1], walks with i
+            auto xPrev = x.begin();     // x[i-1], lags by one; valid for i >= 1
+            auto rIt = r->begin();
+            auto aboveIt = above.begin();
+            auto belowPrev = below.begin(); // below[i-1], lags by one; valid for i >= 1
+            // i = 0: superdiagonal only (below has no i-1 term yet).
+            *rIt += *aboveIt * *xNext;
+            ++rIt;
+            ++aboveIt;
+            ++xNext;
+            for (int ii = 1; ii < n - 1; ++ii, ++rIt, ++aboveIt, ++belowPrev, ++xNext, ++xPrev) {
+                *rIt += *aboveIt * *xNext;
+                *rIt += *belowPrev * *xPrev;
+            }
+            // i = n-1: subdiagonal only (above has no term at the last row).
+            *rIt += *belowPrev * *xPrev;
         }
 
         Vector_<> TridagBetaInverse(const Vector_<>& diag, const Vector_<>& above, const Vector_<>& below) {

dal-cpp/tests/math/matrix/test_banded.cpp

@@ -108,3 +108,36 @@ TEST(MatrixTest, TestTriDiagonalSolve) {
     for (int i = 0; i < n; ++i)
         ASSERT_NEAR(calculated[i], expected[i], 1e-6);
 }
+
+// Varying x with distinct, asymmetric above/below bands exercises every branch of
+// the fused TriMultiply (interior points combine all three bands; both boundaries
+// drop one term). Hand-computed expected values lock the fused contract.
+TEST(MatrixTest, TestTriDiagonalMultiplyAsymmetricFused) {
+    const auto n = 5;
+    Sparse::TriDiagonal_ trig(n);
+    trig.Set(0, 0, 2.0);
+    trig.Set(0, 1, 3.0); // above[0]
+    trig.Set(1, 0, 4.0); // below[0]
+    trig.Set(1, 1, 5.0);
+    trig.Set(1, 2, 6.0); // above[1]
+    trig.Set(2, 1, 7.0); // below[1]
+    trig.Set(2, 2, 8.0);
+    trig.Set(2, 3, 9.0); // above[2]
+    trig.Set(3, 2, 10.0); // below[2]
+    trig.Set(3, 3, 11.0);
+    trig.Set(3, 4, 12.0); // above[3]
+    trig.Set(4, 3, 13.0); // below[3]
+    trig.Set(4, 4, 14.0);
+
+    Vector_<> x = {1.0, 2.0, 3.0, 4.0, 5.0};
+    Vector_<> calculated(n);
+    trig.MultiplyLeft(x, &calculated);
+    // r[0] = diag[0]*x[0] + above[0]*x[1]
+    // r[1] = diag[1]*x[1] + above[1]*x[2] + below[0]*x[0]
+    // r[2] = diag[2]*x[2] + above[2]*x[3] + below[1]*x[1]
+    // r[3] = diag[3]*x[3] + above[3]*x[4] + below[2]*x[2]
+    // r[4] = diag[4]*x[4] + below[3]*x[3]
+    Vector_<> expected = {2 * 1 + 3 * 2, 5 * 2 + 6 * 3 + 4 * 1, 8 * 3 + 9 * 4 + 7 * 2, 11 * 4 + 12 * 5 + 10 * 3, 14 * 5 + 13 * 4};
+    for (int i = 0; i < n; ++i)
+        ASSERT_DOUBLE_EQ(calculated[i], expected[i]);
+}

dal-cpp/tests/math/matrix/test_bcg.cpp

@@ -3,6 +3,8 @@
 //
 
 #include <gtest/gtest.h>
+#include <algorithm>
+#include <cmath>
 #include <dal/math/matrix/banded.hpp>
 #include <dal/math/matrix/sparse.hpp>
 #include <dal/math/matrix/bcg.hpp>
@@ -43,3 +45,67 @@ TEST(MatrixTest, TestBCGSolve) {
     ASSERT_NEAR(result[9], 0.07446809, 1e-8);
 }
 
+// Tri-diagonal systems with distinct band values exercise every branch of the
+// fused Krylov sweeps. CG requires a symmetric positive-definite matrix, so it
+// gets a symmetric system; BCG handles the asymmetric case (and its shadow path).
+// The contract is residual ||Ax - b||_inf near machine precision after a tight solve.
+namespace {
+    void BuildSymmetricTridiag(Sparse::Square_* mat, int n) {
+        for (int i = 0; i < n; ++i) {
+            mat->Set(i, i, 4.0 + 0.1 * static_cast<double>(i));
+            if (i > 0) {
+                const double off = -0.5 - 0.01 * static_cast<double>(i);
+                mat->Set(i, i - 1, off);
+                mat->Set(i - 1, i, off);
+            }
+        }
+    }
+
+    void BuildAsymmetricTridiag(Sparse::Square_* mat, int n) {
+        for (int i = 0; i < n; ++i) {
+            mat->Set(i, i, 4.0 + 0.1 * static_cast<double>(i));
+            if (i > 0)
+                mat->Set(i, i - 1, -0.5 - 0.01 * static_cast<double>(i));
+            if (i + 1 < n)
+                mat->Set(i, i + 1, -0.3 + 0.02 * static_cast<double>(i));
+        }
+    }
+
+    double ResidualInfNorm(const Sparse::Square_& A, const Vector_<>& x, const Vector_<>& b) {
+        Vector_<> ax(x.size());
+        A.MultiplyLeft(x, &ax);
+        double worst = 0.0;
+        for (int i = 0; i < static_cast<int>(ax.size()); ++i)
+            worst = std::max(worst, std::fabs(ax[i] - b[i]));
+        return worst;
+    }
+} // namespace
+
+TEST(MatrixTest, TestCGSolveSymmetricLowResidual) {
+    const int n = 40;
+    Sparse::Square_* mat = Sparse::NewBandDiagonal(n, 1, 1);
+    BuildSymmetricTridiag(mat, n);
+
+    Vector_<> b(n);
+    for (int i = 0; i < n; ++i)
+        b[i] = 1.0 + 0.1 * static_cast<double>(i % 7);
+    Vector_<> x(n, 0.0);
+
+    Sparse::CGSolve(*mat, b, 1e-12, 1e-14, 500, &x);
+    ASSERT_LT(ResidualInfNorm(*mat, x, b), 1e-8);
+}
+
+TEST(MatrixTest, TestBCGSolveAsymmetricLowResidual) {
+    const int n = 40;
+    Sparse::Square_* mat = Sparse::NewBandDiagonal(n, 1, 1);
+    BuildAsymmetricTridiag(mat, n);
+
+    Vector_<> b(n);
+    for (int i = 0; i < n; ++i)
+        b[i] = 1.0 + 0.1 * static_cast<double>(i % 7);
+    Vector_<> x(n, 0.0);
+
+    Sparse::BCGSolve(*mat, b, 1e-12, 1e-14, 500, &x);
+    ASSERT_LT(ResidualInfNorm(*mat, x, b), 1e-8);
+}
+