From 8a879e6e25dd904f05055118ffa5090b8f4ad506 Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Sun, 30 Aug 2026 12:03:06 -0400 Subject: [PATCH] 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 ea934073d0..7315658613 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 2ff047d6c2..59baa946f0 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -3141,6 +3141,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()), @@ -3165,7 +3169,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); } @@ -3219,7 +3223,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 @@ -3326,7 +3330,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 @@ -3426,7 +3439,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); @@ -3434,10 +3447,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 @@ -3506,7 +3521,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`. @@ -3590,7 +3612,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); @@ -3598,10 +3620,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 def1ae28d1..ae6569dbf6 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 66939bcf3e..7e4e04a76f 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()