Tpetra_Details_crsUtils.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_CRSUTILS_HPP
#define TPETRA_DETAILS_CRSUTILS_HPP
#include <numeric>
#include <type_traits>

#include "TpetraCore_config.h"
#include "Kokkos_Core.hpp"
#include "Tpetra_Details_Behavior.hpp"
#include "Tpetra_Details_CrsPadding.hpp"
#include <iostream>
#include <memory>


namespace Tpetra {
namespace Details {

namespace impl {

template<class ViewType>
ViewType
make_uninitialized_view(
  const std::string& name,
  const size_t size,
  const bool verbose,
  const std::string* const prefix)
{
  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Allocate Kokkos::View " << name
       << ": " << size << std::endl;
    std::cerr << os.str();
  }
  using Kokkos::view_alloc;
  using Kokkos::WithoutInitializing;
  return ViewType(view_alloc(name, WithoutInitializing), size);
}

template<class ViewType>
ViewType
make_initialized_view(
  const std::string& name,
  const size_t size,
  const bool verbose,
  const std::string* const prefix)
{
  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Allocate & initialize Kokkos::View "
       << name << ": " << size << std::endl;
    std::cerr << os.str();
  }
  return ViewType(name, size);
}

template<class OutViewType, class InViewType>
void
assign_to_view(OutViewType& out,
               const InViewType& in,
               const char viewName[],
               const bool verbose,
               const std::string* const prefix)
{
  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Assign to Kokkos::View " << viewName
       << ": Old size: " << out.extent(0)
       << ", New size: " << in.extent(0) << std::endl;
    std::cerr << os.str();
  }
  out = in;
}

template<class MemorySpace, class ViewType>
auto create_mirror_view(
  const MemorySpace& memSpace,
  const ViewType& view,
  const bool verbose,
  const std::string* const prefix) ->
  decltype(Kokkos::create_mirror_view(memSpace, view))
{
  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Create mirror view: "
       << "view.extent(0): " << view.extent(0) << std::endl;
    std::cerr << os.str();
  }
  return Kokkos::create_mirror_view(memSpace, view);
}

enum class PadCrsAction {
  INDICES_ONLY,
  INDICES_AND_VALUES
};

template<class RowPtr, class Indices, class Values, class Padding>
void
pad_crs_arrays(
  const PadCrsAction action,
  const RowPtr& row_ptr_beg,
  const RowPtr& row_ptr_end,
  Indices& indices,
  Values& values,
  const Padding& padding,
  const int my_rank,
  const bool verbose)
{
  using Kokkos::view_alloc;
  using Kokkos::WithoutInitializing;
  using std::endl;
  std::unique_ptr<std::string> prefix;

  const size_t maxNumToPrint = verbose ?
    Behavior::verbosePrintCountThreshold() : size_t(0);
  if (verbose) {
    std::ostringstream os;
    os << "Proc " << my_rank << ": Tpetra::...::pad_crs_arrays: ";
    prefix = std::unique_ptr<std::string>(new std::string(os.str()));
    os << "Start" << endl;
    std::cerr << os.str();
  }
  Kokkos::HostSpace hostSpace;

  if (verbose) {
    std::ostringstream os;
    os << *prefix << "On input: ";
    auto row_ptr_beg_h =
      Kokkos::create_mirror_view(hostSpace, row_ptr_beg);
    Kokkos::deep_copy(row_ptr_beg_h, row_ptr_beg);
    verbosePrintArray(os, row_ptr_beg_h, "row_ptr_beg before scan",
                      maxNumToPrint);
    os << ", ";
    auto row_ptr_end_h =
      Kokkos::create_mirror_view(hostSpace, row_ptr_end);
    Kokkos::deep_copy(row_ptr_end_h, row_ptr_end);
    verbosePrintArray(os, row_ptr_end_h, "row_ptr_end before scan",
                      maxNumToPrint);
    os << ", indices.extent(0): " << indices.extent(0)
       << ", values.extent(0): " << values.extent(0)
       << ", padding: ";
    padding.print(os);
    os << endl;
    std::cerr << os.str();
  }

  if (row_ptr_beg.size() == 0) {
    if (verbose) {
      std::ostringstream os;
      os << *prefix << "Done; local matrix has no rows" << endl;
      std::cerr << os.str();
    }
    return; // nothing to do
  }

  const size_t lclNumRows(row_ptr_beg.size() - 1);
  RowPtr newAllocPerRow =
    make_uninitialized_view<RowPtr>("newAllocPerRow", lclNumRows,
                                    verbose, prefix.get());
  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Fill newAllocPerRow & compute increase" << endl;
    std::cerr << os.str();
  }
  size_t increase = 0;
  {
    auto row_ptr_end_h = create_mirror_view(
      hostSpace, row_ptr_end, verbose, prefix.get());
    Kokkos::deep_copy(row_ptr_end_h, row_ptr_end);
    auto row_ptr_beg_h = create_mirror_view(
      hostSpace, row_ptr_beg, verbose, prefix.get());
    Kokkos::deep_copy(row_ptr_beg_h, row_ptr_beg);

    auto newAllocPerRow_h = create_mirror_view(
      hostSpace, newAllocPerRow, verbose, prefix.get());
    using host_range_type = Kokkos::RangePolicy<
      Kokkos::DefaultHostExecutionSpace, size_t>;
    Kokkos::parallel_reduce
      ("Tpetra::CrsGraph: Compute new allocation size per row",
       host_range_type(0, lclNumRows),
       [&] (const size_t lclRowInd, size_t& lclIncrease) {
         const size_t start = row_ptr_beg_h[lclRowInd];
         const size_t end   = row_ptr_beg_h[lclRowInd+1];
         TEUCHOS_ASSERT( end >= start );
         const size_t oldAllocSize = end - start;
         const size_t oldNumEnt = row_ptr_end_h[lclRowInd] - start;
         TEUCHOS_ASSERT( oldNumEnt <= oldAllocSize );

         // This is not a pack routine.  Do not shrink!  Shrinking now
         // to fit the number of entries would ignore users' hint for
         // the max number of entries in each row.  Also, CrsPadding
         // only counts entries and ignores any available free space.

         auto result = padding.get_result(lclRowInd);
         const size_t newNumEnt = oldNumEnt + result.numInSrcNotInTgt;
         if (newNumEnt > oldAllocSize) {
           lclIncrease += (newNumEnt - oldAllocSize);
           newAllocPerRow_h[lclRowInd] = newNumEnt;
         }
         else {
           newAllocPerRow_h[lclRowInd] = oldAllocSize;
         }
       }, increase);

    if (verbose) {
      std::ostringstream os;
      os << *prefix << "increase: " << increase << ", ";
      verbosePrintArray(os, newAllocPerRow_h, "newAllocPerRow",
                        maxNumToPrint);
      os << endl;
      std::cerr << os.str();
    }

    if (increase == 0) {
      return;
    }
    Kokkos::deep_copy(newAllocPerRow, newAllocPerRow_h);
  }

  using inds_value_type = typename Indices::non_const_value_type;
  using vals_value_type = typename Values::non_const_value_type;

  const size_t newIndsSize = size_t(indices.size()) + increase;
  auto indices_new = make_uninitialized_view<Indices>(
    "Tpetra::CrsGraph column indices", newIndsSize, verbose,
    prefix.get());

  Values values_new;
  if (action == PadCrsAction::INDICES_AND_VALUES) {
    const size_t newValsSize = newIndsSize;
    // NOTE (mfh 10 Feb 2020) If we don't initialize values_new here,
    // then the CrsMatrix tests fail.
    values_new = make_initialized_view<Values>(
      "Tpetra::CrsMatrix values", newValsSize, verbose, prefix.get());
  }

  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Repack" << endl;
    std::cerr << os.str();
  }
  using execution_space = typename Indices::execution_space;
  using range_type = Kokkos::RangePolicy<execution_space, size_t>;
  Kokkos::parallel_scan(
    "Tpetra::CrsGraph or CrsMatrix repack",
    range_type(0, lclNumRows+1),
    KOKKOS_LAMBDA (const size_t lclRow, size_t& newRowBeg,
                   const bool finalPass)
    {
      // row_ptr_beg    has lclNumRows + 1 entries.
      // row_ptr_end    has lclNumRows     entries.
      // newAllocPerRow has lclNumRows     entries.
      const size_t row_beg = row_ptr_beg[lclRow];
      const size_t row_end =
        lclRow < lclNumRows ? row_ptr_end[lclRow] : row_beg;
      const size_t numEnt = row_end - row_beg;
      const size_t newRowAllocSize =
        lclRow < lclNumRows ? newAllocPerRow[lclRow] : size_t(0);
      if (finalPass) {
        if (lclRow < lclNumRows) {
          const Kokkos::pair<size_t, size_t> oldRange(
            row_beg, row_beg + numEnt);
          const Kokkos::pair<size_t, size_t> newRange(
            newRowBeg, newRowBeg + numEnt);
          auto oldColInds = subview(indices, oldRange);
          auto newColInds = subview(indices_new, newRange);
          // memcpy works fine on device; the next step is to
          // introduce two-level parallelism and use team copy.
          memcpy(newColInds.data(), oldColInds.data(),
                 numEnt * sizeof(inds_value_type));
          if (action == PadCrsAction::INDICES_AND_VALUES) {
            auto oldVals = subview(values, oldRange);
            auto newVals = subview(values_new, newRange);
            memcpy(newVals.data(), oldVals.data(),
                   numEnt * sizeof(vals_value_type));
          }
        }
        // It's the final pass, so we can modify these arrays.
        row_ptr_beg[lclRow] = newRowBeg;
        if (lclRow < lclNumRows) {
          row_ptr_end[lclRow] = newRowBeg + numEnt;
        }
      }
      newRowBeg += newRowAllocSize;
    });

  if (verbose) {
    std::ostringstream os;

    os << *prefix;
    auto row_ptr_beg_h =
      Kokkos::create_mirror_view(hostSpace, row_ptr_beg);
    Kokkos::deep_copy(row_ptr_beg_h, row_ptr_beg);
    verbosePrintArray(os, row_ptr_beg_h, "row_ptr_beg after scan",
                      maxNumToPrint);
    os << endl;

    os << *prefix;
    auto row_ptr_end_h =
      Kokkos::create_mirror_view(hostSpace, row_ptr_end);
    Kokkos::deep_copy(row_ptr_end_h, row_ptr_end);
    verbosePrintArray(os, row_ptr_end_h, "row_ptr_end after scan",
                      maxNumToPrint);
    os << endl;

    std::cerr << os.str();
  }

  assign_to_view(indices, indices_new,
                 "Tpetra::CrsGraph column indices",
                 verbose, prefix.get());
  assign_to_view(values, values_new,
                 "Tpetra::CrsMatrix values",
                 verbose, prefix.get());

  if (verbose) {
    auto indices_h = Kokkos::create_mirror_view(hostSpace, indices);
    Kokkos::deep_copy(indices_h, indices);
    auto values_h = Kokkos::create_mirror_view(hostSpace, values);
    Kokkos::deep_copy(values_h, values);
    std::ostringstream os;
    os << "On output: ";
    verbosePrintArray(os, indices_h, "indices", maxNumToPrint);
    os << ", ";
    verbosePrintArray(os, values_h, "values", maxNumToPrint);
    os << ", padding: ";
    padding.print(os);
    os << endl;
  }

  if (verbose) {
    std::ostringstream os;
    os << *prefix << "Done" << endl;
    std::cerr << os.str();
  }
}

template <class Pointers, class InOutIndices, class InIndices, class IndexMap>
size_t
insert_crs_indices(
    typename Pointers::value_type const row,
    Pointers const& row_ptrs,
    InOutIndices& cur_indices,
    size_t& num_assigned,
    InIndices const& new_indices,
    IndexMap&& map,
    std::function<void(size_t const, size_t const, size_t const)> cb)
{
  if (new_indices.size() == 0) {
    return 0;
  }

  if (cur_indices.size() == 0) {
    // No room to insert new indices
    return Teuchos::OrdinalTraits<size_t>::invalid();
  }

  using offset_type = typename std::decay<decltype (row_ptrs[0])>::type;
  using ordinal_type = typename std::decay<decltype (cur_indices[0])>::type;

  const offset_type start = row_ptrs[row];
  offset_type end = start + static_cast<offset_type> (num_assigned);
  const size_t num_avail = (row_ptrs[row + 1] < end) ? size_t (0) :
    row_ptrs[row + 1] - end;
  const size_t num_new_indices = static_cast<size_t> (new_indices.size ());
  size_t num_inserted = 0;

  for (size_t k = 0; k < num_new_indices; ++k) {
    const ordinal_type idx = std::forward<IndexMap>(map)(new_indices[k]);
    offset_type row_offset = start;
    for (; row_offset < end; ++row_offset) {
      if (idx == cur_indices[row_offset]) {
        break;
      }
    }

    if (row_offset == end) {
      if (num_inserted >= num_avail) { // not enough room
        return Teuchos::OrdinalTraits<size_t>::invalid();
      }
      // This index is not yet in indices
      cur_indices[end++] = idx;
      num_inserted++;
    }
    if (cb) {
      cb(k, start, row_offset - start);
    }
  }
  num_assigned += num_inserted;
  return num_inserted;
}

template <class Pointers, class Indices1, class Indices2, class IndexMap, class Callback>
size_t
find_crs_indices(
    typename Pointers::value_type const row,
    Pointers const& row_ptrs,
    const size_t curNumEntries,
    Indices1 const& cur_indices,
    Indices2 const& new_indices,
    IndexMap&& map,
    Callback&& cb)
{
  if (new_indices.size() == 0)
    return 0;

  using ordinal = typename Indices1::value_type;
  auto invalid_ordinal = Teuchos::OrdinalTraits<ordinal>::invalid();

  const size_t start = static_cast<size_t> (row_ptrs[row]);
  const size_t end = start + curNumEntries;
  size_t num_found = 0;
  for (size_t k = 0; k < new_indices.size(); k++)
  {
    auto row_offset = start;
    auto idx = std::forward<IndexMap>(map)(new_indices[k]);
    if (idx == invalid_ordinal)
      continue;
    for (; row_offset < end; row_offset++)
    {
      if (idx == cur_indices[row_offset])
      {
        std::forward<Callback>(cb)(k, start, row_offset - start);
        num_found++;
      }
    }
  }
  return num_found;
}

} // namespace impl


template<class RowPtr, class Indices, class Padding>
void
padCrsArrays(
    const RowPtr& rowPtrBeg,
    const RowPtr& rowPtrEnd,
    Indices& indices,
    const Padding& padding,
    const int my_rank,
    const bool verbose)
{
  using impl::pad_crs_arrays;
  // send empty values array
  Indices values;
  pad_crs_arrays<RowPtr, Indices, Indices, Padding>(
    impl::PadCrsAction::INDICES_ONLY, rowPtrBeg, rowPtrEnd,
    indices, values, padding, my_rank, verbose);
}

template<class RowPtr, class Indices, class Values, class Padding>
void
padCrsArrays(
    const RowPtr& rowPtrBeg,
    const RowPtr& rowPtrEnd,
    Indices& indices,
    Values& values,
    const Padding& padding,
    const int my_rank,
    const bool verbose)
{
  using impl::pad_crs_arrays;
  pad_crs_arrays<RowPtr, Indices, Values, Padding>(
    impl::PadCrsAction::INDICES_AND_VALUES, rowPtrBeg, rowPtrEnd,
    indices, values, padding, my_rank, verbose);
}

template <class Pointers, class InOutIndices, class InIndices>
size_t
insertCrsIndices(
    typename Pointers::value_type const row,
    Pointers const& rowPtrs,
    InOutIndices& curIndices,
    size_t& numAssigned,
    InIndices const& newIndices,
    std::function<void(const size_t, const size_t, const size_t)> cb =
        std::function<void(const size_t, const size_t, const size_t)>())
{
  static_assert(std::is_same<typename std::remove_const<typename InOutIndices::value_type>::type,
                             typename std::remove_const<typename InIndices::value_type>::type>::value,
    "Expected views to have same value type");

  // Provide a unit map for the more general insert_indices
  using ordinal = typename InOutIndices::value_type;
  auto numInserted = impl::insert_crs_indices(row, rowPtrs, curIndices,
    numAssigned, newIndices, [](ordinal const idx) { return idx; }, cb);
  return numInserted;
}

template <class Pointers, class InOutIndices, class InIndices>
size_t
insertCrsIndices(
    typename Pointers::value_type const row,
    Pointers const& rowPtrs,
    InOutIndices& curIndices,
    size_t& numAssigned,
    InIndices const& newIndices,
    std::function<typename InOutIndices::value_type(const typename InIndices::value_type)> map,
    std::function<void(const size_t, const size_t, const size_t)> cb =
      std::function<void(const size_t, const size_t, const size_t)>())
{
  auto numInserted = impl::insert_crs_indices(row, rowPtrs, curIndices,
    numAssigned, newIndices, map, cb);
  return numInserted;
}


template <class Pointers, class Indices1, class Indices2, class Callback>
size_t
findCrsIndices(
    typename Pointers::value_type const row,
    Pointers const& rowPtrs,
    const size_t curNumEntries,
    Indices1 const& curIndices,
    Indices2 const& newIndices,
    Callback&& cb)
{
  static_assert(std::is_same<typename std::remove_const<typename Indices1::value_type>::type,
                             typename std::remove_const<typename Indices2::value_type>::type>::value,
    "Expected views to have same value type");
  // Provide a unit map for the more general find_crs_indices
  using ordinal = typename Indices2::value_type;
  auto numFound = impl::find_crs_indices(row, rowPtrs, curNumEntries, curIndices, newIndices,
    [=](ordinal ind){ return ind; }, cb);
  return numFound;
}

template <class Pointers, class Indices1, class Indices2, class IndexMap, class Callback>
size_t
findCrsIndices(
    typename Pointers::value_type const row,
    Pointers const& rowPtrs,
    const size_t curNumEntries,
    Indices1 const& curIndices,
    Indices2 const& newIndices,
    IndexMap&& map,
    Callback&& cb)
{
  return impl::find_crs_indices(row, rowPtrs, curNumEntries, curIndices, newIndices, map, cb);
}

} // namespace Details
} // namespace Tpetra

#endif // TPETRA_DETAILS_CRSUTILS_HPP