block_cyclic.h
Go to the documentation of this file.
1 /*
2  * This file is a part of TiledArray.
3  * Copyright (C) 2020 Virginia Tech
4  *
5  * This program is free software: you can redistribute it and/or modify
6  * it under the terms of the GNU General Public License as published by
7  * the Free Software Foundation, either version 3 of the License, or
8  * (at your option) any later version.
9  *
10  * This program is distributed in the hope that it will be useful,
11  * but WITHOUT ANY WARRANTY; without even the implied warranty of
12  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13  * GNU General Public License for more details.
14  *
15  * You should have received a copy of the GNU General Public License
16  * along with this program. If not, see <http://www.gnu.org/licenses/>.
17  *
18  * David Williams-Young
19  * Computational Research Division, Lawrence Berkeley National Laboratory
20  *
21  * block_cyclic.h
22  * Created: 7 Feb, 2020
23  * Edited: 13 May, 2020 (DBWY)
24  *
25  */
26 
27 #ifndef TILEDARRAY_MATH_LINALG_SCALAPACK_TO_BLOCKCYCLIC_H__INCLUDED
28 #define TILEDARRAY_MATH_LINALG_SCALAPACK_TO_BLOCKCYCLIC_H__INCLUDED
29 
30 #include <TiledArray/config.h>
31 #if TILEDARRAY_HAS_SCALAPACK
32 
34 #include <TiledArray/dist_array.h>
35 #include <TiledArray/error.h>
36 #include <TiledArray/tensor.h>
37 #include <TiledArray/tiled_range.h>
38 
39 #include <blacspp/grid.hpp>
40 #include <blacspp/information.hpp>
41 
42 #include <scalapackpp/block_cyclic.hpp>
43 #include <scalapackpp/util/sfinae.hpp>
44 
46 
47 template <typename T,
48  typename = scalapackpp::detail::enable_if_scalapack_supported_t<T>>
49 class BlockCyclicMatrix : public madness::WorldObject<BlockCyclicMatrix<T>> {
50  using world_base_t = madness::WorldObject<BlockCyclicMatrix<T>>;
51  using col_major_mat_t =
52  Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;
53 
54  std::shared_ptr<blacspp::Grid> internal_grid_ = nullptr;
55 
56  scalapackpp::BlockCyclicDist2D
57  bc_dist_;
58  col_major_mat_t local_mat_;
59  std::pair<size_t, size_t> dims_;
60 
61  template <typename Tile,
62  typename = std::enable_if_t<
63  TiledArray::detail::is_contiguous_tensor_v<Tile>>>
64  void put_tile(const Tile& tile) {
65  // Extract Tile information
66  const auto* lo = tile.range().lobound_data();
67  const auto* up = tile.range().upbound_data();
68  auto tile_map = eigen_map(tile);
69 
70  // Extract distribution information
71  const auto mb = bc_dist_.mb();
72  const auto nb = bc_dist_.nb();
73 
74  const auto m = dims_.first;
75  const auto n = dims_.second;
76 
77  // Loop over 2D BC compatible blocks
78  long i_extent, j_extent, i_t = 0;
79  for (auto i = lo[0]; i < up[0]; i += i_extent, i_t += i_extent) {
80  long j_t = 0;
81  for (auto j = lo[1]; j < up[1]; j += j_extent, j_t += j_extent) {
82  // Determine indices of start of BC owning block
83  const auto i_block_begin = (i / mb) * mb;
84  const auto j_block_begin = (j / nb) * nb;
85 
86  // Determine indices of end of BC owning block
87  const auto i_block_end =
88  std::min(static_cast<decltype(i)>(m), i_block_begin + mb);
89  const auto j_block_end =
90  std::min(static_cast<decltype(j)>(n), j_block_begin + nb);
91 
92  // Cut block if necessacary to adhere to tile dimensions
93  const auto i_last = std::min(i_block_end, up[0]);
94  const auto j_last = std::min(j_block_end, up[1]);
95 
96  // Calculate extents of the block to be copied
97  i_extent = i_last - i;
98  j_extent = j_last - j;
99 
100  if (bc_dist_.i_own(i, j)) {
101  // Calculate local indices from BC distribution
102  auto [i_local, j_local] = bc_dist_.local_indx(i, j);
103 
104  // Copy the block into local storage
105  local_mat_.block(i_local, j_local, i_extent, j_extent) =
106  tile_map.block(i_t, j_t, i_extent, j_extent);
107 
108  } else {
109  // Send the subblock to a remote rank for processing
110  Tensor<T> subblock;
111  if constexpr (TiledArray::detail::is_ta_tensor_v<Tile>)
112  subblock = tile.block({i, j}, {i_last, j_last});
113  else {
114  auto tile_blk_range = TiledArray::BlockRange(
116  {i_last, j_last});
117  using std::data;
118  auto tile_blk_view =
119  TiledArray::make_const_map(data(tile), tile_blk_range);
120  subblock = tile_blk_view;
121  }
122  world_base_t::send(
123  owner(i, j),
124  &BlockCyclicMatrix<T>::template put_tile<decltype(subblock)>,
125  subblock);
126  }
127 
128  } // for (j)
129  } // for (i)
130 
131  } // put_tile
132 
133  template <typename Tile,
134  typename = std::enable_if_t<
135  TiledArray::detail::is_contiguous_tensor_v<Tile>>>
136  Tile extract_submatrix(std::vector<size_t> lo, std::vector<size_t> up) {
137  assert(bc_dist_.i_own(lo[0], lo[1]));
138 
139  auto [i_st, j_st] = bc_dist_.local_indx(lo[0], lo[1]);
140 
141  auto i_extent = up[0] - lo[0];
142  auto j_extent = up[1] - lo[1];
143 
144  Range range(lo, up);
145  Tile tile(range);
146 
147  auto tile_map = eigen_map(tile);
148 
149  tile_map = local_mat_.block(i_st, j_st, i_extent, j_extent);
150 
151  return tile;
152 
153  } // extract_submatrix
154 
155  public:
166  BlockCyclicMatrix(madness::World& world, const blacspp::Grid& grid, size_t M,
167  size_t N, size_t MB, size_t NB)
168  : world_base_t(world), bc_dist_(grid, MB, NB), dims_{M, N} {
169  // TODO: Check if world / grid are compatible
170 
171  // Determine size of local BC buffer
172  auto [Mloc, Nloc] = bc_dist_.get_local_dims(M, N);
173  local_mat_.resize(Mloc, Nloc);
174  local_mat_.fill(0);
175 
176  world_base_t::process_pending();
177  };
178 
187  template <typename Tile, typename Policy>
189  const blacspp::Grid& grid, size_t MB, size_t NB)
190  : BlockCyclicMatrix(array.world(), grid, array.trange().dim(0).extent(),
191  array.trange().dim(1).extent(), MB, NB) {
192  TA_ASSERT(array.trange().rank() == 2);
193 
194  for (auto it = array.begin(); it != array.end(); ++it) put_tile(it->get());
195  world_base_t::process_pending();
196  }
197 
200 
203 
204  const auto& dist() const { return bc_dist_; }
205  const auto& dims() const { return dims_; }
206  const auto& local_mat() const { return local_mat_; }
207 
208  auto& dist() { return bc_dist_; }
209  auto& dims() { return dims_; }
210  auto& local_mat() { return local_mat_; }
211 
212  inline size_t owner(size_t I, size_t J) const noexcept {
213  return blacspp::coordinate_rank(bc_dist_.grid(),
214  bc_dist_.owner_coordinate(I, J));
215  }
216 
217  template <typename Array>
218  Array tensor_from_matrix(const TiledRange& trange) const {
219  using Tile = typename Array::value_type;
220  auto construct_tile = [&](Tile& tile, const Range& range) {
221  tile = Tile(range);
222 
223  // Extract Tile information
224  const auto* lo = tile.range().lobound_data();
225  const auto* up = tile.range().upbound_data();
226  auto tile_map = eigen_map(tile);
227 
228  // Extract distribution information
229  const auto mb = bc_dist_.mb();
230  const auto nb = bc_dist_.nb();
231 
232  const auto m = dims_.first;
233  const auto n = dims_.second;
234 
235  // Loop over 2D BC compatible blocks
236  size_t i_extent, j_extent;
237  for (size_t i = lo[0], i_t = 0ul; i < up[0];
238  i += i_extent, i_t += i_extent)
239  for (size_t j = lo[1], j_t = 0ul; j < up[1];
240  j += j_extent, j_t += j_extent) {
241  // Determine indices of start of BC owning block
242  const decltype(m) i_block_begin = (i / mb) * mb;
243  const decltype(n) j_block_begin = (j / nb) * nb;
244 
245  // Determine indices of end of BC owning block
246  const auto i_block_end = std::min(m, i_block_begin + mb);
247  const auto j_block_end = std::min(n, j_block_begin + nb);
248 
249  // Cut block if necessacary to adhere to tile dimensions
250  const auto i_last = std::min(i_block_end, static_cast<size_t>(up[0]));
251  const auto j_last = std::min(j_block_end, static_cast<size_t>(up[1]));
252 
253  // Calculate extents of the block to be copied
254  i_extent = i_last - i;
255  j_extent = j_last - j;
256 
257  if (bc_dist_.i_own(i, j)) {
258  // Calculate local indices from BC distribution
259  auto [i_local, j_local] = bc_dist_.local_indx(i, j);
260 
261  // Copy the block into local storage
262  tile_map.block(i_t, j_t, i_extent, j_extent) =
263  local_mat_.block(i_local, j_local, i_extent, j_extent);
264 
265  } else {
266  std::vector<size_t> lo{i, j};
267  std::vector<size_t> up{i_last, j_last};
268  madness::Future<Tensor<T>> remtile_fut = world_base_t::send(
269  owner(i, j),
270  &BlockCyclicMatrix<T>::template extract_submatrix<Tensor<T>>,
271  lo, up);
272 
273  if constexpr (TiledArray::detail::is_ta_tensor_v<Tile>)
274  tile.block(lo, up) = remtile_fut.get();
275  else {
276  auto tile_blk_range = TiledArray::BlockRange(
277  TiledArray::detail::make_ta_range(tile.range()), lo, up);
278  using std::data;
279  auto tile_blk_view =
280  TiledArray::make_map(data(tile), tile_blk_range);
281  tile_blk_view = remtile_fut.get();
282  }
283  }
284  }
285 
286  return norm(tile);
287  };
288 
289  return make_array<Array>(world_base_t::get_world(), trange, construct_tile);
290  }
291 
292 }; // class BlockCyclicMatrix
293 
306 template <typename Array>
307 BlockCyclicMatrix<typename std::remove_cv_t<Array>::element_type>
308 array_to_block_cyclic(const Array& array, const blacspp::Grid& grid, size_t MB,
309  size_t NB) {
311  array, grid, MB, NB);
312 }
313 
324 template <typename Array>
325 std::remove_cv_t<Array> block_cyclic_to_array(
326  const BlockCyclicMatrix<typename std::remove_cv_t<Array>::element_type>&
327  matrix,
328  const TiledRange& trange) {
329  return matrix.template tensor_from_matrix<std::remove_cv_t<Array>>(trange);
330 }
331 
332 } // namespace TiledArray::math::linalg::scalapack
333 
334 #endif // TILEDARRAY_HAS_SCALAPACK
335 #endif // TILEDARRAY_MATH_LINALG_SCALAPACK_TO_BLOCKCYCLIC_H__INCLUDED
impl_type::value_type value_type
Tile type.
Definition: dist_array.h:86
BlockCyclicMatrix(const DistArray< Tile, Policy > &array, const blacspp::Grid &grid, size_t MB, size_t NB)
Construct a BlockCyclic matrix from a DistArray.
Definition: block_cyclic.h:188
BlockCyclicMatrix & operator=(BlockCyclicMatrix &&)=default
Eigen::Map< const Eigen::Matrix< typename T::value_type, Eigen::Dynamic, Eigen::Dynamic, Storage >, Eigen::AutoAlign > eigen_map(const T &tensor, const std::size_t m, const std::size_t n)
Construct a const Eigen::Map object for a given Tensor object.
Definition: eigen.h:77
KroneckerDeltaTile< _N >::numeric_type min(const KroneckerDeltaTile< _N > &arg)
TensorConstMap< T, Range_, OpResult > make_const_map(const T *const data, const Index &lower_bound, const Index &upper_bound)
Definition: tensor_map.h:98
decltype(auto) range() const
Range accessor.
Definition: tile.h:213
std::remove_cv_t< Array > block_cyclic_to_array(const BlockCyclicMatrix< typename std::remove_cv_t< Array >::element_type > &matrix, const TiledRange &trange)
Convert a block-cyclic matrix to DistArray.
Definition: block_cyclic.h:325
TensorMap< T, Range_, OpResult > make_map(T *const data, const Index &lower_bound, const Index &upper_bound)
Definition: tensor_map.h:53
BlockCyclicMatrix(madness::World &world, const blacspp::Grid &grid, size_t M, size_t N, size_t MB, size_t NB)
Construct and allocate memory for a BlockCyclic matrix.
Definition: block_cyclic.h:166
decltype(auto) norm(const Tile< Arg > &arg)
Vector 2-norm of a tile.
Definition: tile.h:1527
Range that references a subblock of another range.
Definition: block_range.h:34
decltype(auto) block(const Index1 &lower_bound, const Index2 &upper_bound)
Constructs a view of the block defined by lower_bound and upper_bound.
Definition: tile.h:423
iterator begin()
Begin iterator factory function.
Definition: dist_array.h:484
BlockCyclicMatrix & operator=(const BlockCyclicMatrix &)=default
#define TA_ASSERT(EXPR,...)
Definition: error.h:39
Range data of a tiled array.
Definition: tiled_range.h:32
const trange_type & trange() const
Tiled range accessor.
Definition: dist_array.h:917
const TiledArray::Range & make_ta_range(const TiledArray::Range &range)
Definition: btas.h:58
Forward declarations.
Definition: dist_array.h:57
iterator end()
End iterator factory function.
Definition: dist_array.h:498
size_t owner(size_t I, size_t J) const noexcept
Definition: block_cyclic.h:212
Array tensor_from_matrix(const TiledRange &trange) const
Definition: block_cyclic.h:218
BlockCyclicMatrix< typename std::remove_cv_t< Array >::element_type > array_to_block_cyclic(const Array &array, const blacspp::Grid &grid, size_t MB, size_t NB)
Convert a dense DistArray to block-cyclic storage format.
Definition: block_cyclic.h:308
An N-dimensional tensor object.
Definition: tensor.h:50
An N-dimensional shallow copy wrapper for tile objects.
Definition: tile.h:82
BlockCyclicMatrix(const BlockCyclicMatrix &)=default
A (hyperrectangular) interval on , space of integer -indices.
Definition: range.h:46