From 84c13e71375f210d7679eb5a34b2618f83311e42 Mon Sep 17 00:00:00 2001 From: "Ryan M. Richard" Date: Wed, 12 Aug 2026 12:45:37 -0500 Subject: [PATCH 1/2] rolls our own GEMM to avoid a bug in Eigen --- .../backends/eigen/eigen_tensor_impl.cpp | 22 ++- docs/source/developer/design/index.rst | 1 + .../developer/design/uq_contraction.rst | 152 ++++++++++++++++++ .../testing/contraction_assignment.hpp | 35 ++++ 4 files changed, 209 insertions(+), 1 deletion(-) create mode 100644 docs/source/developer/design/uq_contraction.rst diff --git a/cxx/src/tensorwrapper/backends/eigen/eigen_tensor_impl.cpp b/cxx/src/tensorwrapper/backends/eigen/eigen_tensor_impl.cpp index 28d13020..6f8157db 100644 --- a/cxx/src/tensorwrapper/backends/eigen/eigen_tensor_impl.cpp +++ b/cxx/src/tensorwrapper/backends/eigen/eigen_tensor_impl.cpp @@ -18,6 +18,7 @@ #include "eigen_tensor_impl.hpp" #include #include +#include namespace tensorwrapper::backends::eigen { @@ -248,7 +249,26 @@ void EIGEN_TENSOR::contraction_assignment_(label_type this_label, map_t rmatrix(new_rhs_buffer.data(), rrows, rcols); map_t omatrix(out_buffer.data(), lrows, rcols); - omatrix = lmatrix * rmatrix; + if constexpr(types::is_uq_type_v) { + // Eigen's dense GEMM seeds each output cell via an internal + // zero-construction path that Eigen::NumTraits::Zero() + // cannot intercept; for UQ scalar types that hidden zero carries the + // compile-time default truncation order, which then clamps every + // accumulated result via operator+=/operator*='s min-order rule. + // Seeding with a genuinely empty FloatType{} instead lets the first + // real term's order propagate through unmodified. See + // TensorWrapper/docs/source/design for the full rationale. + for(std::size_t i = 0; i < lrows; ++i) { + for(std::size_t j = 0; j < rcols; ++j) { + FloatType acc{}; + for(std::size_t k = 0; k < lcols; ++k) + acc += lmatrix(i, k) * rmatrix(k, j); + omatrix(i, j) = acc; + } + } + } else { + omatrix = lmatrix * rmatrix; + } // The last transpose part of TTGT this->permute_assignment(this_label, olabels, *pout_tensor); diff --git a/docs/source/developer/design/index.rst b/docs/source/developer/design/index.rst index 81b3540c..54d73e70 100644 --- a/docs/source/developer/design/index.rst +++ b/docs/source/developer/design/index.rst @@ -33,6 +33,7 @@ Design of TensorWrapper expression op_graph sparse_maps + uq_contraction .. toctree:: :maxdepth: 2 diff --git a/docs/source/developer/design/uq_contraction.rst b/docs/source/developer/design/uq_contraction.rst new file mode 100644 index 00000000..da1f9fc3 --- /dev/null +++ b/docs/source/developer/design/uq_contraction.rst @@ -0,0 +1,152 @@ +.. Copyright 2026 NWChemEx-Project +.. +.. Licensed under the Apache License, Version 2.0 (the "License"); +.. you may not use this file except in compliance with the License. +.. You may obtain a copy of the License at +.. +.. http://www.apache.org/licenses/LICENSE-2.0 +.. +.. Unless required by applicable law or agreed to in writing, software +.. distributed under the License is distributed on an "AS IS" BASIS, +.. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +.. See the License for the specific language governing permissions and +.. limitations under the License. + +.. _tw_designing_uq_contraction: + +##################################### +Contracting Uncertainty-Aware Tensors +##################################### + +The point of this page is to record why tensor contraction for uncertainty +quantification (UQ) scalar types (``sigma::Uncertain``, ``sigma::Interval``, +``sigma::Affine``, ``sigma::ThresholdedAffine``, and ``sigma::TaylorModel``) +does **not** go through the same code path as contraction for ``float``/ +``double``, even though both live in the same +``EigenTensorImpl::contraction_assignment_`` function. + +************************************ +What is UQ-aware tensor contraction? +************************************ + +``EigenTensorImpl::contraction_assignment_`` implements tensor contraction via +TTGT (Transpose-Transpose-GEMM-Transpose): the operands are physically +transposed into a 2D matrix layout, multiplied with a dense GEMM, and the +result is transposed back. This works uniformly for any scalar type ``Eigen`` +knows how to multiply and add -- which includes UQ scalar types, since each has +an ``Eigen::NumTraits`` specialization (see ``sigma/include/sigma/*/eigen_compat.hpp``). + +************************************** +Why do we need to treat it specially? +************************************** + +For UQ scalar types, correctness depends on more than the numeric value: a +type like ``sigma::TaylorModel`` also carries a *truncation order*, and +operations between operands of different orders resolve to +``min(lhs.max_order(), rhs.max_order())`` (documented, intentional behavior -- +see ``sigma::TaylorModel::operator+=``/``operator*=``). + +Eigen's dense GEMM/GEBP kernel seeds each output accumulator cell with an +internal "zero" of the scalar type before accumulating into it. For UQ scalar +types this seed is **not** a genuinely empty/default-constructed value -- it is +a concrete value constructed via an implicit conversion from a numeric literal +(e.g. ``TaylorModel(0)``), which for ``TaylorModel`` resolves to +``default_max_order()`` (compile-time constant, currently 2). Because mixed- +order operations resolve downward via ``min()``, every subsequent accumulation +into that cell is silently clamped to order 2, **regardless of the order the +real operands were built at**. The bug is invisible for ``float``/``double`` +because those types carry no such metadata -- only the numeric value is ever +wrong-or-right, and Eigen's zero seed is numerically correct (``0.0``). + +We confirmed this is genuinely unreachable through the public customization +point: overriding ``Eigen::NumTraits>::Zero()`` to return an +explicitly empty ``TaylorModel()`` had no effect -- instrumenting the override +with a call counter showed it is never invoked by Eigen's GEMM/GEBP kernel for +this scalar type. The zero-seed is constructed through some other, non- +interceptable internal mechanism. + +********************** +Existing Options +********************** + +Patch Eigen directly + Eigen is a vendored third-party dependency several layers below + ``TensorWrapper`` that this project does not own or maintain a patched fork + of. Locating and patching the exact internal zero-construction site inside + ``GeneralMatrixMatrix.h``/the GEBP kernel would be fragile (version-specific) + and outside the maintenance boundary this project wants to take on. + +Override ``Eigen::NumTraits::Zero()`` + The natural customization point for "what does zero look like for this + scalar type." Tried and confirmed ineffective for this code path (see + above) -- Eigen's dense GEMM does not consult it here. + +*************** +Chosen Strategy +*************** + +``contraction_assignment_`` branches at compile time on +``tensorwrapper::types::is_uq_type_v`` (an existing trait, already +used the same way elsewhere in this codebase, e.g. +``generate/generate_eigenvalues.cpp``, and throughout ``SCF``'s eigensolver). +For UQ scalar types, instead of ``omatrix = lmatrix * rmatrix``, a hand-rolled +triple loop performs the same GEMM manually: + +.. code-block:: c++ + + if constexpr(types::is_uq_type_v) { + for(std::size_t i = 0; i < lrows; ++i) { + for(std::size_t j = 0; j < rcols; ++j) { + FloatType acc{}; // genuinely empty/default-constructed + for(std::size_t k = 0; k < lcols; ++k) + acc += lmatrix(i, k) * rmatrix(k, j); + omatrix(i, j) = acc; + } + } + } else { + omatrix = lmatrix * rmatrix; + } + +The key detail is that ``acc`` starts as a *real* ``FloatType{}`` -- for +``TaylorModel`` this is genuinely empty (``empty() == true``), never a +concrete order-2 value. UQ scalar types' ``operator+=`` special-cases an empty +left-hand side (``if(empty()) { return *this = other; }``), so the first term +accumulated into ``acc`` adopts the real operand's order unmodified. Every +later ``+=``/``*=`` in the inner loop then combines two operands that already +carry the correct (matching, in the common case) order, so ``min()`` is a +no-op. The configured truncation order survives the contraction intact. + +This only changes behavior for ``is_uq_type_v``; ``float``/ +``double`` contractions are untouched and keep using Eigen's optimized, +vectorized GEMM. + +********************** +Further Considerations +********************** + +Performance + The hand-rolled loop is a plain, unblocked, unvectorized triple loop -- it + does not have Eigen's cache blocking or SIMD. This is an intentional + tradeoff: correctness for UQ scalar types over GEMM performance for them. + **Do not "simplify" this back to** ``omatrix = lmatrix * rmatrix`` **for UQ + types** -- that regresses this exact bug. If UQ contraction performance + becomes a bottleneck, the fix should be a UQ-aware blocked/vectorized loop + that still seeds accumulators from an empty ``FloatType{}``, not a return + to Eigen's GEMM. + +Other backends + This fix lives in ``EigenTensorImpl`` (the dense Eigen backend) only. Other + backends (e.g. the CUDA/cuTensor backend) would need the equivalent + treatment if/when they are extended to support UQ scalar types; nothing + about this fix propagates automatically to them. + +Non-``TaylorModel`` UQ types + ``sigma::Uncertain``, ``sigma::Interval``, ``sigma::Affine``, and + ``sigma::ThresholdedAffine`` do not carry a truncation order the way + ``TaylorModel`` does, so they were not directly affected by the order- + collapse symptom. They are still routed through the hand-rolled loop + (via the same ``is_uq_type_v`` gate) since they are also + ``RequireInitialization = 1`` custom scalar types subject to the same + unauditable Eigen zero-seed; using a genuinely empty accumulator for them + is at least as correct as, and no more expensive than, whatever Eigen's + internal seed happened to be. diff --git a/tests/cxx/unit_tests/tensorwrapper/backends/testing/contraction_assignment.hpp b/tests/cxx/unit_tests/tensorwrapper/backends/testing/contraction_assignment.hpp index a6226f33..b5e2a02d 100644 --- a/tests/cxx/unit_tests/tensorwrapper/backends/testing/contraction_assignment.hpp +++ b/tests/cxx/unit_tests/tensorwrapper/backends/testing/contraction_assignment.hpp @@ -16,6 +16,7 @@ #pragma once #include +#include #include namespace tensorwrapper::testing { @@ -254,5 +255,39 @@ void contraction_assignment_tests() { REQUIRE(elements_equal(tensor3.get_elem({1, 1, 1}), tensor3_value_type(44.0))); } + + // Regression test: contraction must preserve the operands' Taylor-model + // truncation order instead of collapsing to default_max_order() via + // Eigen's hidden GEMM zero-seed. + if constexpr(types::is_taylor_model_v) { + SECTION("ij,jk->ik preserves non-default Taylor order") { + using order_type = typename matrix_value_type::Order; + const std::size_t order_value = 4; + const auto order = order_type(order_value); + REQUIRE(order_value != + matrix_value_type::default_max_order().value); + + std::vector lhs_data(4); + std::vector rhs_data(4); + for(std::size_t i = 0; i < 4; ++i) { + lhs_data[i] = matrix_value_type(i + 1.0, i + 2.0, order); + rhs_data[i] = matrix_value_type(i + 1.0, i + 2.0, order); + } + std::span lhs_span(lhs_data.data(), + lhs_data.size()); + std::span rhs_span(rhs_data.data(), + rhs_data.size()); + MatrixType lhs(lhs_span, matrix_shape); + MatrixType rhs(rhs_span, matrix_shape); + + label_type o("i,k"); + label_type l("i,j"); + label_type r("j,k"); + matrix.contraction_assignment(o, l, r, lhs, rhs); + + REQUIRE(matrix.get_elem({0, 0}).max_order() == order_value); + REQUIRE(matrix.get_elem({1, 1}).max_order() == order_value); + } + } } } // namespace tensorwrapper::testing From 49efc84190210efce5d5282da2244f5946a3da4d Mon Sep 17 00:00:00 2001 From: "Ryan M. Richard" Date: Tue, 1 Sep 2026 11:04:16 -0500 Subject: [PATCH 2/2] account for uncertainty in log --- .../generate/generate_eigenvalues.cpp | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/cxx/unit_tests/tensorwrapper/generate/generate_eigenvalues.cpp b/tests/cxx/unit_tests/tensorwrapper/generate/generate_eigenvalues.cpp index ebac9fdb..85ed270a 100644 --- a/tests/cxx/unit_tests/tensorwrapper/generate/generate_eigenvalues.cpp +++ b/tests/cxx/unit_tests/tensorwrapper/generate/generate_eigenvalues.cpp @@ -25,6 +25,21 @@ using namespace tensorwrapper::generate; using namespace tensorwrapper::operations; using namespace tensorwrapper::utilities; +namespace { +// Composing two transcendental functions (log then exp) forces a rigorous +// Interval to widen by more than the default 1e-16 tolerance even when +// compared against itself, since an interval's difference with itself is not +// zero-width unless the interval is a single point. This tolerance accounts +// for that unavoidable widening. +template +constexpr double log_spacing_tol = + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v ? + 5e-3 : + 1e-9; +} // namespace + TEMPLATE_LIST_TEST_CASE("generate_eigenvalues", "", types::floating_point_types) { SECTION("n == 1") { @@ -70,7 +85,7 @@ TEMPLATE_LIST_TEST_CASE("generate_eigenvalues", "", auto gen = make_rng(1); auto result = generate_eigenvalues(spec, gen); auto corr = make_tensor({3}, expected); - REQUIRE(approximately_equal(result, corr)); + REQUIRE(approximately_equal(result, corr, log_spacing_tol)); } SECTION("degenerate spacing") {