From f602d22d6d69a95a7037343570de919310f57f6b Mon Sep 17 00:00:00 2001 From: Kshitij Surjuse Date: Sun, 30 Aug 2026 11:17:12 -0400 Subject: [PATCH] ToT: strided arena kernels admit float/double/complex inner storage; complex ToT scale/add/subt/mult The arena strided-DGEMM ToT x ToT kernels (arena_strided_dgemm_ce_e, _ce_ce_right, _ce_ce_left) and their ContEngine install gates were numeric_type == double only, so every ToT contraction with std::complex inner storage silently took the generic per-cell path. - kernels templated on the inner numeric type T, admitted through detail::is_strided_dgemm_numeric_v (float, double, complex, complex); factor/beta typed T, cell pointers T*; all three operands must share T (mixed real x complex ToT products keep the per-cell path) - diagnostic helpers classify_run / classify_operand / gather_rescuable / measure_segments: cell pointers made generic - cont_engine.h ce+e / ce+ce install gates and the hc+e reuse gate: view cells + same numeric type + the trait; static_cast(factor) - Tensor scale/add/subt/mult fill lambdas and ArenaTensor scale_to pick up detail's mixed complex x scalar operator* (as the non-nested scale branch already does), so e.g. 2 * A("i;j") compiles for complex ToTs - tests: complex cases for the three kernels; a complex row in the ToT fixture type list (runs tot_expressions / tot_dist_array for complex ToTs); typedefs case compares scalar_type against scalar_t Measured on MPQC Kramers PNS-CCD (complex ToTs, HSeOH/cc-pVDZ): 0 -> 531 strided ce+e installs, first CC iteration 52.0 -> 20.5 s, energy unchanged. --- src/TiledArray/expressions/cont_engine.h | 41 ++++--- src/TiledArray/tensor/arena_einsum.h | 132 ++++++++++++++--------- src/TiledArray/tensor/arena_tensor.h | 9 +- src/TiledArray/tensor/tensor.h | 9 ++ tests/arena_strided_dgemm.cpp | 119 ++++++++++++++++++++ tests/tot_array_fixture.h | 20 ++-- tests/tot_dist_array_part1.cpp | 5 +- 7 files changed, 256 insertions(+), 79 deletions(-) diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h index a18607af7e..ea934073d0 100644 --- a/src/TiledArray/expressions/cont_engine.h +++ b/src/TiledArray/expressions/cont_engine.h @@ -1366,26 +1366,27 @@ class ContEngine : public BinaryEngine { // non-identity inner result perm is applied downstream and left // to the per-cell path here. // The strided kernel is specialized to view (arena) inner cells - // with double storage, and its static_assert requires that of - // ALL THREE operands (result, left, right). Gate on the same - // 3-operand predicate so a mixed-operand contraction (e.g. a - // view/double result with a non-view or non-double operand, or - // float/complex inner) stays on the generic per-cell path and - // never instantiates the double-view-only kernel (which would - // be a hard compile error rather than a graceful fallback). + // whose numeric type is one of the four BLAS gemm element types + // (is_strided_dgemm_numeric_v), shared by ALL THREE operands + // (result, left, right). Gate on the same 3-operand predicate + // so a mixed-operand contraction (e.g. a view result with a + // non-view operand, a non-BLAS inner numeric type, or a + // real-ToT x complex-ToT product) stays on the generic + // per-cell path and never instantiates the view-only kernel + // (which would be a hard compile error rather than a graceful + // fallback). if constexpr ( TiledArray::is_tensor_view_v && TiledArray::is_tensor_view_v && TiledArray::is_tensor_view_v && - std::is_same_v< - typename result_tile_element_type::numeric_type, - double> && + TiledArray::detail::is_strided_dgemm_numeric_v< + typename result_tile_element_type::numeric_type> && std::is_same_v< typename left_tile_element_type::numeric_type, - double> && + typename result_tile_element_type::numeric_type> && std::is_same_v< typename right_tile_element_type::numeric_type, - double>) { + 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_; @@ -1401,7 +1402,9 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(M), static_cast(N), static_cast(K), gh.left_op(), - gh.right_op(), double(factor)); + gh.right_op(), + static_cast(factor)); }; } // ce+ce (hce+ce): inner CONTRACTION (num_contract_ranks() >= @@ -1517,7 +1520,10 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(Mo), static_cast(No), static_cast(Ko), gh.left_op(), - gh.right_op(), double(factor), left_inner_T); + gh.right_op(), + static_cast(factor), + left_inner_T); }; } else if (left_arm_ok) { const scalar_type factor = this->factor_; @@ -1537,11 +1543,14 @@ class ContEngine : public BinaryEngine { Cc, Lt, Rt, static_cast(Mo), static_cast(No), static_cast(Ko), gh.left_op(), - gh.right_op(), double(factor), right_inner_T); + gh.right_op(), + static_cast(factor), + right_inner_T); }; } // [strided-dgemm] install-decision instrumentation. For each - // ToT contraction reaching this double-view path, report + // ToT contraction reaching this view-cell path, report // whether a strided-DGEMM regime (hce+e / hc+e / hce+ce) // FIRED or the contraction REVERTED to the generic by-cell // evaluation path (with the blocking guard). Gated by diff --git a/src/TiledArray/tensor/arena_einsum.h b/src/TiledArray/tensor/arena_einsum.h index f49b259c4b..e5147ac43b 100644 --- a/src/TiledArray/tensor/arena_einsum.h +++ b/src/TiledArray/tensor/arena_einsum.h @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -37,6 +38,18 @@ namespace TiledArray::detail { +/// Numeric types the arena strided-DGEMM kernels (ce+e, ce+ce) are +/// instantiated for: the four BLAS gemm element types (float, double, +/// std::complex, std::complex). The kernels are layout/stride +/// machinery over math::blas::gemm, so all four share one body; the ContEngine +/// install gates and the hc+e reuse gate consult this trait so any other inner +/// numeric type keeps the per-cell path. +template +inline constexpr bool is_strided_dgemm_numeric_v = + std::is_same_v || std::is_same_v || + std::is_same_v> || + std::is_same_v>; + /// Env-gated (TA_STRIDED_DGEMM_VERBOSE) toggle for the strided-DGEMM install /// logger. Reads the environment once. Set TA_STRIDED_DGEMM_VERBOSE=1 to have /// the ContEngine print, per ToT contraction, whether a strided-DGEMM regime @@ -282,7 +295,7 @@ template inline int classify_run(GetCell getcell, std::size_t n) { if (n == 0) return 0; long s0 = -1; - const double* base = nullptr; + decltype(getcell(std::size_t{0}).data()) base = nullptr; for (std::size_t i = 0; i < n; ++i) { const auto& c = getcell(i); if (!c) return 1; // absent @@ -320,7 +333,7 @@ inline int classify_operand(GetR getR, GetL getL, std::size_t nrun, const auto& lk = getL(k); if (!lk || static_cast(lk.size()) != P * Q) return 13; // single-cell long s0 = -1; - const double* base = nullptr; + decltype(getR(std::size_t{0}, std::size_t{0}).data()) base = nullptr; for (std::size_t i = 0; i < nrun; ++i) { const auto& c = getR(k, i); if (!c) return 13; @@ -397,7 +410,7 @@ inline int gather_rescuable(GetC getC, GetR getR, GetL getL, std::size_t nrun, if (nrun == 0 || nrun > 1024) return 0; std::size_t pres[1024]; std::size_t np = 0; - const double* rbase = nullptr; + decltype(getC(std::size_t{0}).data()) rbase = nullptr; for (std::size_t i = 0; i < nrun; ++i) { const auto& c = getC(i); if (!c) continue; @@ -418,7 +431,7 @@ inline int gather_rescuable(GetC getC, GetR getR, GetL getL, std::size_t nrun, bool any_k = false; for (std::size_t k = 0; k < nK; ++k) { if (!getL(k)) continue; // absent single-cell -> skip k (β=1) - const double* ob = nullptr; + decltype(getR(std::size_t{0}, std::size_t{0}).data()) ob = nullptr; long os = -1; for (std::size_t j = 0; j < np; ++j) { const auto& oc = getR(k, pres[j]); @@ -472,8 +485,8 @@ inline void measure_segments(GetC getC, GetR getR, GetL getL, std::size_t nrun, ++mu; continue; } - const double* cb = c0.data(); - const double* rb = r0.data(); + const auto* cb = c0.data(); + const auto* rb = r0.data(); std::size_t end = mu + 1; long sC = -1, sR = -1; while (end < nrun) { @@ -1166,11 +1179,12 @@ inline std::atomic g_strided_dgemm_ce_e_calls{0}; /// present, uniform inner size, single constant stride); else an inline per-k /// rank-1 fallback for THAT cell only. Orientation-aware (left_op/right_op pick /// per-(m,n,k) offsets). M=left-external, N=right-external, K=outer-contracted. -template +template void arena_strided_dgemm_ce_e(ResultOuter& C, const LeftOuter& L, const RightOuter& R, std::size_t M, std::size_t N, std::size_t K, math::blas::Op left_op, - math::blas::Op right_op, double factor) { + math::blas::Op right_op, T factor) { namespace blas = TiledArray::math::blas; using integer = blas::integer; static_assert(is_tensor_view_v && @@ -1178,10 +1192,13 @@ void arena_strided_dgemm_ce_e(ResultOuter& C, const LeftOuter& L, is_tensor_view_v, "arena_strided_dgemm_ce_e: arena (view) inner cells only"); static_assert( - std::is_same_v && - std::is_same_v && - std::is_same_v, - "arena_strided_dgemm_ce_e: double inner storage only"); + std::is_same_v && + std::is_same_v && + std::is_same_v, + "arena_strided_dgemm_ce_e: inner storages must share numeric type T"); + static_assert(is_strided_dgemm_numeric_v, + "arena_strided_dgemm_ce_e: float/double/complex/" + "complex inner storage only"); if (M == 0 || N == 0 || K == 0) return; const std::size_t nbatch = static_cast(C.nbatch()); if (nbatch == 0) return; @@ -1253,7 +1270,7 @@ void arena_strided_dgemm_ce_e(ResultOuter& C, const LeftOuter& L, /*K=*/static_cast(K), factor, /*A=*/l0.data(), /*lda=*/static_cast(ldA), /*B=*/r0.data(), /*ldb=*/static_cast(ldB), - /*beta=*/1.0, + /*beta=*/T(1), /*C=*/Cc.data(), /*ldc=*/static_cast(Q)); } #ifdef TA_STRIDED_DGEMM_COUNT @@ -1278,7 +1295,7 @@ void arena_strided_dgemm_ce_e(ResultOuter& C, const LeftOuter& L, record_ce_e_fallback(why); } // inline per-k rank-1 fallback for THIS cell (computed once) - double* c = Cc.data(); + T* c = Cc.data(); std::uint64_t _fl = 0; for (std::size_t k = 0; k < K; ++k) { const auto& lk = lc[lbase + a_off(m, k)]; @@ -1287,8 +1304,8 @@ void arena_strided_dgemm_ce_e(ResultOuter& C, const LeftOuter& L, const std::size_t pp = lk.size(), qq = rk.size(); if (static_cast(Cc.size()) != static_cast(pp * qq)) continue; - const double* lp = lk.data(); - const double* rp = rk.data(); + const T* lp = lk.data(); + const T* rp = rk.data(); _fl += 2ull * pp * qq; for (std::size_t p = 0; p < pp; ++p) for (std::size_t q = 0; q < qq; ++q) @@ -1335,12 +1352,13 @@ inline bool& ce_ce_strided_disabled() { /// (each cell once -> no double-count). C must be pre-shaped (a_1-major); the /// result outer is (m, μ̃) row-major (left-then-right concatenation, matching /// make_result_range). Accumulates into C (beta=1). -template +template void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, const RightOuter& R, std::size_t Mo, std::size_t No, std::size_t Ko, math::blas::Op left_op, math::blas::Op right_op, - double factor, + T factor, bool left_inner_transposed = false) { // left_inner_transposed: the external-carrying LEFT inner cell is stored // (a4,a1)=Q x P (matrix_transpose) instead of canonical (a1,a4)=P x Q. Folded @@ -1353,10 +1371,13 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, is_tensor_view_v, "arena_strided_dgemm_ce_ce_right: arena (view) inner cells only"); static_assert( - std::is_same_v && - std::is_same_v && - std::is_same_v, - "arena_strided_dgemm_ce_ce_right: double inner storage only"); + std::is_same_v && + std::is_same_v && + std::is_same_v, + "arena_strided_dgemm_ce_ce_right: inner storages must share numeric type T"); + static_assert(is_strided_dgemm_numeric_v, + "arena_strided_dgemm_ce_ce_right: float/double/complex/" + "complex inner storage only"); const std::size_t Mmu = No; // right outer-external rides BLAS M const std::size_t nK = Ko; // outer-contracted is looped with beta=1 const std::size_t nbatch = static_cast(C.nbatch()); @@ -1444,7 +1465,7 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, const typename LeftOuter::value_type& lk) { ScopedPhaseTimer _fb_timer(g_fallback_ns_ce_ce); std::uint64_t _fl = 0; - const double* l = lk.data(); + const T* l = lk.data(); for (std::size_t mu = 0; mu < Mmu; ++mu) { auto& Cc = cc[cbase + c_off(m, mu)]; const auto& rk = rc[rbase + r_off(k, mu)]; @@ -1453,14 +1474,14 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, const long Ql = static_cast(rk.size()); if (Ql == 0 || static_cast(lk.size()) != Pl * Ql) continue; _fl += 2ull * static_cast(Pl) * Ql; - double* c = Cc.data(); - const double* rr = rk.data(); + T* c = Cc.data(); + const T* rr = rk.data(); for (long a1 = 0; a1 < Pl; ++a1) { - double acc = 0; + T acc = 0; if (left_inner_transposed) { for (long a4 = 0; a4 < Ql; ++a4) acc += l[a4 * Pl + a1] * rr[a4]; } else { - const double* lr = l + a1 * Ql; + const T* lr = l + a1 * Ql; for (long a4 = 0; a4 < Ql; ++a4) acc += lr[a4] * rr[a4]; } c[a1] += factor * acc; @@ -1492,7 +1513,7 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, continue; } - const double* Lk = lk.data(); // P x Q (or Q x P if transposed) + const T* Lk = lk.data(); // P x Q (or Q x P if transposed) std::size_t mu = 0; while (mu < Mmu) { const auto& rc0 = rc[rbase + r_off(k, mu)]; @@ -1503,8 +1524,8 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, ++mu; continue; } - const double* rstart = rc0.data(); // segment μ̃-run base on R, stride sR - double* cstart = cc0.data(); // segment μ̃-run base on C, stride sC + const T* rstart = rc0.data(); // segment μ̃-run base on R, stride sR + T* cstart = cc0.data(); // segment μ̃-run base on C, stride sC // Grow the maximal segment, recomputing the strides locally (never // reuse a run-wide stale stride). std::size_t end = mu + 1; @@ -1548,7 +1569,7 @@ void arena_strided_dgemm_ce_ce_right(ResultOuter& C, const LeftOuter& L, /*A=*/rstart, /*lda=*/static_cast(ldR), /*B=*/Lk, /*ldb=*/static_cast(left_inner_transposed ? P : Q), - /*beta=*/1.0, + /*beta=*/T(1), /*C=*/cstart, /*ldc=*/static_cast(ldC)); } #ifdef TA_STRIDED_DGEMM_COUNT @@ -1581,12 +1602,13 @@ inline std::atomic g_strided_dgemm_ce_ce_left_calls{0}; /// only (each cell once -> no double-count). Orientation-aware (l_off/r_off from /// left_op/right_op of the OUTER GemmHelper, exactly as the right core). C must /// be pre-shaped; the result outer is (m, n) row-major. Accumulates (beta=1). -template +template void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, const RightOuter& R, std::size_t Mo, std::size_t No, std::size_t Ko, math::blas::Op left_op, - math::blas::Op right_op, double factor, + math::blas::Op right_op, T factor, bool right_inner_transposed = false) { // right_inner_transposed: the external-carrying RIGHT inner cell is stored // (b1,a4)=P x Q (matrix_transpose) instead of canonical (a4,b1)=Q x P. Folded @@ -1599,10 +1621,13 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, is_tensor_view_v, "arena_strided_dgemm_ce_ce_left: arena (view) inner cells only"); static_assert( - std::is_same_v && - std::is_same_v && - std::is_same_v, - "arena_strided_dgemm_ce_ce_left: double inner storage only"); + std::is_same_v && + std::is_same_v && + std::is_same_v, + "arena_strided_dgemm_ce_ce_left: inner storages must share numeric type T"); + static_assert(is_strided_dgemm_numeric_v, + "arena_strided_dgemm_ce_ce_left: float/double/complex/" + "complex inner storage only"); const std::size_t nK = Ko; // outer-contracted, looped with beta=1 const std::size_t nbatch = static_cast(C.nbatch()); if (nbatch == 0 || Mo == 0 || nK == 0 || No == 0) return; @@ -1681,7 +1706,7 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, const typename RightOuter::value_type& rk) { ScopedPhaseTimer _fb_timer(g_fallback_ns_ce_ce); std::uint64_t _fl = 0; - const double* bd = rk.data(); // canonical Q x P row-major + const T* bd = rk.data(); // canonical Q x P row-major for (std::size_t m = 0; m < Mo; ++m) { auto& Cc = cc[cbase + c_off(m, n)]; const auto& lk = lc[lbase + l_off(m, k)]; @@ -1690,14 +1715,14 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, const long Ql = static_cast(lk.size()); if (Ql == 0 || static_cast(rk.size()) != Ql * Pl) continue; _fl += 2ull * static_cast(Pl) * Ql; - double* c = Cc.data(); - const double* a = lk.data(); // Ql vector + T* c = Cc.data(); + const T* a = lk.data(); // Ql vector for (long a4 = 0; a4 < Ql; ++a4) { - const double av = a[a4]; + const T av = a[a4]; if (right_inner_transposed) { for (long p = 0; p < Pl; ++p) c[p] += factor * av * bd[p * Ql + a4]; } else { - const double* br = bd + a4 * Pl; + const T* br = bd + a4 * Pl; for (long p = 0; p < Pl; ++p) c[p] += factor * av * br[p]; } } @@ -1728,7 +1753,7 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, continue; } - const double* Rk = rk.data(); // Q x P (or P x Q if transposed) + const T* Rk = rk.data(); // Q x P (or P x Q if transposed) std::size_t m = 0; while (m < Mo) { const auto& lc0 = lc[lbase + l_off(m, k)]; @@ -1739,8 +1764,8 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, ++m; continue; } - const double* lstart = lc0.data(); // segment m-run base on L, stride sA - double* cstart = cc0.data(); // segment m-run base on C, stride sC + const T* lstart = lc0.data(); // segment m-run base on L, stride sA + T* cstart = cc0.data(); // segment m-run base on C, stride sC // Grow the maximal segment, recomputing the strides locally (never // reuse a run-wide stale stride). std::size_t end = m + 1; @@ -1784,7 +1809,7 @@ void arena_strided_dgemm_ce_ce_left(ResultOuter& C, const LeftOuter& L, /*A=*/lstart, /*lda=*/static_cast(ldA), /*B=*/Rk, /*ldb=*/static_cast(right_inner_transposed ? Q : P), - /*beta=*/1.0, + /*beta=*/T(1), /*C=*/cstart, /*ldc=*/static_cast(ldC)); } #ifdef TA_STRIDED_DGEMM_COUNT @@ -2196,16 +2221,19 @@ bool run_regime_a_arena(const Plan& plan, const HIndex& h, std::size_t batch, if constexpr (a_is_tot && b_is_tot) { using IIndex = ::Einsum::index::Index; // hc+e reuse gate: the result/operand inner cells must be the kernel's - // (view + double) inner type; mirror arena_strided_dgemm_ce_e's - // static_assert so non-view / non-double ToT keep the per-cell path. + // (view + BLAS numeric) inner type; mirror arena_strided_dgemm_ce_e's + // static_assert so non-view / non-BLAS-numeric / mixed-type ToT keep the + // per-cell path. using LInnerT = typename ArrayA_t::value_type::value_type; using RInnerT = typename ArrayB_t::value_type::value_type; constexpr bool ce_e_kernel_ok = is_tensor_view_v && is_tensor_view_v && is_tensor_view_v && - std::is_same_v && - std::is_same_v && - std::is_same_v; + is_strided_dgemm_numeric_v && + std::is_same_v && + std::is_same_v; // Inner OUTER-PRODUCT (K_inner==0) is the strided-reusable shape; any // inner contraction (hc+ce) stays per-cell (two-level stride). The // runtime toggle lets tests/benches force the per-cell path. @@ -2285,7 +2313,7 @@ bool run_regime_a_arena(const Plan& plan, const HIndex& h, std::size_t batch, arena_strided_dgemm_ce_e(cview, ai, bi, /*M=*/std::size_t{1}, /*N=*/std::size_t{1}, /*K=*/Kvol, blas::NoTranspose, blas::NoTranspose, - /*factor=*/1.0); + /*factor=*/typename InnerT::numeric_type(1)); continue; // tile-i contribution complete } } diff --git a/src/TiledArray/tensor/arena_tensor.h b/src/TiledArray/tensor/arena_tensor.h index 783e8d5956..4c34a9502c 100644 --- a/src/TiledArray/tensor/arena_tensor.h +++ b/src/TiledArray/tensor/arena_tensor.h @@ -445,7 +445,14 @@ void scale_to(ArenaTensor& dst, Scalar factor) { if (!dst) return; auto* d = dst.data(); const auto n = dst.size(); - for (std::size_t i = 0; i < n; ++i) d[i] *= factor; + if constexpr (requires(T& x) { x *= factor; }) { + for (std::size_t i = 0; i < n; ++i) d[i] *= factor; + } else { + // no compound operator for this element/factor pair (e.g. complex + // *= int): use detail's mixed complex x scalar operator* + using namespace TiledArray::detail; + for (std::size_t i = 0; i < n; ++i) d[i] = d[i] * factor; + } } /// `dst += src`. Asserts both views non-null and shape-compatible. diff --git a/src/TiledArray/tensor/tensor.h b/src/TiledArray/tensor/tensor.h index 998b764ba2..2ff047d6c2 100644 --- a/src/TiledArray/tensor/tensor.h +++ b/src/TiledArray/tensor/tensor.h @@ -1972,6 +1972,9 @@ class Tensor { auto fill = [factor](typename value_type::value_type* dst, const typename value_type::value_type* src, std::size_t n) { + // detail's mixed complex x scalar operator* (e.g. complex * + // int), as in the non-nested branch below + using namespace TiledArray::detail; for (std::size_t i = 0; i < n; ++i) dst[i] = src[i] * factor; }; return detail::arena_trivial_unary(*this, fill); @@ -1979,6 +1982,9 @@ class Tensor { auto fill = [factor](typename value_type::value_type* dst, const typename value_type::value_type* src, std::size_t n) { + // detail's mixed complex x scalar operator* (e.g. complex * + // int), as in the non-nested branch below + using namespace TiledArray::detail; for (std::size_t i = 0; i < n; ++i) dst[i] = src[i] * factor; }; return detail::arena_trivial_unary(*this, fill); @@ -2127,6 +2133,7 @@ class Tensor { using ElemT = typename value_type::value_type; auto fill = [factor](ElemT* dst, const ElemT* l, const ElemT* r, std::size_t n) { + using namespace TiledArray::detail; // mixed complex x scalar operator* for (std::size_t i = 0; i < n; ++i) dst[i] = (l[i] + r[i]) * factor; }; return detail::arena_trivial_binary(*this, right, fill); @@ -2607,6 +2614,7 @@ class Tensor { using ElemT = typename value_type::value_type; auto fill = [factor](ElemT* dst, const ElemT* l, const ElemT* r, std::size_t n) { + using namespace TiledArray::detail; // mixed complex x scalar operator* for (std::size_t i = 0; i < n; ++i) dst[i] = (l[i] - r[i]) * factor; }; return detail::arena_trivial_binary(*this, right, fill); @@ -2862,6 +2870,7 @@ class Tensor { using ElemT = typename value_type::value_type; auto fill = [factor](ElemT* dst, const ElemT* l, const ElemT* r, std::size_t n) { + using namespace TiledArray::detail; // mixed complex x scalar operator* for (std::size_t i = 0; i < n; ++i) dst[i] = (l[i] * r[i]) * factor; }; return detail::arena_trivial_binary(*this, right, fill); diff --git a/tests/arena_strided_dgemm.cpp b/tests/arena_strided_dgemm.cpp index bd99b12f43..3b1a1049e9 100644 --- a/tests/arena_strided_dgemm.cpp +++ b/tests/arena_strided_dgemm.cpp @@ -5,6 +5,7 @@ #include "TiledArray/math/blas.h" #include "tiledarray.h" #include "unit_test_config.h" +#include #include #include #include @@ -2109,4 +2110,122 @@ BOOST_AUTO_TEST_CASE(ce_ce_seg_killswitch_matches_left) { } } + +// --------------------------------------------------------------------------- +// std::complex inner storage: the three strided kernels are templated +// on the inner numeric type (is_strided_dgemm_numeric_v), so complex ToT +// contractions (e.g. Kramers/relativistic CSV amplitudes) take the same +// zero-copy strided path. Each case fabricates complex arena tiles inline and +// checks against a naive complex reference, including a complex factor. +using ZInner = TA::ArenaTensor, TA::Range>; +using ZOuter = TA::Tensor; + +namespace { +ZOuter make_filled_z(const TA::Range& r, + const std::function& shape_fn, + double base) { + ZOuter t = TA::detail::arena_outer_init(r, 1, shape_fn); + for (std::size_t o = 0; o < t.range().volume(); ++o) { + ZInner& c = t.data()[o]; + if (!c) continue; + for (std::size_t e = 0; e < c.size(); ++e) + c.data()[e] = std::complex(base + 0.01 * o + e, 0.5 * e - 0.1 * o); + } + return t; +} +void check_close_z(const std::complex& got, + const std::complex& ref) { + BOOST_CHECK_SMALL(std::abs(got - ref), 1e-10 * std::max(1.0, std::abs(ref))); +} +} // namespace + +BOOST_AUTO_TEST_CASE(ce_e_complex_matches_reference) { + namespace blas = TA::math::blas; + using Z = std::complex; + const std::size_t M = 2, N = 3, K = 4, P = 3, Q = 5; + const Z factor(0.5, -1.25); + ZOuter L = make_filled_z(TA::Range{M, K}, [&](std::size_t) { return TA::Range{P}; }, 1.0); + ZOuter R = make_filled_z(TA::Range{N, K}, [&](std::size_t) { return TA::Range{Q}; }, 2.0); + ZOuter C = TA::detail::arena_outer_init( + TA::Range{M, N}, 1, [&](std::size_t) { return TA::Range{P, Q}; }); // zero-init + TA::detail::arena_strided_dgemm_ce_e(C, L, R, M, N, K, blas::NoTranspose, + blas::Transpose, factor); + for (std::size_t m = 0; m < M; ++m) + for (std::size_t n = 0; n < N; ++n) { + std::vector ref(P * Q, Z{}); + for (std::size_t k = 0; k < K; ++k) { + const Z* lp = L.data()[m * K + k].data(); + const Z* rp = R.data()[n * K + k].data(); + for (std::size_t p = 0; p < P; ++p) + for (std::size_t q = 0; q < Q; ++q) ref[p * Q + q] += factor * lp[p] * rp[q]; + } + const Z* got = C.data()[m * N + n].data(); + for (std::size_t e = 0; e < P * Q; ++e) check_close_z(got[e], ref[e]); + } +} + +BOOST_AUTO_TEST_CASE(ce_ce_right_complex_matches_reference) { + namespace blas = TA::math::blas; + using Z = std::complex; + const std::size_t Mo = 2, Mmu = 3, nK = 2, P = 4, Q = 5; + const Z factor(-0.75, 0.3); + // L outer (Mo,nK) row-major (m slow, k fast), inner {P,Q}. + ZOuter L = make_filled_z(TA::Range{Mo, nK}, [&](std::size_t) { return TA::Range{P, Q}; }, 1.0); + // R outer (Mmu,nK) canonical (mu slow, k fast), inner {Q}. + ZOuter R = make_filled_z(TA::Range{Mmu, nK}, [&](std::size_t) { return TA::Range{Q}; }, 2.0); + // C outer (Mo,Mmu) row-major (m slow, mu fast), inner {P}. + ZOuter C = TA::detail::arena_outer_init( + TA::Range{Mo, Mmu}, 1, [&](std::size_t) { return TA::Range{P}; }); + TA::detail::arena_strided_dgemm_ce_ce_right(C, L, R, /*Mo=*/Mo, /*No=*/Mmu, + /*Ko=*/nK, blas::NoTranspose, + blas::Transpose, factor); + for (std::size_t m = 0; m < Mo; ++m) + for (std::size_t mu = 0; mu < Mmu; ++mu) { + std::vector ref(P, Z{}); + for (std::size_t k = 0; k < nK; ++k) { + const Z* l = L.data()[m * nK + k].data(); // P x Q row-major + const Z* r = R.data()[mu * nK + k].data(); // Q + for (std::size_t a1 = 0; a1 < P; ++a1) { + Z acc{}; + for (std::size_t a4 = 0; a4 < Q; ++a4) acc += l[a1 * Q + a4] * r[a4]; + ref[a1] += factor * acc; + } + } + const Z* got = C.data()[m * Mmu + mu].data(); + for (std::size_t a1 = 0; a1 < P; ++a1) check_close_z(got[a1], ref[a1]); + } +} + +BOOST_AUTO_TEST_CASE(ce_ce_left_complex_matches_reference) { + namespace blas = TA::math::blas; + using Z = std::complex; + const std::size_t Mo = 2, No = 3, nK = 2, P = 4, Q = 5; + const Z factor(1.5, 2.0); + // L (clean) outer (Mo,nK) row-major (m slow, k fast), inner {Q}. + ZOuter L = make_filled_z(TA::Range{Mo, nK}, [&](std::size_t) { return TA::Range{Q}; }, 1.0); + // R (matrix) outer (nK,No) canonical (k slow, n fast), inner {Q,P} row-major. + ZOuter R = make_filled_z(TA::Range{nK, No}, [&](std::size_t) { return TA::Range{Q, P}; }, 2.0); + // C outer (Mo,No) row-major (m slow, n fast), inner {P}. + ZOuter C = TA::detail::arena_outer_init( + TA::Range{Mo, No}, 1, [&](std::size_t) { return TA::Range{P}; }); + TA::detail::arena_strided_dgemm_ce_ce_left(C, L, R, /*Mo=*/Mo, /*No=*/No, + /*Ko=*/nK, blas::NoTranspose, + blas::NoTranspose, factor); + for (std::size_t m = 0; m < Mo; ++m) + for (std::size_t n = 0; n < No; ++n) { + std::vector ref(P, Z{}); + for (std::size_t k = 0; k < nK; ++k) { + const Z* l = L.data()[m * nK + k].data(); // Q vector + const Z* r = R.data()[k * No + n].data(); // Q x P row-major + for (std::size_t p = 0; p < P; ++p) { + Z acc{}; + for (std::size_t a4 = 0; a4 < Q; ++a4) acc += l[a4] * r[a4 * P + p]; + ref[p] += factor * acc; + } + } + const Z* got = C.data()[m * No + n].data(); + for (std::size_t p = 0; p < P; ++p) check_close_z(got[p], ref[p]); + } +} + BOOST_AUTO_TEST_SUITE_END() diff --git a/tests/tot_array_fixture.h b/tests/tot_array_fixture.h index 2c27824961..47c3a58a4f 100644 --- a/tests/tot_array_fixture.h +++ b/tests/tot_array_fixture.h @@ -23,6 +23,7 @@ #include #include #include +#include #include "unit_test_config.h" #ifdef TILEDARRAY_HAS_BTAS #include @@ -44,15 +45,16 @@ using namespace TiledArray; // These are all of the template parameters we are going to test over -using test_params = - boost::mpl::list>>, - std::tuple>>, - std::tuple>> +using test_params = boost::mpl::list< + std::tuple>>, + std::tuple>>, + std::tuple>>, + std::tuple, Tensor>>> #ifdef TILEDARRAY_HAS_BTAS - , - std::tuple>>, - std::tuple>>, - std::tuple>> + , + std::tuple>>, + std::tuple>>, + std::tuple>> // ,std::tuple, Range>>, // std::tuple, Range>>, // std::tuple, Range>> @@ -60,7 +62,7 @@ using test_params = // std::tuple, Range>>>, // std::tuple, Range>>> #endif - >; + >; // These typedefs unpack the unit test template parameter //{ diff --git a/tests/tot_dist_array_part1.cpp b/tests/tot_dist_array_part1.cpp index d95bb050a2..dff37f526d 100644 --- a/tests/tot_dist_array_part1.cpp +++ b/tests/tot_dist_array_part1.cpp @@ -39,8 +39,11 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(typedefs, TestParam, test_params) { } { + // scalar_type is the REAL type underlying the element type (complex + // -> T), so compare against the fixture's element type stripped likewise. constexpr bool is_same = - std::is_same_v; + std::is_same_v>; BOOST_TEST(is_same); }