gemm_helper.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
 *
 *  gemm_helper.h
 *  Jan 20, 2014
 *
 */
 
#ifndef TILEDARRAY_MATH_GEMM_HELPER_H__INCLUDED
#define TILEDARRAY_MATH_GEMM_HELPER_H__INCLUDED
 
#include <TiledArray/error.h>
#include <TiledArray/math/blas.h>
 
#include <vector>
 
namespace TiledArray::math {
 
 
class GemmHelper {
 private:
  blas::Op left_op_;
  blas::Op right_op_;
  unsigned int result_rank_;  
 
 
  struct ContractArg {
    unsigned int inner[2];  
    unsigned int outer[2];  
    unsigned int rank;      
  } left_,                  
      right_;               
 
 public:
  GemmHelper(const blas::Op left_op, const blas::Op right_op,
             const unsigned int result_rank, const unsigned int left_rank,
             const unsigned int right_rank)
      : left_op_(left_op),
        right_op_(right_op),
        result_rank_(result_rank),
        left_(),
        right_() {
    // Compute the number of contracted dimensions in left and right.
    TA_ASSERT(((left_rank + right_rank - result_rank) % 2u) == 0u);
 
    left_.rank = left_rank;
    right_.rank = right_rank;
    const unsigned int contract_size = num_contract_ranks();
 
    // Store the inner and outer dimension ranges for the left-hand argument.
    if (left_op == blas::NoTranspose) {
      left_.outer[0] = 0u;
      left_.outer[1] = left_.inner[0] = left_rank - contract_size;
      left_.inner[1] = left_rank;
    } else {
      left_.inner[0] = 0ul;
      left_.inner[1] = left_.outer[0] = contract_size;
      left_.outer[1] = left_rank;
    }
 
    // Store the inner and outer dimension ranges for the right-hand argument.
    if (right_op == blas::NoTranspose) {
      right_.inner[0] = 0u;
      right_.inner[1] = right_.outer[0] = contract_size;
      right_.outer[1] = right_rank;
    } else {
      right_.outer[0] = 0u;
      right_.outer[1] = right_.inner[0] = right_rank - contract_size;
      right_.inner[1] = right_rank;
    }
  }
 
 
  GemmHelper(const GemmHelper& other)
      : left_op_(other.left_op_),
        right_op_(other.right_op_),
        result_rank_(other.result_rank_),
        left_(other.left_),
        right_(other.right_) {}
 
 
  GemmHelper& operator=(const GemmHelper& other) {
    left_op_ = other.left_op_;
    right_op_ = other.right_op_;
    result_rank_ = other.result_rank_;
    left_ = other.left_;
    right_ = other.right_;
 
    return *this;
  }
 
 
  unsigned int num_contract_ranks() const {
    return (left_.rank + right_.rank - result_rank_) >> 1;
  }
 
 
  unsigned int result_rank() const { return result_rank_; }
 
 
  unsigned int left_rank() const { return left_.rank; }
 
 
  unsigned int right_rank() const { return right_.rank; }
 
  unsigned int left_inner_begin() const { return left_.inner[0]; }
  unsigned int left_inner_end() const { return left_.inner[1]; }
  unsigned int left_outer_begin() const { return left_.outer[0]; }
  unsigned int left_outer_end() const { return left_.outer[1]; }
 
  unsigned int right_inner_begin() const { return right_.inner[0]; }
  unsigned int right_inner_end() const { return right_.inner[1]; }
  unsigned int right_outer_begin() const { return right_.outer[0]; }
  unsigned int right_outer_end() const { return right_.outer[1]; }
 
 
  template <typename R, typename Left, typename Right>
  R make_result_range(const Left& left, const Right& right) const {
    // Get pointers to lower and upper bounds of left and right.
    const auto* MADNESS_RESTRICT const left_lower = left.lobound_data();
    const auto* MADNESS_RESTRICT const left_upper = left.upbound_data();
    const auto* MADNESS_RESTRICT const right_lower = right.lobound_data();
    const auto* MADNESS_RESTRICT const right_upper = right.upbound_data();
 
    // Create the start and finish indices
    std::vector<std::size_t> lower, upper;
    lower.reserve(result_rank_);
    upper.reserve(result_rank_);
 
    // Copy left-hand argument outer dimensions to start and finish
    for (unsigned int i = left_.outer[0]; i < left_.outer[1]; ++i) {
      lower.push_back(left_lower[i]);
      upper.push_back(left_upper[i]);
    }
 
    // Copy right-hand argument outer dimensions to start and finish
    for (unsigned int i = right_.outer[0]; i < right_.outer[1]; ++i) {
      lower.push_back(right_lower[i]);
      upper.push_back(right_upper[i]);
    }
 
    // Construct the result tile range
    return R(lower, upper);
  }
 
 
  template <typename Left, typename Result>
  bool left_result_congruent(const Left& left, const Result& result) const {
    return std::equal(left + left_.outer[0], left + left_.outer[1], result);
  }
 
 
  template <typename Right, typename Result>
  bool right_result_congruent(const Right& right, const Result& result) const {
    return std::equal(right + right_.outer[0], right + right_.outer[1],
                      result + (left_.outer[1] - left_.outer[0]));
  }
 
 
  template <typename Left, typename Right>
  bool left_right_congruent(const Left& left, const Right& right) const {
    return std::equal(left + left_.inner[0], left + left_.inner[1],
                      right + right_.inner[0]);
  }
 
 
  template <typename Left, typename Right>
  void compute_matrix_sizes(blas::integer& m, blas::integer& n,
                            blas::integer& k, const Left& left,
                            const Right& right) const {
    // Check that the arguments are not empty and have the correct ranks
    TA_ASSERT(left.rank() == left_.rank);
    TA_ASSERT(right.rank() == right_.rank);
    const auto* MADNESS_RESTRICT const left_extent = left.extent_data();
    const auto* MADNESS_RESTRICT const right_extent = right.extent_data();
 
    // Compute fused dimension sizes
    m = 1;
    for (unsigned int i = left_.outer[0]; i < left_.outer[1]; ++i)
      m *= left_extent[i];
    k = 1;
    for (unsigned int i = left_.inner[0]; i < left_.inner[1]; ++i)
      k *= left_extent[i];
    n = 1;
    for (unsigned int i = right_.outer[0]; i < right_.outer[1]; ++i)
      n *= right_extent[i];
  }
 
  blas::Op left_op() const { return left_op_; }
  blas::Op right_op() const { return right_op_; }
};  // class GemmHelper
 
}  // namespace TiledArray::math
 
#endif  // TILEDARRAY_MATH_GEMM_HELPER_H__INCLUDED