Tpetra_Details_computeOffsets.hpp

Go to the documentation of this file.

/*
// @HEADER
// ***********************************************************************
//
//          Tpetra: Templated Linear Algebra Services Package
//                 Copyright (2008) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// ************************************************************************
// @HEADER
*/

#ifndef TPETRA_DETAILS_COMPUTEOFFSETS_HPP
#define TPETRA_DETAILS_COMPUTEOFFSETS_HPP


#include "TpetraCore_config.h"
#include "Tpetra_Details_getEntryOnHost.hpp"
#include <limits>
#include <type_traits>

namespace Tpetra {
namespace Details {

//
// Implementation details for computeOffsetsFromCounts (see below).
// Users should skip over this anonymous namespace.
//
namespace { // (anonymous)

template<class OffsetType,
         class CountType,
         class SizeType>
class ComputeOffsetsFromCounts {
public:
  static_assert (std::is_integral<OffsetType>::value,
                 "OffsetType must be a built-in integer.");
  static_assert (std::is_integral<CountType>::value,
                 "CountType must be a built-in integer.");
  static_assert (std::is_integral<SizeType>::value,
                 "SizeType must be a built-in integer.");

  using offsets_view_type =
    Kokkos::View<OffsetType*, Kokkos::AnonymousSpace>;
  using counts_view_type =
    Kokkos::View<const CountType*, Kokkos::AnonymousSpace>;

  ComputeOffsetsFromCounts (const offsets_view_type& offsets,
                            const counts_view_type& counts) :
    offsets_ (offsets),
    counts_ (counts),
    size_ (counts.extent (0))
  {}

  KOKKOS_INLINE_FUNCTION void
  operator () (const SizeType i, OffsetType& update,
               const bool finalPass) const
  {
    const auto curVal = (i < size_) ? counts_[i] : OffsetType ();
    if (finalPass) {
      offsets_[i] = update;
    }
    update += (i < size_) ? curVal : OffsetType ();
  }

  template<class ExecutionSpace>
  static OffsetType
  run (const ExecutionSpace& execSpace,
       const offsets_view_type& offsets,
       const counts_view_type& counts)
  {
    const SizeType numCounts (counts.extent (0));
    using range_type = Kokkos::RangePolicy<ExecutionSpace, SizeType>;
    range_type range (execSpace, 0, numCounts + SizeType (1));
    using functor_type =
      ComputeOffsetsFromCounts<OffsetType, CountType, SizeType>;
    functor_type functor (offsets, counts);
    OffsetType total (0);
    const char funcName[] = "Tpetra::Details::computeOffsetsFromCounts";
    Kokkos::parallel_scan (range, functor, total, funcName);
    return total;
  }

private:
  offsets_view_type offsets_;
  counts_view_type counts_;
  SizeType size_;
};

template<class OffsetType,
         class CountType,
         class SizeType>
class ComputeOffsetsFromConstantCount {
public:
  static_assert (std::is_integral<OffsetType>::value,
                 "OffsetType must be a built-in integer.");
  static_assert (std::is_integral<CountType>::value,
                 "CountType must be a built-in integer.");
  static_assert (std::is_integral<SizeType>::value,
                 "SizeType must be a built-in integer.");

  using offsets_view_type =
    Kokkos::View<OffsetType*, Kokkos::AnonymousSpace>;

  ComputeOffsetsFromConstantCount (const offsets_view_type& offsets,
                                   const CountType count) :
    offsets_ (offsets),
    count_ (count)
  {}

  KOKKOS_INLINE_FUNCTION void
  operator () (const SizeType i) const
  {
    offsets_[i] = count_*i;
  }

  template<class ExecutionSpace>
  static OffsetType
  run (const ExecutionSpace& execSpace,
       const offsets_view_type& offsets,
       const CountType count)
  {
    const SizeType numOffsets (offsets.extent (0));
    using range_type = Kokkos::RangePolicy<ExecutionSpace, SizeType>;
    range_type range (execSpace, 0, numOffsets);
    using functor_type =
      ComputeOffsetsFromConstantCount<OffsetType, CountType, SizeType>;
    functor_type functor (offsets, count);
    const OffsetType total = numOffsets*count;
    const char funcName[] =
      "Tpetra::Details::computeOffsetsFromConstantCount";
    Kokkos::parallel_for (range, functor, funcName);
    return total;
  }

private:
  offsets_view_type offsets_;
  CountType count_;
};

} // namespace (anonymous)

template<class ExecutionSpace,
         class OffsetsViewType,
         class CountsViewType,
         class SizeType = typename OffsetsViewType::size_type>
typename OffsetsViewType::non_const_value_type
computeOffsetsFromCounts (const ExecutionSpace& execSpace,
                          const OffsetsViewType& ptr,
                          const CountsViewType& counts)
{
  static_assert (Kokkos::is_execution_space<ExecutionSpace>::value,
                 "ExecutionSpace must be a Kokkos execution space.");
  static_assert (Kokkos::Impl::is_view<OffsetsViewType>::value,
                 "OffsetsViewType (the type of ptr) must be a Kokkos::View.");
  static_assert (Kokkos::Impl::is_view<CountsViewType>::value,
                 "CountsViewType (the type of counts) must be a Kokkos::View.");
  static_assert (std::is_same<typename OffsetsViewType::value_type,
                   typename OffsetsViewType::non_const_value_type>::value,
                 "OffsetsViewType (the type of ptr) must be a nonconst Kokkos::View.");
  static_assert (static_cast<int> (OffsetsViewType::rank) == 1,
                 "OffsetsViewType (the type of ptr) must be a rank-1 Kokkos::View.");
  static_assert (static_cast<int> (CountsViewType::rank) == 1,
                 "CountsViewType (the type of counts) must be a rank-1 Kokkos::View.");

  using offset_type = typename OffsetsViewType::non_const_value_type;
  static_assert (std::is_integral<offset_type>::value,
                 "The entries of ptr must be built-in integers.");
  using count_type = typename CountsViewType::non_const_value_type;
  static_assert (std::is_integral<count_type>::value,
                 "The entries of counts must be built-in integers.");
  static_assert (std::is_integral<SizeType>::value,
                 "SizeType must be a built-in integer type.");

  const char funcName[] = "Tpetra::Details::computeOffsetsFromCounts";

  const auto numOffsets = ptr.size ();
  const auto numCounts = counts.size ();
  offset_type total (0);

  if (numOffsets != 0) {
    TEUCHOS_TEST_FOR_EXCEPTION
      (numCounts >= numOffsets, std::invalid_argument, funcName <<
       ": counts.size() = " << numCounts << " >= ptr.size() = " <<
       numOffsets << ".");

#ifdef TPETRA_AVOID_PARALLEL_SCAN
    // KDD 11/25/19  Temporary workaround to the one scan whose behavior
    // makes a difference for GPU runs.  On the GPU platforms, serializing
    // this scan (via Kokkos' magic) makes the code perform correctly.  
    // This code serializes the scan, without the Kokkos magic, for debugging
    // on stria.  #6345
    // To use this code, add -DTPETRA_AVOID_PARALLEL_SCAN to CMAKE_CXX_FLAGS

    Kokkos::View<offset_type*,Kokkos::HostSpace> 
            ptrAPS(Kokkos::ViewAllocateWithoutInitializing("APS_offsets"),
                   numOffsets);
    Kokkos::View<count_type*,Kokkos::HostSpace> 
            countsAPS(Kokkos::ViewAllocateWithoutInitializing("APS_counts"),
                      numCounts);

    Kokkos::deep_copy(countsAPS, counts);

    ptrAPS(0) = 0;
    for (size_t i = 0; i < numCounts; i++)
      ptrAPS(i+1) = ptrAPS(i) + countsAPS(i);

    total = ptrAPS(numCounts);

    Kokkos::deep_copy(ptr, ptrAPS);

#else
    using Kokkos::AnonymousSpace;
    using Kokkos::View;
    View<offset_type*, AnonymousSpace> ptr_a = ptr;
    View<const count_type*, AnonymousSpace> counts_a;

    using offsets_device_type = typename OffsetsViewType::device_type;
    using counts_copy_type = View<count_type*, offsets_device_type>;
    counts_copy_type counts_copy;

    using offsets_memory_space =
      typename offsets_device_type::memory_space;
    using counts_memory_space = typename CountsViewType::memory_space;
    constexpr bool countsAccessibleFromOffsetsExecSpace =
      Kokkos::Impl::VerifyExecutionCanAccessMemorySpace<
        offsets_memory_space, counts_memory_space>::value;
    if (countsAccessibleFromOffsetsExecSpace) {
      // NOTE (mfh 21 Aug 2019) Some compilers have trouble deducing
      // that operator= works if more than one template argument
      // differ.  If that should happen, introduce an intermediate
      // type here.
      counts_a = counts;
    }
    else {
      using Kokkos::view_alloc;
      using Kokkos::WithoutInitializing;
      counts_copy = counts_copy_type
        (view_alloc ("counts_copy", WithoutInitializing), numCounts);
      Kokkos::deep_copy (counts_copy, counts);
      counts_a = counts_copy;
    }

    using functor_type =
      ComputeOffsetsFromCounts<offset_type, count_type, SizeType>;
    total = functor_type::run (execSpace, ptr_a, counts_a);
#endif
  }

  return total;
}

template<class OffsetsViewType,
         class CountsViewType,
         class SizeType = typename OffsetsViewType::size_type>
typename OffsetsViewType::non_const_value_type
computeOffsetsFromCounts (const OffsetsViewType& ptr,
                          const CountsViewType& counts)
{
  using execution_space = typename OffsetsViewType::execution_space;
  return computeOffsetsFromCounts (execution_space (), ptr, counts);
}

template<class OffsetsViewType,
         class CountType,
         class SizeType = typename OffsetsViewType::size_type>
typename OffsetsViewType::non_const_value_type
computeOffsetsFromConstantCount (const OffsetsViewType& ptr,
                                 const CountType count)
{
  static_assert (Kokkos::Impl::is_view<OffsetsViewType>::value,
                 "ptr must be a Kokkos::View.");
  static_assert (std::is_same<typename OffsetsViewType::value_type,
                   typename OffsetsViewType::non_const_value_type>::value,
                 "ptr must be a nonconst Kokkos::View.");
  static_assert (static_cast<int> (OffsetsViewType::rank) == 1,
                 "ptr must be a rank-1 Kokkos::View.");

  using offset_type = typename OffsetsViewType::non_const_value_type;
  static_assert (std::is_integral<offset_type>::value,
                 "The type of each entry of ptr must be a "
                 "built-in integer.");
  static_assert (std::is_integral<CountType>::value,
                 "CountType must be a built-in integer.");
  static_assert (std::is_integral<SizeType>::value,
                 "SizeType must be a built-in integer.");

  using device_type = typename OffsetsViewType::device_type;
  using execution_space = typename device_type::execution_space;

  offset_type total (0);
  if (ptr.extent (0) != 0) {
    using CT = CountType;
    using functor_type =
      ComputeOffsetsFromConstantCount<offset_type, CT, SizeType>;
    execution_space execSpace;
    functor_type::run (execSpace, ptr, count);
  }
  return total;
}

} // namespace Details
} // namespace Tpetra

#endif // TPETRA_DETAILS_COMPUTEOFFSETS_HPP