From 2ff70600e19cd1aeb98ca415026bb2d792bfa78a Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Sun, 30 Aug 2026 12:03:06 -0400 Subject: [PATCH 1/7] ToT: conj() on nested-tile contractions; real-plain x complex-ToT products via one real gemm conj(A * B) on tensor-of-tensor operands did not compile: ContEngine's per-cell multiply-add ops static_cast the ComplexConjugate contraction factor to the element type, Tensor::gemm forwarded it as the BLAS alpha, the ContractReduce<..., ComplexConjugate<...>> specializations named their result type through a value-returning gemm nested tiles do not have, and their nested-tile accumulate was an abort() stub. - detail::elem_factor(factor): the per-element multiplier for a contraction factor -- the factor for numeric factors, 1 for ComplexConjugate<...> (conjugation and scale are applied to the finished result by ContractReduce's finalization, as for non-nested tiles); used by the ToT per-cell ops in ContEngine and by Tensor::gemm's alpha - ContractReduce ComplexConjugate specializations: result_type is Result; the nested-tile accumulate and the arena-aware partial-result reduce are shared with the primary template (ContractReduceBase::accumulate_nested / reduce_results) - ToT x real plain-tensor products: when the plain element type is the real part of the inner element type, the complex slabs are viewed as real matrices with the inner extent doubled (re,im interleaved) and the strided GEMM fast path runs in real arithmetic with alpha = beta = 1 - tests: tot_conj suite in tot_expressions.cpp -- conj(a), permuted conj(a), conj(a)*b, a*conj(b), conj(a*b), conj(a)*b with an inner contraction, and the mixed-type ToT x real-plain product, against explicit references --- src/TiledArray/expressions/cont_engine.h | 24 ++- src/TiledArray/tensor/complex.h | 23 ++ src/TiledArray/tensor/tensor.h | 48 +++-- src/TiledArray/tile_op/contract_reduce.h | 124 ++++++----- tests/tot_expressions.cpp | 261 +++++++++++++++++++++++ 5 files changed, 400 insertions(+), 80 deletions(-) diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h index f73179c6b8..ed43db9b47 100644 --- a/src/TiledArray/expressions/cont_engine.h +++ b/src/TiledArray/expressions/cont_engine.h @@ -1171,7 +1171,9 @@ class ContEngine : public BinaryEngine { const auto* rp = right.data(); result_tile_element_type acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; - return static_cast(factor) * acc; + return TiledArray::detail::elem_factor( + factor) * + acc; }; this->element_nonreturn_op_ = [flat_dot]( result_tile_element_type& result, @@ -1295,7 +1297,8 @@ class ContEngine : public BinaryEngine { Numeric acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; // result cell is pre-shaped [1] by the unit_range plan. - result.data()[0] += static_cast(factor) * acc; + result.data()[0] += + TiledArray::detail::elem_factor(factor) * acc; }; if (this->outer_product_uses_summa()) { this->arena_plan_ = @@ -1403,8 +1406,9 @@ class ContEngine : public BinaryEngine { static_cast(N), static_cast(K), gh.left_op(), gh.right_op(), - static_cast(factor)); + TiledArray::detail::elem_factor< + typename result_tile_element_type:: + numeric_type>(factor)); }; } // ce+ce (hce+ce): inner CONTRACTION (num_contract_ranks() >= @@ -1521,8 +1525,9 @@ class ContEngine : public BinaryEngine { static_cast(No), static_cast(Ko), gh.left_op(), gh.right_op(), - static_cast(factor), + TiledArray::detail::elem_factor< + typename result_tile_element_type:: + numeric_type>(factor), left_inner_T); }; } else if (left_arm_ok) { @@ -1544,8 +1549,9 @@ class ContEngine : public BinaryEngine { static_cast(No), static_cast(Ko), gh.left_op(), gh.right_op(), - static_cast(factor), + TiledArray::detail::elem_factor< + typename result_tile_element_type:: + numeric_type>(factor), right_inner_T); }; } @@ -1702,7 +1708,7 @@ class ContEngine : public BinaryEngine { const auto* rp = right.data(); Numeric acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; - acc *= static_cast(factor); + acc *= TiledArray::detail::elem_factor(factor); if (TA::empty(result)) { using R = typename result_tile_element_type::range_type; TiledArray::container::svector ext( diff --git a/src/TiledArray/tensor/complex.h b/src/TiledArray/tensor/complex.h index fa60a2c39b..25e94e2e03 100644 --- a/src/TiledArray/tensor/complex.h +++ b/src/TiledArray/tensor/complex.h @@ -291,6 +291,29 @@ TILEDARRAY_FORCE_INLINE L& operator*=(L& value, return value; } +template +struct is_complex_conjugate : std::false_type {}; +template +struct is_complex_conjugate> : std::true_type {}; +/// true if \c T is a ComplexConjugate<...> contraction factor +template +inline constexpr bool is_complex_conjugate_v = + is_complex_conjugate>::value; + +/// The numeric multiplier to bake into a per-element (per-cell) multiply-add +/// op for a contraction with factor \c factor: the factor itself (converted +/// to \c Numeric) for a numeric factor, and \c Numeric(1) for a +/// ComplexConjugate<...> factor -- the conjugation (and, for +/// ComplexConjugate, the scale) of such a factor is applied to the +/// finished result by ContractReduce's finalization step, not per element. +template +TILEDARRAY_FORCE_INLINE Numeric elem_factor(const Scalar& factor) { + if constexpr (is_complex_conjugate_v) + return Numeric(1); + else + return static_cast(factor); +} + template inline auto abs(const ComplexConjugate& a) { return std::abs(a.factor()); diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index bc7a4de9d3..6e09e69583 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -3382,6 +3382,10 @@ class Tensor { Tensor& gemm(const Tensor& A, const Tensor& B, const W alpha, const math::GemmHelper& gemm_helper) { numeric_type beta = 1; + // A ComplexConjugate<...> alpha (conj(A*B) at the expression level) is + // applied to the finished result by ContractReduce's finalization step; + // the gemm itself runs with alpha = 1 (see detail::elem_factor). + const numeric_type alpha_n = detail::elem_factor(alpha); if (this->empty()) { *this = Tensor(gemm_helper.make_result_range(A.range_, B.range()), @@ -3406,7 +3410,7 @@ class Tensor { } for (size_t i = 0; i < this->nbatch(); ++i) { auto Ci = this->batch(i); - TiledArray::gemm(alpha, A.batch(i), B.batch(i), + TiledArray::gemm(alpha_n, A.batch(i), B.batch(i), twostep ? numeric_type(0) : numeric_type(1), Ci, gemm_helper); } @@ -3460,7 +3464,7 @@ class Tensor { #else // TA_ENABLE_TILE_OPS_LOGGING for (size_t i = 0; i < this->nbatch(); ++i) { auto Ci = this->batch(i); - TiledArray::detail::gemm(alpha, A.batch(i), B.batch(i), beta, Ci, + TiledArray::detail::gemm(alpha_n, A.batch(i), B.batch(i), beta, Ci, gemm_helper); } #endif // TA_ENABLE_TILE_OPS_LOGGING @@ -3567,7 +3571,16 @@ class Tensor { if constexpr (detail::is_numeric_v && is_tensor_view_v && is_tensor_view_v) { using Real = std::remove_cv_t; - if constexpr (std::is_same_v, Real>) { + using Vr = std::remove_cv_t; + // Same element type: one gemm in that type. Real plain scalars (Vr) + // against complex inner cells: the complex slabs are viewed as real + // matrices with the inner extent doubled (re,im interleaved), so one + // real gemm with alpha = beta = 1 accumulates both parts exactly. + constexpr bool same_type = std::is_same_v; + constexpr bool interleaved = + !same_type && std::is_same_v, Real>; + constexpr integer cw = interleaved ? 2 : 1; // reals per inner element + if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { // kernel-total timer: destroyed at `return *this;` below, so it @@ -3667,7 +3680,7 @@ class Tensor { detail::g_scale[0].gemm_flop.fetch_add( 2ull * static_cast(K) * static_cast(N) * - static_cast(A), + static_cast(A) * cw, std::memory_order_relaxed); } const integer Ai = static_cast(A); @@ -3675,10 +3688,12 @@ class Tensor { TiledArray::math::blas::gemm( TiledArray::math::blas::Transpose, TiledArray::math::blas::NoTranspose, - /*M=*/N, /*N=*/Ai, /*K=*/K, Real(1), + /*M=*/N, /*N=*/Ai * cw, /*K=*/K, Vr(1), /*A=*/right_data, /*lda=*/N, - /*B=*/lc0[0].data(), /*ldb=*/ldb, Real(1), - /*C=*/rc0[0].data(), /*ldc=*/ldc); + /*B=*/reinterpret_cast(lc0[0].data()), + /*ldb=*/ldb * cw, Vr(1), + /*C=*/reinterpret_cast(rc0[0].data()), + /*ldc=*/ldc * cw); } else { // per-cell AXPY fallback for this row if (detail::scale_gemm_timing_enabled()) { // classify fallback reason (re-scan; observation only, does @@ -3747,7 +3762,14 @@ class Tensor { if constexpr (detail::is_numeric_v && is_tensor_view_v && is_tensor_view_v) { using Real = std::remove_cv_t; - if constexpr (std::is_same_v, Real>) { + using Ur = std::remove_cv_t; + // see the tot_x_t block: same type, or real plain x complex inner + // cells via the re,im-interleaved real gemm + constexpr bool same_type = std::is_same_v; + constexpr bool interleaved = + !same_type && std::is_same_v, Real>; + constexpr integer cw = interleaved ? 2 : 1; // reals per inner element + if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { // kernel-total timer (see tot_x_t block); destroyed at `return`. @@ -3831,7 +3853,7 @@ class Tensor { detail::g_scale[1].gemm_flop.fetch_add( 2ull * static_cast(M) * static_cast(K) * - static_cast(A), + static_cast(A) * cw, std::memory_order_relaxed); } const integer Ai = static_cast(A); @@ -3839,10 +3861,12 @@ class Tensor { TiledArray::math::blas::gemm( TiledArray::math::blas::NoTranspose, TiledArray::math::blas::NoTranspose, - /*M=*/M, /*N=*/Ai, /*K=*/K, Real(1), + /*M=*/M, /*N=*/Ai * cw, /*K=*/K, Ur(1), /*A=*/left_data, /*lda=*/K, - /*B=*/right_data[n].data(), /*ldb=*/ldb, Real(1), - /*C=*/this_data[n].data(), /*ldc=*/ldc); + /*B=*/reinterpret_cast(right_data[n].data()), + /*ldb=*/ldb * cw, Ur(1), + /*C=*/reinterpret_cast(this_data[n].data()), + /*ldc=*/ldc * cw); } else { // per-cell AXPY fallback for this column if (detail::scale_gemm_timing_enabled()) { // classify fallback reason (re-scan; observation only) + diff --git a/src/TiledArray/tile_op/contract_reduce.h b/src/TiledArray/tile_op/contract_reduce.h index c36e630c77..9e0cd1e77f 100644 --- a/src/TiledArray/tile_op/contract_reduce.h +++ b/src/TiledArray/tile_op/contract_reduce.h @@ -250,6 +250,57 @@ class ContractReduceBase { strided_oprod_op() const { return pimpl_->strided_oprod_op_; } + /// Reduce two partial results: \c result += \c arg. Arena ToT partials + /// reduced from disjoint K-panel subsets can carry different inner-cell + /// sparsity, so their shapes are unioned before accumulating. + template + void reduce_results(R& result, const R& arg) const { + if constexpr ( + detail::is_contraction_arena_tot_v< + R, std::remove_cv_t>, + std::remove_cv_t>>) { + detail::arena_tot_add_to(result, arg); + } else { + using TiledArray::add_to; + add_to(result, arg); + } + } + + /// Nested-tile accumulate: \c result += \c left * \c right through the + /// per-cell multiply-add op (via the arena plan and the strided + /// outer-product op when installed). Shared by the primary ContractReduce + /// and its ComplexConjugate specializations -- a ComplexConjugate factor is + /// applied to the finished result by those specializations' finalization + /// step, so the accumulate itself is identical. + template + void accumulate_nested(R& result, const L& left, const Rt& right) const { + using TiledArray::empty; + using TiledArray::gemm; + TA_ASSERT(this->elem_muladd_op()); + if constexpr (detail::is_contraction_arena_tot_v< + R, std::remove_cv_t>, + std::remove_cv_t>>) { + // The result tile is shaped from operand inner cells. A SUMMA + // reduction streams K-panels one at a time: the first panel sizes the + // result; a later panel of a contracted-dimension-sparse ToT operand + // can touch inner cells the first panel left null, so each subsequent + // panel extends the result to cover its own cells. + if (this->arena_plan().has_value()) { + if (empty(result)) + result = this->arena_plan()->reserve_and_construct( + left, right, this->gemm_helper()); + else + this->arena_plan()->grow_to_cover(result, left, right, + this->gemm_helper()); + } + if (this->strided_oprod_op()) { + this->strided_oprod_op()(result, left, right, this->gemm_helper()); + return; + } + } + gemm(result, left, right, this->gemm_helper(), this->elem_muladd_op()); + } + void set_strided_oprod_op( TiledArray::function_ref op) { pimpl_->strided_oprod_op_ = op; @@ -384,19 +435,7 @@ class ContractReduce : public ContractReduceBase { /// target /// \param[in] arg The argument that will be added to \c result void operator()(result_type& result, const result_type& arg) const { - if constexpr ( - detail::is_contraction_arena_tot_v< - result_type, - std::remove_cv_t>, - std::remove_cv_t>>) { - // Two partial contraction results reduced from disjoint K-panel - // subsets can carry different inner-cell sparsity; union their shapes - // before accumulating. - detail::arena_tot_add_to(result, arg); - } else { - using TiledArray::add_to; - add_to(result, arg); - } + this->reduce_results(result, arg); } /// Contract a pair of tiles and add to a target tile @@ -413,34 +452,7 @@ class ContractReduce : public ContractReduceBase { if (empty(left) || empty(right)) return; if constexpr (!ContractReduceBase_::plain_tensors) { - TA_ASSERT(this->elem_muladd_op()); - if constexpr (detail::is_contraction_arena_tot_v< - result_type, - std::remove_cv_t< - std::remove_reference_t>, - std::remove_cv_t< - std::remove_reference_t>>) { - // The result tile is shaped from operand inner cells. A SUMMA - // reduction streams K-panels one at a time: the first panel sizes the - // result; a later panel of a contracted-dimension-sparse ToT operand - // can touch inner cells the first panel left null, so each subsequent - // panel extends the result to cover its own cells. - if (this->arena_plan().has_value()) { - if (empty(result)) - result = this->arena_plan()->reserve_and_construct( - left, right, this->gemm_helper()); - else - this->arena_plan()->grow_to_cover(result, left, right, - this->gemm_helper()); - } - if (this->strided_oprod_op()) { - this->strided_oprod_op()(result, left, right, - ContractReduceBase_::gemm_helper()); - return; - } - } - gemm(result, left, right, ContractReduceBase_::gemm_helper(), - this->elem_muladd_op()); + this->accumulate_nested(result, left, right); } else { // plain tensors TA_ASSERT(!this->elem_muladd_op()); if (empty(result)) @@ -476,10 +488,8 @@ class ContractReduce(), std::declval(), 1, - std::declval())) - result_type; ///< The result tile type. + second_argument_type; ///< The right tile type + typedef Result result_type; ///< The result tile type. typedef TiledArray::detail::ComplexConjugate scalar_type; using typename ContractReduceBase_::elem_muladd_op_type; @@ -555,8 +565,7 @@ class ContractReducereduce_results(result, arg); } /// Contract a pair of tiles and add to a target tile @@ -568,10 +577,10 @@ class ContractReduceelem_muladd_op()); - // not yet implemented - abort(); + this->accumulate_nested(result, left, right); } else { TA_ASSERT(!this->elem_muladd_op()); using TiledArray::empty; @@ -608,10 +617,8 @@ class ContractReduce(), std::declval(), 1, - std::declval())) - result_type; ///< The result tile type. + second_argument_type; ///< The right tile type + typedef Result result_type; ///< The result tile type. typedef TiledArray::detail::ComplexConjugate scalar_type; using typename ContractReduceBase_::elem_muladd_op_type; @@ -687,8 +694,7 @@ class ContractReducereduce_results(result, arg); } /// Contract a pair of tiles and add to a target tile @@ -700,10 +706,10 @@ class ContractReduceelem_muladd_op()); - // not yet implemented - abort(); + this->accumulate_nested(result, left, right); } else { TA_ASSERT(!this->elem_muladd_op()); using TiledArray::empty; diff --git a/tests/tot_expressions.cpp b/tests/tot_expressions.cpp index cb8fbe5f52..e4472b5ca7 100644 --- a/tests/tot_expressions.cpp +++ b/tests/tot_expressions.cpp @@ -1,3 +1,4 @@ +#include #include "tot_array_fixture.h" template @@ -4605,3 +4606,263 @@ BOOST_AUTO_TEST_CASE(ik_mn_eq_ij_mn_times_kj_mn) { } BOOST_AUTO_TEST_SUITE_END() + +//------------------------------------------------------------------------------ +// conj() on ToT expressions, and the mixed-type ToT x real-plain-tensor +// product. Every case is checked against an explicit reference computed from +// single-tile arrays. +//------------------------------------------------------------------------------ + +namespace { + +// Owning nested tiles only (the btas inner rows of test_params are not +// exercised here). +using conj_test_params = boost::mpl::list< + std::tuple>>, + std::tuple, Tensor>>>>; + +template +E mk(double re, double im) { + if constexpr (TiledArray::detail::is_complex_v) + return E(re, im); + else + return E(re); +} + +template +E cj(const E& x) { + return TiledArray::detail::conj(x); +} + +template +void check_close(const E& got, const E& ref) { + BOOST_CHECK_SMALL(std::abs(got - ref), + 1e-10 * std::max(1.0, double(std::abs(ref)))); +} + +// Single-tile rank-2-outer ToT with rank-1 inner cells of extent na: +// A(i,j)(a) = gen(i, j, a) +template +Array make_tot_1(World& world, std::size_t ni, std::size_t nj, std::size_t na, + Gen gen) { + using inner_t = typename Array::value_type::value_type; + TiledRange tr{TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { + inner_t t(Range{static_cast(na)}); + for (std::size_t a = 0; a < na; ++a) t.at_ordinal(a) = gen(idx[0], idx[1], a); + return t; + }); + world.gop.fence(); + return arr; +} + +// Single-tile rank-2-outer ToT with rank-2 inner cells (na x nb): +// A(i,j)(a,b) = gen(i, j, a, b) +template +Array make_tot_2(World& world, std::size_t ni, std::size_t nj, std::size_t na, + std::size_t nb, Gen gen) { + using inner_t = typename Array::value_type::value_type; + TiledRange tr{TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { + inner_t t(Range{static_cast(na), static_cast(nb)}); + for (std::size_t a = 0; a < na; ++a) + for (std::size_t b = 0; b < nb; ++b) t(a, b) = gen(idx[0], idx[1], a, b); + return t; + }); + world.gop.fence(); + return arr; +} + +template +auto single_tile(const Array& arr) { + return arr.find({0, 0}).get(); +} + +} // namespace + +BOOST_FIXTURE_TEST_SUITE(tot_conj, ToTArrayFixture) + +// c(i,j;a) = conj(a(i,j;a)) +BOOST_AUTO_TEST_CASE_TEMPLATE(unary, TestParam, conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, na = 4; + auto gen = [](auto i, auto j, auto a) { + return mk(1.0 + i + 2.0 * j + 0.5 * a, 0.3 * i - j + a); + }; + array_t a = make_tot_1(m_world, ni, nj, na, gen); + array_t c; + c("i,j;a") = conj(a("i,j;a")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t x = 0; x < na; ++x) + check_close(tile(i, j).at_ordinal(x), cj(gen(i, j, x))); +} + +// c(j,i;a) = conj(a(i,j;a)) (outer permutation + conj) +BOOST_AUTO_TEST_CASE_TEMPLATE(unary_permuted, TestParam, conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, na = 4; + auto gen = [](auto i, auto j, auto a) { + return mk(1.0 + i + 2.0 * j + 0.5 * a, 0.3 * i - j + a); + }; + array_t a = make_tot_1(m_world, ni, nj, na, gen); + array_t c; + c("j,i;a") = conj(a("i,j;a")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t x = 0; x < na; ++x) + check_close(tile(j, i).at_ordinal(x), cj(gen(i, j, x))); +} + +// c(i,k;a,b) = sum_j conj(a(i,j;a)) * b(j,k;b) (outer contraction, inner +// outer product, conj on the left operand) +BOOST_AUTO_TEST_CASE_TEMPLATE(conj_left_outer_product, TestParam, + conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + auto ga = [](auto i, auto j, auto a) { + return mk(1.0 + i - j + 0.5 * a, 0.25 * i + j - a); + }; + auto gb = [](auto j, auto k, auto b) { + return mk(2.0 - j + k + 0.1 * b, 0.5 * j - k + 0.2 * b); + }; + array_t a = make_tot_1(m_world, ni, nj, na, ga); + array_t b = make_tot_1(m_world, nj, nk, nb, gb); + array_t c; + c("i,k;a,b") = conj(a("i,j;a")) * b("j,k;b"); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + E ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += cj(ga(i, j, x)) * gb(j, k, y); + check_close(tile(i, k)(x, y), ref); + } +} + +// c(i,k;a,b) = sum_j a(i,j;a) * conj(b(j,k;b)) +BOOST_AUTO_TEST_CASE_TEMPLATE(conj_right_outer_product, TestParam, + conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + auto ga = [](auto i, auto j, auto a) { + return mk(1.0 + i - j + 0.5 * a, 0.25 * i + j - a); + }; + auto gb = [](auto j, auto k, auto b) { + return mk(2.0 - j + k + 0.1 * b, 0.5 * j - k + 0.2 * b); + }; + array_t a = make_tot_1(m_world, ni, nj, na, ga); + array_t b = make_tot_1(m_world, nj, nk, nb, gb); + array_t c; + c("i,k;a,b") = a("i,j;a") * conj(b("j,k;b")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + E ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += ga(i, j, x) * cj(gb(j, k, y)); + check_close(tile(i, k)(x, y), ref); + } +} + +// c(i,k;a,b) = conj( sum_j a(i,j;a) * b(j,k;b) ) +BOOST_AUTO_TEST_CASE_TEMPLATE(conj_of_product, TestParam, conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + auto ga = [](auto i, auto j, auto a) { + return mk(1.0 + i - j + 0.5 * a, 0.25 * i + j - a); + }; + auto gb = [](auto j, auto k, auto b) { + return mk(2.0 - j + k + 0.1 * b, 0.5 * j - k + 0.2 * b); + }; + array_t a = make_tot_1(m_world, ni, nj, na, ga); + array_t b = make_tot_1(m_world, nj, nk, nb, gb); + array_t c; + c("i,k;a,b") = conj(a("i,j;a") * b("j,k;b")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + E ref{}; + for (std::size_t j = 0; j < nj; ++j) ref += ga(i, j, x) * gb(j, k, y); + check_close(tile(i, k)(x, y), cj(ref)); + } +} + +// c(i,k;a) = sum_j sum_b conj(a(i,j;a,b)) * b(j,k;b) (outer + inner +// contraction, conj on the left operand) +BOOST_AUTO_TEST_CASE_TEMPLATE(conj_left_inner_contraction, TestParam, + conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + auto ga = [](auto i, auto j, auto a, auto b) { + return mk(1.0 + i - j + 0.5 * a - 0.3 * b, 0.25 * i + j - a + 0.1 * b); + }; + auto gb = [](auto j, auto k, auto b) { + return mk(2.0 - j + k + 0.1 * b, 0.5 * j - k + 0.2 * b); + }; + array_t a = make_tot_2(m_world, ni, nj, na, nb, ga); + array_t b = make_tot_1(m_world, nj, nk, nb, gb); + array_t c; + c("i,k;a") = conj(a("i,j;a,b")) * b("j,k;b"); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + E ref{}; + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t y = 0; y < nb; ++y) + ref += cj(ga(i, j, x, y)) * gb(j, k, y); + check_close(tile(i, k).at_ordinal(x), ref); + } +} + +// c(i,k;a) = sum_j a(i,j;a) * t(j,k) with a REAL plain array t: the +// ToT x plain-tensor product with different element types (complex ToT, real +// plain tensor). For the real row this is the same-type product. +BOOST_AUTO_TEST_CASE_TEMPLATE(tot_times_real_plain, TestParam, + conj_test_params) { + using array_t = tensor_type; + using E = typename inner_type::value_type; + using plain_t = DistArray, policy_type>; + const std::size_t ni = 2, nj = 3, nk = 4, na = 3; + auto ga = [](auto i, auto j, auto a) { + return mk(1.0 + i - j + 0.5 * a, 0.25 * i + j - a); + }; + auto gt = [](auto j, auto k) { return 0.5 + j - 0.25 * k; }; + array_t a = make_tot_1(m_world, ni, nj, na, ga); + TiledRange tr{TiledRange1{0, static_cast(nj)}, + TiledRange1{0, static_cast(nk)}}; + plain_t t(m_world, tr); + t.init_elements([=](const auto& idx) { return gt(idx[0], idx[1]); }); + m_world.gop.fence(); + array_t c; + c("i,k;a") = a("i,j;a") * t("j,k"); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + E ref{}; + for (std::size_t j = 0; j < nj; ++j) ref += ga(i, j, x) * gt(j, k); + check_close(tile(i, k).at_ordinal(x), ref); + } +} + +BOOST_AUTO_TEST_SUITE_END() From 59cc2e8701ff3ac8bc404d8e0dab564238e7cad0 Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Mon, 7 Sep 2026 13:12:13 -0400 Subject: [PATCH 2/7] Tensor: static_assert the complex == Real[2] layout the interleaved gemm relies on The interleaved real-gemm path reinterprets a std::complex/ slab as a real matrix with twice the column count. The standard guarantees that layout (array-oriented access, [complex.numbers.general]/4); pin it with a static_assert at both interleaved sites so an exotic ABI fails to compile instead of silently producing garbage. Also reflow one over-long line in the tot_expressions make_tot_1 helper (clang-format). --- src/TiledArray/tensor/tensor.h | 11 +++++++++++ tests/tot_expressions.cpp | 3 ++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 6e09e69583..6727cf7136 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -3580,6 +3580,13 @@ class Tensor { constexpr bool interleaved = !same_type && std::is_same_v, Real>; constexpr integer cw = interleaved ? 2 : 1; // reals per inner element + // The interleaved view reinterprets the complex slab as Vr[2*n]: the + // standard guarantees array-oriented access to std::complex + // ([complex.numbers.general]/4), pinned here so an exotic ABI cannot + // turn the reinterpret_cast below into silent undefined behavior. + static_assert(!interleaved || (sizeof(Real) == 2 * sizeof(Vr) && + alignof(Real) == alignof(Vr)), + "interleaved real gemm needs std::complex == Vr[2]"); if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { @@ -3769,6 +3776,10 @@ class Tensor { constexpr bool interleaved = !same_type && std::is_same_v, Real>; constexpr integer cw = interleaved ? 2 : 1; // reals per inner element + // see the tot_x_t block: std::complex must be exactly Ur[2] + static_assert(!interleaved || (sizeof(Real) == 2 * sizeof(Ur) && + alignof(Real) == alignof(Ur)), + "interleaved real gemm needs std::complex == Ur[2]"); if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { diff --git a/tests/tot_expressions.cpp b/tests/tot_expressions.cpp index e4472b5ca7..0c8c613dff 100644 --- a/tests/tot_expressions.cpp +++ b/tests/tot_expressions.cpp @@ -4651,7 +4651,8 @@ Array make_tot_1(World& world, std::size_t ni, std::size_t nj, std::size_t na, Array arr(world, tr); arr.init_elements([=](const auto& idx) { inner_t t(Range{static_cast(na)}); - for (std::size_t a = 0; a < na; ++a) t.at_ordinal(a) = gen(idx[0], idx[1], a); + for (std::size_t a = 0; a < na; ++a) + t.at_ordinal(a) = gen(idx[0], idx[1], a); return t; }); world.gop.fence(); From 0a8dea7b607fdcbcbf26a5b1b2e8529420a5906e Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Tue, 8 Sep 2026 16:43:52 -0400 Subject: [PATCH 3/7] ToT conj: apply a ComplexConjugate factor exactly once, in the outer op's finalization A conj(A*B) factor cannot be applied per cell (conj does not distribute into a sum of products), so ContEngine now hands every consumer one of two things: elem_scale(), the numeric per-cell multiplier (the factor itself, 1 for a ComplexConjugate<...>), and outer_factor(), the factor of the outer tile op (1 for a numeric factor, whose scale the cells absorb; the ComplexConjugate<...> itself otherwise, so its finalization conjugates AND scales the finished tile). No site reads factor_ directly any more. Before, six of sixteen consumers used elem_factor() and the rest the raw factor, so S*conj(A*B) and -conj(A*B) dropped the scale/sign, the inner-Hadamard form double-conjugated, a complex plain operand came out un-conjugated, and every arena-cell conj(A*B) failed to compile. - ContractReduceBase: a ComplexConjugate alpha may accompany a per-cell op (the "non-unit alpha must be absorbed" assertion applies to numeric factors only). - BatchedContractReduce: the finalization forwards to the wrapped op's, so batched (fused-mode) conj(A*B) is conjugated like the unbatched one (this was wrong for plain tiles too). - ArenaTensor: operator*= overloads for ComplexConjugate factors (in-place, via the free scale_to kernel); the gemm overload rejects a ComplexConjugate factor at compile time. - Tensor::scale_to: a ComplexConjugate factor on owning nested cells conjugates each cell in place instead of `cell = conj(cell) * S` (allocation + copy per cell); real elements skip the conjugation; arena cells go through the free kernel for every factor type. - Tensor::gemm: a ComplexConjugate alpha is a static_assert again (it was silently mapped to 1). - ContractReduce: the two ComplexConjugate specializations (120 lines apart in two finalization lines) are one specialization on ComplexConjugate with an if constexpr on void. - complex.h: the mixed-scalar operator* templates are declared before the ComplexConjugate operators, so an integer-literal scale (2 * conj(a*b)) compiles. - [scale-timing]: the interleaved real-gemm path counts FLOPs in element multiply-adds, the unit of the other paths. - tests (tot_conj): scaled / negated / integer-scaled conj of a product, scaled conj of an inner contraction and of an inner Hadamard, conj of a ToT x complex plain product (with an inner permutation), multi-tile scaled conj, fused-mode conj (ToT and plain), arena-cell conj forms, and the arena ToT x real plain products that exercise the interleaved real-gemm fast path. --- src/TiledArray/expressions/cont_engine.h | 104 ++-- src/TiledArray/tensor/arena_tensor.h | 34 +- src/TiledArray/tensor/complex.h | 57 ++- src/TiledArray/tensor/tensor.h | 63 ++- .../tile_op/batched_contract_reduce.h | 8 +- src/TiledArray/tile_op/contract_reduce.h | 170 ++----- tests/tot_expressions.cpp | 444 ++++++++++++++++++ 7 files changed, 661 insertions(+), 219 deletions(-) diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h index ed43db9b47..d4faea643f 100644 --- a/src/TiledArray/expressions/cont_engine.h +++ b/src/TiledArray/expressions/cont_engine.h @@ -113,6 +113,43 @@ class ContEngine : public BinaryEngine { protected: scalar_type factor_; ///< Contraction scaling factor + /// The numeric type of the tile's elements (for a tensor-of-tensors, of the + /// inner cells): the type of the per-cell multiplier below. + using elem_scalar_type = + typename TiledArray::detail::numeric_type::type; + + /// A contraction's factor is applied in exactly one of two places. A + /// numeric factor is absorbed into the per-cell (per-element) multiply-add + /// ops and the outer tile op runs with factor 1. A ComplexConjugate<...> + /// factor (`conj(A*B)`, `S*conj(A*B)`, `-conj(A*B)`) cannot be applied per + /// cell -- conj does not distribute into a sum of products -- so the + /// per-cell ops run with multiplier 1 and the outer op's finalization + /// conjugates AND scales the finished tile (ContractReduce's ComplexConjugate + /// specializations). Every consumer of the factor goes through these two + /// accessors; none may read factor_ directly. + + /// \return the multiplier for the per-cell ops + template + Numeric elem_scale() const { + return TiledArray::detail::elem_factor(factor_); + } + + /// \return true if the per-cell ops need no multiplier + bool elem_scale_is_one() const { + if constexpr (TiledArray::detail::is_complex_conjugate_v) + return true; + else + return factor_ == scalar_type(1); + } + + /// \return the factor handed to the outer (tile-level) op + scalar_type outer_factor() const { + if constexpr (TiledArray::detail::is_complex_conjugate_v) + return factor_; + else + return scalar_type(1); + } + protected: op_type op_; ///< Tile operation @@ -431,8 +468,9 @@ class ContEngine : public BinaryEngine { auto total_perm = make_total_perm(); - // factor_ is absorbed into inner_tile_nonreturn_op_ - op_ = op_type(left_op, right_op, scalar_type(1), outer_size(indices_), + // a numeric factor_ is absorbed into the per-cell op; a + // ComplexConjugate one is applied by op_'s finalization + op_ = op_type(left_op, right_op, outer_factor(), outer_size(indices_), outer_size(left_indices_), outer_size(right_indices_), total_perm, this->element_nonreturn_op_, std::move(this->arena_plan_)); @@ -494,8 +532,9 @@ class ContEngine : public BinaryEngine { auto total_perm = make_total_perm(); - // factor_ is absorbed into inner_tile_nonreturn_op_ - op_ = op_type(left_op, right_op, scalar_type(1), outer_size(indices_), + // a numeric factor_ is absorbed into the per-cell op; a + // ComplexConjugate one is applied by op_'s finalization + op_ = op_type(left_op, right_op, outer_factor(), outer_size(indices_), outer_size(left_indices_), outer_size(right_indices_), total_perm, this->element_nonreturn_op_, std::move(this->arena_plan_)); @@ -779,8 +818,9 @@ class ContEngine : public BinaryEngine { "annotation of the result"); } - // factor_ is absorbed into element_nonreturn_op_ - op_ = op_type(left_op, right_op, scalar_type(1), + // a numeric factor_ is absorbed into the per-cell op; a + // ComplexConjugate one is applied by op_'s finalization + op_ = op_type(left_op, right_op, outer_factor(), outer_size(indices_) - nh + u + u_right, outer_size(left_indices_) - nh + u, outer_size(right_indices_) - nh + u_right, @@ -1156,7 +1196,7 @@ class ContEngine : public BinaryEngine { // ContractReduce (built in init_struct) routes the !plain_tensors case to // gemm(result, left, right, helper, elem_muladd_op), invoking this op per // outer cell. - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); // shared flat (non-conjugating) scalar dot of two inner cells, scaled by // factor; returns the contribution for one outer cell (0 if either // operand cell is empty). The numerically-sensitive accumulation lives @@ -1245,7 +1285,7 @@ class ContEngine : public BinaryEngine { "nested Hadamard on view inner tiles: the arena fast path " "was inactive (arena disabled, or a non-identity inner " "result permutation -- not yet supported on view cells)"); - if (this->factor_ == scalar_type{1}) { + if (this->elem_scale_is_one()) { this->element_nonreturn_op_ = TiledArray::detail::make_fused_hadamard_lambda< result_tile_element_type, left_tile_element_type, @@ -1254,7 +1294,7 @@ class ContEngine : public BinaryEngine { this->element_nonreturn_op_ = TiledArray::detail::make_fused_hadamard_scaled_lambda< result_tile_element_type, left_tile_element_type, - right_tile_element_type>(this->factor_); + right_tile_element_type>(this->elem_scale()); } } // element_return_op_ left null: a view cell cannot be @@ -1282,7 +1322,7 @@ class ContEngine : public BinaryEngine { break; } if (result_inner_all_phantom) { - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); this->element_nonreturn_op_ = [factor](result_tile_element_type& result, const left_tile_element_type& left, @@ -1326,9 +1366,9 @@ class ContEngine : public BinaryEngine { } else { using op_type = TiledArray::detail::ContractReduce< result_tile_element_type, left_tile_element_type, - right_tile_element_type, scalar_type>; - // The inner op is built *perm-free* on purpose. factor_ is - // absorbed into element_nonreturn_op_; operand inner transposes + right_tile_element_type, elem_scalar_type>; + // The inner op is built *perm-free* on purpose. It carries the + // per-cell multiplier (elem_scale()); operand inner transposes // are folded into the inner GEMM via left_/right_inner_permtype_. // A non-identity inner *result* permutation is NOT placed on this // op (make_fused_contraction_lambda asserts a perm-free op); it @@ -1338,7 +1378,7 @@ class ContEngine : public BinaryEngine { // Hadamard outer product. auto contrreduce_op = op_type( to_cblas_op(this->left_inner_permtype_), - to_cblas_op(this->right_inner_permtype_), this->factor_, + to_cblas_op(this->right_inner_permtype_), this->elem_scale(), inner_size(this->indices_), inner_size(this->left_indices_), inner_size(this->right_indices_)); // perm-free per-cell in-place contraction; used by both outer @@ -1392,7 +1432,7 @@ class ContEngine : public BinaryEngine { typename result_tile_element_type::numeric_type>) { if (contrreduce_op.gemm_helper().num_contract_ranks() == 0 && !bool(inner(this->perm_))) { - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); this->arena_strided_gemm_ce_e_tile_op_ = [factor](result_tile_type& Cc, const left_tile_type& Lt, const right_tile_type& Rt, @@ -1507,7 +1547,7 @@ class ContEngine : public BinaryEngine { TiledArray::expressions::PermutationType::identity && inner_pt_ok(this->right_inner_permtype_); if (right_arm_ok) { - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); const bool left_inner_T = this->left_inner_permtype_ == TiledArray::expressions::PermutationType:: @@ -1531,7 +1571,7 @@ class ContEngine : public BinaryEngine { left_inner_T); }; } else if (left_arm_ok) { - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); const bool right_inner_T = this->right_inner_permtype_ == TiledArray::expressions::PermutationType:: @@ -1694,7 +1734,7 @@ class ContEngine : public BinaryEngine { } if (result_inner_all_phantom) { const std::size_t phantom_rank = result_inner.size(); - const scalar_type factor = this->factor_; + const auto factor = this->elem_scale(); this->element_nonreturn_op_ = [phantom_rank, factor](result_tile_element_type& result, const left_tile_element_type& left, @@ -1720,21 +1760,21 @@ class ContEngine : public BinaryEngine { } else { using op_type = TiledArray::detail::ContractReduce< result_tile_element_type, left_tile_element_type, - right_tile_element_type, scalar_type>; - // factor_ is absorbed into inner_tile_nonreturn_op_ + right_tile_element_type, elem_scalar_type>; + // the inner op carries the per-cell multiplier (elem_scale()) auto contrreduce_op = (inner_target_indices != inner(this->indices_)) ? op_type( to_cblas_op(this->left_inner_permtype_), to_cblas_op(this->right_inner_permtype_), - this->factor_, inner_size(this->indices_), + this->elem_scale(), inner_size(this->indices_), inner_size(this->left_indices_), inner_size(this->right_indices_), (!this->implicit_permute_inner_ ? inner(this->perm_) : Permutation{})) : op_type(to_cblas_op(this->left_inner_permtype_), to_cblas_op(this->right_inner_permtype_), - this->factor_, inner_size(this->indices_), + this->elem_scale(), inner_size(this->indices_), inner_size(this->left_indices_), inner_size(this->right_indices_)); constexpr bool arena_eligible = @@ -1792,7 +1832,7 @@ class ContEngine : public BinaryEngine { // is contract then inner must implement (ternary) multiply-add; // if the outer is hadamard then the inner is binary multiply const bool outer_uses_summa = this->outer_product_uses_summa(); - if (this->factor_ == scalar_type{1}) { + if (this->elem_scale_is_one()) { using base_op_type = TiledArray::detail::Mult { } else { using base_op_type = TiledArray::detail::ScalMult< result_tile_element_type, left_tile_element_type, - right_tile_element_type, scalar_type, false, false>; + right_tile_element_type, elem_scalar_type, false, false>; using op_type = TiledArray::detail::BinaryWrapper< base_op_type>; // can't consume inputs if they are used // multiple times, e.g. when outer op is gemm auto mult_op = (inner_target_indices != inner(this->indices_)) - ? op_type(base_op_type(this->factor_), + ? op_type(base_op_type(this->elem_scale()), !this->implicit_permute_inner_ ? inner(this->perm_) : Permutation{}) - : op_type(base_op_type(this->factor_)); + : op_type(base_op_type(this->elem_scale())); constexpr bool arena_eligible_h_scaled = TiledArray::detail::is_contraction_arena_tot_v< result_tile_type, left_tile_type, right_tile_type>; @@ -1896,7 +1936,7 @@ class ContEngine : public BinaryEngine { this->element_nonreturn_op_ = TiledArray::detail::make_fused_hadamard_scaled_lambda< result_tile_element_type, left_tile_element_type, - right_tile_element_type>(this->factor_); + right_tile_element_type>(this->elem_scale()); } else { this->element_nonreturn_op_ = [mult_op, outer_uses_summa]( @@ -1962,8 +2002,7 @@ class ContEngine : public BinaryEngine { // the fused arena scale ops are factor-free; a non-unit // expression-level prefactor (ScalMult) takes the fallback op, // which absorbs it - if (this->outer_product_uses_summa() && - this->factor_ == scalar_type(1)) { + if (this->outer_product_uses_summa() && this->elem_scale_is_one()) { // The inner perm handed to the plan must match how the inner // *result* permutation is applied for this result cell type -- // and the two cell types apply it in different places: @@ -2003,9 +2042,10 @@ class ContEngine : public BinaryEngine { [perm = !this->implicit_permute_inner_ ? inner(this->perm_) : Permutation{}, outer_uses_summa = this->outer_product_uses_summa(), - factor = this->factor_](result_tile_element_type& result, - const left_tile_element_type& left, - const right_tile_element_type& right) { + factor = this->elem_scale()]( + result_tile_element_type& result, + const left_tile_element_type& left, + const right_tile_element_type& right) { if (outer_uses_summa) { using TiledArray::axpy_to; if constexpr (tot_x_t) { diff --git a/src/TiledArray/tensor/arena_tensor.h b/src/TiledArray/tensor/arena_tensor.h index 7d2bba61c9..d4895e802d 100644 --- a/src/TiledArray/tensor/arena_tensor.h +++ b/src/TiledArray/tensor/arena_tensor.h @@ -439,18 +439,45 @@ void fill(ArenaTensor& dst, const U& value) { std::fill_n(dst.data(), dst.size(), static_cast(value)); } -/// `dst *= factor`. No-op on a null view. +/// `dst *= factor`. No-op on a null view. A ComplexConjugate<...> factor +/// conjugates (and scales) every arena scalar in place. template void scale_to(ArenaTensor& dst, Scalar factor) { if (!dst) return; auto* d = dst.data(); const auto n = dst.size(); // operator*= is the permissive one (std::complex::operator*=(const T&) - // accepts any arithmetic factor), so no mixed-scalar detail operator is + // accepts any arithmetic factor; detail::operator*=(T&, ComplexConjugate) + // conjugates a scalar in place), so no mixed-scalar detail operator is // needed here for (std::size_t i = 0; i < n; ++i) d[i] *= factor; } +/// `dst *= conj_op()`, `dst *= conj_op(S)`, `dst *= -conj_op()`: in-place +/// conjugation (and scale) of every arena scalar. These overloads exist so +/// that a view is never routed to the generic +/// detail::operator*=(L&, ComplexConjugate), which needs a value-returning +/// conj(L) that a non-owning view cannot provide. +template +ArenaTensor& operator*=(ArenaTensor& dst, + const detail::ComplexConjugate& factor) { + scale_to(dst, factor); + return dst; +} +template +ArenaTensor& operator*=( + ArenaTensor& dst, + const detail::ComplexConjugate& factor) { + scale_to(dst, factor); + return dst; +} +template +ArenaTensor& operator*=(ArenaTensor& dst, + const detail::ComplexConjugate& factor) { + scale_to(dst, factor); + return dst; +} + /// `dst += src`. Asserts shape compatibility. /// A null `src` is an implicit zero and a no-op; a null `dst` with a populated /// `src` throws (a view has no storage to write into and cannot allocate, so @@ -599,6 +626,9 @@ auto gemm(ArenaTensor& result, const ArenaTensor& left, (gemm_helper.right_op() == math::blas::NoTranspose) ? N : K; const integer ldc = N; + static_assert(!detail::is_complex_conjugate_v, + "ArenaTensor gemm: a ComplexConjugate<...> factor cannot be " + "applied inside the gemm; conjugate the finished result"); math::blas::gemm(gemm_helper.left_op(), gemm_helper.right_op(), M, N, K, static_cast(factor), left.data(), lda, right.data(), ldb, T(1), result.data(), ldc); diff --git a/src/TiledArray/tensor/complex.h b/src/TiledArray/tensor/complex.h index 25e94e2e03..55467ff4e7 100644 --- a/src/TiledArray/tensor/complex.h +++ b/src/TiledArray/tensor/complex.h @@ -220,6 +220,37 @@ inline ComplexConjugate ComplexConjugate::operator-() template struct is_numeric> : public std::true_type {}; +// Mixed-scalar products (an integral or floating scalar with a std::complex of +// a different scalar type). Declared BEFORE the ComplexConjugate operators +// below: those evaluate `conj(value) * op.factor()`, and an integer-literal +// scale (`2 * conj(a*b)` -> ComplexConjugate) resolves only if these are +// visible at the template definition (they are not found by ADL). +template >> +TILEDARRAY_FORCE_INLINE auto operator*(const L l, const std::complex r) { + return static_cast(l) * r; +} + +template >> +TILEDARRAY_FORCE_INLINE auto operator*(const std::complex l, const R r) { + return l * static_cast(r); +} + +template +TILEDARRAY_FORCE_INLINE + std::enable_if_t, std::complex> + operator*(const L l, const std::complex r) { + return std::complex(l, 0.) * r; +} + +template +TILEDARRAY_FORCE_INLINE + std::enable_if_t, std::complex> + operator*(const std::complex l, const R r) { + return l * std::complex(r, 0.); +} + /// ComplexConjugate operator factory function /// \tparam S The scalar type @@ -321,32 +352,6 @@ inline auto abs(const ComplexConjugate& a) { inline int abs(const ComplexConjugate& a) { return 1; } -template >> -TILEDARRAY_FORCE_INLINE auto operator*(const L l, const std::complex r) { - return static_cast(l) * r; -} - -template >> -TILEDARRAY_FORCE_INLINE auto operator*(const std::complex l, const R r) { - return l * static_cast(r); -} - -template -TILEDARRAY_FORCE_INLINE - std::enable_if_t, std::complex> - operator*(const L l, const std::complex r) { - return std::complex(l, 0.) * r; -} - -template -TILEDARRAY_FORCE_INLINE - std::enable_if_t, std::complex> - operator*(const std::complex l, const R r) { - return l * std::complex(r, 0.); -} - } // namespace detail namespace conversions { diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 6727cf7136..8c0f63dae7 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -2057,14 +2057,41 @@ class Tensor { // early exit for empty this if (empty()) return *this; - if constexpr (is_arena_tensor_v) { - // Arena inner cells: route through each cell's own in-place scale_to (the - // free arena kernel), which handles a ComplexConjugate factor by - // conjugating each arena scalar in place. Going through `cell *= factor` - // would instead select the generic operator*=(.., ComplexConjugate) -> - // detail::conj(cell), which has no value-returning conj for ArenaTensor. - return inplace_unary( - [factor](value_type& MADNESS_RESTRICT c) { c.scale_to(factor); }); + if constexpr (detail::is_complex_conjugate_v && + !detail::is_complex_v) { + // real elements: conj is the identity, only the scale (if any) remains + if constexpr (std::is_same_v>) + return *this; + else if constexpr (std::is_same_v>) + return neg_to(); + else + return scale_to(factor.factor()); + } else if constexpr (is_arena_tensor_v) { + // Arena inner cells: the free arena kernel scales -- or, for a + // ComplexConjugate factor, conjugates and scales -- every arena scalar + // in place. (The member ArenaTensor::scale_to takes numeric factors + // only; `cell *= factor` is the same kernel via the ArenaTensor + // operator*= overloads.) + return inplace_unary([factor](value_type& MADNESS_RESTRICT c) { + ::TiledArray::scale_to(c, factor); + }); + } else if constexpr (detail::is_complex_conjugate_v && + detail::is_ta_tensor_v) { + // Owning nested cells with a ComplexConjugate factor: conjugate (and + // scale) each cell IN PLACE. `cell *= factor` would resolve to + // `cell = conj(cell) * S`, i.e. a fresh allocation, copy and free of + // every cell. + return inplace_unary([factor](value_type& MADNESS_RESTRICT res) { + if constexpr (std::is_same_v>) + res.conj_to(); + else if constexpr (std::is_same_v>) { + res.conj_to(); + res.neg_to(); + } else + res.conj_to(factor.factor()); + }); } else { return inplace_unary( [factor](value_type& MADNESS_RESTRICT res) { res *= factor; }); @@ -3382,10 +3409,14 @@ class Tensor { Tensor& gemm(const Tensor& A, const Tensor& B, const W alpha, const math::GemmHelper& gemm_helper) { numeric_type beta = 1; - // A ComplexConjugate<...> alpha (conj(A*B) at the expression level) is - // applied to the finished result by ContractReduce's finalization step; - // the gemm itself runs with alpha = 1 (see detail::elem_factor). - const numeric_type alpha_n = detail::elem_factor(alpha); + // A ComplexConjugate<...> factor is not a gemm alpha: conj does not + // distribute into the sum of products, so the expression layer applies it + // to the finished result (ContractReduce's finalization; see + // ContEngine::outer_factor) and hands the gemm a numeric alpha. + static_assert(!detail::is_complex_conjugate_v, + "Tensor::gemm: a ComplexConjugate<...> alpha cannot be " + "applied inside the gemm; conjugate the finished result"); + const numeric_type alpha_n = static_cast(alpha); if (this->empty()) { *this = Tensor(gemm_helper.make_result_range(A.range_, B.range()), @@ -3684,10 +3715,13 @@ class Tensor { if (detail::scale_gemm_timing_enabled()) { detail::g_scale[0].gemm_runs.fetch_add( 1, std::memory_order_relaxed); + // element multiply-adds (A in inner elements), the same + // unit as fb_flop and the same-type path, so the + // [scale-timing] coverage ratios compare like with like detail::g_scale[0].gemm_flop.fetch_add( 2ull * static_cast(K) * static_cast(N) * - static_cast(A) * cw, + static_cast(A), std::memory_order_relaxed); } const integer Ai = static_cast(A); @@ -3861,10 +3895,11 @@ class Tensor { if (detail::scale_gemm_timing_enabled()) { detail::g_scale[1].gemm_runs.fetch_add( 1, std::memory_order_relaxed); + // element multiply-adds, the unit of fb_flop (see above) detail::g_scale[1].gemm_flop.fetch_add( 2ull * static_cast(M) * static_cast(K) * - static_cast(A) * cw, + static_cast(A), std::memory_order_relaxed); } const integer Ai = static_cast(A); diff --git a/src/TiledArray/tile_op/batched_contract_reduce.h b/src/TiledArray/tile_op/batched_contract_reduce.h index 78ec5b8e42..345200606d 100644 --- a/src/TiledArray/tile_op/batched_contract_reduce.h +++ b/src/TiledArray/tile_op/batched_contract_reduce.h @@ -119,11 +119,13 @@ class BatchedContractReduce { /// Create a new, empty result object result_type operator()() const { return result_type(); } - /// Post processing step (no result permutation supported) - result_type operator()(const result_type& temp) const { + /// Post processing step (no result permutation supported): the wrapped + /// op's own finalization -- a no-op for a numeric factor, conjugation (and + /// scale) of the finished tile for a ComplexConjugate<...> one + result_type operator()(result_type& temp) const { using TiledArray::empty; TA_ASSERT(!empty(temp)); - return temp; + return op_(temp); } /// Reduce two result objects (both carry the full fused range) diff --git a/src/TiledArray/tile_op/contract_reduce.h b/src/TiledArray/tile_op/contract_reduce.h index 9e0cd1e77f..94b2c5174f 100644 --- a/src/TiledArray/tile_op/contract_reduce.h +++ b/src/TiledArray/tile_op/contract_reduce.h @@ -124,8 +124,12 @@ class ContractReduceBase { perm_(std::forward(perm)), elem_muladd_op_(std::forward(elem_muladd_op)), arena_plan_(std::forward(arena_plan_in)) { - // non-unit alpha must be absorbed into elem_muladd_op - if (elem_muladd_op_) TA_ASSERT(alpha == scalar_type(1)); + // a numeric factor must be absorbed into elem_muladd_op (alpha is then + // 1); a ComplexConjugate<...> factor cannot be applied per cell and is + // carried here for the finalization (conj and scale of the result) + if constexpr (!TiledArray::detail::is_complex_conjugate_v) { + if (elem_muladd_op_) TA_ASSERT(alpha == scalar_type(1)); + } } math::GemmHelper gemm_helper_; ///< Gemm helper object @@ -466,142 +470,17 @@ class ContractReduce : public ContractReduceBase { }; // class ContractReduce -/// Contract and (sum) reduce operation +/// Contract and (sum) reduce operation with a ComplexConjugate factor -/// This encodes a binary tensor contraction mapped to a GEMM, as well as the -/// sum reduction and post-processing. +/// The contraction of \c conj(A*B) (\c Scalar = \c void), \c S*conj(A*B) +/// (numeric \c Scalar) or \c -conj(A*B) (\c Scalar = \c ComplexNegTag). +/// Conjugation does not distribute into the sum of products, so the +/// contraction and reduction run with unit factor and the finalization +/// conjugates (and scales) the finished result. /// \tparam Result The result tile type /// \tparam Left The left-hand tile type /// \tparam Right The right-hand tile type -template -class ContractReduce> - : public ContractReduceBase> { - public: - typedef ContractReduce> - ContractReduce_; ///< This class type - typedef ContractReduceBase> - ContractReduceBase_; ///< This class type - typedef typename ContractReduceBase_::first_argument_type - first_argument_type; ///< The left tile type - typedef typename ContractReduceBase_::second_argument_type - second_argument_type; ///< The right tile type - typedef Result result_type; ///< The result tile type. - typedef TiledArray::detail::ComplexConjugate scalar_type; - - using typename ContractReduceBase_::elem_muladd_op_type; - using typename ContractReduceBase_::left_value_type; - using typename ContractReduceBase_::result_value_type; - using typename ContractReduceBase_::right_value_type; - - // Compiler generated defaults are fine. N.B. This has shallow copy semantics. - - ContractReduce() = default; - ContractReduce(const ContractReduce_&) = default; - ContractReduce(ContractReduce_&&) = default; - ~ContractReduce() = default; - ContractReduce_& operator=(const ContractReduce_&) = default; - ContractReduce_& operator=(ContractReduce_&&) = default; - - /// Construct contract/reduce functor - - /// \tparam ElemeElemMultAddOpntOp a callable with signature - /// elem_muladd_op_type \param left_op The left-hand BLAS matrix operation - /// \param right_op The right-hand BLAS matrix operation - /// \param alpha The scaling factor applied to the contracted tiles - /// \param result_rank The rank of the result tensor - /// \param left_rank The rank of the left-hand tensor - /// \param right_rank The rank of the right-hand tensor - /// \param perm The permutation to be applied to the result tensor - /// (default = no permute) - /// \param elem_muladd_op The element multiply-add op - template < - typename Perm = BipartitePermutation, - typename ElemMultAddOp = TiledArray::function_ref, - typename Plan = typename ContractReduceBase_::arena_plan_storage_t, - typename = std::enable_if_t< - TiledArray::detail::is_permutation_v> && - std::is_invocable_r_v, - result_value_type&, const left_value_type&, - const right_value_type&> && - std::is_same_v, - typename ContractReduceBase_::arena_plan_storage_t>>> - ContractReduce(const math::blas::Op left_op, const math::blas::Op right_op, - const scalar_type alpha, const unsigned int result_rank, - const unsigned int left_rank, const unsigned int right_rank, - Perm&& perm = {}, ElemMultAddOp&& elem_muladd_op = {}, - Plan&& arena_plan_in = {}) - : ContractReduceBase_(left_op, right_op, alpha, result_rank, left_rank, - right_rank, std::forward(perm), - std::forward(elem_muladd_op), - std::forward(arena_plan_in)) {} - - /// Create a result type object - - /// Initialize a result object for subsequent reductions - result_type operator()() const { return result_type(); } - - /// Post processing step - result_type operator()(result_type& temp) const { - using TiledArray::empty; - TA_ASSERT(!empty(temp)); - - if (!ContractReduceBase_::perm()) { - using TiledArray::conj_to; - return conj_to(temp); - } - - using TiledArray::conj; - return conj(temp, ContractReduceBase_::perm()); - } - - /// Reduce two result objects - - /// Add \c arg to \c result . - /// \param[in,out] result The result object that will be the reduction - /// target - /// \param[in] arg The argument that will be added to \c result - void operator()(result_type& result, const result_type& arg) const { - this->reduce_results(result, arg); - } - - /// Contract a pair of tiles and add to a target tile - - /// Contract \c left and \c right and add the result to \c result. - /// \param[in,out] result The result object that will be the reduction - /// target - /// \param[in] left The left-hand tile to be contracted - /// \param[in] right The right-hand tile to be contracted - void operator()(result_type& result, const first_argument_type& left, - const second_argument_type& right) const { - using TiledArray::empty; - if (empty(left) || empty(right)) return; - if constexpr (!ContractReduceBase_::plain_tensors) { - this->accumulate_nested(result, left, right); - } else { - TA_ASSERT(!this->elem_muladd_op()); - using TiledArray::empty; - using TiledArray::gemm; - if (empty(result)) - result = gemm(left, right, 1, ContractReduceBase_::gemm_helper()); - else - gemm(result, left, right, 1, ContractReduceBase_::gemm_helper()); - } - } - -}; // class ContractReduce - -/// Contract and reduce operation - -/// This object uses a tile contraction operation to form a pair reduction -/// operation. -/// \tparam Result The result tile type -/// \tparam Left The left-hand tile type -/// \tparam Right The right-hand tile type -/// \tparam Scalar The scaling factor type +/// \tparam Scalar The ComplexConjugate parameter (see above) template class ContractReduce> @@ -677,14 +556,22 @@ class ContractReduce) { + if (!ContractReduceBase_::perm()) { + using TiledArray::conj_to; + return conj_to(temp); + } + using TiledArray::conj; + return conj(temp, ContractReduceBase_::perm()); + } else { + if (!ContractReduceBase_::perm()) { + using TiledArray::conj_to; + return conj_to(temp, ContractReduceBase_::factor().factor()); + } + using TiledArray::conj; + return conj(temp, ContractReduceBase_::factor().factor(), + ContractReduceBase_::perm()); } - - using TiledArray::conj; - return conj(temp, ContractReduceBase_::factor().factor(), - ContractReduceBase_::perm()); } /// Reduce two result objects @@ -712,7 +599,6 @@ class ContractReduceaccumulate_nested(result, left, right); } else { TA_ASSERT(!this->elem_muladd_op()); - using TiledArray::empty; using TiledArray::gemm; if (empty(result)) result = gemm(left, right, 1, ContractReduceBase_::gemm_helper()); diff --git a/tests/tot_expressions.cpp b/tests/tot_expressions.cpp index 0c8c613dff..1e926693bb 100644 --- a/tests/tot_expressions.cpp +++ b/tests/tot_expressions.cpp @@ -1,4 +1,8 @@ #include + +#include +#include + #include "tot_array_fixture.h" template @@ -4683,6 +4687,162 @@ auto single_tile(const Array& arr) { return arr.find({0, 0}).get(); } +// rank-3-outer ToT (leading fused mode h) with rank-1 inner cells, one tile +// per mode, so the outer products below run in batched (fused) mode +template +Array make_tot_h1(World& world, std::size_t nh, std::size_t ni, std::size_t nj, + std::size_t na, Gen gen) { + using inner_t = typename Array::value_type::value_type; + TiledRange tr{TiledRange1{0, static_cast(nh)}, + TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { + inner_t t(Range{static_cast(na)}); + for (std::size_t a = 0; a < na; ++a) + t.at_ordinal(a) = gen(idx[0], idx[1], idx[2], a); + return t; + }); + world.gop.fence(); + return arr; +} + +// single-tile rank-2 plain array +template +Array make_plain_2(World& world, std::size_t n0, std::size_t n1, Gen gen) { + TiledRange tr{TiledRange1{0, static_cast(n0)}, + TiledRange1{0, static_cast(n1)}}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { return gen(idx[0], idx[1]); }); + world.gop.fence(); + return arr; +} + +// single-tile rank-3 plain array +template +Array make_plain_3(World& world, std::size_t n0, std::size_t n1, std::size_t n2, + Gen gen) { + TiledRange tr{TiledRange1{0, static_cast(n0)}, + TiledRange1{0, static_cast(n1)}, + TiledRange1{0, static_cast(n2)}}; + Array arr(world, tr); + arr.init_elements( + [=](const auto& idx) { return gen(idx[0], idx[1], idx[2]); }); + world.gop.fence(); + return arr; +} + +// multi-tile (tile extent ts along both outer modes) ToT with rank-1 cells +template +Array make_tot_1_tiled(World& world, std::size_t ni, std::size_t nj, + std::size_t na, std::size_t ts, Gen gen) { + using inner_t = typename Array::value_type::value_type; + auto tr1 = [ts](std::size_t n) { + std::vector b; + for (std::size_t x = 0; x < n; x += ts) b.push_back(x); + b.push_back(n); + return TiledRange1(b.begin(), b.end()); + }; + TiledRange tr{tr1(ni), tr1(nj)}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { + inner_t t(Range{static_cast(na)}); + for (std::size_t a = 0; a < na; ++a) + t.at_ordinal(a) = gen(idx[0], idx[1], a); + return t; + }); + world.gop.fence(); + return arr; +} + +// gather a (possibly multi-tile) array into one dense tile over its element +// range, so a multi-tile result can be checked like a single tile +template +auto gather_tiles(const Array& arr) { + using tile_t = typename Array::value_type; + tile_t out(arr.trange().elements_range()); + for (auto it = arr.begin(); it != arr.end(); ++it) { + const tile_t& t = it->get(); + for (const auto& idx : t.range()) out(idx) = t(idx); + } + return out; +} + +// single-tile rank-2-outer ToT with ARENA (view) inner cells of extent na -- +// the cell type MPQC's ToT arrays use, and the only one that reaches the +// interleaved real-gemm fast path of Tensor::gemm +template +ArenaArr make_arena_tot_1(World& world, std::size_t ni, std::size_t nj, + std::size_t na, Gen gen) { + using ArenaOuter = typename ArenaArr::value_type; + using ArenaInner = typename ArenaOuter::value_type; + TiledRange tr{TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + ArenaArr arr(world, tr); + const long NA = static_cast(na); + arr.init_tiles([=](const TiledArray::Range& r) { + ArenaOuter t = TiledArray::detail::arena_outer_init( + r, 1, [=](std::size_t) { return TiledArray::Range{NA}; }); + for (std::size_t o = 0; o < t.range().volume(); ++o) { + ArenaInner& c = t.data()[o]; + if (!c) continue; + const std::size_t i = o / nj, j = o % nj; + for (std::size_t a = 0; a < na; ++a) c.data()[a] = gen(i, j, a); + } + return t; + }); + world.gop.fence(); + return arr; +} + +// rank-2 (na x nb) arena-backed inner cells +template +ArenaArr make_arena_tot_2(World& world, std::size_t ni, std::size_t nj, + std::size_t na, std::size_t nb, Gen gen) { + using ArenaOuter = typename ArenaArr::value_type; + using ArenaInner = typename ArenaOuter::value_type; + TiledRange tr{TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + ArenaArr arr(world, tr); + const long NA = static_cast(na), NB = static_cast(nb); + arr.init_tiles([=](const TiledArray::Range& r) { + ArenaOuter t = TiledArray::detail::arena_outer_init( + r, 1, [=](std::size_t) { return TiledArray::Range{NA, NB}; }); + for (std::size_t o = 0; o < t.range().volume(); ++o) { + ArenaInner& c = t.data()[o]; + if (!c) continue; + const std::size_t i = o / nj, j = o % nj; + for (std::size_t a = 0; a < na; ++a) + for (std::size_t b = 0; b < nb; ++b) + c.data()[a * nb + b] = gen(i, j, a, b); + } + return t; + }); + world.gop.fence(); + return arr; +} + +using Z = std::complex; +using arena_inner_z = TiledArray::ArenaTensor; +using arena_tot_z = DistArray, DensePolicy>; +using plain_z = DistArray, DensePolicy>; +using plain_d = DistArray, DensePolicy>; +using own_tot_z = DistArray>, DensePolicy>; + +auto gaz1 = [](auto i, auto j, auto a) { + return Z(1.0 + i - j + 0.5 * a, 0.25 * i + j - a); +}; +auto gbz1 = [](auto j, auto k, auto b) { + return Z(2.0 - j + k + 0.1 * b, 0.5 * j - k + 0.2 * b); +}; +auto gaz2 = [](auto i, auto j, auto a, auto b) { + return Z(1.0 + i - j + 0.5 * a - 0.3 * b, 0.25 * i + j - a + 0.1 * b); +}; +auto gtd = [](auto j, auto k) { return 0.5 + j - 0.25 * k; }; +auto gtz = [](auto j, auto k) { + return Z(0.5 + j - 0.25 * k, 0.3 * j + 0.7 * k); +}; + } // namespace BOOST_FIXTURE_TEST_SUITE(tot_conj, ToTArrayFixture) @@ -4866,4 +5026,288 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(tot_times_real_plain, TestParam, } } +// ---- scaled / negated conj(a*b): the ComplexConjugate factor must be +// applied exactly once (conj AND scale), whichever inner product the +// annotation selects + +// c(i,k;a,b) = 2 conj( sum_j a(i,j;a) b(j,k;b) ) with a floating and an +// integer scale, -conj(...), and the outer-permuted scaled form +BOOST_AUTO_TEST_CASE(scaled_conj_of_product) { + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + own_tot_z a = make_tot_1(m_world, ni, nj, na, gaz1); + own_tot_z b = make_tot_1(m_world, nj, nk, nb, gbz1); + auto ref = [=](auto i, auto k, auto x, auto y) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) r += gaz1(i, j, x) * gbz1(j, k, y); + return r; + }; + auto check = [&](const own_tot_z& c, Z scale, const char* what) { + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + BOOST_TEST_CONTEXT(what << " i=" << i << " k=" << k << " x=" << x + << " y=" << y) + check_close(tile(i, k)(x, y), scale * std::conj(ref(i, k, x, y))); + } + }; + own_tot_z c; + c("i,k;a,b") = 2.0 * conj(a("i,j;a") * b("j,k;b")); + check(c, Z(2.0), "2.0*conj(a*b)"); + c("i,k;a,b") = 2 * conj(a("i,j;a") * b("j,k;b")); + check(c, Z(2.0), "2*conj(a*b)"); + c("i,k;a,b") = -conj(a("i,j;a") * b("j,k;b")); + check(c, Z(-1.0), "-conj(a*b)"); + c("k,i;b,a") = 2.0 * conj(a("i,j;a") * b("j,k;b")); + { + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) + check_close(tile(k, i)(y, x), 2.0 * std::conj(ref(i, k, x, y))); + } +} + +// c(i,k;a) = 2 conj( sum_j sum_b a(i,j;a,b) b(j,k;b) ): scaled conj of an +// inner contraction +BOOST_AUTO_TEST_CASE(scaled_conj_of_inner_contraction) { + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + own_tot_z a = make_tot_2(m_world, ni, nj, na, nb, gaz2); + own_tot_z b = make_tot_1(m_world, nj, nk, nb, gbz1); + own_tot_z c; + c("i,k;a") = 2.0 * conj(a("i,j;a,b") * b("j,k;b")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t y = 0; y < nb; ++y) + ref += gaz2(i, j, x, y) * gbz1(j, k, y); + check_close(tile(i, k).at_ordinal(x), 2.0 * std::conj(ref)); + } +} + +// c(i,k;a) = 2 conj( sum_j a(i,j;a) b(j,k;a) ) and -conj(...): scaled conj +// with an inner Hadamard product (outer contraction) +BOOST_AUTO_TEST_CASE(scaled_conj_of_inner_hadamard) { + const std::size_t ni = 2, nj = 3, nk = 2, na = 3; + own_tot_z a = make_tot_1(m_world, ni, nj, na, gaz1); + own_tot_z b = make_tot_1(m_world, nj, nk, na, gbz1); + auto ref = [=](auto i, auto k, auto x) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) r += gaz1(i, j, x) * gbz1(j, k, x); + return r; + }; + own_tot_z c; + c("i,k;a") = 2.0 * conj(a("i,j;a") * b("j,k;a")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + check_close(tile(i, k).at_ordinal(x), 2.0 * std::conj(ref(i, k, x))); + c("i,k;a") = -conj(a("i,j;a") * b("j,k;a")); + tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + check_close(tile(i, k).at_ordinal(x), -std::conj(ref(i, k, x))); +} + +// conj of a ToT x COMPLEX plain product: c(i,k;a) = 2 conj( sum_j a(i,j;a) +// t(j,k) ) and, with an inner permutation, c(i,k;b,a) = conj( sum_j +// a(i,j;a,b) t(j,k) ); a real t would hide a conjugation applied to the +// wrong operand +BOOST_AUTO_TEST_CASE(conj_of_tot_times_complex_plain) { + const std::size_t ni = 2, nj = 3, nk = 4, na = 3, nb = 2; + own_tot_z a = make_tot_1(m_world, ni, nj, na, gaz1); + plain_z t = make_plain_2(m_world, nj, nk, gtz); + own_tot_z c; + c("i,k;a") = 2.0 * conj(a("i,j;a") * t("j,k")); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) ref += gaz1(i, j, x) * gtz(j, k); + check_close(tile(i, k).at_ordinal(x), 2.0 * std::conj(ref)); + } + own_tot_z a2 = make_tot_2(m_world, ni, nj, na, nb, gaz2); + c("i,k;b,a") = conj(a2("i,j;a,b") * t("j,k")); + tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += gaz2(i, j, x, y) * gtz(j, k); + check_close(tile(i, k)(y, x), std::conj(ref)); + } +} + +// multi-tile outer ranges (K-panel reduce over several j tiles): +// c(i,k;a,b) = 2 conj( sum_j a(i,j;a) b(j,k;b) ) +BOOST_AUTO_TEST_CASE(multitile_scaled_conj_of_product) { + const std::size_t ni = 3, nj = 5, nk = 3, na = 2, nb = 2, ts = 2; + own_tot_z a = make_tot_1_tiled(m_world, ni, nj, na, ts, gaz1); + own_tot_z b = make_tot_1_tiled(m_world, nj, nk, nb, ts, gbz1); + own_tot_z c; + c("i,k;a,b") = 2.0 * conj(a("i,j;a") * b("j,k;b")); + auto tile = gather_tiles(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += gaz1(i, j, x) * gbz1(j, k, y); + check_close(tile(i, k)(x, y), 2.0 * std::conj(ref)); + } +} + +// batched (fused leading mode h) contraction, ToT and plain twins: +// c(h,i,k;a,b) = 2 conj( sum_j a(h,i,j;a) b(h,j,k;b) ) and +// c(h,i,k) = conj( sum_j a(h,i,j) b(h,j,k) ): the batched reduce must run +// the same conj/scale finalization as the unbatched one +BOOST_AUTO_TEST_CASE(fused_conj_of_product) { + const std::size_t nh = 2, ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + auto g1 = [](auto h, auto i, auto j, auto a) { + return Z(1.0 + h + i - j + 0.5 * a, 0.25 * i + j - a - h); + }; + auto g2 = [](auto h, auto j, auto k, auto b) { + return Z(2.0 - j + k + 0.1 * b - h, 0.5 * j - k + 0.2 * b + h); + }; + own_tot_z a = make_tot_h1(m_world, nh, ni, nj, na, g1); + own_tot_z b = make_tot_h1(m_world, nh, nj, nk, nb, g2); + own_tot_z c; + c("h,i,k;a,b") = 2.0 * conj(a("h,i,j;a") * b("h,j,k;b")); + auto tile = c.find({0, 0, 0}).get(); + for (std::size_t h = 0; h < nh; ++h) + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += g1(h, i, j, x) * g2(h, j, k, y); + check_close(tile(h, i, k)(x, y), 2.0 * std::conj(ref)); + } + auto p1 = [](auto h, auto i, auto j) { + return Z(1.0 + h + i - 0.5 * j, 0.25 * i + j - h); + }; + auto p2 = [](auto h, auto j, auto k) { + return Z(0.5 + j - 0.25 * k + h, 0.3 * j + k - 0.5 * h); + }; + plain_z pa = make_plain_3(m_world, nh, ni, nj, p1); + plain_z pb = make_plain_3(m_world, nh, nj, nk, p2); + plain_z pc; + pc("h,i,k") = conj(pa("h,i,j") * pb("h,j,k")); + auto ptile = pc.find({0, 0, 0}).get(); + for (std::size_t h = 0; h < nh; ++h) + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) ref += p1(h, i, j) * p2(h, j, k); + check_close(ptile(h, i, k), std::conj(ref)); + } +} + +// ---- ARENA (view) inner cells, MPQC's ToT cell type: the conj forms must +// compile and agree with the owning-cell results; ToT x real plain is the one +// form that reaches the interleaved real-gemm fast path + +// c(i,k;a,b) = conj / 2 conj / -conj( sum_j a(i,j;a) b(j,k;b) ), the +// outer-permuted variant, and the inner contraction +// c(i,k;a) = conj( sum_j sum_b a(i,j;a,b) b(j,k;b) ) +BOOST_AUTO_TEST_CASE(arena_conj_of_product) { + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + arena_tot_z a = make_arena_tot_1(m_world, ni, nj, na, gaz1); + arena_tot_z b = make_arena_tot_1(m_world, nj, nk, nb, gbz1); + auto ref = [=](auto i, auto k, auto x, auto y) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) r += gaz1(i, j, x) * gbz1(j, k, y); + return r; + }; + auto check = [&](const arena_tot_z& c, Z scale, bool permuted, + const char* what) { + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + const Z got = permuted ? tile(k, i).data()[y * na + x] + : tile(i, k).data()[x * nb + y]; + BOOST_TEST_CONTEXT(what << " i=" << i << " k=" << k << " x=" << x + << " y=" << y) + check_close(got, scale * std::conj(ref(i, k, x, y))); + } + }; + arena_tot_z c; + c("i,k;a,b") = conj(a("i,j;a") * b("j,k;b")); + check(c, Z(1.0), false, "arena conj(a*b)"); + c("i,k;a,b") = 2.0 * conj(a("i,j;a") * b("j,k;b")); + check(c, Z(2.0), false, "arena 2.0*conj(a*b)"); + c("i,k;a,b") = -conj(a("i,j;a") * b("j,k;b")); + check(c, Z(-1.0), false, "arena -conj(a*b)"); + c("k,i;b,a") = conj(a("i,j;a") * b("j,k;b")); + check(c, Z(1.0), true, "arena conj(a*b) permuted"); + + arena_tot_z a2 = make_arena_tot_2(m_world, ni, nj, na, nb, gaz2); + arena_tot_z d; + d("i,k;a") = conj(a2("i,j;a,b") * b("j,k;b")); + auto tile = single_tile(d); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t y = 0; y < nb; ++y) + r += gaz2(i, j, x, y) * gbz1(j, k, y); + check_close(tile(i, k).data()[x], std::conj(r)); + } +} + +// interleaved real-gemm fast path: complex arena ToT x real plain and real +// plain x complex arena ToT, rank-1 and rank-2 inner cells +BOOST_AUTO_TEST_CASE(arena_tot_times_real_plain) { + const std::size_t ni = 2, nj = 3, nk = 4, na = 3, nb = 2; + arena_tot_z a = make_arena_tot_1(m_world, ni, nj, na, gaz1); + plain_d t = make_plain_2(m_world, nj, nk, gtd); + arena_tot_z c; + c("i,k;a") = a("i,j;a") * t("j,k"); + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) ref += gaz1(i, j, x) * gtd(j, k); + check_close(tile(i, k).data()[x], ref); + } + plain_d u = make_plain_2(m_world, nk, ni, gtd); + c("k,j;a") = u("k,i") * a("i,j;a"); + tile = single_tile(c); + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t x = 0; x < na; ++x) { + Z ref{}; + for (std::size_t i = 0; i < ni; ++i) ref += gtd(k, i) * gaz1(i, j, x); + check_close(tile(k, j).data()[x], ref); + } + arena_tot_z a2 = make_arena_tot_2(m_world, ni, nj, na, nb, gaz2); + c("i,k;a,b") = a2("i,j;a,b") * t("j,k"); + tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + Z ref{}; + for (std::size_t j = 0; j < nj; ++j) + ref += gaz2(i, j, x, y) * gtd(j, k); + check_close(tile(i, k).data()[x * nb + y], ref); + } +} + BOOST_AUTO_TEST_SUITE_END() From 586adc7ec502780828991310b8273daa5edcd5c7 Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Tue, 8 Sep 2026 17:45:55 -0400 Subject: [PATCH 4/7] tests: tot_conj gather_tiles fetches every tile with find() (rank-safe) The multi-tile check gathered tiles through the local-tile iterators, so under two MPI ranks the other rank's tiles stayed empty and the Debug assertion in Tensor::operator() fired (CI run-np-2; run-np-1 passed). --- tests/tot_expressions.cpp | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/tot_expressions.cpp b/tests/tot_expressions.cpp index 1e926693bb..19495eb880 100644 --- a/tests/tot_expressions.cpp +++ b/tests/tot_expressions.cpp @@ -4756,13 +4756,15 @@ Array make_tot_1_tiled(World& world, std::size_t ni, std::size_t nj, } // gather a (possibly multi-tile) array into one dense tile over its element -// range, so a multi-tile result can be checked like a single tile +// range, so a multi-tile result can be checked like a single tile; every +// tile is fetched with find(), so this is rank-safe (the local-tile +// iterators would leave the other ranks' tiles empty) template auto gather_tiles(const Array& arr) { using tile_t = typename Array::value_type; tile_t out(arr.trange().elements_range()); - for (auto it = arr.begin(); it != arr.end(); ++it) { - const tile_t& t = it->get(); + for (const auto& tidx : arr.trange().tiles_range()) { + const tile_t t = arr.find(tidx).get(); for (const auto& idx : t.range()) out(idx) = t(idx); } return out; From c45aeb3fe5c96ca5927c89f3752181202c47eb08 Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Tue, 8 Sep 2026 18:57:05 -0400 Subject: [PATCH 5/7] Review round 2: interleaved-gemm layout as a predicate term; factor note scoped to ToT-aware ops - Tensor::gemm (tot_x_t and t_x_tot): the std::complex == Vr[2] layout requirement moves from a static_assert into the `interleaved` predicate (detail::is_interleaved_real_view_v, shared by both blocks; alignment needs only >=), so an ABI that breaks it falls back to the per-cell loop instead of failing to compile. The ToT operand's inner scalar must match the result's for either GEMM branch (both slabs are viewed through the same scalar pointer). - ContEngine: the factor comment states what the code does -- plain-tensor contractions hand factor_ to ContractReduce (GEMM alpha, or the ComplexConjugate finalization) and the shape GEMMs take it too; only the ToT-aware ops go through elem_scale()/outer_factor(). --- src/TiledArray/expressions/cont_engine.h | 25 +++++++++----- src/TiledArray/tensor/complex.h | 12 +++++++ src/TiledArray/tensor/tensor.h | 43 +++++++++++------------- 3 files changed, 48 insertions(+), 32 deletions(-) diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h index d4faea643f..31ccc6291d 100644 --- a/src/TiledArray/expressions/cont_engine.h +++ b/src/TiledArray/expressions/cont_engine.h @@ -118,15 +118,22 @@ class ContEngine : public BinaryEngine { using elem_scalar_type = typename TiledArray::detail::numeric_type::type; - /// A contraction's factor is applied in exactly one of two places. A - /// numeric factor is absorbed into the per-cell (per-element) multiply-add - /// ops and the outer tile op runs with factor 1. A ComplexConjugate<...> - /// factor (`conj(A*B)`, `S*conj(A*B)`, `-conj(A*B)`) cannot be applied per - /// cell -- conj does not distribute into a sum of products -- so the - /// per-cell ops run with multiplier 1 and the outer op's finalization - /// conjugates AND scales the finished tile (ContractReduce's ComplexConjugate - /// specializations). Every consumer of the factor goes through these two - /// accessors; none may read factor_ directly. + /// Where the contraction's factor is applied. + /// + /// Plain-tensor contractions (the `!tot_aware_op` branches) hand factor_ to + /// the outer ContractReduce as is: a numeric factor is its GEMM alpha, a + /// ComplexConjugate<...> one is applied by its finalization. The shape + /// GEMMs likewise take factor_ (only its magnitude enters the norms). + /// + /// ToT-aware contractions (nested tiles, and dot_inner) split it in two, + /// through the accessors below, and no such consumer may read factor_ + /// directly. A numeric factor is absorbed into the per-cell (per-element) + /// multiply-add ops and the outer tile op runs with factor 1. A + /// ComplexConjugate<...> factor (`conj(A*B)`, `S*conj(A*B)`, `-conj(A*B)`) + /// cannot be applied per cell -- conj does not distribute into a sum of + /// products -- so the per-cell ops run with multiplier 1 and the outer op's + /// finalization conjugates AND scales the finished tile (ContractReduce's + /// ComplexConjugate specialization). /// \return the multiplier for the per-cell ops template diff --git a/src/TiledArray/tensor/complex.h b/src/TiledArray/tensor/complex.h index 55467ff4e7..6c9ae91450 100644 --- a/src/TiledArray/tensor/complex.h +++ b/src/TiledArray/tensor/complex.h @@ -331,6 +331,18 @@ template inline constexpr bool is_complex_conjugate_v = is_complex_conjugate>::value; +/// Whether inner cells of scalar type \c Inner can be viewed as arrays of +/// \c Real for the interleaved real GEMM on complex cells: \c Inner is +/// std::complex laid out as two contiguous \c Real, the standard's +/// array-oriented access guarantee ([complex.numbers.general]/4). The size +/// and alignment terms are part of the predicate rather than an assertion, so +/// an ABI that breaks the layout falls back to the per-cell loop instead of a +/// reinterpret_cast into undefined behavior. +template +inline constexpr bool is_interleaved_real_view_v = + std::is_same_v, Inner> && + sizeof(Inner) == 2 * sizeof(Real) && alignof(Inner) >= alignof(Real); + /// The numeric multiplier to bake into a per-element (per-cell) multiply-add /// op for a contraction with factor \c factor: the factor itself (converted /// to \c Numeric) for a numeric factor, and \c Numeric(1) for a diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 8c0f63dae7..a5d6849d6d 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -3603,21 +3603,19 @@ class Tensor { is_tensor_view_v) { using Real = std::remove_cv_t; using Vr = std::remove_cv_t; - // Same element type: one gemm in that type. Real plain scalars (Vr) - // against complex inner cells: the complex slabs are viewed as real - // matrices with the inner extent doubled (re,im interleaved), so one - // real gemm with alpha = beta = 1 accumulates both parts exactly. - constexpr bool same_type = std::is_same_v; - constexpr bool interleaved = - !same_type && std::is_same_v, Real>; + // Both slabs are viewed through Vr* below, so the left ToT's inner + // scalar must be the result's. Same element type: one gemm in that + // type. Real plain scalars (Vr) against complex inner cells: the + // complex slabs are viewed as real matrices with the inner extent + // doubled (re,im interleaved), so one real gemm with alpha = beta = 1 + // accumulates both parts exactly (detail::is_interleaved_real_view_v + // carries the layout terms; a mismatch falls back to the per-cell loop). + constexpr bool left_matches = + std::is_same_v, Real>; + constexpr bool same_type = left_matches && std::is_same_v; + constexpr bool interleaved = left_matches && !std::is_same_v && + detail::is_interleaved_real_view_v; constexpr integer cw = interleaved ? 2 : 1; // reals per inner element - // The interleaved view reinterprets the complex slab as Vr[2*n]: the - // standard guarantees array-oriented access to std::complex - // ([complex.numbers.general]/4), pinned here so an exotic ABI cannot - // turn the reinterpret_cast below into silent undefined behavior. - static_assert(!interleaved || (sizeof(Real) == 2 * sizeof(Vr) && - alignof(Real) == alignof(Vr)), - "interleaved real gemm needs std::complex == Vr[2]"); if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { @@ -3804,16 +3802,15 @@ class Tensor { is_tensor_view_v) { using Real = std::remove_cv_t; using Ur = std::remove_cv_t; - // see the tot_x_t block: same type, or real plain x complex inner - // cells via the re,im-interleaved real gemm - constexpr bool same_type = std::is_same_v; - constexpr bool interleaved = - !same_type && std::is_same_v, Real>; + // see the tot_x_t block: the right ToT's inner scalar must be the + // result's; then same type, or real plain x complex inner cells via + // the re,im-interleaved real gemm + constexpr bool right_matches = + std::is_same_v, Real>; + constexpr bool same_type = right_matches && std::is_same_v; + constexpr bool interleaved = right_matches && !std::is_same_v && + detail::is_interleaved_real_view_v; constexpr integer cw = interleaved ? 2 : 1; // reals per inner element - // see the tot_x_t block: std::complex must be exactly Ur[2] - static_assert(!interleaved || (sizeof(Real) == 2 * sizeof(Ur) && - alignof(Real) == alignof(Ur)), - "interleaved real gemm needs std::complex == Ur[2]"); if constexpr (same_type || interleaved) { if (gemm_helper.left_op() == TiledArray::math::blas::NoTranspose && gemm_helper.right_op() == TiledArray::math::blas::NoTranspose) { From 8187a19332b38087418ba40eaf2763d0b0830a63 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Thu, 10 Sep 2026 22:35:13 -0400 Subject: [PATCH 6/7] Review cleanups: correct the ComplexNegTag claim, drop the redundant factor cast Follow-up to the factor-normalization round; no behavior change except the two new assertions. - ContractReduce's ComplexConjugate specialization documented -conj(A*B) as arriving with Scalar = ComplexNegTag. It does not: operator-(ConjMultExpr) (expressions/mult_expr.h) spells the negation conj_op(-1), and ComplexConjugate would not compile in the finalization, which reads factor().factor() -- a member that specialization does not have. Doc corrected and a static_assert added, so the constraint fails with a sentence instead of "no member named 'factor'". - ContEngine: elem_scale() already returns the per-cell multiplier as the inner numeric type, so re-applying detail::elem_factor to it was a no-op cast. Dropped at the six sites that did it. - BatchedContractReduce's finalization now forwards to the wrapped op's, which also applies that op's permutation; assert the batched op is perm-free, the invariant ContEngine maintains by construction. - Two comments still named factor_ where the code reads elem_scale(); the one on the scale-inner fallback op now also records why the factor must not be applied there for a ComplexConjugate (op_'s finalization does it, and doing both conjugates the plain operand twice). - ArenaTensor: note that the ComplexConjugate and operator*= overloads are not redundant with the ComplexConjugate one -- that one alone is ambiguous against detail::operator*= under partial ordering (each is more specialized in one parameter). --- src/TiledArray/expressions/cont_engine.h | 35 +++++++------------ src/TiledArray/tensor/arena_tensor.h | 7 ++++ .../tile_op/batched_contract_reduce.h | 4 +++ src/TiledArray/tile_op/contract_reduce.h | 16 +++++++-- 4 files changed, 37 insertions(+), 25 deletions(-) diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h index 31ccc6291d..a7ff78c1d4 100644 --- a/src/TiledArray/expressions/cont_engine.h +++ b/src/TiledArray/expressions/cont_engine.h @@ -1218,9 +1218,7 @@ class ContEngine : public BinaryEngine { const auto* rp = right.data(); result_tile_element_type acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; - return TiledArray::detail::elem_factor( - factor) * - acc; + return factor * acc; }; this->element_nonreturn_op_ = [flat_dot]( result_tile_element_type& result, @@ -1269,7 +1267,7 @@ class ContEngine : public BinaryEngine { // Mirror the owning-tile path (init_inner_tile_op_owning_): the // SUMMA shapes each result cell from a non-empty left inner cell // (left_range plan), and the per-cell op accumulates `r += l * rr` - // -- or `r += (l * rr) * factor_` when scaled -- via + // -- or `r += (l * rr) * elem_scale()` when scaled -- via // fused_hadamard_inplace into the pre-shaped view cell. No // value-returning per-cell op is needed, so this works for view // cells; non-identity inner result permutation is rejected here @@ -1344,8 +1342,7 @@ class ContEngine : public BinaryEngine { Numeric acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; // result cell is pre-shaped [1] by the unit_range plan. - result.data()[0] += - TiledArray::detail::elem_factor(factor) * acc; + result.data()[0] += factor * acc; }; if (this->outer_product_uses_summa()) { this->arena_plan_ = @@ -1452,10 +1449,7 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(M), static_cast(N), static_cast(K), gh.left_op(), - gh.right_op(), - TiledArray::detail::elem_factor< - typename result_tile_element_type:: - numeric_type>(factor)); + gh.right_op(), factor); }; } // ce+ce (hce+ce): inner CONTRACTION (num_contract_ranks() >= @@ -1571,11 +1565,7 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(Mo), static_cast(No), static_cast(Ko), gh.left_op(), - gh.right_op(), - TiledArray::detail::elem_factor< - typename result_tile_element_type:: - numeric_type>(factor), - left_inner_T); + gh.right_op(), factor, left_inner_T); }; } else if (left_arm_ok) { const auto factor = this->elem_scale(); @@ -1595,11 +1585,7 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(Mo), static_cast(No), static_cast(Ko), gh.left_op(), - gh.right_op(), - TiledArray::detail::elem_factor< - typename result_tile_element_type:: - numeric_type>(factor), - right_inner_T); + gh.right_op(), factor, right_inner_T); }; } // [strided-gemm] install-decision instrumentation. For each @@ -1755,7 +1741,7 @@ class ContEngine : public BinaryEngine { const auto* rp = right.data(); Numeric acc{0}; for (std::size_t j = 0; j < n; ++j) acc += lp[j] * rp[j]; - acc *= TiledArray::detail::elem_factor(factor); + acc *= factor; if (TA::empty(result)) { using R = typename result_tile_element_type::range_type; TiledArray::container::svector ext( @@ -2043,8 +2029,11 @@ class ContEngine : public BinaryEngine { // cells. The Hadamard outer product is an assignment // `result = (perm ^ tot) * scalar`, which needs value-returning // `scale`; only owning inner cells support it. - // N.B. the expression-level scalar prefactor (factor_, != 1 for - // ScalMult expressions) multiplies the plain operand's element + // N.B. the expression-level scalar prefactor (elem_scale(), != 1 + // for a scaled ScalMult expression) multiplies the plain operand's + // element. A ComplexConjugate factor_ leaves it 1 and is applied by + // op_'s finalization instead -- applying it here as well would + // conjugate the plain operand a second time. auto fallback_op = [perm = !this->implicit_permute_inner_ ? inner(this->perm_) : Permutation{}, diff --git a/src/TiledArray/tensor/arena_tensor.h b/src/TiledArray/tensor/arena_tensor.h index d4895e802d..8bd2d372a1 100644 --- a/src/TiledArray/tensor/arena_tensor.h +++ b/src/TiledArray/tensor/arena_tensor.h @@ -458,6 +458,13 @@ void scale_to(ArenaTensor& dst, Scalar factor) { /// that a view is never routed to the generic /// detail::operator*=(L&, ComplexConjugate), which needs a value-returning /// conj(L) that a non-owning view cannot provide. +/// +/// N.B. all three are needed; the ComplexConjugate one does NOT subsume +/// the other two. Against detail's own +/// operator*=(L&, const ComplexConjugate&) it is ambiguous: this +/// overload is more specialized in the first parameter, detail's in the +/// second, so neither wins partial ordering. The and +/// overloads are more specialized in BOTH parameters and therefore win. template ArenaTensor& operator*=(ArenaTensor& dst, const detail::ComplexConjugate& factor) { diff --git a/src/TiledArray/tile_op/batched_contract_reduce.h b/src/TiledArray/tile_op/batched_contract_reduce.h index 345200606d..782f1c758c 100644 --- a/src/TiledArray/tile_op/batched_contract_reduce.h +++ b/src/TiledArray/tile_op/batched_contract_reduce.h @@ -125,6 +125,10 @@ class BatchedContractReduce { result_type operator()(result_type& temp) const { using TiledArray::empty; TA_ASSERT(!empty(temp)); + // the wrapped op's finalization also applies ITS permutation; the batched + // op is built perm-free (ContEngine), and a perm leaking in here would + // permute the fused result silently + TA_ASSERT(!op_.perm()); return op_(temp); } diff --git a/src/TiledArray/tile_op/contract_reduce.h b/src/TiledArray/tile_op/contract_reduce.h index 94b2c5174f..2cf556c298 100644 --- a/src/TiledArray/tile_op/contract_reduce.h +++ b/src/TiledArray/tile_op/contract_reduce.h @@ -472,8 +472,12 @@ class ContractReduce : public ContractReduceBase { /// Contract and (sum) reduce operation with a ComplexConjugate factor -/// The contraction of \c conj(A*B) (\c Scalar = \c void), \c S*conj(A*B) -/// (numeric \c Scalar) or \c -conj(A*B) (\c Scalar = \c ComplexNegTag). +/// The contraction of \c conj(A*B) (\c Scalar = \c void) or of a scaled +/// conjugate -- \c S*conj(A*B) and \c -conj(A*B), both numeric \c Scalar +/// (the expression layer spells the negation \c conj_op(-1), +/// see \c operator-(ConjMultExpr) in expressions/mult_expr.h; it never +/// produces \c ComplexConjugate here, which the finalization +/// below could not read a factor from). /// Conjugation does not distribute into the sum of products, so the /// contraction and reduction run with unit factor and the finalization /// conjugates (and scales) the finished result. @@ -500,6 +504,14 @@ class ContractReduce scalar_type; + // the finalization below reads factor().factor() for a non-void Scalar, + // which ComplexConjugate does not have; see the class doc + static_assert( + !std::is_same_v, + "ContractReduce: a negated conjugate contraction factor is spelled " + "ComplexConjugate(-1), not " + "ComplexConjugate"); + using typename ContractReduceBase_::elem_muladd_op_type; using typename ContractReduceBase_::left_value_type; using typename ContractReduceBase_::result_value_type; From ef1ecb06a5f6c31e41bee21ed5dec1ef6c981a40 Mon Sep 17 00:00:00 2001 From: Eduard Valeyev Date: Fri, 11 Sep 2026 08:46:42 -0400 Subject: [PATCH 7/7] tests: conj() on BTAS nested cells, the third supported inner-tile family conj(A*B) on Tensor, Range>> fails to compile on master -- the per-cell op static_casts the ComplexConjugate factor, tripping the "gemm without custom element op is only applicable to plain tensors" assertion in kernels.h. It works on this branch: the factor reaches the cells as 1 and the finalization conjugates the finished tile, through btas's own value-returning conj / conj_to CPOs (external/btas.h). Nothing in the branch was aimed at BTAS, so pin the behavior with a test. Covers conj(a*b), 2*conj(a*b), -conj(a*b), unary conj(a) and conj(a)*b against explicit references. The fixture's btas rows are real (int/float/double), so this is the first complex BTAS ToT case in the suite; make_tot_1_data fills through data() because btas::Tensor has no at_ordinal(). Raised by a review comment that read the opposite way -- that the branch left BTAS ill-formed. --- tests/tot_expressions.cpp | 84 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 83 insertions(+), 1 deletion(-) diff --git a/tests/tot_expressions.cpp b/tests/tot_expressions.cpp index 19495eb880..73b852241a 100644 --- a/tests/tot_expressions.cpp +++ b/tests/tot_expressions.cpp @@ -4809,7 +4809,9 @@ ArenaArr make_arena_tot_2(World& world, std::size_t ni, std::size_t nj, const long NA = static_cast(na), NB = static_cast(nb); arr.init_tiles([=](const TiledArray::Range& r) { ArenaOuter t = TiledArray::detail::arena_outer_init( - r, 1, [=](std::size_t) { return TiledArray::Range{NA, NB}; }); + r, 1, [=](std::size_t) { + return TiledArray::Range{NA, NB}; + }); for (std::size_t o = 0; o < t.range().volume(); ++o) { ArenaInner& c = t.data()[o]; if (!c) continue; @@ -5312,4 +5314,84 @@ BOOST_AUTO_TEST_CASE(arena_tot_times_real_plain) { } } +// ---- BTAS inner cells, the third supported nested-tile family. conj(A*B) +// on Tensor> does not compile on master (the per-cell +// op static_casts the ComplexConjugate factor); it works here because the +// factor reaches the cells as 1 and the finalization conjugates the finished +// tile through btas's own conj/conj_to CPOs (external/btas.h). + +using btas_inner_z = btas::Tensor; +using btas_tot_z = DistArray, DensePolicy>; + +// btas::Tensor has no at_ordinal(); fill through data() so the same generator +// serves the owning, arena and btas rows. +template +Array make_tot_1_data(World& world, std::size_t ni, std::size_t nj, + std::size_t na, Gen gen) { + using inner_t = typename Array::value_type::value_type; + TiledRange tr{TiledRange1{0, static_cast(ni)}, + TiledRange1{0, static_cast(nj)}}; + Array arr(world, tr); + arr.init_elements([=](const auto& idx) { + inner_t t(Range{static_cast(na)}); + for (std::size_t a = 0; a < na; ++a) t.data()[a] = gen(idx[0], idx[1], a); + return t; + }); + world.gop.fence(); + return arr; +} + +// c(i,k;a,b) = conj / 2 conj / -conj( sum_j a(i,j;a) b(j,k;b) ) on btas cells +BOOST_AUTO_TEST_CASE(btas_conj_of_product) { + const std::size_t ni = 2, nj = 3, nk = 2, na = 2, nb = 3; + btas_tot_z a = make_tot_1_data(m_world, ni, nj, na, gaz1); + btas_tot_z b = make_tot_1_data(m_world, nj, nk, nb, gbz1); + auto ref = [=](auto i, auto k, auto x, auto y) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) r += gaz1(i, j, x) * gbz1(j, k, y); + return r; + }; + auto check = [&](const btas_tot_z& c, Z scale, const char* what) { + auto tile = single_tile(c); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + BOOST_TEST_CONTEXT(what << " i=" << i << " k=" << k << " x=" << x + << " y=" << y) + check_close(tile(i, k).data()[x * nb + y], + scale * std::conj(ref(i, k, x, y))); + } + }; + btas_tot_z c; + c("i,k;a,b") = conj(a("i,j;a") * b("j,k;b")); + check(c, Z(1.0), "btas conj(a*b)"); + c("i,k;a,b") = 2.0 * conj(a("i,j;a") * b("j,k;b")); + check(c, Z(2.0), "btas 2.0*conj(a*b)"); + c("i,k;a,b") = -conj(a("i,j;a") * b("j,k;b")); + check(c, Z(-1.0), "btas -conj(a*b)"); + + // unary conj, and conj on one operand of the product + btas_tot_z u; + u("i,j;a") = conj(a("i,j;a")); + auto utile = single_tile(u); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t j = 0; j < nj; ++j) + for (std::size_t x = 0; x < na; ++x) + check_close(utile(i, j).data()[x], std::conj(gaz1(i, j, x))); + + btas_tot_z l; + l("i,k;a,b") = conj(a("i,j;a")) * b("j,k;b"); + auto ltile = single_tile(l); + for (std::size_t i = 0; i < ni; ++i) + for (std::size_t k = 0; k < nk; ++k) + for (std::size_t x = 0; x < na; ++x) + for (std::size_t y = 0; y < nb; ++y) { + Z r{}; + for (std::size_t j = 0; j < nj; ++j) + r += std::conj(gaz1(i, j, x)) * gbz1(j, k, y); + check_close(ltile(i, k).data()[x * nb + y], r); + } +} + BOOST_AUTO_TEST_SUITE_END()