expressions: native no-external general products via a synthetic unit left-external mode

A general product whose every outer index is 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
(this shape used to segfault through wild stride reads). Evaluate it
with a synthetic unit left-external mode instead: the folded product
becomes (1,K) x (K) -> (1), the exact shape of the already-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
with a unit extent. Replaces the interim gate.

Tests: dense ToT (incl. the no-external root fed by a general T x ToT
inner node), plain dense (the Hadamard-reduction shape), and
block-sparse (exercising the gemm_batched unit handling), all
differential-tested against legacy einsum.

Author: evaleev

src/TiledArray/expressions/cont_engine.h

@@ -666,19 +666,6 @@ class ContEngine : public BinaryEngine<Derived> {
     TA_ASSERT(nh > 0u);  // else this is a pure contraction
     n_fused_modes_ = nh;
 
-    // a general product of tensors-of-tensors 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")) is not supported by the
-    // batched tile op yet (its folded GEMM has no free modes);
-    // einsum() evaluates this shape natively
-    if constexpr (TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
-      if (outer_size(indices_) == nh)
-        TA_EXCEPTION(
-            "general products of tensors-of-tensors without external (free) "
-            "outer indices are not yet supported in the expression layer; "
-            "use TiledArray::einsum() for this contraction");
-    }
-
     // initialize perm_; a target that differs from the canonical (fused...,
     // left-free..., right-free...) result layout cannot be folded into the
     // batched tile op (BatchedContractReduce must be perm-free), so the
@@ -687,12 +674,27 @@ class ContEngine : public BinaryEngine<Derived> {
     this->init_perm(target_indices);
     general_repermute_ = (outer(target_indices) != outer(indices_));
 
-    // the tile op operates on the folded (fused-mode-free) shapes
-    const auto left_op = to_cblas_op(left_outer_permtype_);
+    // 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
+    // with a unit extent.
+    const unsigned int u = (outer_size(indices_) == nh) ? 1u : 0u;
+
+    // the tile op operates on the folded (fused-mode-free) shapes; the
+    // synthetic unit mode leads the folded left operand, so it is NoTrans
+    const auto left_op =
+        u ? math::blas::NoTranspose : to_cblas_op(left_outer_permtype_);
     const auto right_op = to_cblas_op(right_outer_permtype_);
     if constexpr (!TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
-      op_ = op_type(left_op, right_op, factor_, outer_size(indices_) - nh,
-                    outer_size(left_indices_) - nh,
+      op_ = op_type(left_op, right_op, factor_, outer_size(indices_) - nh + u,
+                    outer_size(left_indices_) - nh + u,
                     outer_size(right_indices_) - nh);
     } else {
       // the batched tile op must be perm-free (BatchedContractReduce cannot
@@ -710,7 +712,8 @@ class ContEngine : public BinaryEngine<Derived> {
 
       // factor_ is absorbed into element_nonreturn_op_
       op_ = op_type(left_op, right_op, scalar_type(1),
-                    outer_size(indices_) - nh, outer_size(left_indices_) - nh,
+                    outer_size(indices_) - nh + u,
+                    outer_size(left_indices_) - nh + u,
                     outer_size(right_indices_) - nh, BipartitePermutation{},
                     this->element_nonreturn_op_, std::move(this->arena_plan_));
       // ce+e, ce+ce_right and ce+ce_left are mutually exclusive; at most one
@@ -750,7 +753,11 @@ class ContEngine : public BinaryEngine<Derived> {
   trange_type make_trange_general() const {
     const unsigned int nh = n_fused_modes_;
     const unsigned int nc = op_.gemm_helper().num_contract_ranks();
-    const unsigned int neA = op_.gemm_helper().left_rank() - nc;
+    // 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;
+    const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
     const unsigned int neB = op_.gemm_helper().right_rank() - nc;
 
     typename trange_type::Ranges ranges(nh + neA + neB);
@@ -802,7 +809,11 @@ 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();
-    const unsigned int neA = op_.gemm_helper().left_rank() - nc;
+    // 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;
+    const unsigned int neA = op_.gemm_helper().left_rank() - nc - u;
     const unsigned int neB = op_.gemm_helper().right_rank() - nc;
 
     // Get pointers to the argument sizes

src/TiledArray/sparse_shape.h

@@ -1720,42 +1720,54 @@ class SparseShape {
     const auto* left_extent = tile_norms_.range().extent_data();
     const auto* right_extent = other.tile_norms_.range().extent_data();
 
+    // a no-external product carries a SYNTHETIC unit left-external mode in
+    // the GemmHelper only (see ContEngine::init_struct_general); detect it
+    // from the one-rank mismatch with the actual norm tensor and pad the
+    // folded left/result views with a unit extent
+    const bool unit_external =
+        (tile_norms_.range().rank() + 1u == nfused + gemm_helper.left_rank());
+    const unsigned int u = unit_external ? 1u : 0u;
+
     // check that the ranks match the folded gemm ranks plus the fused modes,
     // and that the fused and contracted mode extents of the two shapes are
     // congruent
-    TA_ASSERT(tile_norms_.range().rank() == nfused + gemm_helper.left_rank());
+    TA_ASSERT(tile_norms_.range().rank() + u ==
+              nfused + gemm_helper.left_rank());
     TA_ASSERT(other.tile_norms_.range().rank() ==
               nfused + gemm_helper.right_rank());
     for (unsigned int d = 0u; d < nfused; ++d)
       TA_ASSERT(left_extent[d] == right_extent[d]);
     for (unsigned int i = gemm_helper.left_inner_begin(),
                       j = gemm_helper.right_inner_begin();
          i < gemm_helper.left_inner_end(); ++i, ++j)
-      TA_ASSERT(left_extent[nfused + i] == right_extent[nfused + j]);
+      TA_ASSERT(left_extent[nfused + i - u] == right_extent[nfused + j]);
 
     integer H = 1, M = 1, N = 1, K = 1;
     for (unsigned int d = 0u; d < nfused; ++d) H *= left_extent[d];
-    for (unsigned int i = gemm_helper.left_outer_begin();
-         i < gemm_helper.left_outer_end(); ++i)
-      M *= left_extent[nfused + i];
+    if (!unit_external)
+      for (unsigned int i = gemm_helper.left_outer_begin();
+           i < gemm_helper.left_outer_end(); ++i)
+        M *= left_extent[nfused + i];
     for (unsigned int i = gemm_helper.left_inner_begin();
          i < gemm_helper.left_inner_end(); ++i)
-      K *= left_extent[nfused + i];
+      K *= left_extent[nfused + i - u];
     for (unsigned int i = gemm_helper.right_outer_begin();
          i < gemm_helper.right_outer_end(); ++i)
       N *= right_extent[nfused + i];
 
     // result size vectors: fused modes (from this), then the left and right
-    // outer modes
-    const unsigned int result_rank = nfused + gemm_helper.result_rank();
+    // outer modes (the synthetic unit left-external mode is absent from the
+    // actual result)
+    const unsigned int result_rank = nfused + gemm_helper.result_rank() - u;
     std::shared_ptr<vector_type> result_size_vectors(
         new vector_type[result_rank], std::default_delete<vector_type[]>());
     unsigned int x = 0ul;
     for (unsigned int i = 0u; i < nfused; ++i, ++x)
       result_size_vectors.get()[x] = size_vectors_.get()[i];
-    for (unsigned int i = gemm_helper.left_outer_begin();
-         i < gemm_helper.left_outer_end(); ++i, ++x)
-      result_size_vectors.get()[x] = size_vectors_.get()[nfused + i];
+    if (!unit_external)
+      for (unsigned int i = gemm_helper.left_outer_begin();
+           i < gemm_helper.left_outer_end(); ++i, ++x)
+        result_size_vectors.get()[x] = size_vectors_.get()[nfused + i];
     for (unsigned int i = gemm_helper.right_outer_begin();
          i < gemm_helper.right_outer_end(); ++i, ++x)
       result_size_vectors.get()[x] = other.size_vectors_.get()[nfused + i];
@@ -1770,11 +1782,12 @@ class SparseShape {
       lobounds.push_back(tile_norms_.range().lobound_data()[d]);
       upbounds.push_back(tile_norms_.range().upbound_data()[d]);
     }
-    for (unsigned int i = gemm_helper.left_outer_begin();
-         i < gemm_helper.left_outer_end(); ++i) {
-      lobounds.push_back(tile_norms_.range().lobound_data()[nfused + i]);
-      upbounds.push_back(tile_norms_.range().upbound_data()[nfused + i]);
-    }
+    if (!unit_external)
+      for (unsigned int i = gemm_helper.left_outer_begin();
+           i < gemm_helper.left_outer_end(); ++i) {
+        lobounds.push_back(tile_norms_.range().lobound_data()[nfused + i]);
+        upbounds.push_back(tile_norms_.range().upbound_data()[nfused + i]);
+      }
     for (unsigned int i = gemm_helper.right_outer_begin();
          i < gemm_helper.right_outer_end(); ++i) {
       lobounds.push_back(other.tile_norms_.range().lobound_data()[nfused + i]);
@@ -1784,10 +1797,13 @@ class SparseShape {
 
     // the range spanned by modes [nfused, rank) of \p r, rebased to zero
     // lobounds (scratch view for the slab-batched norm GEMM)
-    auto fold_range = [nfused](const range_type& r) {
+    auto fold_range = [nfused](const range_type& r,
+                               const bool prepend_unit = false) {
       const auto* extent = r.extent_data();
-      container::svector<index1_type> extents(extent + nfused,
-                                              extent + r.rank());
+      container::svector<index1_type> extents;
+      extents.reserve(r.rank() - nfused + (prepend_unit ? 1u : 0u));
+      if (prepend_unit) extents.push_back(1);
+      extents.insert(extents.end(), extent + nfused, extent + r.rank());
       return range_type(extents);
     };
 
@@ -1797,10 +1813,11 @@ class SparseShape {
 
     if (k_rank > 0u) {
       // the contracted-mode size vector; identical for every slab since the
-      // contracted modes follow the fused modes
+      // contracted modes follow the fused modes (helper coordinates carry
+      // the synthetic unit mode, the actual size vectors do not)
       const vector_type k_sizes = recursive_outer_product(
-          size_vectors_.get() + nfused + gemm_helper.left_inner_begin(), k_rank,
-          [](const vector_type& size_vector) -> const vector_type& {
+          size_vectors_.get() + nfused + gemm_helper.left_inner_begin() - u,
+          k_rank, [](const vector_type& size_vector) -> const vector_type& {
             return size_vector;
           });
 
@@ -1828,10 +1845,11 @@ class SparseShape {
       // slab-batched norm GEMM: fold the fused modes into the tensor batch
       // dimension by zero-copy reshape; result_folded shares result_norms'
       // buffer, so the accumulation lands in place
-      auto left_folded = left.reshape(fold_range(left.range()), H);
+      auto left_folded =
+          left.reshape(fold_range(left.range(), unit_external), H);
       auto right_folded = right.reshape(fold_range(right.range()), H);
-      auto result_folded =
-          result_norms.reshape(fold_range(result_norms.range()), H);
+      auto result_folded = result_norms.reshape(
+          fold_range(result_norms.range(), unit_external), H);
       result_folded.gemm(left_folded, right_folded, abs_factor, gemm_helper);
 
       // Hard zero tiles that are below the zero threshold.

src/TiledArray/tile_op/batched_contract_reduce.h

@@ -62,12 +62,16 @@ class BatchedContractReduce {
 
   /// \return the range spanned by modes [nfused_, rank) of \p r, rebased to
   /// zero lobounds (the folded view is a GEMM scratch view; only extents
-  /// matter)
+  /// matter); \p prepend_unit prepends a unit extent (the synthetic
+  /// left-external mode of a no-external product, see
+  /// ContEngine::init_struct_general)
   template <typename Range_>
-  Range_ fold_range(const Range_& r) const {
+  Range_ fold_range(const Range_& r, const bool prepend_unit = false) const {
     const auto* extent = r.extent_data();
-    container::svector<typename Range_::index1_type> extents(extent + nfused_,
-                                                             extent + r.rank());
+    container::svector<typename Range_::index1_type> extents;
+    extents.reserve(r.rank() - nfused_ + (prepend_unit ? 1u : 0u));
+    if (prepend_unit) extents.push_back(1);
+    extents.insert(extents.end(), extent + nfused_, extent + r.rank());
     return Range_(extents);
   }
 
@@ -150,7 +154,13 @@ class BatchedContractReduce {
 
     const auto& gh = op_.gemm_helper();
     const unsigned int nc = gh.num_contract_ranks();
-    const unsigned int neA = gh.left_rank() - nc;
+    // a no-external product carries a SYNTHETIC unit left-external mode in
+    // the GemmHelper only (see ContEngine::init_struct_general); detect it
+    // from the one-rank mismatch with the actual left tile and pad the
+    // folded left/result views with a unit extent
+    const bool unit_external =
+        (left.range().rank() + 1u == nfused_ + gh.left_rank());
+    const unsigned int neA = gh.left_rank() - nc - (unit_external ? 1u : 0u);
     const unsigned int neB = gh.right_rank() - nc;
 
     // both args must carry the fused modes as their leading modes, with
@@ -163,7 +173,8 @@ class BatchedContractReduce {
     TA_ASSERT(batch == fused_volume(right.range()));
 
     // folded, zero-copy argument views
-    auto left_folded = left.reshape(fold_range(left.range()), batch);
+    auto left_folded =
+        left.reshape(fold_range(left.range(), unit_external), batch);
     auto right_folded = right.reshape(fold_range(right.range()), batch);
 
     if (empty(result)) {
@@ -195,7 +206,8 @@ class BatchedContractReduce {
     } else {
       // accumulate through a folded, zero-copy view of the result
       const auto full_range = result.range();
-      auto result_folded = result.reshape(fold_range(full_range), batch);
+      auto result_folded =
+          result.reshape(fold_range(full_range, unit_external), batch);
       op_(result_folded, left_folded, right_folded);
       // the wrapped op may REBIND the result instead of writing in place:
       // the arena grow-to-cover path (a later K-panel touching inner cells

tests/general_product.cpp

@@ -731,20 +731,64 @@ BOOST_AUTO_TEST_CASE(expression_general_kitchen_sink) {
   BOOST_CHECK_SMALL(tot_max_abs_diff(w, w_ref), 1e-10);
 }
 
-BOOST_AUTO_TEST_CASE(expression_general_product_tot_no_externals_gated) {
+BOOST_AUTO_TEST_CASE(expression_general_product_tot_no_externals) {
   // a ToT x ToT general product with NO external (free) outer indices --
-  // every outer index fused or contracted -- is not supported by the
-  // batched tile op (and used to segfault in the folded GEMM); the engine
-  // must reject it with an informative error (einsum() handles this shape
-  // natively via its no-external regime)
+  // every outer index fused or contracted. The folded product has no free
+  // modes, so it is evaluated with a synthetic unit left-external mode
+  // carried by the tile op's GemmHelper only (this shape used to segfault).
   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, {3}, 2.0);
 
   TArrayToT c;
-  BOOST_CHECK_THROW(c("i,j;a,b") = a("x,i,j;a") * b("x,i,j;b"),
-                    TiledArray::Exception);
+  BOOST_REQUIRE_NO_THROW(c("i,j;a,b") = a("x,i,j;a") * b("x,i,j;b"));
+  auto c_ref = TA::einsum(a("x,i,j;a"), b("x,i,j;b"), "i,j;a,b");
+  BOOST_CHECK_SMALL(tot_max_abs_diff(c, c_ref), 1e-10);
+
+  // the same no-external root with the left operand produced by a general
+  // T x ToT product at an INNER node (the original motivating expression)
+  TA::TiledRange tr_g{{0, 3, 5}, {0, 2, 4}};   // x, i
+  TA::TiledRange tr_cv{{0, 3, 5}, {0, 2, 3}};  // x, j
+  auto g = make_patterned_array(world, tr_g, 1.0);
+  auto cv = make_patterned_tot_array(world, tr_cv, {2}, 2.0);
+  TArrayToT i1, w, w_ref;
+  i1("x,i,j;a") = g("x,i") * cv("x,j;a");
+  w_ref("i,j;a,b") = i1("x,i,j;a") * b("x,i,j;b");
+  BOOST_REQUIRE_NO_THROW(w("i,j;a,b") =
+                             (g("x,i") * cv("x,j;a")) * b("x,i,j;b"));
+  BOOST_CHECK_SMALL(tot_max_abs_diff(w, w_ref), 1e-10);
+}
+
+BOOST_AUTO_TEST_CASE(expression_general_product_no_externals) {
+  // the plain-tensor analogue of the no-external general product (the
+  // Hadamard-reduction shape): C("i") = A("i,j") * B("i,j")
+  auto& world = TA::get_default_world();
+  ForceLegacyEinsum legacy_oracle;  // keep the einsum reference independent
+  TA::TiledRange tr{{0, 2, 4}, {0, 3, 5}};  // i, 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);
+}
+
+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
+  auto& world = TA::get_default_world();
+  ForceLegacyEinsum legacy_oracle;  // keep the einsum reference independent
+  TA::TiledRange tr{{0, 2, 5}, {0, 3, 4}, {0, 2, 6, 7}};  // b, i, j
+  auto a = make_patterned_sparse_array(world, tr, 1.0, 3);
+  auto b = make_patterned_sparse_array(world, tr, 2.0, 4);
+
+  TA::TSpArrayD c;
+  BOOST_REQUIRE_NO_THROW(c("b,i") = a("b,i,j") * b("b,i,j"));
+  auto c_ref = TA::einsum(a("b,i,j"), b("b,i,j"), "b,i");
+  BOOST_CHECK_SMALL(diff_norm_sp(c, c_ref, "b,i"), 1e-10);
 }
 
 BOOST_AUTO_TEST_CASE(expression_general_product_tot_inner_outer_product) {