Tpetra_Details_crsUtils.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @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 "Tpetra_Details_WrappedDualView.hpp"
#include <iostream>
#include <memory>
#include <unordered_map>


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_wdv,
  Values& values_wdv,
  const Padding& padding,
  const int my_rank,
  const bool verbose)
{
  using execution_space = typename Indices::t_dev::execution_space;
  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);
    // DEEP_COPY REVIEW - NOT TESTED
    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);
    // DEEP_COPY REVIEW - NOT TESTED
    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_wdv.extent(0)
       << ", values.extent(0): " << values_wdv.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;
  {
    // Must do on host because padding uses std::map
    execution_space exec_space_instance = execution_space();
    auto row_ptr_end_h = create_mirror_view(
      hostSpace, row_ptr_end, verbose, prefix.get());
    // DEEP_COPY REVIEW - DEVICE-TO-HOSTMIRROR
    Kokkos::deep_copy(exec_space_instance, row_ptr_end_h, row_ptr_end);
    auto row_ptr_beg_h = create_mirror_view(
      hostSpace, row_ptr_beg, verbose, prefix.get());
    // DEEP_COPY REVIEW - DEVICE-TO-HOSTMIRROR
    Kokkos::deep_copy(exec_space_instance, row_ptr_beg_h, row_ptr_beg);

    // lbv 03/15/23: The execution space deep_copy does an asynchronous
    // copy so we really want to fence that space before touching the
    // host data as it is not guarenteed to have arrived by the time we
    // start the parallel_reduce below which might use a different
    // execution space, see:
    // https://kokkos.github.io/kokkos-core-wiki/API/core/view/deep_copy.html#semantics
    exec_space_instance.fence();

    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;
    }
    // DEEP_COPY REVIEW - HOSTMIRROR-TO-DEVICE
    Kokkos::deep_copy(execution_space(), newAllocPerRow, newAllocPerRow_h);
  }

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

  {
    auto indices_old = indices_wdv.getDeviceView(Access::ReadOnly);
    const size_t newIndsSize = size_t(indices_old.size()) + increase;
    auto indices_new = make_uninitialized_view<typename Indices::t_dev>(
      "Tpetra::CrsGraph column indices", newIndsSize, verbose,
      prefix.get());

    typename Values::t_dev values_new;
    auto values_old = values_wdv.getDeviceView(Access::ReadOnly);
    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<typename Values::t_dev>(
        "Tpetra::CrsMatrix values", newValsSize, verbose, prefix.get());
    }

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

    using range_type = Kokkos::RangePolicy<execution_space, size_t>;
    Kokkos::parallel_scan(
      "Tpetra::CrsGraph or CrsMatrix repack",
      range_type(size_t(0), size_t(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 = Kokkos::subview(indices_old, oldRange);
            auto newColInds = Kokkos::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 = 
                  Kokkos::subview(values_old, oldRange);
              auto newVals = Kokkos::subview(values_new, newRange);
              memcpy((void*) 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);
      // DEEP_COPY REVIEW - NOT TESTED
      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);
      // DEEP_COPY REVIEW - NOT TESTED
      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::cout << os.str();
    }

    indices_wdv = Indices(indices_new);
    values_wdv = Values(values_new);
  }
  
  if (verbose) {
    auto indices_h = indices_wdv.getHostView(Access::ReadOnly);
    auto values_h = values_wdv.getHostView(Access::ReadOnly);
    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;

  size_t numIndicesLookup = num_assigned + num_new_indices;

  // Threshold determined from test/Utils/insertCrsIndicesThreshold.cpp
  const size_t useLookUpTableThreshold = 400; 

  if (numIndicesLookup <= useLookUpTableThreshold || num_new_indices == 1) {
    // For rows with few nonzeros, can use a serial search to find duplicates
    // Or if inserting only one index, serial search is as fast as anything else
    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);
      }
    }
  }
  else {
    // For rows with many nonzeros, use a lookup table to find duplicates
    std::unordered_map<ordinal_type, offset_type> idxLookup(numIndicesLookup);

    // Put existing indices into the lookup table
    for (size_t k = 0; k < num_assigned; k++) {
      idxLookup[cur_indices[start+k]] = start+k;
    }

    // Check for new indices in table; insert if not there yet
    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;

      auto it = idxLookup.find(idx);
      if (it == idxLookup.end()) {
        if (num_inserted >= num_avail) { // not enough room
          return Teuchos::OrdinalTraits<size_t>::invalid();
        }
        // index not found; insert it
        row_offset = end;
        cur_indices[end++] = idx;
        idxLookup[idx] = row_offset;
        num_inserted++;
      }
      else {
        // index found; note its position
        row_offset = it->second;
      }
      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 std::remove_const<typename Indices1::value_type>::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_wdv,
    const Padding& padding,
    const int my_rank,
    const bool verbose)
{
  using impl::pad_crs_arrays;
  // send empty values array
  Indices values_null; 
  pad_crs_arrays<RowPtr, Indices, Indices, Padding>( 
    impl::PadCrsAction::INDICES_ONLY, rowPtrBeg, rowPtrEnd,
    indices_wdv, values_null, padding, my_rank, verbose);
}

template<class RowPtr, class Indices, class Values, class Padding>
void
padCrsArrays(
    const RowPtr& rowPtrBeg,
    const RowPtr& rowPtrEnd,
    Indices& indices_wdv,
    Values& values_wdv,
    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_wdv, values_wdv, 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