Merge pull request #555 from ValeevGroup/evaleev/fix/arena-tot-permuted-binary-ops

arena: support permuted Hadamard add/subt/mult on Tensor<ArenaTensor> ToT

Author: evaleev

src/TiledArray/expressions/cont_engine.h

@@ -305,9 +305,18 @@ class ContEngine : public BinaryEngine<Derived> {
           // this->product_type() is Tensor::Contraction, and,
           // this->implicit_permute_inner_ is false
 
-          return this->inner_product_type() == TensorProduct::Scale
-                     ? BipartitePermutation(outer(this->perm_))
-                     : this->perm_;
+          if (this->inner_product_type() == TensorProduct::Scale) {
+            // Owning inner cells apply the inner result permutation in the
+            // per-cell scale op, so they carry only the outer perm here. View
+            // (arena) cells instead use a perm-free per-cell op + an unpermuted
+            // arena plan and rely on op_'s post-processing permute for the
+            // inner perm -- so they carry the full perm, like inner
+            // Contraction.
+            if constexpr (!TiledArray::is_tensor_view_v<
+                              result_tile_element_type>)
+              return BipartitePermutation(outer(this->perm_));
+          }
+          return this->perm_;
         };
 
         auto total_perm = make_total_perm();
@@ -341,9 +350,18 @@ class ContEngine : public BinaryEngine<Derived> {
           // this->product_type() is Tensor::Contraction, and,
           // this->implicit_permute_inner_ is false
 
-          return this->inner_product_type() == TensorProduct::Scale
-                     ? BipartitePermutation(outer(this->perm_))
-                     : this->perm_;
+          if (this->inner_product_type() == TensorProduct::Scale) {
+            // Owning inner cells apply the inner result permutation in the
+            // per-cell scale op, so they carry only the outer perm here. View
+            // (arena) cells instead use a perm-free per-cell op + an unpermuted
+            // arena plan and rely on op_'s post-processing permute for the
+            // inner perm -- so they carry the full perm, like inner
+            // Contraction.
+            if constexpr (!TiledArray::is_tensor_view_v<
+                              result_tile_element_type>)
+              return BipartitePermutation(outer(this->perm_));
+          }
+          return this->perm_;
         };
 
         auto total_perm = make_total_perm();
@@ -949,10 +967,30 @@ class ContEngine : public BinaryEngine<Derived> {
                   result_tile_type, left_tile_type, right_tile_type>;
           if constexpr (arena_eligible_scale) {
             if (this->product_type() == TensorProduct::Contraction) {
+              // The inner perm handed to the plan must match how the inner
+              // *result* permutation is applied for this result cell type --
+              // and the two cell types apply it in different places:
+              //
+              //   * View (arena) cells: pass an identity inner perm so the
+              //     plan is always built (pre-shaping result cells in the
+              //     unpermuted operand inner layout) and the perm-free fused
+              //     scale op is selected; the inner result perm is applied
+              //     downstream by op_'s post-processing permute (carried in
+              //     make_total_perm for view cells).
+              //
+              //   * Owning cells: pass the inner result perm so the plan bails
+              //     (nullopt) on a non-identity inner perm, falling back to the
+              //     per-cell op that applies the inner perm itself -- matching
+              //     the outer-only total_perm make_total_perm carries here.
+              //     (A trivial inner perm still lets the plan + fused op run.)
+              Permutation plan_inner_perm;
+              if constexpr (!TiledArray::is_tensor_view_v<
+                                result_tile_element_type>)
+                plan_inner_perm = inner(this->perm_);
               this->arena_plan_ =
                   TiledArray::detail::make_contraction_arena_plan<
                       result_tile_type, left_tile_type, right_tile_type>(
-                      kind, std::nullopt, inner(this->perm_));
+                      kind, std::nullopt, plan_inner_perm);
             }
           }
           // Fallback per-element op for the scale inner-product when no

src/TiledArray/tensor/tensor.h

@@ -1953,33 +1953,28 @@ class Tensor {
       return !static_cast<bool>(perm) || perm.is_identity();
   }
 
-  /// Permuted add for `Tensor<ArenaTensor>` ToT operands. A non-trivial
-  /// permutation of arena ToT tiles is not yet supported; an identity (or
-  /// null) permutation falls through to the plain element-wise add.
+  /// Permuted add for `Tensor<ArenaTensor>` ToT operands. The operands are
+  /// congruent by the time a permuted product reaches a tile op, so the
+  /// elementwise `add(right)` is valid and `perm` is the result permutation;
+  /// `permute` applies it (shallow outer reindex + inner-slab rewrite).
   template <typename Right, typename Perm>
     requires(is_arena_tensor_v<value_type> &&
              is_arena_tensor_v<typename Right::value_type> &&
              detail::is_permutation_v<Perm>)
   Tensor add(const Right& right, const Perm& perm) const {
-    if (!arena_perm_is_trivial(perm))
-      TA_EXCEPTION(
-          "TA::Tensor<ArenaTensor>::add: permuted add of a tensor-of-tensors "
-          "is not yet supported");
-    return add(right);
+    auto result = add(right);
+    return arena_perm_is_trivial(perm) ? result : result.permute(perm);
   }
 
   /// Permuted scaled add for `Tensor<ArenaTensor>` ToT operands; see the
-  /// permuted-add overload above for the permutation restriction.
+  /// permuted-add overload above for the congruent-operand rationale.
   template <typename Right, typename Scalar, typename Perm>
     requires(is_arena_tensor_v<value_type> &&
              is_arena_tensor_v<typename Right::value_type> &&
              detail::is_numeric_v<Scalar> && detail::is_permutation_v<Perm>)
   Tensor add(const Right& right, const Scalar factor, const Perm& perm) const {
-    if (!arena_perm_is_trivial(perm))
-      TA_EXCEPTION(
-          "TA::Tensor<ArenaTensor>::add: permuted scaled add of a "
-          "tensor-of-tensors is not yet supported");
-    return add(right, factor);
+    auto result = add(right, factor);
+    return arena_perm_is_trivial(perm) ? result : result.permute(perm);
   }
 
   /// Add this and \c other to construct a new tensor
@@ -2382,8 +2377,15 @@ class Tensor {
       typename std::enable_if<is_tensor<Right>::value &&
                               detail::is_permutation_v<Perm>>::type* = nullptr>
   Tensor subt(const Right& right, const Perm& perm) const {
-    if constexpr (is_tensor_view_v<value_type>) {
-      // Permutation isn't supported for view inner cells (fixed storage
+    if constexpr (is_arena_tensor_v<value_type> &&
+                  is_arena_tensor_v<typename Right::value_type>) {
+      // arena ToT x arena ToT: operands are congruent at tile-op time, so the
+      // elementwise `subt(right)` is valid; apply the result permutation as a
+      // post-pass (shallow outer reindex + inner-slab rewrite).
+      auto result = subt(right);
+      return arena_perm_is_trivial(perm) ? result : result.permute(perm);
+    } else if constexpr (is_tensor_view_v<value_type>) {
+      // Permutation isn't supported for other view inner cells (fixed storage
       // layout). Subt+permute would require materialization.
       TA_EXCEPTION(
           "Tensor<View>::subt(right, perm): permutation is not "
@@ -2443,11 +2445,10 @@ class Tensor {
   Tensor subt(const Right& right, const Scalar factor, const Perm& perm) const {
     if constexpr (is_arena_tensor_v<value_type> &&
                   is_arena_tensor_v<typename Right::value_type>) {
-      if (!arena_perm_is_trivial(perm))
-        TA_EXCEPTION(
-            "TA::Tensor<ArenaTensor>::subt: permuted scaled subt of a "
-            "tensor-of-tensors is not yet supported");
-      return subt(right, factor);
+      // arena ToT x arena ToT scaled subtraction; see the unscaled permuted
+      // subt overload above for the congruent-operand rationale.
+      auto result = subt(right, factor);
+      return arena_perm_is_trivial(perm) ? result : result.permute(perm);
     } else {
       return binary(
           right,
@@ -2622,11 +2623,15 @@ class Tensor {
   decltype(auto) mult(const Right& right, const Perm& perm) const {
     if constexpr (is_arena_tensor_v<value_type> &&
                   is_arena_tensor_v<typename Right::value_type>) {
-      if (!arena_perm_is_trivial(perm))
-        TA_EXCEPTION(
-            "TA::Tensor<ArenaTensor>::mult: permuted mult of a "
-            "tensor-of-tensors is not yet supported");
-      return mult(right);
+      // arena ToT x arena ToT Hadamard product. By the time a permuted product
+      // reaches a tile op, the engine has already brought both operands to a
+      // common (congruent) layout, so the elementwise `mult(right)` is valid;
+      // `perm` is the result permutation (common layout -> target). Apply it
+      // as a post-pass: `permute` reindexes the outer cells shallowly
+      // (arena_permute_shallow) and rewrites the inner slab if the inner part
+      // of the permutation is non-trivial (arena_inner_permute).
+      auto result = mult(right);
+      return arena_perm_is_trivial(perm) ? result : result.permute(perm);
     } else if constexpr (detail::is_numeric_v<value_type> &&
                          is_arena_tensor_v<typename Right::value_type>) {
       // t x tot: a plain scalar tile times an arena ToT tile. The 2-arg
@@ -2697,11 +2702,11 @@ class Tensor {
                       const Perm& perm) const {
     if constexpr (is_arena_tensor_v<value_type> &&
                   is_arena_tensor_v<typename Right::value_type>) {
-      if (!arena_perm_is_trivial(perm))
-        TA_EXCEPTION(
-            "TA::Tensor<ArenaTensor>::mult: permuted scaled mult of a "
-            "tensor-of-tensors is not yet supported");
-      return mult(right, factor);
+      // arena ToT x arena ToT scaled Hadamard product; see the unscaled
+      // permuted mult overload above for the congruent-operand rationale.
+      // Scale during the elementwise product, then permute the result.
+      auto result = mult(right, factor);
+      return arena_perm_is_trivial(perm) ? result : result.permute(perm);
     } else {
       return binary(
           right,