btas.h

Go to the documentation of this file.

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2018  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/>.
 *
 *  btas.h
 *  January 19, 2018
 *
 */

#ifndef TILEDARRAY_CONVERSIONS_BTAS_H__INCLUDED
#define TILEDARRAY_CONVERSIONS_BTAS_H__INCLUDED

#include <limits>

#include <TiledArray/block_range.h>
#include <TiledArray/external/btas.h>
#include <TiledArray/tensor.h>
#include <TiledArray/tensor/tensor_map.h>
#include <TiledArray/dense_shape.h>
#include <TiledArray/policies/dense_policy.h>
#include <TiledArray/sparse_shape.h>
#include <TiledArray/policies/sparse_policy.h>
#include <TiledArray/pmap/replicated_pmap.h>

namespace TiledArray {


  template <typename T, typename Range_, typename Storage_, typename Allocator_>
  inline void btas_subtensor_to_tensor(const btas::Tensor<T,Range_,Storage_>& src, Tensor<T, Allocator_>& dst) {
    TA_ASSERT(dst.range().rank() == src.range().rank());

    const auto& src_range = src.range();
    const auto& dst_range = dst.range();
    auto src_blk_range = TiledArray::BlockRange(detail::make_ta_range(src_range), dst_range.lobound(), dst_range.upbound());
    using std::data;
    auto src_view = TiledArray::make_const_map(data(src), src_blk_range);

    dst = src_view;
  }


  template <typename T, typename Allocator_, typename Range_, typename Storage_>
  inline void tensor_to_btas_subtensor(const Tensor<T, Allocator_>& src, btas::Tensor<T,Range_,Storage_>& dst) {
    TA_ASSERT(dst.range().rank() == src.range().rank());

    const auto& src_range = src.range();
    const auto& dst_range = dst.range();
    auto dst_blk_range = TiledArray::BlockRange(detail::make_ta_range(dst_range), src_range.lobound(), src_range.upbound());
    using std::data;
    auto dst_view = TiledArray::make_map(data(dst), dst_blk_range);

    dst_view = src;
  }

namespace detail {


template <typename DistArray_, typename BTAS_Tensor_>
void counted_btas_subtensor_to_tensor(const BTAS_Tensor_* src,
                                      DistArray_* dst, const typename DistArray_::size_type i, madness::AtomicInt* counter)
{
  typename DistArray_::value_type tensor(dst->trange().make_tile_range(i));
  btas_subtensor_to_tensor(*src, tensor);
  dst->set(i, tensor);
  (*counter)++;
}


template <typename TA_Tensor_, typename BTAS_Tensor_>
void counted_tensor_to_btas_subtensor(const TA_Tensor_& src,
                                      BTAS_Tensor_* dst, madness::AtomicInt* counter)
{
  tensor_to_btas_subtensor(src, *dst);
  (*counter)++;
}

template <bool sparse>
auto make_shape(World& world, const TiledArray::TiledRange& trange);

template <>
inline auto make_shape<true>(World& world, const TiledArray::TiledRange& trange) {
  TiledArray::Tensor<float> tile_norms(trange.tiles_range(), std::numeric_limits<float>::max());
  return TiledArray::SparseShape<float>(world, tile_norms, trange);
}

template <>
inline auto make_shape<false>(World&, const TiledArray::TiledRange&) {
  return TiledArray::DenseShape{};
}

} // namespace detail


template <typename DistArray_,
    typename T,
    typename Range,
    typename Storage>
DistArray_
btas_tensor_to_array(World& world, const TiledArray::TiledRange& trange,
                     const btas::Tensor<T, Range, Storage>& src, bool replicated = false)
{
  // Test preconditions
  const auto rank = trange.tiles_range().rank();
  TA_USER_ASSERT(rank == src.range().rank(),
                 "TiledArray::btas_tensor_to_array(): rank of destination trange does not match the rank of source BTAS tensor.");
  auto dst_range_extents = trange.elements_range().extent();
  for(std::remove_const_t<decltype(rank)> d=0; d!=rank; ++d) {
    TA_USER_ASSERT(dst_range_extents[d] == src.range().extent(d),
                   "TiledArray::btas_tensor_to_array(): source dimension does not match destination dimension.");
  }

  using Tensor_ = btas::Tensor<T,Range,Storage>;
  using Policy_ = typename DistArray_::policy_type;
  const auto is_sparse = std::is_same<Policy_,TiledArray::SparsePolicy>::value;

  // Check that this is not a distributed computing environment
  if(! replicated)
    TA_USER_ASSERT(world.size() == 1,
                   "An array can be created from a btas::Tensor if the number of World ranks is greater than 1 only when replicated=true.");

  // Make a shape, only used if making a sparse array
  using Shape_ = typename DistArray_::shape_type;
  Shape_ shape = detail::make_shape<is_sparse>(world, trange);

  // Create a new tensor
  DistArray_ array = (replicated && (world.size() > 1)
                     ? DistArray_(world, trange, shape,
                     std::static_pointer_cast<typename DistArray_::pmap_interface>(
                     std::make_shared<detail::ReplicatedPmap>(
                         world, trange.tiles_range().volume())))
             : DistArray_(world, trange, shape));

  // Spawn copy tasks
  madness::AtomicInt counter;
  counter = 0;
  std::int64_t n = 0;
  for(std::size_t i = 0; i < array.size(); ++i) {
    world.taskq.add(& detail::counted_btas_subtensor_to_tensor<DistArray_, Tensor_>,
                    &src, &array, i, &counter);
    ++n;
  }

  // Wait until the write tasks are complete
  array.world().await([&counter,n] () { return counter == n; });

  // Analyze tiles norms and truncate based on sparse policy
  if (is_sparse) truncate(array);

  return array;
}


template <typename Tile, typename Policy>
btas::Tensor<typename Tile::value_type>
array_to_btas_tensor(const TiledArray::DistArray<Tile,Policy>& src)
{
  // Test preconditions
  if(! src.pmap()->is_replicated())
    TA_USER_ASSERT(src.world().size() == 1,
                   "TiledArray::array_to_btas_tensor(): a non-replicated array can only be converted to a btas::Tensor if the number of World ranks is 1.");

  // Construct the result
  using result_type = btas::Tensor<typename TiledArray::DistArray<Tile,Policy>::element_type>;
  using result_range_type = typename result_type::range_type;
  // if array is sparse must initialize to zero
  result_type result(result_range_type(src.trange().elements_range().extent()), 0.0);

  // Spawn tasks to copy array tiles to btas::Tensor
  madness::AtomicInt counter;
  counter = 0;
  int n = 0;
  for(std::size_t i = 0; i < src.size(); ++i) {
    if(! src.is_zero(i)) {
      src.world().taskq.add(
          & detail::counted_tensor_to_btas_subtensor<Tile, result_type>,
          src.find(i), &result, &counter);
      ++n;
    }
  }

  // Wait until the write tasks are complete
  src.world().await([&counter,n] () { return counter == n; });

  return result;
}

}  // namespace TiledArray

#endif // TILEDARRAY_CONVERSIONS_BTAS_H__INCLUDED