CA-Red2Band: Communication-Avoiding Reduction To Band

include/dlaf/communication/broadcast_panel.h

@@ -18,11 +18,13 @@
 #include <pika/execution.hpp>
 
 #include <dlaf/common/index2d.h>
+#include <dlaf/common/range2d.h>
 #include <dlaf/communication/communicator_pipeline.h>
 #include <dlaf/communication/index.h>
 #include <dlaf/communication/kernels/broadcast.h>
 #include <dlaf/communication/message.h>
 #include <dlaf/matrix/copy_tile.h>
+#include <dlaf/matrix/index.h>
 #include <dlaf/matrix/panel.h>
 #include <dlaf/matrix/tile.h>
 #include <dlaf/types.h>
@@ -89,7 +91,6 @@ auto& get_taskchain(comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row
     return col_task_chain;
   }
 }
-}  // namespace internal
 
 /// Broadcast
 ///
@@ -105,9 +106,6 @@ auto& get_taskchain(comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row
 /// - linking as external tile, if the tile is already available locally for the rank
 /// - receiving the tile from the owning rank (via a broadcast)
 ///
-/// Be aware that the last tile will just be available on @p panel, but it won't be transposed to
-/// @p panelT.
-///
 /// @param rank_root specifies on which rank the @p panel is the source of the data
 /// @param panel
 ///   on rank_root it is the source panel (a)
@@ -125,7 +123,8 @@ template <class T, Device D, Coord axis, matrix::StoreTransposed storage,
 void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& panel,
                matrix::Panel<orthogonal(axis), T, D, storageT>& panelT,
                comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row_task_chain,
-               comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain) {
+               comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain,
+               common::IterableRange2D<SizeType, matrix::LocalTile_TAG> range) {
   constexpr Coord axisT = orthogonal(axis);
 
   constexpr Coord coord = std::decay_t<decltype(panel)>::coord;
@@ -183,13 +182,6 @@ void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& p
 
   auto& chain_step2 = internal::get_taskchain<comm_dir_step2>(row_task_chain, col_task_chain);
 
-  const SizeType last_tile = std::max(panelT.rangeStart(), panelT.rangeEnd() - 1);
-  const auto owner = dist.template rankGlobalTile<coordT>(last_tile);
-  const auto range = dist.rankIndex().get(coordT) == owner
-                         ? common::iterate_range2d(*panelT.iteratorLocal().begin(),
-                                                   LocalTileIndex(coordT, panelT.rangeEndLocal() - 1, 1))
-                         : panelT.iteratorLocal();
-
   for (const auto& indexT : range) {
     auto [index_diag, owner_diag] = internal::transposedOwner<coordT>(dist, indexT);
 
@@ -208,6 +200,39 @@ void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& p
     }
   }
 }
+}  // namespace internal
+
+template <class T, Device D, Coord axis, matrix::StoreTransposed storage,
+          matrix::StoreTransposed storageT, class = std::enable_if_t<!std::is_const_v<T>>>
+void broadcast(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& panel,
+               matrix::Panel<orthogonal(axis), T, D, storageT>& panelT,
+               comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row_task_chain,
+               comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain) {
+  constexpr Coord coordT = std::decay_t<decltype(panelT)>::coord;
 
+  const auto& dist = panel.parentDistribution();
+
+  const SizeType last_tile = std::max(panelT.rangeStart(), panelT.rangeEnd() - 1);
+
+  if (panel.rangeStart() == panel.rangeEnd())
+    return;
+
+  const auto owner = dist.template rankGlobalTile<coordT>(last_tile);
+  const auto range = dist.rankIndex().get(coordT) == owner
+                         ? common::iterate_range2d(*panelT.iteratorLocal().begin(),
+                                                   LocalTileIndex(coordT, panelT.rangeEndLocal() - 1, 1))
+                         : panelT.iteratorLocal();
+
+  internal::broadcast(rank_root, panel, panelT, row_task_chain, col_task_chain, range);
+}
+
+template <class T, Device D, Coord axis, matrix::StoreTransposed storage,
+          matrix::StoreTransposed storageT, class = std::enable_if_t<!std::is_const_v<T>>>
+void broadcast_all(comm::IndexT_MPI rank_root, matrix::Panel<axis, T, D, storage>& panel,
+                   matrix::Panel<orthogonal(axis), T, D, storageT>& panelT,
+                   comm::CommunicatorPipeline<comm::CommunicatorType::Row>& row_task_chain,
+                   comm::CommunicatorPipeline<comm::CommunicatorType::Col>& col_task_chain) {
+  internal::broadcast(rank_root, panel, panelT, row_task_chain, col_task_chain, panelT.iteratorLocal());
+}
 }
 }

include/dlaf/eigensolver/reduction_to_band.h

@@ -119,4 +119,18 @@ Matrix<T, Device::CPU> reduction_to_band(comm::CommunicatorGrid& grid, Matrix<T,
 
   return ReductionToBand<B, D, T>::call(grid, mat_a, band_size);
 }
+
+template <Backend B, Device D, class T>
+internal::CARed2BandResult<T, D> ca_reduction_to_band(comm::CommunicatorGrid& grid, Matrix<T, D>& mat_a,
+                                                      const SizeType band_size) {
+  DLAF_ASSERT(matrix::square_size(mat_a), mat_a);
+  DLAF_ASSERT(matrix::square_blocksize(mat_a), mat_a);
+  DLAF_ASSERT(matrix::single_tile_per_block(mat_a), mat_a);
+  DLAF_ASSERT(matrix::equal_process_grid(mat_a, grid), mat_a, grid);
+
+  DLAF_ASSERT(band_size >= 2, band_size);
+  DLAF_ASSERT(mat_a.blockSize().rows() % band_size == 0, mat_a.blockSize().rows(), band_size);
+
+  return CAReductionToBand<B, D, T>::call(grid, mat_a, band_size);
+}
 }

include/dlaf/eigensolver/reduction_to_band/api.h

@@ -24,9 +24,24 @@ struct ReductionToBand {
                                      const SizeType band_size);
 };
 
+template <class T, Device D>
+struct CARed2BandResult {
+  Matrix<T, Device::CPU> taus_1st;
+  // hh_1st are stored in-place
+  Matrix<T, Device::CPU> taus_2nd;
+  Matrix<T, D> hh_2nd;
+};
+
+template <Backend B, Device D, class T>
+struct CAReductionToBand {
+  static CARed2BandResult<T, D> call(comm::CommunicatorGrid& grid, Matrix<T, D>& mat_a,
+                                     const SizeType band_size);
+};
+
 // ETI
 #define DLAF_EIGENSOLVER_REDUCTION_TO_BAND_ETI(KWORD, BACKEND, DEVICE, DATATYPE) \
-  KWORD template struct ReductionToBand<BACKEND, DEVICE, DATATYPE>;
+  KWORD template struct ReductionToBand<BACKEND, DEVICE, DATATYPE>;              \
+  KWORD template struct CAReductionToBand<BACKEND, DEVICE, DATATYPE>;
 
 DLAF_EIGENSOLVER_REDUCTION_TO_BAND_ETI(extern, Backend::MC, Device::CPU, float)
 DLAF_EIGENSOLVER_REDUCTION_TO_BAND_ETI(extern, Backend::MC, Device::CPU, double)