einsum: delegate nesting-preserving products to the expression layer

evaleev · evaleev · commit 9602d4f3d8e7 · 2026-06-12T16:51:43.000-04:00
Two general-product shapes that einsum handled with bespoke machinery are
nesting-preserving (the result keeps the operand nesting), so the general-
product expression engine can own them; only the type-mutating ToT*ToT-&gt;T
denest reduction remains einsum's by default.

* Fused broadcast C(h,e) = A(h) * B(h,e) (one operand entirely fused, no
  contraction): replaces the replicate_array() expand-then-Hadamard with a
  direct delegation. The engine evaluates it natively once the fully-fused
  operand leads -- its folded left tile is rank-0, restored by a synthetic
  unit left-external mode. ContEngine::synthetic_unit_left_external() now
  fires for a rank-0 LEFT operand (broadcast), not only a rank-0 RESULT
  (no-external). einsum orders the fully-fused operand first and applies any
  inner result permutation afterward.

* No-external products C(h;..) = A(h,i;..) * B(h,i;..) (fused + contracted,
  no free index) now default to the expression layer (the existing
  no-external synthetic-unit mode). The legacy local kernel is retained only
  behind einsum_legacy_subworld() as the cross-check oracle, mirroring the
  generalized-subworld path.

Tests: dense fused broadcast (both operand orders), ToT no-external with
inner Hadamard and inner contraction.
diff --git a/src/TiledArray/einsum/tiledarray.h b/src/TiledArray/einsum/tiledarray.h
@@ -673,23 +673,38 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
     auto range_map =
         (RangeMap(a, A.array().trange()) | RangeMap(b, B.array().trange()));
 
-    // special Hadamard
+    // Fused broadcast: one operand is ENTIRELY fused (h == its index set) and
+    // there is no contraction -- C(h,e) = A * B, a per-fused-block scale. The
+    // general-product expression engine evaluates this natively when the
+    // fully-fused operand LEADS (it folds to a rank-0 left tile, restored by a
+    // synthetic unit left-external mode; see ContEngine::
+    // synthetic_unit_left_external). Order the fully-fused operand first and
+    // delegate -- no physical replication needed.
     if (h.size() == a.size() || h.size() == b.size()) {
       TA_ASSERT(!i && e);
-      bool const small_a = h.size() == a.size();
-      auto const delta_trng = make_trange(range_map, e);
+      _ein_call.branch = "fused-broadcast-expression";
+      const bool a_leads = (h.size() == a.size());  // A entirely fused?
+      // the engine produces the result in its canonical inner layout
+      // (inner fused, then left-then-right inner externals, in the order the
+      // operands are handed to it) and rejects a non-identity inner result
+      // permutation; evaluate canonically, then apply the inner permutation
+      // (cf. the generalized-inner-perm-recurse path below). The canonical
+      // inner-external order tracks the operand order, so mirror the swap.
+      std::string canon_inner;
+      if constexpr (IsArrayToT<ArrayC>) {
+        auto inner_e = a_leads ? (inner.A ^ inner.B) : (inner.B ^ inner.A);
+        canon_inner = ";" + (std::string)(inner.h + inner_e);
+      }
+      std::string canon_layout = std::string(c) + canon_inner;
+      ArrayC C0;
+      if (a_leads)  // A fully fused -> already leads
+        C0(canon_layout) = A * B;
+      else  // B fully fused -> put it first (multiplication commutes)
+        C0(canon_layout) = B * A;
       std::string target_layout = std::string(c) + inner.c;
+      if (target_layout == canon_layout) return C0;
       ArrayC C;
-      if (small_a) {
-        auto temp = replicate_array(A.array(), delta_trng);
-        std::string temp_layout = std::string(e) + "," + A.annotation();
-        C(target_layout) = temp(temp_layout) * B;
-      } else {
-        auto temp = replicate_array(B.array(), delta_trng);
-        std::string temp_layout = std::string(e) + "," + B.annotation();
-        C(target_layout) = A * temp(temp_layout);
-      }
-
+      C(target_layout) = C0(canon_layout);  // inner-permute to requested layout
       return C;
     }
 
@@ -791,7 +806,28 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
       }
     }
 
-    if (!e) {  // hadamard reduction
+    if (!e) {  // no external outer index (fused + contracted): a nesting-
+               // preserving reduction. By default the expression layer handles
+               // it natively (synthetic unit left-external mode, see
+               // ContEngine::synthetic_unit_left_external); the legacy local
+               // kernel below is retained only as the opt-in legacy
+               // cross-check oracle (cf. the generalized-subworld path).
+      if (!detail::einsum_legacy_subworld()) {
+        _ein_call.branch = "no-external-expression";
+        // evaluate in the engine's canonical inner layout, then apply any
+        // inner result permutation (cf. the broadcast / inner-perm-recurse)
+        std::string canon_inner;
+        if constexpr (IsArrayToT<ArrayC>)
+          canon_inner = ";" + (std::string)(inner.h + inner.e);
+        std::string canon_layout = std::string(c) + canon_inner;
+        ArrayC C0;
+        C0(canon_layout) = A * B;
+        std::string target_layout = std::string(c) + inner.c;
+        if (target_layout == canon_layout) return C0;
+        ArrayC C;
+        C(target_layout) = C0(canon_layout);
+        return C;
+      }
 
       _ein_call.branch = "hadamard-reduction-local";
       const auto _ein_he_t0 = _ein_call.active ? now() : time_point{};
diff --git a/src/TiledArray/expressions/cont_engine.h b/src/TiledArray/expressions/cont_engine.h
@@ -677,18 +677,20 @@ class ContEngine : public BinaryEngine<Derived> {
     this->init_perm(target_indices);
     general_repermute_ = (outer(target_indices) != outer(indices_));
 
-    // A product with NO external (free) outer indices (every outer index
-    // fused or contracted, e.g. C("i,j;a,b") = A("x,i,j;a") * B("x,i,j;b"))
-    // folds to a GEMM with no free modes, i.e. rank-0 tensors, which the
-    // tile kernels do not support. Evaluate it with a SYNTHETIC UNIT
-    // left-external mode instead: the folded product becomes
-    // (1,K) x (K) -> (1), the exact shape of the (supported) one-sided
-    // neB == 0 case. The unit mode lives only in the tile op's GemmHelper;
-    // tranges, shapes and tiles carry the true (external-free) ranks, and
-    // BatchedContractReduce / SparseShape::gemm_batched detect the
-    // synthetic mode from the one-rank mismatch and pad their folded views
+    // Some degenerate folded shapes would carry a rank-0 tensor, which the
+    // tile kernels do not support (see synthetic_unit_left_external()):
+    //  - a NO-EXTERNAL product (every outer index fused or contracted, e.g.
+    //    C("i,j;a,b") = A("x,i,j;a") * B("x,i,j;b")) folds to a rank-0 RESULT;
+    //  - a FUSED BROADCAST (the left operand is entirely fused, no contraction,
+    //    e.g. C("b,k") = A("b") * B("b,k")) folds to a rank-0 LEFT operand.
+    // Evaluate both with a SYNTHETIC UNIT left-external mode: the folded
+    // product becomes (1,K) x (K,N) -> (1,N), a supported shape (the no-
+    // external case has N == 1, the broadcast has K == 1). The unit mode lives
+    // only in the tile op's GemmHelper; tranges, shapes and tiles carry the
+    // true ranks, and BatchedContractReduce / SparseShape::gemm_batched detect
+    // the synthetic mode from the one-rank mismatch and pad their folded views
     // with a unit extent.
-    const unsigned int u = (outer_size(indices_) == nh) ? 1u : 0u;
+    const unsigned int u = synthetic_unit_left_external();
 
     // the tile op operates on the folded (fused-mode-free) shapes; the
     // synthetic unit mode leads the folded left operand, so it is NoTrans
@@ -749,17 +751,34 @@ class ContEngine : public BinaryEngine<Derived> {
     }
   }
 
+  /// \return 1 if the folded general product needs a SYNTHETIC unit
+  /// left-external mode, else 0. The folded (fused-mode-free) GEMM cannot host
+  /// a rank-0 tensor, which arises in two degenerate cases:
+  ///  - rank-0 RESULT: every outer index is fused or contracted (no external),
+  ///    i.e. outer_size(indices_) == n_fused_modes_;
+  ///  - rank-0 LEFT operand: the left argument is entirely fused with no
+  ///    contraction (a fused broadcast / per-fused-block scale), i.e.
+  ///    outer_size(left_indices_) == n_fused_modes_.
+  /// A unit left-external mode (carried only in the GemmHelper) restores a
+  /// supported (1,K) x (K,N) -> (1,N) shape in both cases.
+  unsigned int synthetic_unit_left_external() const {
+    return (outer_size(indices_) == n_fused_modes_ ||
+            outer_size(left_indices_) == n_fused_modes_)
+               ? 1u
+               : 0u;
+  }
+
   /// Tiled range factory function for a general product
 
   /// \return The result tiled range: the fused mode ranges followed by the
   /// left- and right-external mode ranges
   trange_type make_trange_general() const {
     const unsigned int nh = n_fused_modes_;
     const unsigned int nc = op_.gemm_helper().num_contract_ranks();
-    // the no-external case carries a synthetic unit left-external mode in
-    // the GemmHelper only (see init_struct_general); the actual tranges do
-    // not have it
-    const unsigned int u = (outer_size(indices_) == n_fused_modes_) ? 1u : 0u;
+    // degenerate folds carry a synthetic unit left-external mode in the
+    // GemmHelper only (see synthetic_unit_left_external()); the actual tranges
+    // do not have it
+    const unsigned int u = synthetic_unit_left_external();
     const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
     const unsigned int neB = op_.gemm_helper().right_rank() - nc;
 
@@ -812,10 +831,10 @@ class ContEngine : public BinaryEngine<Derived> {
                                  std::shared_ptr<const pmap_interface> pmap) {
     const unsigned int nh = n_fused_modes_;
     const unsigned int nc = op_.gemm_helper().num_contract_ranks();
-    // the no-external case carries a synthetic unit left-external mode in
-    // the GemmHelper only (see init_struct_general); the actual tranges do
-    // not have it
-    const unsigned int u = (outer_size(indices_) == nh) ? 1u : 0u;
+    // degenerate folds carry a synthetic unit left-external mode in the
+    // GemmHelper only (see synthetic_unit_left_external()); the actual tranges
+    // do not have it
+    const unsigned int u = synthetic_unit_left_external();
     const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
     const unsigned int neB = op_.gemm_helper().right_rank() - nc;
 
diff --git a/tests/general_product.cpp b/tests/general_product.cpp
@@ -237,6 +237,37 @@ BOOST_AUTO_TEST_CASE(expression_general_product_dense_batched_outer) {
   BOOST_CHECK_SMALL(diff_norm(c, c_ref, "b,i,k"), 1e-10);
 }
 
+BOOST_AUTO_TEST_CASE(expression_general_product_dense_broadcast) {
+  // fused broadcast: one operand is ENTIRELY fused, no contraction, the other
+  // carries an extra external -- C(b,k) = A(b) * B(b,k). Per fused block b this
+  // is scalar * vector (a scale). The general-product engine handles it via a
+  // synthetic unit left-external mode; einsum no longer needs replicate_array.
+  auto& world = TA::get_default_world();
+
+  TA::TiledRange tr_s{{0, 2, 5}};             // b
+  TA::TiledRange tr_v{{0, 2, 5}, {0, 4, 5}};  // b, k
+  auto s = make_patterned_array(world, tr_s, 1.0);
+  auto v = make_patterned_array(world, tr_v, 2.0);
+
+  // independent oracle: physically replicate s over k (NOT through the
+  // general-product engine), then a plain Hadamard multiply
+  TA::TiledRange tr_k{{0, 4, 5}};                     // k
+  auto s_rep = TA::Einsum::replicate_array(s, tr_k);  // s_rep(k, b)
+  TA::TArrayD ref;
+  ref("b,k") = s_rep("k,b") * v("b,k");
+
+  // (a) LEFT operand fully fused, straight through the expression layer
+  TA::TArrayD cc;
+  BOOST_REQUIRE_NO_THROW(cc("b,k") = s("b") * v("b,k"));
+  BOOST_CHECK_SMALL(diff_norm(cc, ref, "b,k"), 1e-10);
+
+  // (b) einsum, both operand orders (RIGHT operand fully fused too)
+  auto c_lr = TA::einsum(s("b"), v("b,k"), "b,k");
+  BOOST_CHECK_SMALL(diff_norm(c_lr, ref, "b,k"), 1e-10);
+  auto c_rl = TA::einsum(v("b,k"), s("b"), "b,k");
+  BOOST_CHECK_SMALL(diff_norm(c_rl, ref, "b,k"), 1e-10);
+}
+
 BOOST_AUTO_TEST_CASE(expression_general_product_noncanonical_root_target) {
   // a root-level general product with a NON-canonical target layout: the
   // product evaluates canonically (r1,p,q) and is re-permuted to the target
@@ -778,6 +809,42 @@ BOOST_AUTO_TEST_CASE(expression_general_product_no_externals) {
   BOOST_CHECK_SMALL(diff_norm(c, c_ref, "i"), 1e-10);
 }
 
+BOOST_AUTO_TEST_CASE(
+    expression_general_product_tot_no_externals_inner_hadamard) {
+  // ToT no-external product with an INNER HADAMARD: outer i,j fused, x
+  // contracted, no free outer index; inner m fused. This is the einsum
+  // "hadamard-reduction-local" shape with an inner-Hadamard cell op -- verify
+  // the expression layer (synthetic unit left-external + inner Hadamard)
+  // reproduces it.
+  auto& world = TA::get_default_world();
+  ForceLegacyEinsum legacy_oracle;  // keep the einsum reference independent
+  TA::TiledRange tr{{0, 3, 5}, {0, 2, 4}, {0, 2, 3}};  // x, i, j
+  auto a = make_patterned_tot_array(world, tr, {2}, 1.0);
+  auto b = make_patterned_tot_array(world, tr, {2}, 2.0);
+
+  TArrayToT c;
+  BOOST_REQUIRE_NO_THROW(c("i,j;m") = a("x,i,j;m") * b("x,i,j;m"));
+  auto c_ref = TA::einsum(a("x,i,j;m"), b("x,i,j;m"), "i,j;m");
+  BOOST_CHECK_SMALL(tot_max_abs_diff(c, c_ref), 1e-10);
+}
+
+BOOST_AUTO_TEST_CASE(
+    expression_general_product_tot_no_externals_inner_contraction) {
+  // ToT no-external product with an INNER CONTRACTION: outer i,j fused, x
+  // contracted, no free outer index; inner c contracted, m/n free. The einsum
+  // "hadamard-reduction-local" shape with an inner-contraction cell op.
+  auto& world = TA::get_default_world();
+  ForceLegacyEinsum legacy_oracle;  // keep the einsum reference independent
+  TA::TiledRange tr{{0, 3, 5}, {0, 2, 4}, {0, 2, 3}};         // x, i, j
+  auto a = make_patterned_tot_array(world, tr, {2, 3}, 1.0);  // inner m, c
+  auto b = make_patterned_tot_array(world, tr, {3, 2}, 2.0);  // inner c, n
+
+  TArrayToT c;
+  BOOST_REQUIRE_NO_THROW(c("i,j;m,n") = a("x,i,j;m,c") * b("x,i,j;c,n"));
+  auto c_ref = TA::einsum(a("x,i,j;m,c"), b("x,i,j;c,n"), "i,j;m,n");
+  BOOST_CHECK_SMALL(tot_max_abs_diff(c, c_ref), 1e-10);
+}
+
 BOOST_AUTO_TEST_CASE(expression_general_product_sparse_no_externals) {
   // block-sparse no-external general product: exercises the synthetic
   // unit-mode handling in SparseShape::gemm_batched
@@ -1241,13 +1308,8 @@ BOOST_AUTO_TEST_CASE(expression_general_product_csv_like) {
   ArenaSpArr c1, c2;
   {
     ScopedEinsumRoute expression_route(false);
-    try {
-      c1 = TA::einsum(a("i2,i1,m;a"), b("m,i2,K"), "i1,i2,K;a");
-      c2 = TA::einsum(a("i2,i1,m;a"), b("m,i2,K"), "i1,i2,K;a");
-    } catch (std::exception& ex) {
-      std::cerr << "EXPRESSION ROUTE THREW: " << ex.what() << std::endl;
-      throw;
-    }
+    c1 = TA::einsum(a("i2,i1,m;a"), b("m,i2,K"), "i1,i2,K;a");
+    c2 = TA::einsum(a("i2,i1,m;a"), b("m,i2,K"), "i1,i2,K;a");
   }
 
   // BISECT: same shape, CANONICAL target, direct expression (no einsum
@@ -1278,7 +1340,6 @@ BOOST_AUTO_TEST_CASE(expression_general_product_csv_like) {
         }
       }
     }
-    std::cerr << "CANONICAL-TARGET DIRECT-EXPR max_diff = " << md << std::endl;
     BOOST_CHECK_SMALL(md, 1e-10);
   }
 
