cont_engine.h

Go to the documentation of this file.

/*
 *  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
 *
 *  cont_engine.h
 *  Mar 31, 2014
 *
 */
 
#ifndef TILEDARRAY_EXPRESSIONS_CONT_ENGINE_H__INCLUDED
#define TILEDARRAY_EXPRESSIONS_CONT_ENGINE_H__INCLUDED
 
#include <TiledArray/dist_eval/contraction_eval.h>
#include <TiledArray/expressions/binary_engine.h>
#include <TiledArray/expressions/permopt.h>
#include <TiledArray/proc_grid.h>
#include <TiledArray/tensor/utility.h>
#include <TiledArray/tile_op/contract_reduce.h>
#include <TiledArray/tile_op/mult.h>
 
namespace TiledArray {
namespace expressions {
 
// Forward declarations
template <typename, typename>
class MultExpr;
template <typename, typename, typename>
class ScalMultExpr;
 
 
template <typename Derived>
class ContEngine : public BinaryEngine<Derived> {
 public:
  // Class hierarchy typedefs
  typedef ContEngine<Derived> ContEngine_;      
  typedef BinaryEngine<Derived> BinaryEngine_;  
  typedef ExprEngine<Derived>
      ExprEngine_;  
 
  // Argument typedefs
  typedef typename EngineTrait<Derived>::left_type
      left_type;  
  typedef typename EngineTrait<Derived>::right_type
      right_type;  
 
  // Operational typedefs
  typedef typename EngineTrait<Derived>::value_type
      value_type;  
  typedef typename EngineTrait<Derived>::scalar_type
      scalar_type;  
  typedef TiledArray::detail::ContractReduce<
      value_type, typename eval_trait<typename left_type::value_type>::type,
      typename eval_trait<typename right_type::value_type>::type,
      scalar_type>
      op_type;  
  typedef
      typename EngineTrait<Derived>::policy policy;  
  typedef typename EngineTrait<Derived>::dist_eval_type
      dist_eval_type;  
 
  // Meta data typedefs
  typedef typename EngineTrait<Derived>::size_type size_type;  
  typedef typename EngineTrait<Derived>::trange_type
      trange_type;  
  typedef typename EngineTrait<Derived>::shape_type shape_type;  
  typedef typename EngineTrait<Derived>::pmap_interface
      pmap_interface;  
 
 protected:
  // Import base class variables to this scope
  using BinaryEngine_::left_;
  using BinaryEngine_::left_indices_;
  using BinaryEngine_::left_inner_permtype_;
  using BinaryEngine_::left_outer_permtype_;
  using BinaryEngine_::right_;
  using BinaryEngine_::right_indices_;
  using BinaryEngine_::right_inner_permtype_;
  using BinaryEngine_::right_outer_permtype_;
  using ExprEngine_::indices_;
  using ExprEngine_::perm_;
  using ExprEngine_::permute_tiles_;
  using ExprEngine_::pmap_;
  using ExprEngine_::shape_;
  using ExprEngine_::trange_;
  using ExprEngine_::world_;
 
 protected:
  scalar_type factor_;  
 
 protected:
  op_type op_;  
  using tile_element_type = typename value_type::value_type;
  std::function<void(tile_element_type&, const tile_element_type&,
                     const tile_element_type&)>
      inner_tile_nonreturn_op_;  
  std::function<tile_element_type(const tile_element_type&,
                                  const tile_element_type&)>
      inner_tile_return_op_;  
  TiledArray::detail::ProcGrid
      proc_grid_;    
  size_type K_ = 1;  
 
  static unsigned int find(const BipartiteIndexList& indices,
                           const std::string& index_label, unsigned int i,
                           const unsigned int n) {
    for (; i < n; ++i) {
      if (indices[i] == index_label) break;
    }
 
    return i;
  }
 
  TensorProduct product_type_ = TensorProduct::Invalid;
  TensorProduct inner_product_type_ = TensorProduct::Invalid;
 
  TensorProduct product_type() const {
    TA_ASSERT(product_type_ !=
              TensorProduct::Invalid);  // init_indices() must initialize this
    TA_ASSERT(product_type_ == TensorProduct::Hadamard ||
              product_type_ == TensorProduct::Contraction);
    return product_type_;
  }
 
  TensorProduct inner_product_type() const {
    TA_ASSERT(inner_product_type_ !=
              TensorProduct::Invalid);  // init_indices() must initialize this
    TA_ASSERT(inner_product_type_ == TensorProduct::Hadamard ||
              inner_product_type_ == TensorProduct::Contraction);
    return inner_product_type_;
  }
 
 public:
 
  template <typename L, typename R>
  ContEngine(const MultExpr<L, R>& expr) : BinaryEngine_(expr), factor_(1) {}
 
 
  template <typename L, typename R, typename S>
  ContEngine(const ScalMultExpr<L, R, S>& expr)
      : BinaryEngine_(expr), factor_(expr.factor()) {}
 
  // Pull base class functions into this class.
  using ExprEngine_::derived;
  using ExprEngine_::indices;
 
 
  void perm_indices(const BipartiteIndexList& target_indices) {
    // assert that init_indices has been called
    TA_ASSERT(left_.indices() && right_.indices());
    if (permute_tiles_) {
      this->template init_indices_<TensorProduct::Contraction>(target_indices);
 
      // propagate the indices down the tree, if needed
      if (left_indices_ != left_.indices()) {
        left_.perm_indices(left_indices_);
      }
      if (right_indices_ != right_.indices()) {
        right_.perm_indices(right_indices_);
      }
    }
  }
 
 
  void init_indices(bool children_initialized = false) {
    if (!children_initialized) {
      left_.init_indices();
      right_.init_indices();
    }
 
    this->template init_indices_<TensorProduct::Contraction>();
  }
 
 
  void init_indices(const BipartiteIndexList& target_indices) {
    init_indices();
    perm_indices(target_indices);
  }
 
 
  void init_struct(const BipartiteIndexList& target_indices) {
    // precondition checks
    // 1. if ToT inner tile op has been initialized
    if constexpr (TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
      TA_ASSERT(inner_tile_nonreturn_op_);
      TA_ASSERT(inner_tile_return_op_);
    }
 
    // Initialize children
    left_.init_struct(left_indices_);
    right_.init_struct(right_indices_);
 
    // Initialize the tile operation in this function because it is used to
    // evaluate the tiled range and shape.
 
    const math::blas::Op left_op =
        (left_outer_permtype_ == PermutationType::matrix_transpose
             ? math::blas::Transpose
             : math::blas::NoTranspose);
    const math::blas::Op right_op =
        (right_outer_permtype_ == PermutationType::matrix_transpose
             ? math::blas::Transpose
             : math::blas::NoTranspose);
 
    if (outer(target_indices) != outer(indices_)) {
      // Initialize permuted structure
      perm_ = ExprEngine_::make_perm(target_indices);
      if constexpr (!TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
        op_ = op_type(left_op, right_op, factor_, outer_size(indices_),
                      outer_size(left_indices_), outer_size(right_indices_),
                      (permute_tiles_ ? perm_ : BipartitePermutation{}));
      } else {
        // factor_ is absorbed into inner_tile_nonreturn_op_
        op_ = op_type(left_op, right_op, scalar_type(1), outer_size(indices_),
                      outer_size(left_indices_), outer_size(right_indices_),
                      (permute_tiles_ ? perm_ : BipartitePermutation{}),
                      this->inner_tile_nonreturn_op_);
      }
      trange_ = ContEngine_::make_trange(outer(perm_));
      shape_ = ContEngine_::make_shape(outer(perm_));
    } else {
      // Initialize non-permuted structure
      if constexpr (!TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
        op_ = op_type(left_op, right_op, factor_, outer_size(indices_),
                      outer_size(left_indices_), outer_size(right_indices_));
      } else {
        // factor_ is absorbed into inner_tile_nonreturn_op_
        op_ = op_type(left_op, right_op, scalar_type(1), outer_size(indices_),
                      outer_size(left_indices_), outer_size(right_indices_),
                      BipartitePermutation{}, this->inner_tile_nonreturn_op_);
      }
      trange_ = ContEngine_::make_trange();
      shape_ = ContEngine_::make_shape();
    }
 
    if (ExprEngine_::override_ptr_ && ExprEngine_::override_ptr_->shape) {
      shape_ = shape_.mask(*ExprEngine_::override_ptr_->shape);
    }
  }
 
 
  void init_distribution(World* world, std::shared_ptr<pmap_interface> pmap) {
    const unsigned int inner_rank = op_.gemm_helper().num_contract_ranks();
    const unsigned int left_rank = op_.gemm_helper().left_rank();
    const unsigned int right_rank = op_.gemm_helper().right_rank();
    const unsigned int left_outer_rank = left_rank - inner_rank;
 
    // Get pointers to the argument sizes
    const auto* MADNESS_RESTRICT const left_tiles_size =
        left_.trange().tiles_range().extent_data();
    const auto* MADNESS_RESTRICT const left_element_size =
        left_.trange().elements_range().extent_data();
    const auto* MADNESS_RESTRICT const right_tiles_size =
        right_.trange().tiles_range().extent_data();
    const auto* MADNESS_RESTRICT const right_element_size =
        right_.trange().elements_range().extent_data();
 
    // Compute the fused sizes of the contraction
    size_type M = 1ul, m = 1ul, N = 1ul, n = 1ul;
    unsigned int i = 0u;
    for (; i < left_outer_rank; ++i) {
      M *= left_tiles_size[i];
      m *= left_element_size[i];
    }
    for (; i < left_rank; ++i) K_ *= left_tiles_size[i];
    for (i = inner_rank; i < right_rank; ++i) {
      N *= right_tiles_size[i];
      n *= right_element_size[i];
    }
 
    // Construct the process grid.
    proc_grid_ = TiledArray::detail::ProcGrid(*world, M, N, m, n);
 
    // Initialize children
    left_.init_distribution(world, proc_grid_.make_row_phase_pmap(K_));
    right_.init_distribution(world, proc_grid_.make_col_phase_pmap(K_));
 
    // Initialize the process map in not already defined
    if (!pmap) pmap = proc_grid_.make_pmap();
    ExprEngine_::init_distribution(world, pmap);
  }
 
 
  trange_type make_trange(const Permutation& perm = {}) const {
    // Compute iteration limits
    const unsigned int left_rank = op_.gemm_helper().left_rank();
    const unsigned int right_rank = op_.gemm_helper().right_rank();
    const unsigned int inner_rank = op_.gemm_helper().num_contract_ranks();
    const unsigned int left_outer_rank = left_rank - inner_rank;
 
    // Construct the trange input and compute the gemm sizes
    typename trange_type::Ranges ranges(op_.gemm_helper().result_rank());
    unsigned int i = 0ul;
    for (unsigned int x = 0ul; x < left_outer_rank; ++x, ++i) {
      const unsigned int pi = (perm ? perm[i] : i);
      ranges[pi] = left_.trange().data()[x];
    }
    for (unsigned int x = inner_rank; x < right_rank; ++x, ++i) {
      const unsigned int pi = (perm ? perm[i] : i);
      ranges[pi] = right_.trange().data()[x];
    }
 
#ifndef NDEBUG
 
    // Check that the contracted dimensions have congruent tilings
    for (unsigned int l = left_outer_rank, r = 0ul; l < left_rank; ++l, ++r) {
      if (!is_congruent(left_.trange().data()[l], right_.trange().data()[r])) {
        if (TiledArray::get_default_world().rank() == 0) {
          TA_USER_ERROR_MESSAGE(
              "The contracted dimensions of the left- "
              "and right-hand arguments are not congruent:"
              << "\n    left  = " << left_.trange()
              << "\n    right = " << right_.trange());
 
          TA_EXCEPTION(
              "The contracted dimensions of the left- and "
              "right-hand expressions are not congruent.");
        }
 
        TA_EXCEPTION(
            "The contracted dimensions of the left- and "
            "right-hand expressions are not congruent.");
      }
    }
#endif  // NDEBUG
 
    return trange_type(ranges.begin(), ranges.end());
  }
 
 
  shape_type make_shape() const {
    const TiledArray::math::GemmHelper shape_gemm_helper(
        math::blas::NoTranspose, math::blas::NoTranspose,
        op_.gemm_helper().result_rank(), op_.gemm_helper().left_rank(),
        op_.gemm_helper().right_rank());
    return left_.shape().gemm(right_.shape(), factor_, shape_gemm_helper);
  }
 
 
  shape_type make_shape(const Permutation& perm) const {
    const TiledArray::math::GemmHelper shape_gemm_helper(
        math::blas::NoTranspose, math::blas::NoTranspose,
        op_.gemm_helper().result_rank(), op_.gemm_helper().left_rank(),
        op_.gemm_helper().right_rank());
    return left_.shape().gemm(right_.shape(), factor_, shape_gemm_helper, perm);
  }
 
  dist_eval_type make_dist_eval() const {
    // Define the impl type
    typedef TiledArray::detail::Summa<typename left_type::dist_eval_type,
                                      typename right_type::dist_eval_type,
                                      op_type, typename Derived::policy>
        impl_type;
 
    typename left_type::dist_eval_type left = left_.make_dist_eval();
    typename right_type::dist_eval_type right = right_.make_dist_eval();
 
    std::shared_ptr<impl_type> pimpl =
        std::make_shared<impl_type>(left, right, *world_, trange_, shape_,
                                    pmap_, perm_, op_, K_, proc_grid_);
 
    return dist_eval_type(pimpl);
  }
 
 
  std::string make_tag() const {
    std::stringstream ss;
    ss << "[*]";
    if (factor_ != scalar_type(1)) ss << "[" << factor_ << "]";
    return ss.str();
  }
 
 
  void print(ExprOStream os, const BipartiteIndexList& target_indices) const {
    ExprEngine_::print(os, target_indices);
    os.inc();
    left_.print(os, left_indices_);
    right_.print(os, right_indices_);
    os.dec();
  }
 
 protected:
  void init_inner_tile_op(const IndexList& inner_target_indices) {
    if constexpr (TiledArray::detail::is_tensor_of_tensor_v<value_type>) {
      using inner_tile_type = typename value_type::value_type;
      const auto inner_prod = this->inner_product_type();
      TA_ASSERT(inner_prod == TensorProduct::Contraction ||
                inner_prod == TensorProduct::Hadamard);
      if (inner_prod == TensorProduct::Contraction) {
        using inner_tile_type = typename value_type::value_type;
        using contract_inner_tile_type =
            TiledArray::detail::ContractReduce<inner_tile_type, inner_tile_type,
                                               inner_tile_type, scalar_type>;
        // factor_ is absorbed into inner_tile_nonreturn_op_
        auto contrreduce_op =
            (inner_target_indices != inner(this->indices_))
                ? contract_inner_tile_type(
                      to_cblas_op(this->left_inner_permtype_),
                      to_cblas_op(this->right_inner_permtype_), this->factor_,
                      inner_size(this->indices_),
                      inner_size(this->left_indices_),
                      inner_size(this->right_indices_),
                      (this->permute_tiles_ ? inner(this->perm_)
                                            : Permutation{}))
                : contract_inner_tile_type(
                      to_cblas_op(this->left_inner_permtype_),
                      to_cblas_op(this->right_inner_permtype_), this->factor_,
                      inner_size(this->indices_),
                      inner_size(this->left_indices_),
                      inner_size(this->right_indices_));
        this->inner_tile_nonreturn_op_ = [contrreduce_op](
                                             inner_tile_type& result,
                                             const inner_tile_type& left,
                                             const inner_tile_type& right) {
          contrreduce_op(result, left, right);
        };
      } else if (inner_prod == TensorProduct::Hadamard) {
        // inner tile op depends on the outer op ... e.g. if outer op
        // is contract then inner must implement (ternary) multiply-add;
        // if the outer is hadamard then the inner is binary multiply
        const auto outer_prod = this->product_type();
        if (this->factor_ == 1) {
          using base_op_type =
              TiledArray::detail::Mult<inner_tile_type, inner_tile_type,
                                       inner_tile_type, false, false>;
          using op_type = TiledArray::detail::BinaryWrapper<
              base_op_type>;  // can't consume inputs if they are used multiple
                              // times, e.g. when outer op is gemm
          auto mult_op = (inner_target_indices != inner(this->indices_))
                             ? op_type(base_op_type(), this->permute_tiles_
                                                           ? inner(this->perm_)
                                                           : Permutation{})
                             : op_type(base_op_type());
          this->inner_tile_nonreturn_op_ = [mult_op, outer_prod](
                                               inner_tile_type& result,
                                               const inner_tile_type& left,
                                               const inner_tile_type& right) {
            if (outer_prod == TensorProduct::Hadamard)
              result = mult_op(left, right);
            else {
              TA_ASSERT(outer_prod == TensorProduct::Hadamard ||
                        outer_prod == TensorProduct::Contraction);
              // there is currently no fused MultAdd ternary Op, only Add and
              // Mult thus implement this as 2 separate steps
              // TODO optimize by implementing (ternary) MultAdd
              if (empty(result))
                result = mult_op(left, right);
              else {
                auto result_increment = mult_op(left, right);
                add_to(result, result_increment);
              }
            }
          };
        } else {
          using base_op_type =
              TiledArray::detail::ScalMult<inner_tile_type, inner_tile_type,
                                           inner_tile_type, scalar_type, false,
                                           false>;
          using op_type = TiledArray::detail::BinaryWrapper<
              base_op_type>;  // can't consume inputs if they are used multiple
                              // times, e.g. when outer op is gemm
          auto mult_op = (inner_target_indices != inner(this->indices_))
                             ? op_type(base_op_type(this->factor_),
                                       this->permute_tiles_ ? inner(this->perm_)
                                                            : Permutation{})
                             : op_type(base_op_type(this->factor_));
          this->inner_tile_nonreturn_op_ = [mult_op, outer_prod](
                                               inner_tile_type& result,
                                               const inner_tile_type& left,
                                               const inner_tile_type& right) {
            TA_ASSERT(outer_prod == TensorProduct::Hadamard ||
                      outer_prod == TensorProduct::Contraction);
            if (outer_prod == TensorProduct::Hadamard)
              result = mult_op(left, right);
            else {
              // there is currently no fused MultAdd ternary Op, only Add and
              // Mult thus implement this as 2 separate steps
              // TODO optimize by implementing (ternary) MultAdd
              if (empty(result))
                result = mult_op(left, right);
              else {
                auto result_increment = mult_op(left, right);
                add_to(result, result_increment);
              }
            }
          };
        }
      } else
        abort();  // unsupported TensorProduct type
      TA_ASSERT(inner_tile_nonreturn_op_);
      this->inner_tile_return_op_ =
          [inner_tile_nonreturn_op = this->inner_tile_nonreturn_op_](
              const inner_tile_type& left, const inner_tile_type& right) {
            inner_tile_type result;
            inner_tile_nonreturn_op(result, left, right);
            return result;
          };
    }
  }
 
};  // class ContEngine
 
}  // namespace expressions
}  // namespace TiledArray
 
#endif  // TILEDARRAY_EXPRESSIONS_CONT_ENGINE_H__INCLUDED