vector_of_arrays.h

Go to the documentation of this file.

//
// Created by Chong Peng on 2019-05-01.
//
 
#ifndef TILEDARRAY_CONVERSIONS_VECTOR_OF_ARRAYS_H_
#define TILEDARRAY_CONVERSIONS_VECTOR_OF_ARRAYS_H_
 
#include <tiledarray.h>
 
namespace TiledArray {
 
namespace detail {
 
 
 
inline TA::TiledRange prepend_dim_to_trange(
    std::size_t array_rank, const TiledArray::TiledRange& array_trange,
    std::size_t block_size = 1) {
  TA::TiledRange1 new_trange1;
  {
    std::vector<std::size_t> new_trange1_v;
    auto range_size = array_rank;
    new_trange1_v.push_back(0);
    for (decltype(range_size) i = block_size; i < range_size; i += block_size) {
      new_trange1_v.push_back(i);
    }
    new_trange1_v.push_back(range_size);
    new_trange1 = TA::TiledRange1(new_trange1_v.begin(), new_trange1_v.end());
  }
 
  TA::TiledRange new_trange;
  {
    auto old_trange1s = array_trange.data();
    old_trange1s.insert(old_trange1s.begin(), new_trange1);
    new_trange = TA::TiledRange(old_trange1s.begin(), old_trange1s.end());
  }
 
  return new_trange;
}
 
template <typename Tile>
TA::SparseShape<float> fuse_vector_of_shapes_tiles(
    madness::World& global_world,
    const std::vector<TA::DistArray<Tile, TA::SparsePolicy>>& arrays,
    const std::size_t array_rank, const TA::TiledRange& fused_trange) {
  if (arrays.size() == 0) {
    TA::Tensor<float> fused_tile_norms(fused_trange.tiles_range(), 0.f);
    return TA::SparseShape<float>(global_world, fused_tile_norms, fused_trange,
                                  true);
  }
  const std::size_t rank = global_world.rank();
  auto size = global_world.size();
  auto first_tile_in_mode0 = *fused_trange.dim(0).begin();
  const auto block_size =
      first_tile_in_mode0.second - first_tile_in_mode0.first;
 
  std::size_t ntiles_per_array = arrays[0].trange().tiles_range().volume();
  // precompute tile volumes for later repeated use
  std::vector<size_t> tile_volumes(ntiles_per_array);
  {
    const auto& tiles_range = arrays[0].trange().tiles_range();
    for (auto&& tile_idx : tiles_range) {
      const auto tile_ord = tiles_range.ordinal(tile_idx);
      tile_volumes[tile_ord] =
          arrays[0].trange().make_tile_range(tile_idx).volume();
    }
  }
 
  TA::Tensor<float> fused_tile_norms(
      fused_trange.tiles_range(), 0.f);  // use nonzeroes for local tiles only
 
  // compute norms of fused tiles
  // N.B. tile norms are stored in scaled format, unscale in order to compute
  // norms of fused tiles
  std::size_t narrays = array_rank;
  size_t fused_tile_ord = 0;
  auto element_offset_in_owner = 0;
  for (size_t vidx = 0, fused_vidx = 0; vidx < narrays;
       vidx += block_size, ++fused_vidx) {
    bool have_rank = (rank == fused_vidx % size);
    // how many arrays actually constribute to this fused tile ... last fused
    // tile may have fewer than block_size
    if (have_rank) {
      const auto vblk_size =
          (narrays - vidx) >= block_size ? block_size : narrays - vidx;
      for (size_t tile_ord = 0; tile_ord != ntiles_per_array;
           ++tile_ord, ++fused_tile_ord) {
        auto array_ptr = arrays.begin() + element_offset_in_owner * vblk_size;
        float unscaled_fused_tile_norm2 = 0;
        const auto tile_volume = tile_volumes[tile_ord];
        for (size_t v = 0, vv = vidx; v != vblk_size; ++v, ++vv) {
          const auto unscaled_tile_norm =
              (*(array_ptr)).shape().data()[tile_ord] * tile_volume;
          unscaled_fused_tile_norm2 += unscaled_tile_norm * unscaled_tile_norm;
          ++array_ptr;
        }
        const auto fused_tile_volume = tile_volume * vblk_size;
        const auto fused_tile_norm =
            std::sqrt(unscaled_fused_tile_norm2) / fused_tile_volume;
 
        *(fused_tile_norms.data() + fused_tile_ord) = fused_tile_norm;
      }
      element_offset_in_owner += 1;
    } else {
      fused_tile_ord += ntiles_per_array;
    }
  }
 
  auto fused_shapes = TA::SparseShape<float>(global_world, fused_tile_norms,
                                             fused_trange, true);
 
  return fused_shapes;
}
 
template <typename Tile>
TA::DenseShape fuse_vector_of_shapes_tiles(
    madness::World&,
    const std::vector<TA::DistArray<Tile, TA::DensePolicy>>& arrays,
    const std::size_t array_rank, const TA::TiledRange& fused_trange) {
  return TA::DenseShape(1, fused_trange);
}
 
 
// TODO rename to tilewise_slice_of_fused_shape
inline TA::SparseShape<float> subshape_from_fused_tile(
    const TA::TiledRange& split_trange,
    const TA::SparsePolicy::shape_type& shape, const std::size_t tile_idx,
    const std::size_t split_ntiles, const std::size_t tile_size) {
  TA_ASSERT(split_ntiles == split_trange.tiles_range().volume());
  TA::Tensor<float> split_tile_norms(split_trange.tiles_range());
 
  // map element i to its tile index
  std::size_t offset = tile_idx * split_ntiles;
 
  // note that unlike fusion we cannot compute exact norm of the split tile
  // to guarantee upper bound we have to multiply the norms by the number of
  // split tiles in the fused tile; to see why multiplication is necessary think
  // of a tile obtained by fusing 1 nonzero tile with one or more zero tiles.
  const auto* split_tile_begin = shape.data().data() + offset;
  std::transform(split_tile_begin, split_tile_begin + split_ntiles,
                 split_tile_norms.data(),
                 [tile_size](const float& elem) { return elem * tile_size; });
 
  auto split_shape =
      TA::SparseShape<float>(split_tile_norms, split_trange, true);
  return split_shape;
}
 
 
inline TA::DenseShape subshape_from_fused_tile(
    const TA::TiledRange& split_trange,
    const TA::DensePolicy::shape_type& shape, const std::size_t tile_idx,
    const std::size_t split_ntiles, const std::size_t tile_size) {
  return TA::DenseShape(tile_size, split_trange);
}
}  // namespace detail
 
namespace detail {
template <typename Array>
class dist_subarray_vec
    : public madness::WorldObject<dist_subarray_vec<Array>> {
 public:
  using Tile = typename Array::value_type;
  using Policy = typename Array::policy_type;
 
  dist_subarray_vec(madness::World& world, const std::vector<Array>& array,
                    const std::size_t rank)
      : madness::WorldObject<dist_subarray_vec<Array>>(world),
        split_array(array),
        rank_(rank) {
    this->process_pending();
  }
 
  virtual ~dist_subarray_vec() {}
 
  template <typename Index>
  madness::Future<Tile> get_tile(int r, Index& i) {
    return split_array.at(r).find(i);
  }
 
  const std::vector<Array>& array_accessor() const { return split_array; }
 
  unsigned long size() const { return rank_; }
 
 private:
  const std::vector<Array>& split_array;
  const int rank_;
};
}  // namespace detail
 
 
template <typename Tile, typename Policy>
TA::DistArray<Tile, Policy> fuse_vector_of_arrays_tiles(
    madness::World& global_world,
    const std::vector<TA::DistArray<Tile, Policy>>& array_vec,
    const std::size_t fused_dim_extent,
    const TiledArray::TiledRange& array_trange, std::size_t block_size = 1) {
  auto nproc = global_world.size();
 
  // make instances of array_vec globally accessible
  using Array = TA::DistArray<Tile, Policy>;
  detail::dist_subarray_vec<Array> arrays(global_world, array_vec,
                                          fused_dim_extent);
 
  // make fused tiledrange
  auto fused_trange =
      detail::prepend_dim_to_trange(fused_dim_extent, array_trange, block_size);
  std::size_t ntiles_per_array = array_trange.tiles_range().volume();
 
  // make fused shape
  auto fused_shape = detail::fuse_vector_of_shapes_tiles(
      global_world, arrays.array_accessor(), fused_dim_extent, fused_trange);
 
  // make fused array
  TA::DistArray<Tile, Policy> fused_array(global_world, fused_trange,
                                          fused_shape);
 
  auto make_tile = [](const TA::Range& range,
                      const std::vector<madness::Future<Tile>>& tiles) {
    TA_ASSERT(range.extent(0) == tiles.size());
    Tile result(range);
    const auto volume = range.volume();
    size_t result_volume = 0;
    auto* result_ptr = result.data();
    for (auto&& fut_of_tile : tiles) {
      TA_ASSERT(fut_of_tile.probe());
      const auto& tile = fut_of_tile.get();
      const auto* tile_data = tile.data();
      const auto tile_volume = tile.size();
      std::copy(tile_data, tile_data + tile_volume, result_ptr);
      result_ptr += tile_volume;
      result_volume += tile_volume;
    }
    TA_ASSERT(volume == result_volume);
    return result;
  };
 
  for (auto&& fused_tile_ord : *fused_array.pmap()) {
    if (!fused_array.is_zero(fused_tile_ord)) {
      // convert ordinal of the fused tile to the ordinals of its constituent
      // tiles
      const auto div0 = std::ldiv(fused_tile_ord, ntiles_per_array);
      // index of the 0th mode of this tile
      const auto tile_idx_mode0 = div0.quot;
      // ordinal of the corresponding tile in the unfused array
      const auto tile_ord_array = div0.rem;
 
      const auto div1 = std::ldiv(tile_idx_mode0, nproc);
      const auto tile_idx_on_owner = div1.quot;
      const auto vector_idx_offset_on_owner = tile_idx_on_owner * block_size;
      const auto owner_rank = div1.rem;
 
      auto fused_tile_range =
          fused_array.trange().make_tile_range(fused_tile_ord);
      // make a vector of Futures to the input tiles
      std::vector<madness::Future<Tile>> input_tiles;
      input_tiles.reserve(fused_tile_range.extent(0));
      for (size_t v = 0, vidx = tile_idx_mode0 * block_size;
           v != block_size && vidx < fused_dim_extent; ++v, ++vidx) {
        using Index = decltype(tile_ord_array);
        input_tiles.emplace_back(
            arrays.task(owner_rank,
                        &detail::dist_subarray_vec<
                            DistArray<Tile, Policy>>::template get_tile<Index>,
                        vector_idx_offset_on_owner + v, tile_ord_array));
      }
      fused_array.set(
          fused_tile_ord,
          global_world.taskq.add(make_tile, std::move(fused_tile_range),
                                 std::move(input_tiles)));
    }
  }
 
  // keep arrays around until everyone is done
  global_world.gop.fence();
 
  return fused_array;
}
 
 
template <typename Tile, typename Policy>
void subarray_from_fused_array(
    madness::World& local_world, const TA::DistArray<Tile, Policy>& fused_array,
    std::size_t tile_idx,
    std::vector<TA::DistArray<Tile, Policy>>& split_arrays,
    const TA::TiledRange& split_trange) {
  TA_ASSERT(tile_idx < fused_array.trange().dim(0).extent());
  auto arrays_size = split_arrays.size();
 
  // calculate the number of elements in the 0th dimension are in this tile
  auto tile_range = fused_array.trange().dim(0).tile(tile_idx);
  auto tile_size = tile_range.second - tile_range.first;
  std::size_t split_ntiles = split_trange.tiles_range().volume();
  auto& shape = fused_array.shape();
 
  // Create tile_size arrays and put them into split_arrays
  for (size_t i = tile_range.first; i < tile_range.second; ++i) {
    auto split_shape = detail::subshape_from_fused_tile(
        split_trange, shape, tile_idx, split_ntiles, tile_size);
    // create split Array object
    TA::DistArray<Tile, Policy> split_array(local_world, split_trange,
                                            split_shape);
    split_arrays.push_back(split_array);
  }
 
  auto make_tile = [](const TA::Range& range, const Tile& fused_tile,
                      const size_t i_offset_in_tile) {
    const auto split_tile_volume = range.volume();
    return Tile(range,
                fused_tile.data() + i_offset_in_tile * split_tile_volume);
  };
 
  auto split_array_ptr = split_arrays.data();
  for (std::size_t index : *(*split_array_ptr).pmap()) {
    std::size_t fused_array_index = tile_idx * split_ntiles + index;
    if (!fused_array.is_zero(fused_array_index)) {
      for (std::size_t i = tile_range.first, tile_count = 0;
           i < tile_range.second; ++i, ++tile_count) {
        auto& array = *(split_array_ptr + arrays_size + tile_count);
        array.set(index, local_world.taskq.add(
                             make_tile, array.trange().make_tile_range(index),
                             fused_array.find(fused_array_index), tile_count));
      }
    }
  }
  return;
}
 
}  // namespace TiledArray
 
#endif  // TILEDARRAY_CONVERSIONS_VECTOR_OF_ARRAYS_H_