From f4a8864fc2968676b49ad120c72e920e56324f3c Mon Sep 17 00:00:00 2001 From: Marclie Date: Tue, 28 May 2024 13:23:16 -0400 Subject: [PATCH 1/3] Add forall function for element-wise operations on dense and sparse DistArray objects that keeps track of the coordinate index for conditional operations. --- src/TiledArray/conversions/foreach.h | 83 ++++++++++++++++++++++++++++ tests/foreach.cpp | 24 ++++++++ 2 files changed, 107 insertions(+) diff --git a/src/TiledArray/conversions/foreach.h b/src/TiledArray/conversions/foreach.h index 20f2d36ec3..a878dba860 100644 --- a/src/TiledArray/conversions/foreach.h +++ b/src/TiledArray/conversions/foreach.h @@ -495,6 +495,89 @@ inline std::enable_if_t, void> foreach_inplace( if (fence) arg.world().gop.fence(); } +/// Modify each element of a dense Array + +/// This function modifies the elements of a \c DistArray object with a const reference to the +/// index of the current element. This allows the user to modify specific elements of the array +/// based on their indices. Users must provide a function/functor that modifies each element. The provided function +/// should take a reference to a \c Tile object and a reference to a \c std::vector +/// representing the indices of the current element within the tile. For example, +/// to copy the upper triangular elements of a nxnxn array to a c++ vector of size n^3: +/// \code +/// std::vector vec(n*n); +/// forall(array, [&vec] (auto& tile, const auto& index) { +/// size_t i = index[0], j = index[1], k = index[2]; +/// if (i <= j && j <= k) { +/// vec[i*n*n+j*n+k] = tile[index]; +/// } else { +/// vec[i*n*n+j*n+k] = 0.0; +/// } +/// }); +/// \endcode +/// Similarly, to set each upper triangular element of a nxnxn array to the square root of values in a c++ vector of size n^3: +/// \code +/// vector vec(n*n*n); +/// std::generate(v.begin(), v.end(), std::rand); +/// forall(array, [&vec] (Tile& tile, index_type& index) { +/// size_t i = index[0], j = index[1], k = index[2]; +/// if (i <= j && j <= k) { +/// tile[index] = std::sqrt(vec[i*n*n+j*n+k]); +/// } else { +/// tile[index] = 0.0; +/// } +/// }); +/// \endcode +/// The expected signature of the element operation is: +/// \code +/// void op(Tile& tile, Range::index_type& index); +/// \endcode +/// \tparam Tile The tile type of \c arg +/// \tparam Policy The policy type of \c arg +/// \tparam Op Mutating element operation +/// \param arg The argument array to be modified +/// \param op The mutating element function +/// \param fence If \c true, this function will fence before and after the data is modified +template ::type>::value>::type> +inline void forall( + DistArray& arg, Op&& op, bool fence = true) { + // Get the rank of the array + const std::size_t rank = arg.trange().rank(); + + // Use foreach_inplace to iterate over tiles and modify elements + foreach_inplace( + arg, + [op = std::forward(op), rank](Tile& tile) mutable { + + // Get the lower and upper bounds of the current tile + const auto& lobound = tile.range().lobound(); + const auto& upbound = tile.range().upbound(); + + // Create a vector to store the current element's coordinate indices + Range::index_type index(rank); + + // Define a recursive function to iterate over the coordinate indices at each dimension + std::function loop_body = [&](std::size_t dim) { + // If all dimensions have been processed, apply the operation + if (dim == rank) { + op(tile, index); + return; + } else { + // Iterate over the elements in the current dimension + for (index[dim] = lobound[dim]; index[dim] < upbound[dim]; ++index[dim]) { + // Recursively call loop_body for the next dimension of the array + loop_body(dim + 1); + } + } + }; + + // Start the iteration from the first dimension + loop_body(0); + }, + fence); // Fence before and after the data is modified +} + /// Apply a function to each tile of a sparse Array /// This function uses an \c Array object to generate a new \c Array where the diff --git a/tests/foreach.cpp b/tests/foreach.cpp index 106a855a0a..7be3068f4a 100644 --- a/tests/foreach.cpp +++ b/tests/foreach.cpp @@ -119,6 +119,30 @@ BOOST_AUTO_TEST_CASE(foreach_unary) { } } +BOOST_AUTO_TEST_CASE(forall_unary) { + + TArrayI result = a.clone(); + forall (result, [](TensorI& tile, Range::index_type &coord_idx) { + if (coord_idx[0] < coord_idx[1]) + tile[coord_idx] = coord_idx[0] * tile[coord_idx]; + else + tile[coord_idx] = coord_idx[1] * tile[coord_idx]; + }, true); + + for (auto index : *result.pmap()) { + TensorI tile0 = a.find(index).get(); + TensorI tile = result.find(index).get(); + const Range &range = tile0.range(); + for (std::size_t i = 0; i < tile.size(); ++i) { + const Range::index_type &coord_idx = range.idx(i); + if (coord_idx[0] < coord_idx[1]) + BOOST_CHECK_EQUAL(tile[i], coord_idx[0] * tile0[i]); + else + BOOST_CHECK_EQUAL(tile[i], coord_idx[1] * tile0[i]); + } + } +} + BOOST_AUTO_TEST_CASE(foreach_unary_sparse) { TSpArrayI result = foreach (c, [](TensorI& result, const TensorI& arg) -> float { From 3407a9a3d2bb16e0211381a1ccb9cde86a987006 Mon Sep 17 00:00:00 2001 From: Marcus Dante Liebenthal <32597392+Marclie@users.noreply.github.com> Date: Tue, 28 May 2024 13:45:01 -0400 Subject: [PATCH 2/3] fix comment --- src/TiledArray/conversions/foreach.h | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/TiledArray/conversions/foreach.h b/src/TiledArray/conversions/foreach.h index a878dba860..7be9cc7797 100644 --- a/src/TiledArray/conversions/foreach.h +++ b/src/TiledArray/conversions/foreach.h @@ -495,7 +495,7 @@ inline std::enable_if_t, void> foreach_inplace( if (fence) arg.world().gop.fence(); } -/// Modify each element of a dense Array +/// Modify each element of an Array object /// This function modifies the elements of a \c DistArray object with a const reference to the /// index of the current element. This allows the user to modify specific elements of the array @@ -504,7 +504,7 @@ inline std::enable_if_t, void> foreach_inplace( /// representing the indices of the current element within the tile. For example, /// to copy the upper triangular elements of a nxnxn array to a c++ vector of size n^3: /// \code -/// std::vector vec(n*n); +/// std::vector vec(n*n*n); /// forall(array, [&vec] (auto& tile, const auto& index) { /// size_t i = index[0], j = index[1], k = index[2]; /// if (i <= j && j <= k) { @@ -540,8 +540,7 @@ inline std::enable_if_t, void> foreach_inplace( template ::type>::value>::type> -inline void forall( - DistArray& arg, Op&& op, bool fence = true) { +inline void forall(DistArray& arg, Op&& op, bool fence = true) { // Get the rank of the array const std::size_t rank = arg.trange().rank(); From 4dbede8268612c73c5064cd2a3def757d29044a9 Mon Sep 17 00:00:00 2001 From: Marclie Date: Fri, 31 May 2024 14:19:50 -0400 Subject: [PATCH 3/3] Added type traits for dispatching to tile-wise and element-wise paths based on the traits of Op. --- src/TiledArray/conversions/foreach.h | 55 ++++++++++------------------ tests/foreach.cpp | 16 +++----- 2 files changed, 25 insertions(+), 46 deletions(-) diff --git a/src/TiledArray/conversions/foreach.h b/src/TiledArray/conversions/foreach.h index a878dba860..5538f74fa6 100644 --- a/src/TiledArray/conversions/foreach.h +++ b/src/TiledArray/conversions/foreach.h @@ -480,7 +480,8 @@ inline std::enable_if_t, DistArray> foreach ( /// function will fence before AND after the data is modified template ::type>::value>::type> + typename std::decay::type>::value>::type, + typename = typename std::enable_if::value>::type> inline std::enable_if_t, void> foreach_inplace( DistArray& arg, Op&& op, bool fence = true) { // The tile data is being modified in place, which means we may need to @@ -539,43 +540,22 @@ inline std::enable_if_t, void> foreach_inplace( /// \param fence If \c true, this function will fence before and after the data is modified template ::type>::value>::type> -inline void forall( + typename std::decay::type>::value>::type, + typename = typename std::enable_if::value>::type> +inline void foreach_inplace( DistArray& arg, Op&& op, bool fence = true) { - // Get the rank of the array - const std::size_t rank = arg.trange().rank(); + + // wrap Op into a shallow-copy copyable handle + auto op_shared_handle = make_op_shared_handle(std::forward(op)); // Use foreach_inplace to iterate over tiles and modify elements foreach_inplace( arg, - [op = std::forward(op), rank](Tile& tile) mutable { - - // Get the lower and upper bounds of the current tile - const auto& lobound = tile.range().lobound(); - const auto& upbound = tile.range().upbound(); - - // Create a vector to store the current element's coordinate indices - Range::index_type index(rank); - - // Define a recursive function to iterate over the coordinate indices at each dimension - std::function loop_body = [&](std::size_t dim) { - // If all dimensions have been processed, apply the operation - if (dim == rank) { - op(tile, index); - return; - } else { - // Iterate over the elements in the current dimension - for (index[dim] = lobound[dim]; index[dim] < upbound[dim]; ++index[dim]) { - // Recursively call loop_body for the next dimension of the array - loop_body(dim + 1); - } - } - }; - - // Start the iteration from the first dimension - loop_body(0); - }, - fence); // Fence before and after the data is modified + [op = std::move(op_shared_handle)](Tile& tile) mutable { + for (const Range::index_type& index : tile.range()) + op(tile, index); + }, fence); // Fence before and after the data is modified } /// Apply a function to each tile of a sparse Array @@ -670,7 +650,8 @@ inline std::enable_if_t, DistArray> foreach ( /// function will fence before AND after the data is modified template ::type>::value>::type> + typename std::decay::type>::value>::type, + typename = typename std::enable_if::value>::type> inline std::enable_if_t, void> foreach_inplace( DistArray& arg, Op&& op, bool fence = true) { // The tile data is being modified in place, which means we may need to @@ -712,7 +693,8 @@ inline std:: } /// This function takes two input tiles and put result into the left tile -template +template ::value>::type> inline std::enable_if_t, void> foreach_inplace( DistArray& left, const DistArray& right, Op&& op, bool fence = true) { @@ -758,7 +740,8 @@ inline std:: } /// This function takes two input tiles and put result into the left tile -template +template ::value>::type> inline std::enable_if_t, void> foreach_inplace( DistArray& left, const DistArray& right, Op&& op, diff --git a/tests/foreach.cpp b/tests/foreach.cpp index 7be3068f4a..d4cb1adda6 100644 --- a/tests/foreach.cpp +++ b/tests/foreach.cpp @@ -119,14 +119,12 @@ BOOST_AUTO_TEST_CASE(foreach_unary) { } } -BOOST_AUTO_TEST_CASE(forall_unary) { +BOOST_AUTO_TEST_CASE(foreach_w_idx) { TArrayI result = a.clone(); - forall (result, [](TensorI& tile, Range::index_type &coord_idx) { - if (coord_idx[0] < coord_idx[1]) - tile[coord_idx] = coord_idx[0] * tile[coord_idx]; - else - tile[coord_idx] = coord_idx[1] * tile[coord_idx]; + foreach_inplace(result, [](TensorI& tile, const Range::index_type &coord_idx) { + long fac = (coord_idx[0] < coord_idx[1]) ? coord_idx[0] : coord_idx[1]; + tile[coord_idx] = fac * tile[coord_idx]; }, true); for (auto index : *result.pmap()) { @@ -135,10 +133,8 @@ BOOST_AUTO_TEST_CASE(forall_unary) { const Range &range = tile0.range(); for (std::size_t i = 0; i < tile.size(); ++i) { const Range::index_type &coord_idx = range.idx(i); - if (coord_idx[0] < coord_idx[1]) - BOOST_CHECK_EQUAL(tile[i], coord_idx[0] * tile0[i]); - else - BOOST_CHECK_EQUAL(tile[i], coord_idx[1] * tile0[i]); + long fac = coord_idx[0] < coord_idx[1] ? coord_idx[0] : coord_idx[1]; + BOOST_CHECK_EQUAL(tile[i], fac * tile0[i]); } } }