summa: fix h-grouped result-tile owner; enable the 3-d grid by default

The deadlock that gated the h-grouped (hgroups>1) path: Summa's output
get_tile(i) computed the result tile's source rank from the GROUP-LOCAL
process-grid coordinates (proc_row*proc_cols + proc_col) but never added
the world-rank offset of the tile's slab group. So a non-zero group
recv'd each result tile from the wrong (group-0) rank, which never sent
it -> hang. set_tile already routed correctly via the h-grouped pmap;
only get_tile disagreed.

Fix: result_tile_owner(i) = within-group cyclic owner + (slab_group *
group_size), matching the h-grouped SlabbedPmap's owner() (and thus
set_tile's destination). Store group_size_ for the offset. hgroups==1 is
unchanged (offset 0). With this the heuristic is enabled by default
(spread surplus ranks over the slab dimension when M*N tiles < P), so
no-external products distribute their result across slab groups instead
of piling on one rank.

Validated np=2: all general_product no-external cases (dense/ToT/sparse)
+ the distributed suite (now incl. dist_no_externals_hgrouped, which
asserts the result spans >1 rank). np=1 + full regression unchanged.

Author: evaleev

src/TiledArray/dist_eval/contraction_eval.h

@@ -134,11 +134,33 @@ class Summa
   // constructed over the group's rank interval). hgroups_ == 1 is the
   // ordinary shared-grid batched contraction.
   const ordinal_type hgroups_;     ///< Number of slab (h) groups
+  const ordinal_type group_size_;  ///< World ranks per slab group (the
+                                   ///< group of slab h spans world ranks
+                                   ///< [(h % hgroups_) * group_size_, ...))
   const ordinal_type first_slab_;  ///< This rank's group's first slab (== its
                                    ///< group index), or nh_ if this rank is
                                    ///< in no group (idle for this eval)
   const ordinal_type my_slabs_;    ///< Number of slabs of this rank's group
 
+  /// \return the world rank that owns result tile \p i: the within-group
+  /// owner (from the group-local process grid) shifted by the world-rank
+  /// offset of the group that owns \p i's slab. For hgroups_ == 1 the offset
+  /// is 0 and this is the ordinary cyclic owner.
+  ProcessID result_tile_owner(const ordinal_type i) const {
+    const ordinal_type source_index = DistEvalImpl_::perm_index_to_source(i);
+    // owner is independent of slab index *within a group*
+    const ordinal_type slab_index = source_index % result_slab_size_;
+    const ordinal_type tile_row = slab_index / proc_grid_.cols();
+    const ordinal_type tile_col = slab_index % proc_grid_.cols();
+    const ordinal_type proc_row = tile_row % proc_grid_.proc_rows();
+    const ordinal_type proc_col = tile_col % proc_grid_.proc_cols();
+    const ProcessID within_group = proc_row * proc_grid_.proc_cols() + proc_col;
+    // shift by the offset of the group that owns this tile's slab
+    const ordinal_type slab = source_index / result_slab_size_;
+    const ordinal_type group = (hgroups_ > 1ul) ? (slab % hgroups_) : 0ul;
+    return ProcessID(group * group_size_) + within_group;
+  }
+
   /// \return the slab index of SUMMA step \p s
   ordinal_type step_h(const ordinal_type s) const { return s / k_; }
   /// \return the within-slab inner-dimension index of SUMMA step \p s
@@ -1728,6 +1750,7 @@ class Summa
         right_slab_size_(right.size() / nh),
         result_slab_size_(proc_grid.rows() * proc_grid.cols()),
         hgroups_(hgroups),
+        group_size_(group_size),
         first_slab_(compute_first_slab(world, nh, hgroups, group_size)),
         my_slabs_(first_slab_ < nh
                       ? (nh - first_slab_ + hgroups - 1ul) / hgroups
@@ -1777,18 +1800,11 @@ class Summa
     TA_ASSERT(TensorImpl_::is_local(i));
     TA_ASSERT(!TensorImpl_::is_zero(i));
 
-    const ordinal_type source_index = DistEvalImpl_::perm_index_to_source(i);
-
-    // Compute tile coordinate in tile grid (the owner of a tile is
-    // independent of its slab index)
-    const ordinal_type slab_index = source_index % result_slab_size_;
-    const ordinal_type tile_row = slab_index / proc_grid_.cols();
-    const ordinal_type tile_col = slab_index % proc_grid_.cols();
-    // Compute process coordinate of tile in the process grid
-    const ordinal_type proc_row = tile_row % proc_grid_.proc_rows();
-    const ordinal_type proc_col = tile_col % proc_grid_.proc_cols();
-    // Compute the process that owns tile
-    const ProcessID source = proc_row * proc_grid_.proc_cols() + proc_col;
+    // The process that owns tile i: the within-group cyclic owner shifted by
+    // the world-rank offset of the tile's slab group (see
+    // result_tile_owner). For hgroups_ == 1 this is the ordinary cyclic
+    // owner over the whole world.
+    const ProcessID source = result_tile_owner(i);
 
     const madness::DistributedID key(DistEvalImpl_::id(), i);
     return TensorImpl_::world().gop.template recv<value_type>(source, key);

src/TiledArray/expressions/cont_engine.h

@@ -859,29 +859,19 @@ class ContEngine : public BinaryEngine<Derived> {
       // h % hgroups_ of hgroup_size_ contiguous ranks; the division
       // remainder ranks idle for this evaluation). For a no-external
       // product (M == N == 1) this degenerates to an effectively 1-d grid
-      // over the slabs.
-      //
-      // N.B. EXPERIMENTAL / WIP: the h-grouped (hgroups_ > 1) path is not
-      // yet correct at np > 1 -- a group other than group 0 sets its result
-      // tiles but the reduce-task result future never resolves (operands are
-      // ready, the reduce receives its adds + submit, yet completion never
-      // fires; see docs/superpowers/plans/2026-06-12-ta-summa-3d-grid.md for
-      // the precise deadlock state). Until that is resolved hgroups_ defaults
-      // to 1 (no behavior change: no-external products evaluate via the
-      // existing degenerate single-grid path, correct but not h-parallel).
-      // Set TA_SUMMA_HGROUPS=<n> to opt into the WIP path for debugging.
+      // over the slabs. TA_SUMMA_HGROUPS overrides the heuristic.
       const size_type P = world->size();
       hgroups_ = 1ul;
       if (n_slabs_ > 1ul && P > 1ul) {
+        const size_type p2d_cap = std::min<size_type>(P, M * N);
+        hgroups_ = std::min<size_type>(n_slabs_,
+                                       std::max<size_type>(1ul, P / p2d_cap));
         if (const char* hg = std::getenv("TA_SUMMA_HGROUPS")) {
           const long requested = std::atol(hg);
           if (requested >= 1)
             hgroups_ = std::min<size_type>(
                 n_slabs_, std::min<size_type>(P, size_type(requested)));
         }
-        // default heuristic (currently disabled pending the WIP fix):
-        //   const size_type p2d_cap = std::min<size_type>(P, M * N);
-        //   hgroups_ = std::min(n_slabs_, std::max(1ul, P / p2d_cap));
       }
       hgroup_size_ = P / hgroups_;
 

tests/general_product.cpp

@@ -8,6 +8,8 @@
 #include "tiledarray.h"
 #include "unit_test_config.h"
 
+#include <set>
+
 BOOST_AUTO_TEST_SUITE(general_product_suite, TA_UT_LABEL_SERIAL)
 
 namespace TA = TiledArray;
@@ -1455,4 +1457,29 @@ BOOST_AUTO_TEST_CASE(dist_inner_node_thc) {
   BOOST_CHECK_SMALL(diff_norm(g, g_ref, "p,q,r,s"), 1e-10);
 }
 
+BOOST_AUTO_TEST_CASE(dist_no_externals_hgrouped) {
+  // a no-external product with several slabs: at np>1 the heuristic engages
+  // the 3-d (h-grouped) grid (M=N=1 ⇒ each slab gets its own rank group),
+  // so the result tiles are distributed across the slab groups rather than
+  // piled on one rank. Exercises the cross-group result-tile transfer.
+  auto& world = TA::get_default_world();
+  ForceLegacyEinsum legacy_oracle;
+  TA::TiledRange tr{{0, 2, 4, 6}, {0, 3, 5}};  // i (3 tiles), j
+  auto a = make_patterned_array(world, tr, 1.0);
+  auto b = make_patterned_array(world, tr, 2.0);
+  TA::TArrayD c;
+  BOOST_REQUIRE_NO_THROW(c("i") = a("i,j") * b("i,j"));
+  auto c_ref = TA::einsum(a("i,j"), b("i,j"), "i");
+  BOOST_CHECK_SMALL(diff_norm(c, c_ref, "i"), 1e-10);
+
+  // at np>1 the no-external result must NOT all live on one rank (the
+  // degeneracy the 3-d grid fixes)
+  if (world.size() > 1) {
+    std::set<std::size_t> owners;
+    for (std::size_t o = 0; o < c.trange().tiles_range().volume(); ++o)
+      owners.insert(c.pmap()->owner(o));
+    BOOST_CHECK_GT(owners.size(), 1u);
+  }
+}
+
 BOOST_AUTO_TEST_SUITE_END()