From 179e1d39f0b8635f37d72be6d3c91f315a99eaea Mon Sep 17 00:00:00 2001 From: Bimal Gaudel Date: Mon, 3 Jun 2024 13:58:46 -0400 Subject: [PATCH] Sparse array evaluation support in manual eval from tot_array_fixture. --- tests/tot_array_fixture.h | 307 ++++++++++++++++++++++++++------------ 1 file changed, 209 insertions(+), 98 deletions(-) diff --git a/tests/tot_array_fixture.h b/tests/tot_array_fixture.h index 4710ab79d7..5d0a0ce4dd 100644 --- a/tests/tot_array_fixture.h +++ b/tests/tot_array_fixture.h @@ -93,102 +93,203 @@ using output_archive_type = madness::archive::BinaryFstreamOutputArchive; enum class ShapeComp { True, False }; -template , bool> = true> -auto random_tensor(TA::Range const& rng) { - using NumericT = typename TensorT::numeric_type; - TensorT result{rng}; - - std::generate(/*std::execution::par, */ - result.begin(), result.end(), - TA::detail::MakeRandom::generate_value); - return result; +namespace fixture { +namespace { + +template +constexpr bool maps_index_to_range_v{}; + +template +constexpr bool maps_index_to_range_v< + Invocable, + std::enable_if_t>>>{ + true}; + +using il_range = std::initializer_list; +using il_trange = std::initializer_list; + +} // namespace + +/// +/// \tparam T Non cv-qualified TA::Tensor type. +/// \tparam Rng TA::Range should be constructible from Rng type. +/// Eg. TA::Range, std::initializer_list. +/// \param rng The range of the result tensor. +/// TA::Range(rng) will be called explicitly. +/// \return A TA::Tensor of a numeric type with random elements. +/// +template >, + typename = std::enable_if_t>> +auto random_tensor(Rng rng) { + using numeric_type = typename T::numeric_type; + + auto gen = [](auto&&) { + return detail::MakeRandom::generate_value(); + }; + + return T(TA::Range(rng), gen); } -template -auto random_tensor(std::initializer_list const& extents) { - auto lobounds = TA::container::svector(extents.size(), 0); - return random_tensor(TA::Range{lobounds, extents}); +/// +/// \tparam T Non cv-qualified +/// TA::Tensor,...> type. +/// \tparam RngO TA::Range should be constructible from RngO type. +/// Eg. TA::Range, std::initializer_list. +/// \tparam RngI TA::Range should be constructible from RngI type. +/// Eg. TA::Range, std::initializer_list. +/// \param rngo The range of the result tensor (ie the outer tensor). +/// TA::Range(rngo) will be called explicitly. +/// \param rngi The range of the inner tensors. Note that ALL inner tensors +/// will have an EQUAL range. TA::Range(rngi) will be +/// called explicitly. +/// \return A TA::Tensor of TA::Tensor with random +/// numeric_type elements. +/// +template < + typename T, typename RngO, typename RngI, // + typename = std::enable_if_t>, // + typename = std::enable_if_t>, // + typename = std::enable_if_t>> +auto random_tensor(RngO rngo, RngI rngi) { + using numeric_type = typename T::numeric_type; + using Inner = typename T::value_type; + + auto gen_inner = [](auto&&) { + return detail::MakeRandom::generate_value(); + }; + + auto gen_outer = [gen_inner, rngi](auto&&) { + return Inner(TA::Range(rngi), gen_inner); + }; + + return T(TA::Range(rngo), gen_outer); } -// -// note: all the inner tensors (elements of the outer tensor) -// have the same @c inner_rng -// +/// +/// \tparam T Non cv-qualified +/// TA::Tensor,...> type. +/// \tparam RngO TA::Range should be constructible from RngO type. +/// Eg. TA::Range, std::initializer_list. +/// \tparam IxMap An invocable type that maps the index of an element in the +/// outer tensor to a value, allowing the construction of +/// TA::Range from that value. +/// \param rngo The range of the result tensor (ie the outer tensor). +/// TA::Range(rngo) will be called explicitly. +/// \param ixmap An invocable that maps the index of an element in the +/// outer tensor to a value, allowing the construction of +/// TA::Range from that value. +/// \return A TA::Tensor of TA::Tensor with random +/// numeric_type elements. +/// template < - typename TensorT, - std::enable_if_t, bool> = true> -auto random_tensor(TA::Range const& outer_rng, TA::Range const& inner_rng) { - using InnerTensorT = typename TensorT::value_type; - TensorT result{outer_rng}; + typename T, typename RngO, typename IxMap, // + typename = std::enable_if_t>, // + typename = std::enable_if_t>, // + std::enable_if_t, bool> = true> +auto random_tensor(RngO rngo, IxMap ixmap) { + using numeric_type = typename T::numeric_type; + + auto gen_inner = [](auto&&) { + return TA::detail::MakeRandom::generate_value(); + }; - std::generate(/*std::execution::par,*/ - result.begin(), result.end(), [inner_rng]() { - return random_tensor(inner_rng); - }); + auto gen_outer = [gen_inner, ixmap](auto const& oix) { + auto inner_rng = TA::Range(ixmap(oix)); + return typename T::value_type(inner_rng, gen_inner); + }; - return result; + return T(TA::Range(rngo), gen_outer); } -template -auto random_tensor(TA::Range const& outer_rng, - std::initializer_list const& inner_extents) { - TA::container::svector lobounds(inner_extents.size(), 0); - return random_tensor(outer_rng, TA::Range(lobounds, inner_extents)); +/// +/// \tparam Array Non cv-qualified TA::DistArray type that has non-nested +/// tile type. Eg. TA::DistArray> +/// \tparam Rng TA::TiledRange should be constructible from Rng type. +/// \param rng The TA::TiledRange of the result TA::DistArray. +/// \return A TA::DistArray of non-nested tile type with random elements. +/// +template < + typename Array, typename Rng = il_trange, + typename = std::enable_if_t == 1>, + typename = std::enable_if_t>> +auto random_array(Rng rng) { + using T = typename Array::value_type; + + auto make_tile = [](auto& tile, auto const& rng) { + tile = random_tensor(rng); + return tile.norm(); + }; + + return TA::make_array(TA::get_default_world(), TA::TiledRange(rng), + make_tile); } /// -/// \tparam Array The type of DistArray to be generated. Cannot be cv-qualified -/// or reference type. -/// \tparam Args TA::Range type for inner tensor if the tile type of the result -/// is a tensor-of-tensor. -/// \param trange The TiledRange of the result DistArray. -/// \param args Either exactly one TA::Range type when the tile type of Array is -/// tensor-of-tensor or nothing. -/// \return Returns a DistArray of type Array whose elements are randomly -/// generated. -/// @note: -/// - Although DistArrays with Sparse policy can be generated all of their -/// tiles are initialized with random values -- technically the returned value -/// is dense. -/// - In case of arrays with tensor-of-tensor tiles, all the inner tensors have -/// the same rank and the same extent of corresponding modes. +/// \tparam Array Non cv-qualified TA::DistArray type that has a nested +/// tile type. +/// Eg. TA::DistArray>> +/// \tparam RngO TA::TiledRange should be constructible form RngO type. +/// \tparam RngI TA::Range should be constructible from RngI type. +/// \param rngo The TA::TiledRange of the result TA::DistArray. +/// \param rngi The range of the inner tensors. Note that ALL inner tensors +/// will have an EQUAL range. TA::Range(rngi) will be +/// called explicitly. +/// \return A TA::DistArray of nested tile type with random elements. /// template < - typename Array, typename... Args, - typename = - std::void_t, - std::enable_if_t, - bool> = true> -auto random_array(TA::TiledRange const& trange, Args const&... args) { - static_assert( - (sizeof...(Args) == 0 && - TA::detail::is_tensor_v) || - (sizeof...(Args) == 1) && - (TA::detail::is_tensor_of_tensor_v)); - - using TensorT = typename Array::value_type; - using PolicyT = typename Array::policy_type; - - auto make_tile_meta = [](auto&&... args) { - return [=](TensorT& tile, TA::Range const& rng) { - tile = random_tensor(rng, args...); - if constexpr (std::is_same_v) - return tile.norm(); - }; + typename Array, typename RngO = il_trange, typename RngI = il_range, + typename = std::enable_if_t == 2>, + typename = std::enable_if_t>, + typename = std::enable_if_t>> +auto random_array(RngO rngo, RngI rngi) { + using T = typename Array::value_type; + + auto make_tile = [rngi](auto& tile, auto const& rng) { + tile = random_tensor(rng, rngi); + return tile.norm(); }; - return TA::make_array(TA::get_default_world(), trange, - make_tile_meta(args...)); + return TA::make_array(TA::get_default_world(), TA::TiledRange(rngo), + make_tile); } -template -auto random_array(std::initializer_list> trange, - Args&&... args) { - return random_array(TA::TiledRange(trange), - std::forward(args)...); +/// +/// \tparam Array Non cv-qualified TA::DistArray type that has a nested +/// tile type. +/// Eg. TA::DistArray>> +/// \tparam RngO TA::TiledRange should be constructible form RngO type. +/// \tparam IxMap An invocable type that maps the index of an element in the +/// outer tensor to a value, allowing the construction of +/// TA::Range from that value. +/// \param rngo The TA::TiledRange of the result TA::DistArray. +/// \param ixmap An invocable that maps the index of an element in the +/// outer tensor to a value, allowing the construction of +/// TA::Range from that value. +/// \return A TA::DistArray of nested tile type with random elements. +template < + typename Array, typename RngO, typename IxMap, + typename = std::enable_if_t == 2>, + typename = std::enable_if_t>, + std::enable_if_t, bool> = true> +auto random_array(RngO rngo, IxMap ixmap) { + using T = typename Array::value_type; + + auto make_tile = [ixmap](auto& tile, auto const& rng) { + tile = random_tensor(rng, ixmap); + return tile.norm(); + }; + + return TA::make_array(TA::get_default_world(), TA::TiledRange(rngo), + make_tile); } +} // namespace fixture + +using fixture::random_array; +using fixture::random_tensor; + /// /// Succinctly call TA::detail::tensor_contract /// @@ -404,6 +505,9 @@ Result general_product(TensorA const& A, TensorB const& B, using TA::detail::max_nested_rank; using TA::detail::nested_rank; + // empty tensors + if (A.empty() || B.empty()) return Result{}; + static_assert(std::is_same_v); @@ -502,12 +606,22 @@ Result general_product(TensorA const& A, TensorB const& B, } } else { typename Result::value_type temp{}; - for (auto ix_I : rng_I) { + for (auto const& ix_I : rng_I) { apply_partial_perm(ix_A, ix_I, setup.I_to_A); apply_partial_perm(ix_B, ix_I, setup.I_to_B); - if constexpr (is_tot) - temp += general_product( + if constexpr (is_tot) { + auto temp_ = general_product( A(ix_A), B(ix_B), args...); + if constexpr (TA::detail::is_nested_tensor_v< + typename Result::value_type>) { + if (temp.empty()) + temp = std::move(temp_); + else + temp += temp_; + } else { + temp += temp_; + } + } else { TA_ASSERT(!(ix_A.empty() || ix_B.empty())); temp += A(ix_A) * B(ix_B); @@ -548,26 +662,14 @@ auto general_product(TA::DistArray A, TA::TiledRange result_trange; { - auto const rank = result_tensor.range().rank(); - auto const result_range = result_tensor.range(); - - TA::container::svector> tr1s(rank, {0}); - - TA::container::svector const ix_hi(result_range.upbound()); - for (auto d = 0; d < rank; ++d) { - TA::container::svector ix(result_range.lobound()); - for (auto& i = ix[d]; i < ix_hi[d]; ++i) { - auto const& elem_tensor = result_tensor(ix); - auto& tr1 = tr1s[d]; - tr1.emplace_back(tr1.back() + elem_tensor.range().extent(d)); - } + TA::container::svector tr1s(setup.rank_C); + for (auto [t, f] : setup.C_to_A) { + tr1s.at(t) = A.trange().at(f); } - - TA::container::svector tr1s_explicit; - tr1s_explicit.reserve(tr1s.size()); - for (auto const& v : tr1s) tr1s_explicit.emplace_back(v.begin(), v.end()); - - result_trange = TA::TiledRange(tr1s_explicit); + for (auto [t, f] : setup.C_to_B) { + tr1s.at(t) = B.trange().at(f); + } + result_trange = TiledRange(tr1s); } TA::DistArray C(world, result_trange); @@ -588,6 +690,15 @@ auto general_product(TA::DistArray A, return general_product(A, B, args...); } +template +auto general_product(TA::DistArray A, + TA::DistArray B, + Setups const&... args) { + auto A_dense = to_dense(A); + auto B_dense = to_dense(B); + return TA::to_sparse(general_product(A_dense, B_dense, args...)); +} + template >> auto manual_eval(OuterInnerSetup const& setups, ArrayA A, ArrayB B) {