proc_grid/pmap: rank-subset ProcGrid + 3-d (proc_h) SlabbedPmap

Infrastructure for a 3-d (proc_h x proc_r x proc_c) batched-Summa grid
that distributes the fused/batch (h, slab) dimension of a general product
across process planes:

- ProcGrid gains a rank-subset constructor (tagged rank_subset to avoid
  colliding with the same-arity test-only ctor) that builds a 2-d grid
  over a contiguous interval [rank_offset, rank_offset + nprocs) of the
  world's ranks; map_row/map_col and the row/col group factories emit
  world-correct ranks via the offset. The legacy full-world ctor is
  unchanged (offset 0).
- SlabbedPmap gains a 3-d variant (proc_h, proc_h_stride): slab h belongs
  to plane h % proc_h of proc_h_stride contiguous ranks, and the per-slab
  base map's plane-local owners are offset by the slab's plane. The
  original 3-argument form (proc_h == 1, slab-replicated) is unchanged.

Author: evaleev

src/TiledArray/pmap/slabbed_pmap.h

@@ -44,6 +44,16 @@ namespace detail {
 /// (Hadamard/batch) modes as their leading dimensions, every slab is
 /// distributed identically over the same 2-d process grid, and the slab
 /// index never participates in inter-process communication patterns.
+/// The 3-d-grid variant adds a process dimension \c proc_h along the slab
+/// (h) axis: the world's first `proc_h * proc_h_stride` ranks are partitioned
+/// into `proc_h` contiguous h-planes of `proc_h_stride` ranks each; slab
+/// \f$ h \f$ belongs to plane \f$ h \% proc\_h \f$, and within its plane is
+/// distributed by the base map, whose owners must then be PLANE-LOCAL ranks
+/// in \f$ [0, proc\_h\_stride) \f$: the owner of index \f$ o \f$ is
+/// \f$ (h \% proc\_h) \cdot proc\_h\_stride + base(o \% S) \f$.
+/// `proc_h == 1` reduces to the slab-replicated map above (and is constructed
+/// via the 2-argument-base constructor, which delegates locality to the base
+/// map).
 class SlabbedPmap : public Pmap {
  protected:
   // Import Pmap protected variables
@@ -56,6 +66,14 @@ class SlabbedPmap : public Pmap {
   const std::shared_ptr<const Pmap> base_;  ///< The per-slab base map
   const size_type slab_size_;               ///< The number of indices per slab
   const size_type nslabs_;                  ///< The number of slabs
+  const size_type proc_h_ = 1;         ///< Process-grid extent along the slab
+                                       ///< (h) axis (the number of h-planes)
+  const size_type proc_h_stride_ = 0;  ///< Ranks per h-plane (0 = the whole
+                                       ///< world; base owners are world ranks
+                                       ///< and locality delegates to the base)
+
+  /// \return whether this map distributes the slab axis over a process plane
+  bool hgrouped() const { return proc_h_stride_ != 0ul; }
 
  public:
   typedef Pmap::size_type size_type;  ///< Size type
@@ -76,6 +94,49 @@ class SlabbedPmap : public Pmap {
     this->local_size_ = base_->local_size() * nslabs_;
   }
 
+  /// Construct an h-grouped (3-d grid) slabbed process map
+
+  /// \param world The world where the tiles will be mapped
+  /// \param base The base process map, defined over one slab; its owners
+  ///        must be GROUP-LOCAL ranks in [0, proc_h_stride)
+  /// \param nslabs The number of slabs
+  /// \param proc_h The number of slab groups (slab h -> group h % proc_h)
+  /// \param proc_h_stride The number of (contiguous) world ranks per group
+  SlabbedPmap(World& world, std::shared_ptr<const Pmap> base,
+              const size_type nslabs, const size_type proc_h,
+              const size_type proc_h_stride)
+      : Pmap(world, base->size() * nslabs),
+        base_(std::move(base)),
+        slab_size_(base_->size()),
+        nslabs_(nslabs),
+        proc_h_(proc_h),
+        proc_h_stride_(proc_h_stride) {
+    TA_ASSERT(base_);
+    TA_ASSERT(nslabs_ > 0ul);
+    TA_ASSERT(proc_h_ > 0ul);
+    TA_ASSERT(proc_h_stride_ > 0ul);
+    TA_ASSERT(proc_h_ * proc_h_stride_ <= size_type(world.size()));
+
+    // this rank's group and group-local rank; ranks beyond the grouped
+    // prefix of the world own nothing
+    const size_type rank = world.rank();
+    if (rank < proc_h_ * proc_h_stride_) {
+      const size_type my_group = rank / proc_h_stride_;
+      const size_type my_group_rank = rank % proc_h_stride_;
+      // count of my group's slabs: h in [0, nslabs) with h % proc_h ==
+      // my_group
+      const size_type my_slabs =
+          (nslabs_ / proc_h_) + (my_group < (nslabs_ % proc_h_) ? 1u : 0u);
+      // count of slab indices the base assigns to my group-local rank
+      size_type base_local = 0ul;
+      for (size_type j = 0ul; j < slab_size_; ++j)
+        if (base_->owner(j) == my_group_rank) ++base_local;
+      this->local_size_ = my_slabs * base_local;
+    } else {
+      this->local_size_ = 0ul;
+    }
+  }
+
   virtual ~SlabbedPmap() {}
 
   /// \return the per-slab base map
@@ -84,25 +145,32 @@ class SlabbedPmap : public Pmap {
   size_type slab_size() const { return slab_size_; }
   /// \return the number of slabs
   size_type nslabs() const { return nslabs_; }
+  /// \return the number of slab (h) groups
+  size_type proc_h() const { return proc_h_; }
 
   /// Maps \c tile to the process that owns it
 
   /// \param tile The tile to be queried
   /// \return Process that logically owns \c tile
   virtual size_type owner(const size_type tile) const {
     TA_ASSERT(tile < size_);
-    return base_->owner(tile % slab_size_);
+    if (!hgrouped()) return base_->owner(tile % slab_size_);
+    const size_type group = (tile / slab_size_) % proc_h_;
+    return group * proc_h_stride_ + base_->owner(tile % slab_size_);
   }
 
   /// Check that the tile is owned by this process
 
   /// \param tile The tile to be checked
   /// \return \c true if \c tile is owned by this process, otherwise \c false
   virtual bool is_local(const size_type tile) const {
-    return base_->is_local(tile % slab_size_);
+    if (!hgrouped()) return base_->is_local(tile % slab_size_);
+    return owner(tile) == size_type(rank_);
   }
 
-  virtual bool known_local_size() const { return base_->known_local_size(); }
+  virtual bool known_local_size() const {
+    return hgrouped() ? true : base_->known_local_size();
+  }
 
   virtual const_iterator begin() const {
     return Iterator(*this, 0ul, size_, 0ul, /*checking=*/true);

src/TiledArray/proc_grid.h

@@ -55,6 +55,14 @@ namespace detail {
 /// \f]
 /// where the positive, real root of \f$P_{\rm{row}}\f$ give the optimal
 /// optimal communication time.
+
+/// Tag to disambiguate the rank-subset ProcGrid constructor from the
+/// (same-arity) test-only constructor.
+struct rank_subset_t {
+  explicit rank_subset_t() = default;
+};
+inline constexpr rank_subset_t rank_subset{};
+
 class ProcGrid {
  public:
   typedef uint_fast32_t size_type;
@@ -71,9 +79,12 @@ class ProcGrid {
                         ///<  may be less than the number of processes in world.
   ProcessID rank_row_;  ///< This process's row in the process grid
   ProcessID rank_col_;  ///< This process's column in the process grid
-  size_type local_rows_;  ///< The number of local element rows
-  size_type local_cols_;  ///< The number of local element columns
-  size_type local_size_;  ///< Number of local elements
+  ProcessID rank_offset_ = 0;  ///< World rank of the grid's first process
+                               ///< (nonzero for a grid over a contiguous
+                               ///< subset of the world's ranks)
+  size_type local_rows_;       ///< The number of local element rows
+  size_type local_cols_;       ///< The number of local element columns
+  size_type local_size_;       ///< Number of local elements
 
   /// Compute the number of process rows that minimizes communication
 
@@ -180,14 +191,16 @@ class ProcGrid {
       proc_cols_ = 1u;
       proc_size_ = 1u;
 
-      // Set this process rank
-      rank_row_ = 0;
-      rank_col_ = 0;
+      if (rank < proc_size_) {
+        // Set this process rank
+        rank_row_ = 0;
+        rank_col_ = 0;
 
-      // Set local counts
-      local_rows_ = rows_;
-      local_cols_ = cols_;
-      local_size_ = size_;
+        // Set local counts
+        local_rows_ = rows_;
+        local_cols_ = cols_;
+        local_size_ = size_;
+      }
 
     } else if (size_ <= nprocs) {  // Max one tile per process
 
@@ -291,6 +304,48 @@ class ProcGrid {
     init(world_->rank(), world_->size(), row_size, col_size);
   }
 
+  /// Construct a process grid over a contiguous subset of the world's ranks
+
+  /// The grid spans world ranks [rank_offset, rank_offset + nprocs); ranks
+  /// outside that interval construct a valid "not in the grid" instance
+  /// (zero local sizes, empty groups). Used by the h-grouped (3-d) batched
+  /// Summa, where each fused-index slab group runs its own 2-d grid.
+  /// \param world The world where the process grid will live
+  /// \param rank_offset The world rank of the grid's first process
+  /// \param nprocs The number of processes spanned by the grid
+  /// \param rows The number of tile rows
+  /// \param cols The number of tile columns
+  /// \param row_size The number of element rows
+  /// \param col_size The number of element columns
+  ProcGrid(World& world, rank_subset_t, const ProcessID rank_offset,
+           const size_type nprocs, const size_type rows, const size_type cols,
+           const std::size_t row_size, const std::size_t col_size)
+      : world_(&world),
+        rows_(rows),
+        cols_(cols),
+        size_(rows_ * cols_),
+        proc_rows_(0ul),
+        proc_cols_(0ul),
+        proc_size_(0ul),
+        rank_row_(-1),
+        rank_col_(-1),
+        rank_offset_(rank_offset),
+        local_rows_(0ul),
+        local_cols_(0ul),
+        local_size_(0ul) {
+    TA_ASSERT(rank_offset >= 0);
+    TA_ASSERT(nprocs >= 1u);
+    TA_ASSERT(rank_offset + nprocs <= size_type(world.size()));
+    const auto world_rank = world.rank();
+    // out-of-grid ranks pass rank == nprocs, which every init() branch
+    // treats as "not in the grid"
+    const size_type rank = (world_rank >= rank_offset &&
+                            world_rank < rank_offset + ProcessID(nprocs))
+                               ? world_rank - rank_offset
+                               : nprocs;
+    init(rank, nprocs, row_size, col_size);
+  }
+
 #ifdef TILEDARRAY_ENABLE_TEST_PROC_GRID
   // Note: The following function is here for testing purposes only. It
   // has the same functionality as the default constructor above, except the
@@ -350,6 +405,7 @@ class ProcGrid {
         proc_size_(other.proc_size_),
         rank_row_(other.rank_row_),
         rank_col_(other.rank_col_),
+        rank_offset_(other.rank_offset_),
         local_rows_(other.local_rows_),
         local_cols_(other.local_cols_),
         local_size_(other.local_size_) {}
@@ -367,6 +423,7 @@ class ProcGrid {
     proc_size_ = other.proc_size_;
     rank_row_ = other.rank_row_;
     rank_col_ = other.rank_col_;
+    rank_offset_ = other.rank_offset_;
     local_rows_ = other.local_rows_;
     local_cols_ = other.local_cols_;
     local_size_ = other.local_size_;
@@ -447,7 +504,7 @@ class ProcGrid {
       // Populate the row process list
       size_type p = rank_row_ * proc_cols_;
       const size_type row_end = p + proc_cols_;
-      for (; p < row_end; ++p) proc_list.push_back(p);
+      for (; p < row_end; ++p) proc_list.push_back(p + rank_offset_);
 
       // Construct the group
       group = madness::Group(*world_, proc_list, did);
@@ -472,7 +529,7 @@ class ProcGrid {
 
       // Populate the column process list
       for (size_type p = rank_col_; p < proc_size_; p += proc_cols_)
-        proc_list.push_back(p);
+        proc_list.push_back(p + rank_offset_);
 
       // Construct the group
       if (proc_list.size() != 0)
@@ -489,7 +546,7 @@ class ProcGrid {
   /// (row,rank_col)
   ProcessID map_row(const size_type row) const {
     TA_ASSERT(row < proc_rows_);
-    return rank_col_ + row * proc_cols_;
+    return rank_col_ + row * proc_cols_ + rank_offset_;
   }
 
   /// Map a column to the process in this process's row
@@ -499,7 +556,7 @@ class ProcGrid {
   /// (rank_row,col)
   ProcessID map_col(const size_type col) const {
     TA_ASSERT(col < proc_cols_);
-    return rank_row_ * proc_cols_ + col;
+    return rank_row_ * proc_cols_ + col + rank_offset_;
   }
 
   /// Construct a cyclic process