Merge pull request #512 from ValeevGroup/evaleev/fix/einsum-hadamard-reduction

fix einsum hadamard reduction

Author: evaleev

src/TiledArray/conversions/make_array.h

@@ -242,7 +242,7 @@ inline Array make_array(World& world, const detail::trange_t<Array>& trange,
                            op);
 }
 
-/// a make_array variant that uses a sequence of tiles
+/// a make_array variant that uses a sequence of {tile_index,tile} pairs
 /// to construct a DistArray with default pmap
 template <typename Array, typename Tiles>
 Array make_array(World& world, const detail::trange_t<Array>& tiled_range,

src/TiledArray/einsum/tiledarray.h

@@ -600,6 +600,8 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
     std::invoke(update_perm_and_indices, std::get<0>(AB));
     std::invoke(update_perm_and_indices, std::get<1>(AB));
 
+    // construct result, with "dense" DistArray; the array will be
+    // reconstructred from local tiles later
     ArrayTerm<ArrayC> C = {ArrayC(world, TiledRange(range_map[c])), c};
     for (auto idx : e) {
       C.tiles *= Range(range_map[idx].tiles_range());
@@ -609,6 +611,16 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
     }
     C.expr = e;
 
+    using Index = Einsum::Index<size_t>;
+
+    // this will collect local tiles of C.array, to be used to rebuild C.array
+    std::vector<std::pair<Index, ResultTensor>> C_local_tiles;
+    auto build_C_array = [&]() {
+      C.array = make_array<ArrayC>(world, TiledRange(range_map[c]),
+                                   C_local_tiles.begin(), C_local_tiles.end(),
+                                   /* replicated = */ false);
+    };
+
     std::get<0>(AB).expr += inner.a;
     std::get<1>(AB).expr += inner.b;
 
@@ -627,19 +639,56 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
       }
     }
 
-    using Index = Einsum::Index<size_t>;
-
     if (!e) {  // hadamard reduction
+
       auto &[A, B] = AB;
       TiledRange trange(range_map[i]);
       RangeProduct tiles;
       for (auto idx : i) {
         tiles *= Range(range_map[idx].tiles_range());
       }
+
+      // the inner product can be either hadamard or a contraction
+      using TensorT = typename decltype(A.array)::value_type::value_type;
+      static_assert(
+          std::is_same_v<TensorT,
+                         typename decltype(A.array)::value_type::value_type>);
+      constexpr bool is_tot = detail::is_tensor_v<TensorT>;
+      auto element_hadamard_op =
+          (is_tot && inner.h)
+              ? std::make_optional(
+                    [&inner, plan = detail::TensorHadamardPlan(inner.A, inner.B,
+                                                               inner.C)](
+                        auto const &l, auto const &r) -> TensorT {
+                      if (l.empty() || r.empty()) return TensorT{};
+                      return detail::tensor_hadamard(l, r, plan);
+                    })
+              : std::nullopt;
+      auto element_contract_op =
+          (is_tot && !inner.h)
+              ? std::make_optional(
+                    [&inner, plan = detail::TensorContractionPlan(
+                                 inner.A, inner.B, inner.C)](
+                        auto const &l, auto const &r) -> TensorT {
+                      if (l.empty() || r.empty()) return TensorT{};
+                      return detail::tensor_contract(l, r, plan);
+                    })
+              : std::nullopt;
+      auto element_product_op = [&inner, &element_hadamard_op,
+                                 &element_contract_op](
+                                    auto const &l, auto const &r) -> TensorT {
+        TA_ASSERT(inner.h ? element_hadamard_op.has_value()
+                          : element_contract_op.has_value());
+        return inner.h ? element_hadamard_op.value()(l, r)
+                       : element_contract_op.value()(l, r);
+      };
+
       auto pa = A.permutation;
       auto pb = B.permutation;
       for (Index h : H.tiles) {
-        if (!C.array.is_local(h)) continue;
+        auto const pc = C.permutation;
+        auto const c = apply(pc, h);
+        if (!C.array.is_local(c)) continue;
         size_t batch = 1;
         for (size_t i = 0; i < h.size(); ++i) {
           batch *= H.batch[i].at(h[i]);
@@ -670,16 +719,8 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
               auto &el = tile({k});
               using TensorT = std::remove_reference_t<decltype(el)>;
 
-              auto mult_op = [&inner](auto const &l, auto const &r) -> TensorT {
-                if (l.empty() || r.empty()) return TensorT{};
-                return inner.h ? TA::detail::tensor_hadamard(l, inner.A, r,
-                                                             inner.B, inner.C)
-                               : TA::detail::tensor_contract(l, inner.A, r,
-                                                             inner.B, inner.C);
-              };
-
               for (auto i = 0; i < vol; ++i)
-                el.add_to(mult_op(aik.data()[i], bik.data()[i]));
+                el.add_to(element_product_op(aik.data()[i], bik.data()[i]));
 
             } else if constexpr (!AreArraySame<ArrayA, ArrayB>) {
               auto aik = ai.batch(k);
@@ -702,14 +743,21 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
             }
           }
         }
-        auto pc = C.permutation;
-        auto shape = apply_inverse(pc, C.array.trange().tile(h));
+        // data is stored as h1 h2 ... but all modes folded as 1 batch dim
+        // first reshape to h = (h1 h2 ...)
+        // n.b. can't just use shape = C.array.trange().tile(h)
+        auto shape = apply_inverse(pc, C.array.trange().tile(c));
         tile = tile.reshape(shape);
+        // then permute to target C layout c = (c1 c2 ...)
         if (pc) tile = tile.permute(pc);
-        C.array.set(h, tile);
+        // and move to C_local_tiles
+        C_local_tiles.emplace_back(std::move(c), std::move(tile));
       }
+
+      build_C_array();
+
       return C.array;
-    }
+    }  // end: hadamard reduction
 
     // generalized contraction
 
@@ -740,7 +788,6 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
     std::invoke(update_tr, std::get<1>(AB));
 
     std::vector<std::shared_ptr<World>> worlds;
-    std::vector<std::tuple<Index, ResultTensor>> local_tiles;
 
     // iterates over tiles of hadamard indices
     for (Index h : H.tiles) {
@@ -798,26 +845,13 @@ auto einsum(expressions::TsrExpr<ArrayA_> A, expressions::TsrExpr<ArrayB_> B,
         shape = apply_inverse(P, shape);
         tile = tile.reshape(shape);
         if (P) tile = tile.permute(P);
-        local_tiles.push_back({c, tile});
+        C_local_tiles.emplace_back(std::move(c), std::move(tile));
       }
       // mark for lazy deletion
       C.ei = ArrayC();
     }
 
-    if constexpr (!ResultShape::is_dense()) {
-      TiledRange tiled_range = TiledRange(range_map[c]);
-      std::vector<std::pair<Index, float>> tile_norms;
-      for (auto &[index, tile] : local_tiles) {
-        tile_norms.push_back({index, tile.norm()});
-      }
-      ResultShape shape(world, tile_norms, tiled_range);
-      C.array = ArrayC(world, TiledRange(range_map[c]), shape);
-    }
-
-    for (auto &[index, tile] : local_tiles) {
-      if (C.array.is_zero(index)) continue;
-      C.array.set(index, tile);
-    }
+    build_C_array();
 
     for (auto &w : worlds) {
       w->gop.fence();