Coarsen the SUMMA (non-BLAS) axis in mixed retile

The mixed (plain x arena-ToT -> arena-ToT) auto-retile coarsened only the BLAS
axes (plain external + contracted K). Extend it to also coarsen the leftover
SUMMA external (the arena-ToT operand's external).

mixed_retile_config gains tot_external_target (default 16): the SUMMA external
is coarsened toward it, coarsen-only (an axis already coarser is kept intact;
refine is unsupported). Detection stays automatic and mixed-only, so ToT*ToT
still reverts to the stock SUMMA path.

Adds g_summa_coarse_{m,n,k}_grid witnesses and tests for both orientations, the
default auto path, and the ToT*ToT no-retile guard.

Author: zhihao-deng

src/TiledArray/dist_eval/contraction_eval.h

@@ -81,6 +81,20 @@ inline std::atomic<std::size_t> g_summa_result_merge_count{0};
 /// tests assert it is > 0 to prove the grouped active distribution engaged.
 /// gop.sum it across ranks for np-correctness.
 inline std::atomic<std::size_t> g_summa_proc_h_grouped_calls{0};
+
+/// witness: the COARSE SUMMA-M / -N / -K grid tile counts (the number of T
+/// tiles the operands are retiled into on each role), recorded by
+/// `ContEngine::init_distribution` whenever a retile plan is active. These
+/// directly expose the "retiled operand trange" tile count on each SUMMA axis:
+/// coarsening SUMMA-N from 64 fine tiles to 4 coarse tiles records
+/// `g_summa_coarse_n_grid == 4`, even though the per-row strided BLAS GEMM count
+/// is INVARIANT (a SUMMA external is not a BLAS axis). STORE (not accumulate) --
+/// the last active contraction's coarse grid; reset to 0 before the contraction
+/// under test. They stay at the reset value on the inactive (stock SUMMA) path.
+/// np=1 or symmetric np>1, so the per-rank store agrees across ranks.
+inline std::atomic<std::size_t> g_summa_coarse_m_grid{0};
+inline std::atomic<std::size_t> g_summa_coarse_n_grid{0};
+inline std::atomic<std::size_t> g_summa_coarse_k_grid{0};
 #endif
 
 /// \brief Distributed contraction evaluator implementation

src/TiledArray/expressions/cont_engine.h

@@ -595,19 +595,27 @@ class ContEngine : public BinaryEngine<Derived> {
   /// inactive plan. The plan is consumed by later phases; in this phase it is
   /// only stored and threaded to the Summa ctor.
   /// Synthesize coarse retile targets for the mixed (plain x arena-ToT ->
-  /// arena-ToT) shape with the env-gated auto-retile. Collapse the PLAIN
-  /// operand's external -- role M when the plain operand is LEFT, role N when it
-  /// is RIGHT -- and the contracted (K) axis toward the element sizes in
-  /// detail::mixed_retile_config (0 => single tile / full collapse). The fused
-  /// (H) axes and the arena-ToT operand's external (the other of M/N) are left
-  /// EMPTY, i.e. kept intact at U. The role axes are extracted by the SAME
-  /// positional H/M/N/K partition that make_retile_plan consults (nf leading
-  /// fused modes; the last nc left axes are K; the remaining left axes are M;
-  /// the right axes past nf+nc are N), so the synthesized targets are in
-  /// canonical H/M/N/K space -- identical to what a .retile() caller supplies.
-  /// This reproduces the verified mixed_T_x_ToT_coarsen_MK (plain left ->
-  /// coarsen M+K) and mixed_ToT_x_T_coarsen_NK (plain right -> coarsen N+K)
-  /// target strategy.
+  /// arena-ToT) shape with the env-gated auto-retile. Per
+  /// detail::mixed_retile_config, coarsen THREE roles:
+  ///   - the PLAIN operand's external (role M when the plain operand is LEFT,
+  ///     role N when it is RIGHT) toward plain_external_target (default 0 =>
+  ///     single tile / full collapse);
+  ///   - the contracted (K) axis toward contracted_target (default 0 => collapse);
+  ///   - the arena-ToT operand's external -- the LEFTOVER SUMMA axis (role N
+  ///     when the plain operand is LEFT, role M when it is RIGHT) -- toward
+  ///     tot_external_target (default 16).
+  /// The fused (H) axes are left EMPTY (kept intact at U). Every target is fed
+  /// through coarsen_tr1, which COARSENS ONLY (merges U tiles onto existing U
+  /// boundaries; never refines), so a leftover-SUMMA axis already coarser than
+  /// the target is kept intact rather than refined (refine is unsupported).
+  /// The role axes are extracted by the SAME positional H/M/N/K partition that
+  /// make_retile_plan consults (nf leading fused modes; the last nc left axes
+  /// are K; the remaining left axes are M; the right axes past nf+nc are N), so
+  /// the synthesized targets are in canonical H/M/N/K space -- identical to what
+  /// a .retile() caller supplies. This is the default verified by
+  /// mixed_T_x_ToT_coarsen_MK_retile_N (plain left -> collapse M+K, coarsen N to
+  /// 16) and mixed_ToT_x_T_coarsen_NK_retile_M (plain right -> collapse N+K,
+  /// coarsen M to 16).
   void synthesize_mixed_targets_(bool left_is_plain,
                                  const math::GemmHelper& outer_gh,
                                  std::vector<TiledRange1>& tH,
@@ -624,14 +632,24 @@ class ContEngine : public BinaryEngine<Derived> {
     const std::size_t ext_target =
         detail::mixed_retile_config.plain_external_target;
     const std::size_t k_target = detail::mixed_retile_config.contracted_target;
+    const std::size_t tot_ext_target =
+        detail::mixed_retile_config.tot_external_target;
     if (left_is_plain) {
-      // plain operand is LEFT => its external is role M (left outer axes).
+      // plain operand is LEFT => its external is role M (left outer axes),
+      // collapsed; the arena-ToT external is role N (right outer axes),
+      // coarsened to tot_ext_target.
       for (unsigned int i = nf; i + nc < left_rank; ++i)
         tM.push_back(coarsen_tr1(left_U.dim(i), ext_target));
+      for (unsigned int i = nf + nc; i < right_rank; ++i)
+        tN.push_back(coarsen_tr1(right_U.dim(i), tot_ext_target));
     } else {
-      // plain operand is RIGHT => its external is role N (right outer axes).
+      // plain operand is RIGHT => its external is role N (right outer axes),
+      // collapsed; the arena-ToT external is role M (left outer axes),
+      // coarsened to tot_ext_target.
       for (unsigned int i = nf + nc; i < right_rank; ++i)
         tN.push_back(coarsen_tr1(right_U.dim(i), ext_target));
+      for (unsigned int i = nf; i + nc < left_rank; ++i)
+        tM.push_back(coarsen_tr1(left_U.dim(i), tot_ext_target));
     }
     // K (contracted) axes are the last nc axes of the left operand.
     for (unsigned int i = left_rank - nc; i < left_rank; ++i)
@@ -1098,6 +1116,17 @@ class ContEngine : public BinaryEngine<Derived> {
         // owner is rank 0, identical to the old fine-u_grid behavior. (Refine
         // at np>1 was rejected above; only coarsen/identity reaches here.)
         const size_type K_coarse = coarse_K_(K_);
+#ifdef TA_STRIDED_DGEMM_COUNT
+        // Witness: expose the coarse SUMMA grid (the retiled operand trange tile
+        // counts on each role). A coarsened SUMMA external shows up here even
+        // though the per-row strided BLAS GEMM count is invariant.
+        TiledArray::detail::g_summa_coarse_m_grid.store(
+            static_cast<std::size_t>(M_grid), std::memory_order_relaxed);
+        TiledArray::detail::g_summa_coarse_n_grid.store(
+            static_cast<std::size_t>(N_grid), std::memory_order_relaxed);
+        TiledArray::detail::g_summa_coarse_k_grid.store(
+            static_cast<std::size_t>(K_coarse), std::memory_order_relaxed);
+#endif
         // left U layout = [M-axes..., K-axes...]; phase rows = M_grid, cols =
         // K_coarse.
         auto left_phase = proc_grid_.make_row_phase_pmap(K_coarse);

src/TiledArray/expressions/mixed_retile_config.h

@@ -36,24 +36,39 @@ namespace TiledArray::expressions::detail {
 /// "coarsen this role to a single tile spanning the full axis extent".
 ///
 /// HARDCODED ON PURPOSE: the optimal coarsening tracks the physical system
-/// size, so a maintainer edits these two constants per target system size and
-/// rebuilds. Only the plain-operand BLAS externals and the contracted (K) axis
-/// are coarsened; the fused (H) axes and the arena-ToT operand's externals are
-/// left at the operands' own tiling. The plain operand's external lands on
-/// SUMMA role M when the plain operand is the LEFT argument, and on role N when
-/// it is the RIGHT argument; K is always shared and always coarsened. Hence a
-/// single "plain external" number covers both orientations.
+/// size, so a maintainer edits these constants per target system size and
+/// rebuilds.
+///
+/// The default strategy coarsens THREE roles:
+///   1. the plain-operand BLAS external (SUMMA role M when the plain operand is
+///      LEFT, role N when it is RIGHT) -> a single tile (target 0);
+///   2. the contracted (SUMMA-K) axis (always shared) -> a single tile;
+///   3. the arena-ToT operand's external -- the leftover SUMMA axis (role N when
+///      the plain operand is LEFT, role M when it is RIGHT) -> tile size 16.
+/// The fused (H) axes are left at the operands' own tiling (empty target).
+///
+/// COARSEN-ONLY (never refine): every target is fed through `coarsen_tr1`,
+/// which only ever MERGES consecutive user (U) tiles onto existing U
+/// boundaries; it never splits a tile. So a leftover-SUMMA axis whose tiles are
+/// ALREADY >= the target (16) is kept intact (no refine -- refine is
+/// unsupported here), while a finer axis is merged up toward 16.
 struct MixedRetileConfig {
   /// target tile size for the plain operand's external (SUMMA-M if the plain
   /// operand is LEFT, SUMMA-N if it is RIGHT). 0 => single tile.
   std::size_t plain_external_target = 0;
   /// target tile size for the contracted (SUMMA-K) axis. 0 => single tile.
   std::size_t contracted_target = 0;
+  /// target tile size for the arena-ToT operand's external -- the leftover
+  /// SUMMA axis (role N if the plain operand is LEFT, role M if it is RIGHT).
+  /// Coarsen-only: an axis already coarser than this is left intact (never
+  /// refined). 0 => single tile.
+  std::size_t tot_external_target = 16;
 };
 
 inline constexpr MixedRetileConfig mixed_retile_config{
     /*plain_external_target=*/0,  // collapse M (or N) on the plain operand
     /*contracted_target=*/0,      // collapse K
+    /*tot_external_target=*/16,   // coarsen the leftover SUMMA axis to 16
 };
 
 /// Mutable enable flag for the mixed auto-retile gate. Default-initialized once