mult_engine.h

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2014  Virginia Tech
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  Justus Calvin
 *  Department of Chemistry, Virginia Tech
 *
 *  mult_engine.h
 *  Mar 31, 2014
 *
 */

#ifndef TILEDARRAY_EXPRESSIONS_MULT_ENGINE_H__INCLUDED
#define TILEDARRAY_EXPRESSIONS_MULT_ENGINE_H__INCLUDED

#include <TiledArray/expressions/cont_engine.h>
#include <TiledArray/external/btas.h>
#include <TiledArray/tile_op/binary_wrapper.h>
#include <TiledArray/tile_op/mult.h>
#include <btas/optimize/contract.h>

namespace TiledArray {
namespace expressions {

// Forward declarations
template <typename, typename>
class MultExpr;
template <typename, typename, typename>
class ScalMultExpr;
template <typename, typename, typename>
class MultEngine;
template <typename, typename, typename, typename>
class ScalMultEngine;

template <typename Tile>
inline auto make_tile_contract_op(const IndexList& left_indices,
                                  const IndexList& right_indices,
                                  const IndexList& result_indices) {
  static_assert(TiledArray::detail::is_ta_tensor_v<Tile> ||
                    TiledArray::detail::is_btas_tensor_v<Tile>,
                "only support BTAS and TA tensors as inner tensors of ToT");
  // btas::contract only works for col-major storage
  if constexpr (TiledArray::detail::is_btas_tensor_v<
                    Tile>) {  // can do arbitrary contractions here
                              //    auto indices_to_integers = [&]() {
                              //      // collect all unique indices
                              //      IndexList::container_type all_indices;
    //      all_indices.reserve(left_indices.size() + right_indices.size() +
    //                          result_indices.size());
    //      all_indices.insert(all_indices.end(), left_indices.data().begin(),
    //                         left_indices.data().end());
    //      all_indices.insert(all_indices.end(), right_indices.data().begin(),
    //                         right_indices.data().end());
    //      all_indices.insert(all_indices.end(), result_indices.data().begin(),
    //                         result_indices.data().end());
    //      auto end_of_unique = std::unique(all_indices.begin(),
    //      all_indices.end()); all_indices.erase(end_of_unique,
    //      all_indices.end());
    //
    //      // map each index list to numbers
    //      auto to_integers = [](const container::svector<std::string> indices,
    //                            const container::svector<std::string>
    //                            all_indices) {
    //        TA_ASSERT(!all_indices.empty());
    //        btas::DEFAULT::index<std::int64_t> result;
    //        result.reserve(indices.size());
    //        for (const auto& idx : indices) {
    //          auto pos = std::find(all_indices.begin(), all_indices.end(),
    //          idx) -
    //                     all_indices.begin();
    //          TA_ASSERT(pos < all_indices.size());
    //          result.push_back(pos);
    //        }
    //        return result;
    //      };
    //      auto left_annotation = to_integers(left_indices.data(),
    //      all_indices); auto right_annotation =
    //      to_integers(right_indices.data(), all_indices); auto
    //      result_annotation = to_integers(result_indices.data(), all_indices);
    //      return std::make_tuple(left_annotation, right_annotation,
    //                             result_annotation);
    //    };
    //    auto [left_annotation, right_annotation, result_annotation] =
    //        indices_to_integers();
    //    return [left_annotation = std::move(left_annotation),
    //            right_annotation = std::move(right_annotation),
    //            result_annotation = std::move(result_annotation)](
    //               const Tile& left_tile, const Tile& right_tile) {
    //      Tile result;
    //      btas::contract(1, left_tile, left_annotation, right_tile,
    //                     right_annotation, 0, result, result_annotation);
    //      return result;
    //    };
  }
}

template <typename Left, typename Right, typename Result>
struct EngineTrait<MultEngine<Left, Right, Result>> {
  static_assert(
      std::is_same<typename EngineTrait<Left>::policy,
                   typename EngineTrait<Right>::policy>::value,
      "The left- and right-hand expressions must use the same policy class");

  // Argument typedefs
  typedef Left left_type;    ///< The left-hand expression type
  typedef Right right_type;  ///< The right-hand expression type

  // Operational typedefs
  typedef TiledArray::detail::Mult<
      Result, typename EngineTrait<Left>::eval_type,
      typename EngineTrait<Right>::eval_type, EngineTrait<Left>::consumable,
      EngineTrait<Right>::consumable>
      op_base_type;  ///< The base tile operation type
  typedef TiledArray::detail::BinaryWrapper<op_base_type>
      op_type;  ///< The tile operation type
  typedef typename op_type::result_type value_type;  ///< The result tile type
  typedef typename eval_trait<value_type>::type
      eval_type;  ///< Evaluation tile type
  typedef typename TiledArray::detail::numeric_type<value_type>::type
      scalar_type;                       ///< Tile scalar type
  typedef typename Left::policy policy;  ///< The result policy type
  typedef TiledArray::detail::DistEval<value_type, policy>
      dist_eval_type;  ///< The distributed evaluator type

  // Meta data typedefs
  typedef typename policy::ordinal_type size_type;   ///< Size type
  typedef typename policy::trange_type trange_type;  ///< Tiled range type
  typedef typename policy::shape_type shape_type;    ///< Shape type
  typedef typename policy::pmap_interface
      pmap_interface;  ///< Process map interface type

  static constexpr bool consumable = is_consumable_tile<eval_type>::value;
  static constexpr unsigned int leaves =
      EngineTrait<Left>::leaves + EngineTrait<Right>::leaves;
};

template <typename Left, typename Right, typename Scalar, typename Result>
struct EngineTrait<ScalMultEngine<Left, Right, Scalar, Result>> {
  static_assert(
      std::is_same<typename EngineTrait<Left>::policy,
                   typename EngineTrait<Right>::policy>::value,
      "The left- and right-hand expressions must use the same policy class");

  // Argument typedefs
  typedef Left left_type;    ///< The left-hand expression type
  typedef Right right_type;  ///< The right-hand expression type

  // Operational typedefs
  typedef Scalar scalar_type;  ///< Tile scalar type
  typedef TiledArray::detail::ScalMult<
      Result, typename EngineTrait<Left>::eval_type,
      typename EngineTrait<Right>::eval_type, scalar_type,
      EngineTrait<Left>::consumable, EngineTrait<Right>::consumable>
      op_base_type;  ///< The base tile operation type
  typedef TiledArray::detail::BinaryWrapper<op_base_type>
      op_type;  ///< The tile operation type
  typedef typename op_type::result_type value_type;  ///< The result tile type
  typedef typename eval_trait<value_type>::type
      eval_type;                         ///< Evaluation tile type
  typedef typename Left::policy policy;  ///< The result policy type
  typedef TiledArray::detail::DistEval<value_type, policy>
      dist_eval_type;  ///< The distributed evaluator type

  // Meta data typedefs
  typedef typename policy::ordinal_type size_type;   ///< Size type
  typedef typename policy::trange_type trange_type;  ///< Tiled range type
  typedef typename policy::shape_type shape_type;    ///< Shape type
  typedef typename policy::pmap_interface
      pmap_interface;  ///< Process map interface type

  static constexpr bool consumable = is_consumable_tile<eval_type>::value;
  static constexpr unsigned int leaves =
      EngineTrait<Left>::leaves + EngineTrait<Right>::leaves;
};

/// Multiplication expression engine

/// This implements any expression encoded with the multiplication operator.
/// This includes Hadamard product, e.g.
/// \code (c("i,j")=)a("i,j")*b("i,j") \endcode ,
/// and pure contractions, e.g. \code (c("i,j")=)a("i,k")*b("k,j") \endcode .
/// \internal mixed Hadamard-contraction case, e.g.
/// \code c("i,j,l")=a("i,l,k")*b("j,l,k") \endcode ,
/// is not supported since
/// this requires that the result labels are assigned by user (currently they
/// are computed by this engine)
/// \tparam Left The left-hand engine type
/// \tparam Right The right-hand engine type
/// \tparam Result The result tile type
template <typename Left, typename Right, typename Result>
class MultEngine : public ContEngine<MultEngine<Left, Right, Result>> {
 public:
  // Class hierarchy typedefs
  typedef MultEngine<Left, Right, Result> MultEngine_;  ///< This class type
  typedef ContEngine<MultEngine_>
      ContEngine_;  ///< Contraction engine base class
  typedef BinaryEngine<MultEngine_> BinaryEngine_;  ///< Binary base class type
  typedef BinaryEngine<MultEngine_>
      ExprEngine_;  ///< Expression engine base class type

  // Argument typedefs
  typedef typename EngineTrait<MultEngine_>::left_type
      left_type;  ///< The left-hand expression type
  typedef typename EngineTrait<MultEngine_>::right_type
      right_type;  ///< The right-hand expression type

  // Operational typedefs
  typedef typename EngineTrait<MultEngine_>::value_type
      value_type;  ///< The result tile type
  typedef typename EngineTrait<MultEngine_>::op_base_type
      op_base_type;  ///< The tile operation type
  typedef typename EngineTrait<MultEngine_>::op_type
      op_type;  ///< The tile operation type
  typedef typename EngineTrait<MultEngine_>::policy
      policy;  ///< The result policy type
  typedef typename EngineTrait<MultEngine_>::dist_eval_type
      dist_eval_type;  ///< The distributed evaluator type
  typedef typename EngineTrait<MultEngine_>::scalar_type
      scalar_type;  ///< Tile scalar type

  // Meta data typedefs
  typedef
      typename EngineTrait<MultEngine_>::size_type size_type;  ///< Size type
  typedef typename EngineTrait<MultEngine_>::trange_type
      trange_type;  ///< Tiled range type
  typedef
      typename EngineTrait<MultEngine_>::shape_type shape_type;  ///< Shape type
  typedef typename EngineTrait<MultEngine_>::pmap_interface
      pmap_interface;  ///< Process map interface type

 public:
  /// Constructor

  /// \tparam L The left-hand argument expression type
  /// \tparam R The right-hand argument expression type
  /// \param expr The parent expression
  template <typename L, typename R>
  MultEngine(const MultExpr<L, R>& expr) : ContEngine_(expr) {}

  /// Set the index list for this expression

  /// This function will set the index list for this expression and its
  /// children such that the number of permutations is minimized. The final
  /// index list may not be set to target, which indicates that the
  /// result of this expression will be permuted to match \c target_indices.
  /// \param target_indices The target index list for this expression
  void perm_indices(const BipartiteIndexList& target_indices) {
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::perm_indices(target_indices);
    else if (this->product_type() == TensorProduct::General) {
      // mirror ContEngine_::perm_indices, but lay out via
      // GeneralPermutationOptimizer (the target determines which shared
      // indices are fused vs contracted)
      if (!this->implicit_permute()) {
        BinaryEngine_::template init_indices_<TensorProduct::General>(
            target_indices);
        if (BinaryEngine_::left_indices_ != BinaryEngine_::left_.indices())
          BinaryEngine_::left_.perm_indices(BinaryEngine_::left_indices_);
        if (BinaryEngine_::right_indices_ != BinaryEngine_::right_.indices())
          BinaryEngine_::right_.perm_indices(BinaryEngine_::right_indices_);
      }
    } else {
      BinaryEngine_::perm_indices(target_indices);
    }
  }

  /// Initialize the index list of this expression

  /// \param target_indices The target index list for this expression
  void init_indices(const BipartiteIndexList& target_indices) {
    // Deduce each child's index SET top-down (Phase E). The index set an
    // inner node must produce depends on what its ancestors need: a shared
    // index of the two children is fused iff this node's target carries it,
    // contracted here otherwise; an index of one child that neither the
    // sibling nor the target carries is consumed entirely within that
    // child's subtree and is not demanded of it.
    // Each child's demand is ordered target-kept-first (the indices this
    // node's result keeps, in target order) followed by the
    // contracted-here indices in the child's leaf-availability order --
    // fused indices thus lead, matching the canonical layouts of
    // GeneralPermutationOptimizer (h leading) so the nbatch fold stays
    // zero-copy. E.g. in the THC-like g("p,q,r,s") = X("p,r1") * X("q,r1")
    // * Z("r1,r2") * ... the index r1 is demanded of X*X by its consumer
    // (fused there), and contracted where Z meets the X*X subtree.
    // N.B. expressions consumed WITHOUT a target (reductions, e.g. dot)
    // take the no-target init_indices() overload below, which retains the
    // bottom-up contraction convention -- general products under
    // reductions remain unsupported.
    {
      auto const avail_l = BinaryEngine_::left_.available_indices();
      auto const avail_r = BinaryEngine_::right_.available_indices();
      auto demand = [](auto const& avail, auto const& sibling,
                       auto const& tgt) {
        container::svector<std::string> r;
        for (auto&& idx : tgt)
          if (avail.count(idx)) r.push_back(idx);
        for (auto&& idx : avail) {
          if (tgt.count(idx)) continue;
          if (sibling.count(idx)) r.push_back(idx);
          // else: the index is consumed entirely within the child's own
          // subtree (contracted deeper down) -- not demanded here. A
          // genuinely orphaned index (single occurrence, demanded nowhere =
          // implicit trace, unsupported) surfaces as the leaf-level
          // target-is-not-a-permutation error.
        }
        return r;
      };
      auto bipartite_demand = [&demand](auto const& avail, auto const& sib,
                                        auto const& tgt) {
        auto const out = demand(outer(avail), outer(sib), outer(tgt));
        auto const in = demand(inner(avail), inner(sib), inner(tgt));
        return BipartiteIndexList(IndexList(out.begin(), out.end()),
                                  IndexList(in.begin(), in.end()));
      };
      // each child dictates the ORDER of its demand (preferred_layout): a
      // product child reorders to its canonical (h, eA, eB) layout -- a
      // general product cannot host a result permutation, and contraction
      // consumers absorb any child layout via the GEMM transpose forms
      BinaryEngine_::left_.init_indices(BinaryEngine_::left_.preferred_layout(
          bipartite_demand(avail_l, avail_r, target_indices)));
      BinaryEngine_::right_.init_indices(BinaryEngine_::right_.preferred_layout(
          bipartite_demand(avail_r, avail_l, target_indices)));
    }

    this->product_type_ = compute_product_type(
        outer(BinaryEngine_::left_.indices()),
        outer(BinaryEngine_::right_.indices()), outer(target_indices));
    this->inner_product_type_ = compute_product_type(
        inner(BinaryEngine_::left_.indices()),
        inner(BinaryEngine_::right_.indices()), inner(target_indices));

    if (this->inner_product_type_ == TensorProduct::General)
      TA_EXCEPTION(
          "MultEngine: general products (fused + contracted + free indices) "
          "between the inner (nested) indices of tensors-of-tensors are not "
          "supported");

    // Check the *outer* indices to determine whether the arguments are
    // - Hadamard-multiplied (outer left, right, and target indices are all
    //   related by permutations),
    // - contracted (no index appears in left, right, AND target), or
    // - generally multiplied (fused, contracted, and free indices coexist,
    //   e.g. "b,i,j" * "b,j,k" -> "b,i,k"); the layout then depends on the
    //   target, which determines the role (fused vs contracted) of every
    //   shared index.
    // Only the outer indices matter here since the inner indices only encode
    // the tile op; the type of the tile op does not need to match the type of
    // the operation on the outer indices
    if (this->product_type() == TensorProduct::Hadamard) {
      BinaryEngine_::perm_indices(target_indices);
    } else if (this->product_type() == TensorProduct::General) {
      this->perm_indices(target_indices);
    } else {
      auto children_initialized = true;
      ContEngine_::init_indices(children_initialized);
      ContEngine_::perm_indices(target_indices);
    }
  }

  /// Initialize the index list of this expression
  void init_indices() {
    // to decide what type of product this is must initialize indices down the
    // tree
    BinaryEngine_::left_.init_indices();
    BinaryEngine_::right_.init_indices();
    auto children_initialized = true;
    this->product_type_ =
        compute_product_type(outer(BinaryEngine_::left_.indices()),
                             outer(BinaryEngine_::right_.indices()));
    this->inner_product_type_ =
        compute_product_type(inner(BinaryEngine_::left_.indices()),
                             inner(BinaryEngine_::right_.indices()));

    if (this->product_type() == TensorProduct::Hadamard) {
      BinaryEngine_::init_indices(children_initialized);
    } else {
      ContEngine_::init_indices(children_initialized);
    }
  }

  /// \return the layout this product prefers for producing the index set of
  /// \p demand: the canonical general-product result layout (h, eA, eB) --
  /// indices supplied by both children (fused here) lead, then indices
  /// supplied by the left child only, then by the right child only, each
  /// group in demand order. h-leading is required when this node resolves to
  /// a general product (general results cannot host a result permutation and
  /// the nbatch fold needs the fused modes leading); for a pure contraction
  /// (empty h) the (eA, eB) order matches the GEMM result, and for a pure
  /// Hadamard the demand order is preserved.
  BipartiteIndexList preferred_layout(const BipartiteIndexList& demand) const {
    auto const avail_l = BinaryEngine_::left_.available_indices();
    auto const avail_r = BinaryEngine_::right_.available_indices();
    auto canonical = [](auto const& d, auto const& al, auto const& ar) {
      container::svector<std::string> h, ea, eb;
      for (auto&& idx : d) {
        bool const inl = al.count(idx);
        bool const inr = ar.count(idx);
        if (inl && inr)
          h.push_back(idx);
        else if (inl)
          ea.push_back(idx);
        else
          eb.push_back(idx);
      }
      h.insert(h.end(), ea.begin(), ea.end());
      h.insert(h.end(), eb.begin(), eb.end());
      return h;
    };
    auto const out = canonical(outer(demand), outer(avail_l), outer(avail_r));
    auto const in = canonical(inner(demand), inner(avail_l), inner(avail_r));
    return BipartiteIndexList(IndexList(out.begin(), out.end()),
                              IndexList(in.begin(), in.end()));
  }

  /// Initialize result tensor structure

  /// This function will initialize the permutation, tiled range, and shape
  /// for the result tensor.
  /// \param target_indices The target index list for the result tensor
  void init_struct(const BipartiteIndexList& target_indices) {
    if (this->product_type() == TensorProduct::General) {
      // the inner tile op (for tensors-of-tensors) must be initialized
      // first; init_struct_general consumes element_nonreturn_op_ and the
      // arena plan it builds
      this->init_inner_tile_op(inner(target_indices));
      ContEngine_::init_struct_general(target_indices);
      return;
    }

    this->init_perm(target_indices);

    // for ContEngine_::init_struct need to initialize element op first
    this->init_inner_tile_op(inner(target_indices));
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::init_struct(target_indices);
    else
      BinaryEngine_::init_struct(target_indices);
  }

  /// Initialize result tensor distribution

  /// This function will initialize the world and process map for the result
  /// tensor.
  /// \param world The world were the result will be distributed
  /// \param pmap The process map for the result tensor tiles
  void init_distribution(World* world,
                         std::shared_ptr<const pmap_interface> pmap) {
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::init_distribution(world, pmap);
    else if (this->product_type() == TensorProduct::General)
      ContEngine_::init_distribution_general(world, pmap);
    else
      BinaryEngine_::init_distribution(world, pmap);
  }

  /// Non-permuting tiled range factory function

  /// \return The result tiled range object
  trange_type make_trange() const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_trange();
    else if (this->product_type() == TensorProduct::General)
      return ContEngine_::make_trange_general();
    else
      return BinaryEngine_::make_trange();
  }

  /// Permuting tiled range factory function

  /// \param perm The permutation to be applied to the array
  /// \return The result tiled range object
  trange_type make_trange(const Permutation& perm) const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_trange(perm);
    else
      return BinaryEngine_::make_trange(perm);
  }

  /// Non-permuting shape factory function

  /// \return The result shape
  shape_type make_shape() const {
    return BinaryEngine_::left_.shape().mult(BinaryEngine_::right_.shape());
  }

  /// Permuting shape factory function

  /// \param perm The permutation to be applied to the array
  /// \return The result shape
  shape_type make_shape(const Permutation& perm) const {
    return BinaryEngine_::left_.shape().mult(BinaryEngine_::right_.shape(),
                                             outer(perm));
  }

  /// Non-permuting tile operation factory function

  /// \return The tile operation
  op_type make_tile_op() const {
    if constexpr (TiledArray::detail::is_tensor_of_tensor_v<
                      value_type>) {  // nested tensors
      const auto inner_prod = this->inner_product_type();
      if (inner_prod == TensorProduct::Hadamard) {
        TA_ASSERT(this->product_type() ==
                  inner_prod);  // Hadamard automatically works for inner
                                // dimensions as well
        return op_type(op_base_type());
      } else if (inner_prod == TensorProduct::Contraction) {
        if constexpr (TiledArray::is_tensor_view_v<
                          typename value_type::value_type>) {
          // arena ToT: a view inner cell cannot host a value-returning
          // per-cell op, so delegate the whole tile product to the arena op
          // built in init_inner_tile_op
          return op_type(op_base_type(typename op_base_type::tile_op_tag{},
                                      this->arena_hadamard_tile_op_));
        } else {
          return op_type(op_base_type(this->element_return_op_));
        }
      } else if (inner_prod == TensorProduct::Scale) {
        return op_type(op_base_type());
      } else
        abort();
    } else {  // plain tensors
      return op_type(op_base_type());
    }
    abort();  // unreachable
  }

  /// Permuting tile operation factory function

  /// \param perm The permutation to be applied to the result
  /// \return The tile operation
  template <typename Perm,
            typename = std::enable_if_t<TiledArray::detail::is_permutation_v<
                std::remove_reference_t<Perm>>>>
  op_type make_tile_op(Perm&& perm) const {
    if constexpr (TiledArray::detail::is_tensor_of_tensor_v<
                      value_type>) {  // nested tensors
      const auto inner_prod = this->inner_product_type();
      if (inner_prod == TensorProduct::Hadamard) {
        TA_ASSERT(this->product_type() ==
                  inner_prod);  // Hadamard automatically works for inner
                                // dimensions as well
        return op_type(op_base_type(), std::forward<Perm>(perm));
      } else if (inner_prod == TensorProduct::Contraction) {
        // inner permutation, if needed, was fused into inner op, do not apply
        // inner part of the perm again
        if constexpr (TiledArray::is_tensor_view_v<
                          typename value_type::value_type>) {
          return op_type(op_base_type(typename op_base_type::tile_op_tag{},
                                      this->arena_hadamard_tile_op_),
                         outer(std::forward<Perm>(perm)));
        } else {
          return op_type(op_base_type(this->element_return_op_),
                         outer(std::forward<Perm>(perm)));
        }
      } else if (inner_prod == TensorProduct::Scale) {
        if constexpr (TiledArray::is_tensor_view_v<
                          typename value_type::value_type>) {
          // arena ToT: a view result cell cannot be value-assigned a scaled
          // tensor, so the element_return_op_ path is unusable. Route through
          // the arena-aware mult CPO with the full permutation instead -- it
          // shapes and fills the result tile as a unit and applies the
          // (outer + inner) result permutation in place.
          return op_type(op_base_type(), std::forward<Perm>(perm));
        } else {
          // inner permutation, if needed, was fused into inner op, do not
          // apply inner part of the perm again
          return op_type(op_base_type(this->element_return_op_),
                         outer(std::forward<Perm>(perm)));
        }
      } else
        abort();
    } else {  // plain tensor
      return op_type(op_base_type(), std::forward<Perm>(perm));
    }
    abort();  // unreachable
  }

  /// Construct the distributed evaluator for this expression

  /// \return The distributed evaluator that will evaluate this expression
  dist_eval_type make_dist_eval() const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_dist_eval();
    else if (this->product_type() == TensorProduct::General)
      return ContEngine_::make_dist_eval_general();
    else
      return BinaryEngine_::make_dist_eval();
  }

  /// Expression identification tag

  /// \return An expression tag used to identify this expression
  const char* make_tag() const { return "[*] "; }

  /// Expression print

  /// \param os The output stream
  /// \param target_indices The target index list for this expression
  void print(ExprOStream os, const BipartiteIndexList& target_indices) const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::print(os, target_indices);
    else
      return BinaryEngine_::print(os, target_indices);
  }
};  // class MultEngine

/// Scaled multiplication expression engine

/// Similar to MultEngine but implements the result of MultEngine scaled by a
/// constant. \tparam Left The left-hand engine type \tparam Right The
/// Right-hand engine type \tparam Scalar The scaling factor type \tparam Result
/// The result tile type
template <typename Left, typename Right, typename Scalar, typename Result>
class ScalMultEngine
    : public ContEngine<ScalMultEngine<Left, Right, Scalar, Result>> {
 public:
  // Class hierarchy typedefs
  typedef ScalMultEngine<Left, Right, Scalar, Result>
      ScalMultEngine_;  ///< This class type
  typedef ContEngine<ScalMultEngine_>
      ContEngine_;  ///< Contraction engine base class
  typedef BinaryEngine<ScalMultEngine_>
      BinaryEngine_;  ///< Binary base class type
  typedef BinaryEngine<ScalMultEngine_>
      ExprEngine_;  ///< Expression engine base class type

  // Argument typedefs
  typedef typename EngineTrait<ScalMultEngine_>::left_type
      left_type;  ///< The left-hand expression type
  typedef typename EngineTrait<ScalMultEngine_>::right_type
      right_type;  ///< The right-hand expression type

  // Operational typedefs
  typedef typename EngineTrait<ScalMultEngine_>::value_type
      value_type;  ///< The result tile type
  typedef typename EngineTrait<ScalMultEngine_>::scalar_type
      scalar_type;  ///< Tile scalar type
  typedef typename EngineTrait<ScalMultEngine_>::op_base_type
      op_base_type;  ///< The tile operation type
  typedef typename EngineTrait<ScalMultEngine_>::op_type
      op_type;  ///< The tile operation type
  typedef typename EngineTrait<ScalMultEngine_>::policy
      policy;  ///< The result policy type
  typedef typename EngineTrait<ScalMultEngine_>::dist_eval_type
      dist_eval_type;  ///< The distributed evaluator type

  // Meta data typedefs
  typedef typename EngineTrait<ScalMultEngine_>::size_type
      size_type;  ///< Size type
  typedef typename EngineTrait<ScalMultEngine_>::trange_type
      trange_type;  ///< Tiled range type
  typedef typename EngineTrait<ScalMultEngine_>::shape_type
      shape_type;  ///< Shape type
  typedef typename EngineTrait<ScalMultEngine_>::pmap_interface
      pmap_interface;  ///< Process map interface type

 public:
  /// Constructor

  /// \tparam L The left-hand argument expression type
  /// \tparam R The right-hand argument expression type
  /// \tparam S The expression scalar type
  /// \param expr The parent expression
  template <typename L, typename R, typename S>
  ScalMultEngine(const ScalMultExpr<L, R, S>& expr) : ContEngine_(expr) {}

  /// Set the index list for this expression

  /// This function will set the index list for this expression and its
  /// children such that the number of permutations is minimized. The final
  /// index list may not be set to target, which indicates that the
  /// result of this expression will be permuted to match \c target_indices.
  /// \param target_indices The target index list for this expression
  void perm_indices(const BipartiteIndexList& target_indices) {
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::perm_indices(target_indices);
    else if (this->product_type() == TensorProduct::General) {
      if (!this->implicit_permute()) {
        BinaryEngine_::template init_indices_<TensorProduct::General>(
            target_indices);
        if (BinaryEngine_::left_indices_ != BinaryEngine_::left_.indices())
          BinaryEngine_::left_.perm_indices(BinaryEngine_::left_indices_);
        if (BinaryEngine_::right_indices_ != BinaryEngine_::right_.indices())
          BinaryEngine_::right_.perm_indices(BinaryEngine_::right_indices_);
      }
    } else {
      BinaryEngine_::perm_indices(target_indices);
    }
  }

  /// Initialize the index list of this expression

  /// \param target_indices The target index list for this expression
  void init_indices(const BipartiteIndexList& target_indices) {
    BinaryEngine_::left_.init_indices();
    BinaryEngine_::right_.init_indices();
    this->product_type_ = compute_product_type(
        outer(BinaryEngine_::left_.indices()),
        outer(BinaryEngine_::right_.indices()), outer(target_indices));

    if (this->product_type() == TensorProduct::Hadamard) {
      // since already initialized left and right arg indices assign the target
      // indices
      BinaryEngine_::perm_indices(target_indices);
    } else if (this->product_type() == TensorProduct::General) {
      // layout via GeneralPermutationOptimizer (the target determines which
      // shared indices are fused vs contracted), then propagate to children
      if (!this->implicit_permute()) {
        BinaryEngine_::template init_indices_<TensorProduct::General>(
            target_indices);
        if (BinaryEngine_::left_indices_ != BinaryEngine_::left_.indices())
          BinaryEngine_::left_.perm_indices(BinaryEngine_::left_indices_);
        if (BinaryEngine_::right_indices_ != BinaryEngine_::right_.indices())
          BinaryEngine_::right_.perm_indices(BinaryEngine_::right_indices_);
      }
    } else {
      ContEngine_::init_indices(target_indices);
    }
  }

  /// Initialize the index list of this expression
  void init_indices() {
    BinaryEngine_::left_.init_indices();
    BinaryEngine_::right_.init_indices();
    this->product_type_ =
        compute_product_type(outer(BinaryEngine_::left_.indices()),
                             outer(BinaryEngine_::right_.indices()));
    if (this->product_type() == TensorProduct::Hadamard) {
      auto outer_indices = outer((left_type::leaves <= right_type::leaves)
                                     ? BinaryEngine_::left_.indices()
                                     : BinaryEngine_::right_.indices());
      // assume inner op is also Hadamard
      // TODO compute inner indices using inner_product_type_
      auto inner_indices = inner((left_type::leaves <= right_type::leaves)
                                     ? BinaryEngine_::left_.indices()
                                     : BinaryEngine_::right_.indices());
      ExprEngine_::indices_ = BipartiteIndexList(outer_indices, inner_indices);
    } else {
      ContEngine_::init_indices();
    }
  }

  /// Initialize result tensor structure

  /// This function will initialize the permutation, tiled range, and shape
  /// for the result tensor.
  /// \param target_indices The target index list for the result tensor
  void init_struct(const BipartiteIndexList& target_indices) {
    // TODO Phase B (batched Summa): evaluate general products natively
    if (this->product_type() == TensorProduct::General)
      TA_EXCEPTION(
          "ScalMultEngine: evaluation of general products (fused + contracted "
          "+ free indices) via the expression layer is not yet implemented; "
          "use TiledArray::einsum() instead");

    this->init_perm(target_indices);

    // for ContEngine_::init_struct need to initialize element op first
    this->init_inner_tile_op(inner(target_indices));
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::init_struct(target_indices);
    else
      BinaryEngine_::init_struct(target_indices);
  }

  /// Initialize result tensor distribution

  /// This function will initialize the world and process map for the result
  /// tensor.
  /// \param world The world were the result will be distributed
  /// \param pmap The process map for the result tensor tiles
  void init_distribution(World* world,
                         std::shared_ptr<const pmap_interface> pmap) {
    if (this->product_type() == TensorProduct::Contraction)
      ContEngine_::init_distribution(world, pmap);
    else
      BinaryEngine_::init_distribution(world, pmap);
  }

  /// Construct the distributed evaluator for this expression

  /// \return The distributed evaluator that will evaluate this expression
  dist_eval_type make_dist_eval() const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_dist_eval();
    else
      return BinaryEngine_::make_dist_eval();
  }

  /// Non-permuting tiled range factory function

  /// \return The result tiled range object
  trange_type make_trange() const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_trange();
    else
      return BinaryEngine_::make_trange();
  }

  /// Permuting tiled range factory function

  /// \param perm The permutation to be applied to the array
  /// \return The result tiled range object
  trange_type make_trange(const Permutation& perm) const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::make_trange(perm);
    else
      return BinaryEngine_::make_trange(perm);
  }

  /// Non-permuting shape factory function

  /// \return The result shape
  shape_type make_shape() const {
    return BinaryEngine_::left_.shape().mult(BinaryEngine_::right_.shape(),
                                             ContEngine_::factor_);
  }

  /// Permuting shape factory function

  /// \param perm The permutation to be applied to the array
  /// \return The result shape
  shape_type make_shape(const Permutation& perm) const {
    return BinaryEngine_::left_.shape().mult(BinaryEngine_::right_.shape(),
                                             ContEngine_::factor_, perm);
  }

  /// Non-permuting tile operation factory function

  /// \return The tile operation
  op_type make_tile_op() const {
    return op_type(op_base_type(ContEngine_::factor_));
  }

  /// Permuting tile operation factory function

  /// \param perm The permutation to be applied to tiles
  /// \return The tile operation
  template <typename Perm,
            typename = std::enable_if_t<TiledArray::detail::is_permutation_v<
                std::remove_reference_t<Perm>>>>
  op_type make_tile_op(Perm&& perm) const {
    return op_type(op_base_type(ContEngine_::factor_),
                   std::forward<Perm>(perm));
  }

  /// Expression identification tag

  /// \return An expression tag used to identify this expression
  std::string make_tag() const {
    std::stringstream ss;
    ss << "[*] [" << ContEngine_::factor_ << "] ";
    return ss.str();
  }

  /// Expression print

  /// \param os The output stream
  /// \param target_indices The target index list for this expression
  void print(ExprOStream os, const BipartiteIndexList& target_indices) const {
    if (this->product_type() == TensorProduct::Contraction)
      return ContEngine_::print(os, target_indices);
    else
      return BinaryEngine_::print(os, target_indices);
  }

};  // class ScalMultEngine

}  // namespace expressions
}  // namespace TiledArray

#endif  // TILEDARRAY_EXPRESSIONS_MULT_ENGINE_H__INCLUDED