Kokkos_StaticCrsGraph.hpp

/*
//@HEADER
// ************************************************************************
//
//                        Kokkos v. 3.0
//       Copyright (2020) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// 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 NTESS "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 NTESS 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.
//
// Questions? Contact Christian R. Trott (crtrott@sandia.gov)
//
// ************************************************************************
//@HEADER
*/

#ifndef KOKKOS_STATICCRSGRAPH_HPP
#define KOKKOS_STATICCRSGRAPH_HPP

#include <string>
#include <vector>

#include <Kokkos_View.hpp>
#include <Kokkos_Parallel.hpp>
#include <Kokkos_Parallel_Reduce.hpp>

namespace Kokkos {

namespace Impl {
template <class RowOffsetsType, class RowBlockOffsetsType>
struct StaticCrsGraphBalancerFunctor {
  typedef typename RowOffsetsType::non_const_value_type int_type;
  RowOffsetsType row_offsets;
  RowBlockOffsetsType row_block_offsets;

  int_type cost_per_row, num_blocks;

  StaticCrsGraphBalancerFunctor(RowOffsetsType row_offsets_,
                                RowBlockOffsetsType row_block_offsets_,
                                int_type cost_per_row_, int_type num_blocks_)
      : row_offsets(row_offsets_),
        row_block_offsets(row_block_offsets_),
        cost_per_row(cost_per_row_),
        num_blocks(num_blocks_) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const int_type& iRow) const {
    const int_type num_rows    = row_offsets.extent(0) - 1;
    const int_type num_entries = row_offsets(num_rows);
    const int_type total_cost  = num_entries + num_rows * cost_per_row;

    const double cost_per_workset = 1.0 * total_cost / num_blocks;

    const int_type row_cost =
        row_offsets(iRow + 1) - row_offsets(iRow) + cost_per_row;

    int_type count = row_offsets(iRow + 1) + cost_per_row * iRow;

    if (iRow == num_rows - 1) row_block_offsets(num_blocks) = num_rows;

    if (true) {
      int_type current_block =
          (count - row_cost - cost_per_row) / cost_per_workset;
      int_type end_block = count / cost_per_workset;

      // Handle some corner cases for the last two blocks.
      if (current_block >= num_blocks - 2) {
        if ((current_block == num_blocks - 2) &&
            (count >= (current_block + 1) * cost_per_workset)) {
          int_type row   = iRow;
          int_type cc    = count - row_cost - cost_per_row;
          int_type block = cc / cost_per_workset;
          while ((block > 0) && (block == current_block)) {
            cc    = row_offsets(row) + row * cost_per_row;
            block = cc / cost_per_workset;
            row--;
          }
          if ((count - cc - row_cost - cost_per_row) <
              num_entries - row_offsets(iRow + 1)) {
            row_block_offsets(current_block + 1) = iRow + 1;
          } else {
            row_block_offsets(current_block + 1) = iRow;
          }
        }
      } else {
        if ((count >= (current_block + 1) * cost_per_workset) ||
            (iRow + 2 == int_type(row_offsets.extent(0)))) {
          if (end_block > current_block + 1) {
            int_type num_block                   = end_block - current_block;
            row_block_offsets(current_block + 1) = iRow;
            for (int_type block = current_block + 2; block <= end_block;
                 block++)
              if ((block < current_block + 2 + (num_block - 1) / 2))
                row_block_offsets(block) = iRow;
              else
                row_block_offsets(block) = iRow + 1;
          } else {
            row_block_offsets(current_block + 1) = iRow + 1;
          }
        }
      }
    }
  }
};
}  // namespace Impl

template <class GraphType>
struct GraphRowViewConst {
  typedef const typename GraphType::data_type ordinal_type;

 private:
  ordinal_type* colidx_;
  const ordinal_type stride_;

 public:
  KOKKOS_INLINE_FUNCTION
  GraphRowViewConst(ordinal_type* const colidx_in, const ordinal_type& stride,
                    const ordinal_type& count)
      : colidx_(colidx_in), stride_(stride), length(count) {}

  template <class OffsetType>
  KOKKOS_INLINE_FUNCTION GraphRowViewConst(
      const typename GraphType::entries_type& colidx_in,
      const ordinal_type& stride, const ordinal_type& count,
      const OffsetType& idx,
      const typename std::enable_if<std::is_integral<OffsetType>::value,
                                    int>::type& = 0)
      : colidx_(&colidx_in(idx)), stride_(stride), length(count) {}

  const ordinal_type length;

  KOKKOS_INLINE_FUNCTION
  ordinal_type& colidx(const ordinal_type& i) const {
    return colidx_[i * stride_];
  }

  KOKKOS_INLINE_FUNCTION
  ordinal_type& operator()(const ordinal_type& i) const { return colidx(i); }
};

template <class DataType, class Arg1Type, class Arg2Type = void,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
          typename SizeType =
              typename ViewTraits<DataType*, Arg1Type, Arg2Type>::size_type,
          class Arg3Type = void>
#else
          class Arg3Type    = void,
          typename SizeType = typename ViewTraits<DataType*, Arg1Type, Arg2Type,
                                                  Arg3Type>::size_type>
#endif
class StaticCrsGraph {
 private:
  typedef ViewTraits<DataType*, Arg1Type, Arg2Type, Arg3Type> traits;

 public:
  typedef DataType data_type;
  typedef typename traits::array_layout array_layout;
  typedef typename traits::execution_space execution_space;
  typedef typename traits::device_type device_type;
  typedef typename traits::memory_traits memory_traits;
  typedef SizeType size_type;

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
  typedef StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType, Arg3Type>
      staticcrsgraph_type;
  typedef StaticCrsGraph<data_type, array_layout,
                         typename traits::host_mirror_space, size_type,
                         memory_traits>
      HostMirror;
#else
  typedef StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type, SizeType>
      staticcrsgraph_type;
  typedef StaticCrsGraph<data_type, array_layout,
                         typename traits::host_mirror_space, memory_traits,
                         size_type>
      HostMirror;
#endif

  typedef View<const size_type*, array_layout, device_type, memory_traits>
      row_map_type;
  typedef View<data_type*, array_layout, device_type, memory_traits>
      entries_type;
  typedef View<const size_type*, array_layout, device_type, memory_traits>
      row_block_type;

  entries_type entries;
  row_map_type row_map;
  row_block_type row_block_offsets;

  KOKKOS_INLINE_FUNCTION
  StaticCrsGraph() : entries(), row_map(), row_block_offsets() {}

  KOKKOS_INLINE_FUNCTION
  StaticCrsGraph(const StaticCrsGraph& rhs)
      : entries(rhs.entries),
        row_map(rhs.row_map),
        row_block_offsets(rhs.row_block_offsets) {}

  template <class EntriesType, class RowMapType>
  KOKKOS_INLINE_FUNCTION StaticCrsGraph(const EntriesType& entries_,
                                        const RowMapType& row_map_)
      : entries(entries_), row_map(row_map_), row_block_offsets() {}

  KOKKOS_INLINE_FUNCTION
  StaticCrsGraph& operator=(const StaticCrsGraph& rhs) {
    entries           = rhs.entries;
    row_map           = rhs.row_map;
    row_block_offsets = rhs.row_block_offsets;
    return *this;
  }

  KOKKOS_DEFAULTED_FUNCTION
  ~StaticCrsGraph() = default;

  KOKKOS_INLINE_FUNCTION
  size_type numRows() const {
    return (row_map.extent(0) != 0)
               ? row_map.extent(0) - static_cast<size_type>(1)
               : static_cast<size_type>(0);
  }

  KOKKOS_INLINE_FUNCTION
  GraphRowViewConst<StaticCrsGraph> rowConst(const data_type i) const {
    const size_type start = row_map(i);
    // count is guaranteed to fit in ordinal_type, as long as no row
    // has duplicate entries.
    const data_type count = static_cast<data_type>(row_map(i + 1) - start);

    if (count == 0) {
      return GraphRowViewConst<StaticCrsGraph>(nullptr, 1, 0);
    } else {
      return GraphRowViewConst<StaticCrsGraph>(entries, 1, count, start);
    }
  }

  void create_block_partitioning(size_type num_blocks,
                                 size_type fix_cost_per_row = 4) {
    View<size_type*, array_layout, device_type> block_offsets(
        "StatisCrsGraph::load_balance_offsets", num_blocks + 1);

    Impl::StaticCrsGraphBalancerFunctor<
        row_map_type, View<size_type*, array_layout, device_type> >
        partitioner(row_map, block_offsets, fix_cost_per_row, num_blocks);

    Kokkos::parallel_for("Kokkos::StaticCrsGraph::create_block_partitioning",
                         Kokkos::RangePolicy<execution_space>(0, numRows()),
                         partitioner);
    typename device_type::execution_space().fence();

    row_block_offsets = block_offsets;
  }
};

//----------------------------------------------------------------------------

template <class StaticCrsGraphType, class InputSizeType>
typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
    const std::string& label, const std::vector<InputSizeType>& input);

template <class StaticCrsGraphType, class InputSizeType>
typename StaticCrsGraphType::staticcrsgraph_type create_staticcrsgraph(
    const std::string& label,
    const std::vector<std::vector<InputSizeType> >& input);

//----------------------------------------------------------------------------

template <class DataType, class Arg1Type, class Arg2Type,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
          typename SizeType, class Arg3Type>
typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType,
                        Arg3Type>::HostMirror
create_mirror_view(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType,
                                        Arg3Type>& input);
#else
          class Arg3Type, typename SizeType>
typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
                        SizeType>::HostMirror
create_mirror_view(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
                                        SizeType>& input);
#endif

template <class DataType, class Arg1Type, class Arg2Type,
#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
          typename SizeType, class Arg3Type>
typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType,
                        Arg3Type>::HostMirror
create_mirror_view(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType,
                                        Arg3Type>& input);
#else
          class Arg3Type, typename SizeType>
typename StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
                        SizeType>::HostMirror
create_mirror(const StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type,
                                   SizeType>& input);
#endif

}  // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

#include <impl/Kokkos_StaticCrsGraph_factory.hpp>

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {
namespace Impl {

template <class GraphType>
struct StaticCrsGraphMaximumEntry {
  typedef typename GraphType::execution_space execution_space;
  typedef typename GraphType::data_type value_type;

  const typename GraphType::entries_type entries;

  StaticCrsGraphMaximumEntry(const GraphType& graph) : entries(graph.entries) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const unsigned i, value_type& update) const {
    if (update < entries(i)) update = entries(i);
  }

  KOKKOS_INLINE_FUNCTION
  void init(value_type& update) const { update = 0; }

  KOKKOS_INLINE_FUNCTION
  void join(volatile value_type& update,
            volatile const value_type& input) const {
    if (update < input) update = input;
  }
};

}  // namespace Impl

#ifdef KOKKOS_ENABLE_DEPRECATED_CODE
template <class DataType, class Arg1Type, class Arg2Type, typename SizeType,
          class Arg3Type>
DataType maximum_entry(const StaticCrsGraph<DataType, Arg1Type, Arg2Type,
                                            SizeType, Arg3Type>& graph) {
  typedef StaticCrsGraph<DataType, Arg1Type, Arg2Type, SizeType, Arg3Type>
      GraphType;
#else
template <class DataType, class Arg1Type, class Arg2Type, class Arg3Type,
          typename SizeType>
DataType maximum_entry(const StaticCrsGraph<DataType, Arg1Type, Arg2Type,
                                            Arg3Type, SizeType>& graph) {
  typedef StaticCrsGraph<DataType, Arg1Type, Arg2Type, Arg3Type, SizeType>
      GraphType;
#endif
  typedef Impl::StaticCrsGraphMaximumEntry<GraphType> FunctorType;

  DataType result = 0;
  Kokkos::parallel_reduce("Kokkos::maximum_entry", graph.entries.extent(0),
                          FunctorType(graph), result);
  return result;
}

}  // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

#endif /* #ifndef KOKKOS_CRSARRAY_HPP */