Ifpack2_BlockComputeResidualVector.hpp

// @HEADER
// *****************************************************************************
//       Ifpack2: Templated Object-Oriented Algebraic Preconditioner Package
//
// Copyright 2009 NTESS and the Ifpack2 contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef IFPACK2_BLOCKCOMPUTERES_IMPL_HPP
#define IFPACK2_BLOCKCOMPUTERES_IMPL_HPP

#include "Ifpack2_BlockHelper.hpp"

namespace Ifpack2 {

namespace BlockHelperDetails {

template <typename MatrixType>
struct AmD {
  using impl_type                       = BlockHelperDetails::ImplType<MatrixType>;
  using local_ordinal_type_1d_view      = typename impl_type::local_ordinal_type_1d_view;
  using size_type_1d_view               = typename impl_type::size_type_1d_view;
  using i64_3d_view                     = typename impl_type::i64_3d_view;
  using impl_scalar_type_1d_view_tpetra = Unmanaged<typename impl_type::impl_scalar_type_1d_view_tpetra>;
  // rowptr points to the start of each row of A_colindsub.
  size_type_1d_view rowptr, rowptr_remote;
  // Indices into A's rows giving the blocks to extract. rowptr(i) points to
  // the i'th row. Thus, g.entries(A_colindsub(rowptr(row) : rowptr(row+1))),
  // where g is A's graph, are the columns AmD uses. If seq_method_, then
  // A_colindsub contains all the LIDs and A_colindsub_remote is empty. If !
  // seq_method_, then A_colindsub contains owned LIDs and A_colindsub_remote
  // contains the remote ones.
  local_ordinal_type_1d_view A_colindsub, A_colindsub_remote;
  // Precomputed direct offsets to A,x blocks, for owned entries (OverlapTag case) or all entries (AsyncTag case)
  i64_3d_view A_x_offsets;
  // Precomputed direct offsets to A,x blocks, for non-owned entries (OverlapTag case). For AsyncTag case this is left empty.
  i64_3d_view A_x_offsets_remote;

  // Currently always true.
  bool is_tpetra_block_crs;

  // If is_tpetra_block_crs, then this is a pointer to A_'s value data.
  impl_scalar_type_1d_view_tpetra tpetra_values;

  AmD()             = default;
  AmD(const AmD &b) = default;
};

template <typename MatrixType>
struct PartInterface {
  using local_ordinal_type         = typename BlockHelperDetails::ImplType<MatrixType>::local_ordinal_type;
  using local_ordinal_type_1d_view = typename BlockHelperDetails::ImplType<MatrixType>::local_ordinal_type_1d_view;
  using local_ordinal_type_2d_view = typename BlockHelperDetails::ImplType<MatrixType>::local_ordinal_type_2d_view;

  PartInterface()                       = default;
  PartInterface(const PartInterface &b) = default;

  // Some terms:
  //   The matrix A is split as A = D + R, where D is the matrix of tridiag
  // blocks and R is the remainder.
  //   A part is roughly a synonym for a tridiag. The distinction is that a part
  // is the set of rows belonging to one tridiag and, equivalently, the off-diag
  // rows in R associated with that tridiag. In contrast, the term tridiag is
  // used to refer specifically to tridiag data, such as the pointer into the
  // tridiag data array.
  //   Local (lcl) row are the LIDs. lclrow lists the LIDs belonging to each
  // tridiag, and partptr points to the beginning of each tridiag. This is the
  // LID space.
  //   Row index (idx) is the ordinal in the tridiag ordering. lclrow is indexed
  // by this ordinal. This is the 'index' space.
  //   A flat index is the mathematical index into an array. A pack index
  // accounts for SIMD packing.

  // Local row LIDs. Permutation from caller's index space to tridiag index
  // space.
  local_ordinal_type_1d_view lclrow;
  // partptr_ is the pointer array into lclrow_.
  local_ordinal_type_1d_view partptr;  // np+1
  local_ordinal_type_2d_view partptr_sub;
  local_ordinal_type_1d_view partptr_schur;
  // packptr_(i), for i the pack index, indexes partptr_. partptr_(packptr_(i))
  // is the start of the i'th pack.
  local_ordinal_type_1d_view packptr;  // npack+1
  local_ordinal_type_1d_view packptr_sub;
  local_ordinal_type_1d_view packindices_sub;
  local_ordinal_type_2d_view packindices_schur;
  // part2rowidx0_(i) is the flat row index of the start of the i'th part. It's
  // an alias of partptr_ in the case of no overlap.
  local_ordinal_type_1d_view part2rowidx0;  // np+1
  local_ordinal_type_1d_view part2rowidx0_sub;
  // part2packrowidx0_(i) is the packed row index. If vector_length is 1, then
  // it's the same as part2rowidx0_; if it's > 1, then the value is combined
  // with i % vector_length to get the location in the packed data.
  local_ordinal_type_1d_view part2packrowidx0;  // np+1
  local_ordinal_type_2d_view part2packrowidx0_sub;
  local_ordinal_type part2packrowidx0_back;  // So we don't need to grab the array from the GPU.
  // rowidx2part_ maps the row index to the part index.
  local_ordinal_type_1d_view rowidx2part;  // nr
  local_ordinal_type_1d_view rowidx2part_sub;
  // True if lcl{row|col} is at most a constant away from row{idx|col}. In
  // practice, this knowledge is not particularly useful, as packing for batched
  // processing is done at the same time as the permutation from LID to index
  // space. But it's easy to detect, so it's recorded in case an optimization
  // can be made based on it.
  bool row_contiguous;

  local_ordinal_type max_partsz;
  local_ordinal_type max_subpartsz;
  local_ordinal_type n_subparts_per_part;
  local_ordinal_type nparts;
};

// Precompute offsets of each A and x entry to speed up residual.
// (Applies for hasBlockCrsMatrix == true and OverlapTag/AsyncTag)
// Reading A, x take up to 4, 6 levels of indirection respectively,
// but precomputing the offsets reduces it to 2 for both.
//
// This function allocates and populates these members of AmD:
// A_x_offsets, A_x_offsets_remote
template <typename MatrixType>
void precompute_A_x_offsets(
    AmD<MatrixType> &amd,
    const PartInterface<MatrixType> &interf,
    const Teuchos::RCP<const typename ImplType<MatrixType>::tpetra_crs_graph_type> &g,
    const typename ImplType<MatrixType>::local_ordinal_type_1d_view &dm2cm,
    int blocksize,
    bool ownedRemoteSeparate) {
  using impl_type                 = ImplType<MatrixType>;
  using i64_3d_view               = typename impl_type::i64_3d_view;
  using size_type                 = typename impl_type::size_type;
  using local_ordinal_type        = typename impl_type::local_ordinal_type;
  using execution_space           = typename impl_type::execution_space;
  auto local_graph                = g->getLocalGraphDevice();
  const auto A_block_rowptr       = local_graph.row_map;
  const auto A_colind             = local_graph.entries;
  local_ordinal_type numLocalRows = interf.rowidx2part.extent(0);
  int blocksize_square            = blocksize * blocksize;
  // shallow-copying views to avoid capturing the amd, interf objects in lambdas
  auto lclrow             = interf.lclrow;
  auto A_colindsub        = amd.A_colindsub;
  auto A_colindsub_remote = amd.A_colindsub_remote;
  auto rowptr             = amd.rowptr;
  auto rowptr_remote      = amd.rowptr_remote;
  bool is_dm2cm_active    = dm2cm.extent(0);
  if (ownedRemoteSeparate) {
    // amd.rowptr points to owned entries only, and amd.rowptr_remote points to nonowned.
    local_ordinal_type maxOwnedEntriesPerRow    = 0;
    local_ordinal_type maxNonownedEntriesPerRow = 0;
    Kokkos::parallel_reduce(
        Kokkos::RangePolicy<execution_space>(0, numLocalRows),
        KOKKOS_LAMBDA(local_ordinal_type i, local_ordinal_type & lmaxOwned, local_ordinal_type & lmaxNonowned) {
          const local_ordinal_type lr    = lclrow(i);
          local_ordinal_type rowNumOwned = rowptr(lr + 1) - rowptr(lr);
          if (rowNumOwned > lmaxOwned)
            lmaxOwned = rowNumOwned;
          // rowptr_remote won't be allocated for single-rank problems
          if (rowptr_remote.extent(0)) {
            local_ordinal_type rowNumNonowned = rowptr_remote(lr + 1) - rowptr_remote(lr);
            if (rowNumNonowned > lmaxNonowned)
              lmaxNonowned = rowNumNonowned;
          } else {
            lmaxNonowned = 0;
          }
        },
        Kokkos::Max<local_ordinal_type>(maxOwnedEntriesPerRow), Kokkos::Max<local_ordinal_type>(maxNonownedEntriesPerRow));
    // Allocate the two offsets views now that we know the dimensions
    // For each one, the middle dimension is 0 for A offsets and 1 for x offsets.
    // Packing them together in one view improves cache line utilization
    amd.A_x_offsets         = i64_3d_view("amd.A_x_offsets", numLocalRows, 2, maxOwnedEntriesPerRow);
    amd.A_x_offsets_remote  = i64_3d_view("amd.A_x_offsets_remote", numLocalRows, 2, maxNonownedEntriesPerRow);
    auto A_x_offsets        = amd.A_x_offsets;
    auto A_x_offsets_remote = amd.A_x_offsets_remote;
    // Now, populate all the offsets. Use ArithTraits<int64_t>::min to mark absent entries.
    Kokkos::parallel_for(
        Kokkos::RangePolicy<execution_space>(0, numLocalRows),
        KOKKOS_LAMBDA(local_ordinal_type i) {
          const local_ordinal_type lr = lclrow(i);
          const size_type A_k0        = A_block_rowptr(lr);
          // Owned entries
          size_type rowBegin             = rowptr(lr);
          local_ordinal_type rowNumOwned = rowptr(lr + 1) - rowBegin;
          for (local_ordinal_type entry = 0; entry < maxOwnedEntriesPerRow; entry++) {
            if (entry < rowNumOwned) {
              const size_type j                      = A_k0 + A_colindsub(rowBegin + entry);
              const local_ordinal_type A_colind_at_j = A_colind(j);
              const local_ordinal_type loc           = is_dm2cm_active ? dm2cm(A_colind_at_j) : A_colind_at_j;
              A_x_offsets(i, 0, entry)               = int64_t(j) * blocksize_square;
              A_x_offsets(i, 1, entry)               = int64_t(loc) * blocksize;
            } else {
              A_x_offsets(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
              A_x_offsets(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
            }
          }
          // Nonowned entries
          if (rowptr_remote.extent(0)) {
            rowBegin                          = rowptr_remote(lr);
            local_ordinal_type rowNumNonowned = rowptr_remote(lr + 1) - rowBegin;
            for (local_ordinal_type entry = 0; entry < maxNonownedEntriesPerRow; entry++) {
              if (entry < rowNumNonowned) {
                const size_type j                      = A_k0 + A_colindsub_remote(rowBegin + entry);
                const local_ordinal_type A_colind_at_j = A_colind(j);
                const local_ordinal_type loc           = A_colind_at_j - numLocalRows;
                A_x_offsets_remote(i, 0, entry)        = int64_t(j) * blocksize_square;
                A_x_offsets_remote(i, 1, entry)        = int64_t(loc) * blocksize;
              } else {
                A_x_offsets_remote(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
                A_x_offsets_remote(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
              }
            }
          }
        });
  } else {
    // amd.rowptr points to both owned and nonowned entries, so it tells us how many columns (last dim) A_x_offsets should have.
    local_ordinal_type maxEntriesPerRow = 0;
    Kokkos::parallel_reduce(
        Kokkos::RangePolicy<execution_space>(0, numLocalRows),
        KOKKOS_LAMBDA(local_ordinal_type i, local_ordinal_type & lmax) {
          const local_ordinal_type lr = lclrow(i);
          local_ordinal_type rowNum   = rowptr(lr + 1) - rowptr(lr);
          if (rowNum > lmax)
            lmax = rowNum;
        },
        Kokkos::Max<local_ordinal_type>(maxEntriesPerRow));
    amd.A_x_offsets  = i64_3d_view("amd.A_x_offsets", numLocalRows, 2, maxEntriesPerRow);
    auto A_x_offsets = amd.A_x_offsets;
    // Populate A,x offsets. Use ArithTraits<int64_t>::min to mark absent entries.
    // For x offsets, add a shift blocksize*numLocalRows to represent that it indexes into x_remote instead of x.
    Kokkos::parallel_for(
        Kokkos::RangePolicy<execution_space>(0, numLocalRows),
        KOKKOS_LAMBDA(local_ordinal_type i) {
          const local_ordinal_type lr = lclrow(i);
          const size_type A_k0        = A_block_rowptr(lr);
          // Owned entries
          size_type rowBegin          = rowptr(lr);
          local_ordinal_type rowOwned = rowptr(lr + 1) - rowBegin;
          for (local_ordinal_type entry = 0; entry < maxEntriesPerRow; entry++) {
            if (entry < rowOwned) {
              const size_type j                      = A_k0 + A_colindsub(rowBegin + entry);
              A_x_offsets(i, 0, entry)               = j * blocksize_square;
              const local_ordinal_type A_colind_at_j = A_colind(j);
              if (A_colind_at_j < numLocalRows) {
                const local_ordinal_type loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
                A_x_offsets(i, 1, entry)     = int64_t(loc) * blocksize;
              } else {
                A_x_offsets(i, 1, entry) = int64_t(A_colind_at_j) * blocksize;
              }
            } else {
              A_x_offsets(i, 0, entry) = Kokkos::ArithTraits<int64_t>::min();
              A_x_offsets(i, 1, entry) = Kokkos::ArithTraits<int64_t>::min();
            }
          }
        });
  }
}

static inline int ComputeResidualVectorRecommendedCudaVectorSize(const int blksize,
                                                                 const int team_size) {
  int total_team_size(0);
  if (blksize <= 5)
    total_team_size = 32;
  else if (blksize <= 9)
    total_team_size = 32;  // 64
  else if (blksize <= 12)
    total_team_size = 96;
  else if (blksize <= 16)
    total_team_size = 128;
  else if (blksize <= 20)
    total_team_size = 160;
  else
    total_team_size = 160;
  return total_team_size / team_size;
}

static inline int ComputeResidualVectorRecommendedHIPVectorSize(const int blksize,
                                                                const int team_size) {
  int total_team_size(0);
  if (blksize <= 5)
    total_team_size = 32;
  else if (blksize <= 9)
    total_team_size = 32;  // 64
  else if (blksize <= 12)
    total_team_size = 96;
  else if (blksize <= 16)
    total_team_size = 128;
  else if (blksize <= 20)
    total_team_size = 160;
  else
    total_team_size = 160;
  return total_team_size / team_size;
}

static inline int ComputeResidualVectorRecommendedSYCLVectorSize(const int blksize,
                                                                 const int team_size) {
  int total_team_size(0);
  if (blksize <= 5)
    total_team_size = 32;
  else if (blksize <= 9)
    total_team_size = 32;  // 64
  else if (blksize <= 12)
    total_team_size = 96;
  else if (blksize <= 16)
    total_team_size = 128;
  else if (blksize <= 20)
    total_team_size = 160;
  else
    total_team_size = 160;
  return total_team_size / team_size;
}

template <typename T>
static inline int ComputeResidualVectorRecommendedVectorSize(const int blksize,
                                                             const int team_size) {
  if (is_cuda<T>::value)
    return ComputeResidualVectorRecommendedCudaVectorSize(blksize, team_size);
  if (is_hip<T>::value)
    return ComputeResidualVectorRecommendedHIPVectorSize(blksize, team_size);
  if (is_sycl<T>::value)
    return ComputeResidualVectorRecommendedSYCLVectorSize(blksize, team_size);
  return -1;
}

template <typename MatrixType>
struct ComputeResidualVector {
 public:
  using impl_type        = BlockHelperDetails::ImplType<MatrixType>;
  using node_device_type = typename impl_type::node_device_type;
  using execution_space  = typename impl_type::execution_space;
  using memory_space     = typename impl_type::memory_space;

  using local_ordinal_type  = typename impl_type::local_ordinal_type;
  using size_type           = typename impl_type::size_type;
  using impl_scalar_type    = typename impl_type::impl_scalar_type;
  using magnitude_type      = typename impl_type::magnitude_type;
  using btdm_scalar_type    = typename impl_type::btdm_scalar_type;
  using btdm_magnitude_type = typename impl_type::btdm_magnitude_type;
  using local_ordinal_type_1d_view      = typename impl_type::local_ordinal_type_1d_view;
  using size_type_1d_view               = typename impl_type::size_type_1d_view;
  using tpetra_block_access_view_type   = typename impl_type::tpetra_block_access_view_type;  // block crs (layout right)
  using impl_scalar_type_1d_view        = typename impl_type::impl_scalar_type_1d_view;
  using impl_scalar_type_2d_view_tpetra = typename impl_type::impl_scalar_type_2d_view_tpetra;  // block multivector (layout left)
  using vector_type_3d_view             = typename impl_type::vector_type_3d_view;
  using btdm_scalar_type_4d_view        = typename impl_type::btdm_scalar_type_4d_view;
  using i64_3d_view                     = typename impl_type::i64_3d_view;
  static constexpr int vector_length    = impl_type::vector_length;

  using member_type = typename Kokkos::TeamPolicy<execution_space>::member_type;

  // enum for max blocksize and vector length
  enum : int { max_blocksize = 32 };

 private:
  ConstUnmanaged<impl_scalar_type_2d_view_tpetra> b;
  ConstUnmanaged<impl_scalar_type_2d_view_tpetra> x;  // x_owned
  ConstUnmanaged<impl_scalar_type_2d_view_tpetra> x_remote;
  Unmanaged<impl_scalar_type_2d_view_tpetra> y;
  Unmanaged<vector_type_3d_view> y_packed;
  Unmanaged<btdm_scalar_type_4d_view> y_packed_scalar;

  // AmD information
  const ConstUnmanaged<size_type_1d_view> rowptr, rowptr_remote;
  const ConstUnmanaged<local_ordinal_type_1d_view> colindsub, colindsub_remote;
  const ConstUnmanaged<impl_scalar_type_1d_view> tpetra_values;

  // block crs graph information
  // for cuda (kokkos crs graph uses a different size_type from size_t)
  const ConstUnmanaged<Kokkos::View<size_t *, node_device_type>> A_block_rowptr;
  const ConstUnmanaged<Kokkos::View<size_t *, node_device_type>> A_point_rowptr;
  const ConstUnmanaged<Kokkos::View<local_ordinal_type *, node_device_type>> A_colind;

  // blocksize
  const local_ordinal_type blocksize_requested;

  // part interface
  const ConstUnmanaged<local_ordinal_type_1d_view> part2packrowidx0;
  const ConstUnmanaged<local_ordinal_type_1d_view> part2rowidx0;
  const ConstUnmanaged<local_ordinal_type_1d_view> rowidx2part;
  const ConstUnmanaged<local_ordinal_type_1d_view> partptr;
  const ConstUnmanaged<local_ordinal_type_1d_view> lclrow;
  const ConstUnmanaged<local_ordinal_type_1d_view> dm2cm;

  // block offsets
  const ConstUnmanaged<i64_3d_view> A_x_offsets;
  const ConstUnmanaged<i64_3d_view> A_x_offsets_remote;

  const bool is_dm2cm_active;
  const bool hasBlockCrsMatrix;

 public:
  template <typename LocalCrsGraphType>
  ComputeResidualVector(const AmD<MatrixType> &amd,
                        const LocalCrsGraphType &block_graph,
                        const LocalCrsGraphType &point_graph,
                        const local_ordinal_type &blocksize_requested_,
                        const PartInterface<MatrixType> &interf,
                        const local_ordinal_type_1d_view &dm2cm_,
                        bool hasBlockCrsMatrix_)
    : rowptr(amd.rowptr)
    , rowptr_remote(amd.rowptr_remote)
    , colindsub(amd.A_colindsub)
    , colindsub_remote(amd.A_colindsub_remote)
    , tpetra_values(amd.tpetra_values)
    , A_block_rowptr(block_graph.row_map)
    , A_point_rowptr(point_graph.row_map)
    , A_colind(block_graph.entries)
    , blocksize_requested(blocksize_requested_)
    , part2packrowidx0(interf.part2packrowidx0)
    , part2rowidx0(interf.part2rowidx0)
    , rowidx2part(interf.rowidx2part)
    , partptr(interf.partptr)
    , lclrow(interf.lclrow)
    , dm2cm(dm2cm_)
    , A_x_offsets(amd.A_x_offsets)
    , A_x_offsets_remote(amd.A_x_offsets_remote)
    , is_dm2cm_active(dm2cm_.span() > 0)
    , hasBlockCrsMatrix(hasBlockCrsMatrix_) {}

  inline void
  SerialDot(const local_ordinal_type &blocksize,
            const local_ordinal_type &lclRowID,
            const local_ordinal_type &lclColID,
            const local_ordinal_type &ii,
            const ConstUnmanaged<local_ordinal_type_1d_view> colindsub_,
            const impl_scalar_type *const KOKKOS_RESTRICT xx,
            /* */ impl_scalar_type *KOKKOS_RESTRICT yy) const {
    const size_type Aj_c  = colindsub_(lclColID);
    auto point_row_offset = A_point_rowptr(lclRowID * blocksize + ii) + Aj_c * blocksize;
    impl_scalar_type val  = 0;
#if defined(KOKKOS_ENABLE_PRAGMA_IVDEP)
#pragma ivdep
#endif
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
    for (local_ordinal_type k1 = 0; k1 < blocksize; ++k1)
      val += tpetra_values(point_row_offset + k1) * xx[k1];
    yy[ii] -= val;
  }

  inline void
  SerialGemv(const local_ordinal_type &blocksize,
             const impl_scalar_type *const KOKKOS_RESTRICT AA,
             const impl_scalar_type *const KOKKOS_RESTRICT xx,
             /* */ impl_scalar_type *KOKKOS_RESTRICT yy) const {
    using tlb = BlockHelperDetails::TpetraLittleBlock<Tpetra::Impl::BlockCrsMatrixLittleBlockArrayLayout>;
    for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0) {
      impl_scalar_type val = 0;
#if defined(KOKKOS_ENABLE_PRAGMA_IVDEP)
#pragma ivdep
#endif
#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
#pragma unroll
#endif
      for (local_ordinal_type k1 = 0; k1 < blocksize; ++k1)
        val += AA[tlb::getFlatIndex(k0, k1, blocksize)] * xx[k1];
      yy[k0] -= val;
    }
  }

  template <typename bbViewType, typename yyViewType>
  KOKKOS_INLINE_FUNCTION void
  VectorCopy(const member_type &member,
             const local_ordinal_type &blocksize,
             const bbViewType &bb,
             const yyViewType &yy) const {
    Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize), [&](const local_ordinal_type &k0) {
      yy(k0) = static_cast<typename yyViewType::const_value_type>(bb(k0));
    });
  }

  template <typename xxViewType, typename yyViewType>
  KOKKOS_INLINE_FUNCTION void
  VectorDot(const member_type &member,
            const local_ordinal_type &blocksize,
            const local_ordinal_type &lclRowID,
            const local_ordinal_type &lclColID,
            const local_ordinal_type &ii,
            const ConstUnmanaged<local_ordinal_type_1d_view> colindsub_,
            const xxViewType &xx,
            const yyViewType &yy) const {
    const size_type Aj_c  = colindsub_(lclColID);
    auto point_row_offset = A_point_rowptr(lclRowID * blocksize + ii) + Aj_c * blocksize;
    impl_scalar_type val  = 0;
    Kokkos::parallel_reduce(
        Kokkos::ThreadVectorRange(member, blocksize),
        [&](const local_ordinal_type &k1, impl_scalar_type &update) {
          update += tpetra_values(point_row_offset + k1) * xx(k1);
        },
        val);
    Kokkos::single(Kokkos::PerThread(member),
                   [=]() {
                     Kokkos::atomic_add(&yy(ii), typename yyViewType::const_value_type(-val));
                   });
  }

  // BMK: This version coalesces accesses to AA for LayoutRight blocks.
  template <typename AAViewType, typename xxViewType, typename yyViewType>
  KOKKOS_INLINE_FUNCTION void
  VectorGemv(const member_type &member,
             const local_ordinal_type &blocksize,
             const AAViewType &AA,
             const xxViewType &xx,
             const yyViewType &yy) const {
    for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0) {
      impl_scalar_type val = 0;
      Kokkos::parallel_reduce(
          Kokkos::ThreadVectorRange(member, blocksize),
          [&](const local_ordinal_type &k1, impl_scalar_type &update) {
            update += AA(k0, k1) * xx(k1);
          },
          val);
      Kokkos::single(Kokkos::PerThread(member),
                     [=]() {
                       Kokkos::atomic_add(&yy(k0), -val);
                     });
    }
  }

  struct SeqTag {};

  KOKKOS_INLINE_FUNCTION
  void
  operator()(const SeqTag &, const local_ordinal_type &i) const {
    const local_ordinal_type blocksize        = blocksize_requested;
    const local_ordinal_type blocksize_square = blocksize * blocksize;

    // constants
    const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
    const local_ordinal_type num_vectors = y.extent(1);
    const local_ordinal_type row         = i * blocksize;
    for (local_ordinal_type col = 0; col < num_vectors; ++col) {
      // y := b
      impl_scalar_type *yy             = &y(row, col);
      const impl_scalar_type *const bb = &b(row, col);
      memcpy(yy, bb, sizeof(impl_scalar_type) * blocksize);

      // y -= Rx
      const size_type A_k0 = A_block_rowptr[i];
      for (size_type k = rowptr[i]; k < rowptr[i + 1]; ++k) {
        const size_type j                = A_k0 + colindsub[k];
        const impl_scalar_type *const xx = &x(A_colind[j] * blocksize, col);
        if (hasBlockCrsMatrix) {
          const impl_scalar_type *const AA = &tpetra_values(j * blocksize_square);
          SerialGemv(blocksize, AA, xx, yy);
        } else {
          for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
            SerialDot(blocksize, i, k, k0, colindsub, xx, yy);
        }
      }
    }
  }

  KOKKOS_INLINE_FUNCTION
  void
  operator()(const SeqTag &, const member_type &member) const {
    // constants
    const local_ordinal_type blocksize        = blocksize_requested;
    const local_ordinal_type blocksize_square = blocksize * blocksize;

    const local_ordinal_type lr = member.league_rank();
    const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
    const local_ordinal_type num_vectors = y.extent(1);

    // subview pattern
    auto bb          = Kokkos::subview(b, block_range, 0);
    auto xx          = bb;
    auto yy          = Kokkos::subview(y, block_range, 0);
    auto A_block_cst = ConstUnmanaged<tpetra_block_access_view_type>(NULL, blocksize, blocksize);

    const local_ordinal_type row = lr * blocksize;
    for (local_ordinal_type col = 0; col < num_vectors; ++col) {
      // y := b
      yy.assign_data(&y(row, col));
      bb.assign_data(&b(row, col));
      if (member.team_rank() == 0)
        VectorCopy(member, blocksize, bb, yy);
      member.team_barrier();

      // y -= Rx
      const size_type A_k0 = A_block_rowptr[lr];

      if (hasBlockCrsMatrix) {
        Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr[lr], rowptr[lr + 1]),
                             [&](const local_ordinal_type &k) {
                               const size_type j = A_k0 + colindsub[k];
                               xx.assign_data(&x(A_colind[j] * blocksize, col));
                               A_block_cst.assign_data(&tpetra_values(j * blocksize_square));
                               VectorGemv(member, blocksize, A_block_cst, xx, yy);
                             });
      } else {
        Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr[lr], rowptr[lr + 1]),
                             [&](const local_ordinal_type &k) {
                               const size_type j = A_k0 + colindsub[k];
                               xx.assign_data(&x(A_colind[j] * blocksize, col));

                               for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
                                 VectorDot(member, blocksize, lr, k, k0, colindsub, xx, yy);
                             });
      }
    }
  }

  // * B: block size for compile-time specialization, or 0 for general case (up to max_blocksize)
  // * async: true if using async importer. overlap is not used in this case.
  //          Whether a column is owned or nonowned is decided at runtime.
  // * overlap: true if processing the columns that are not locally owned,
  //            false if processing locally owned columns.
  // * haveBlockMatrix: true if A is a BlockCrsMatrix, false if it's CrsMatrix.
  template <int B, bool async, bool overlap, bool haveBlockMatrix>
  struct GeneralTag {
    static_assert(!(async && overlap),
                  "ComputeResidualVector: async && overlap is not a valid configuration for GeneralTag");
  };

  // Define AsyncTag and OverlapTag in terms of GeneralTag:
  // P == 0 means only compute on owned columns
  // P == 1 means only compute on nonowned columns
  template <int P, int B, bool haveBlockMatrix>
  using OverlapTag = GeneralTag<B, false, P != 0, haveBlockMatrix>;

  template <int B, bool haveBlockMatrix>
  using AsyncTag = GeneralTag<B, true, false, haveBlockMatrix>;

  // CPU implementation for all cases
  template <int B, bool async, bool overlap, bool haveBlockMatrix>
  KOKKOS_INLINE_FUNCTION void
  operator()(const GeneralTag<B, async, overlap, haveBlockMatrix> &, const local_ordinal_type &rowidx) const {
    const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);

    // constants
    const local_ordinal_type partidx = rowidx2part(rowidx);
    const local_ordinal_type pri     = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
    const local_ordinal_type v       = partidx % vector_length;

    const local_ordinal_type num_vectors    = y_packed.extent(2);
    const local_ordinal_type num_local_rows = lclrow.extent(0);

    // temporary buffer for y flat
    impl_scalar_type yy[B == 0 ? max_blocksize : B] = {};

    const local_ordinal_type lr = lclrow(rowidx);

    auto colindsub_used = overlap ? colindsub_remote : colindsub;
    auto rowptr_used    = overlap ? rowptr_remote : rowptr;

    for (local_ordinal_type col = 0; col < num_vectors; ++col) {
      if constexpr (overlap) {
        // y (temporary) := 0
        memset((void *)yy, 0, sizeof(impl_scalar_type) * blocksize);
      } else {
        // y := b
        const local_ordinal_type row = lr * blocksize;
        memcpy(yy, &b(row, col), sizeof(impl_scalar_type) * blocksize);
      }

      // y -= Rx
      const size_type A_k0 = A_block_rowptr[lr];
      for (size_type k = rowptr_used[lr]; k < rowptr_used[lr + 1]; ++k) {
        const size_type j                      = A_k0 + colindsub_used[k];
        const local_ordinal_type A_colind_at_j = A_colind[j];
        if constexpr (haveBlockMatrix) {
          const local_ordinal_type blocksize_square = blocksize * blocksize;
          const impl_scalar_type *const AA          = &tpetra_values(j * blocksize_square);
          if ((!async && !overlap) || (async && A_colind_at_j < num_local_rows)) {
            const auto loc                   = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
            const impl_scalar_type *const xx = &x(loc * blocksize, col);
            SerialGemv(blocksize, AA, xx, yy);
          } else {
            const auto loc                          = A_colind_at_j - num_local_rows;
            const impl_scalar_type *const xx_remote = &x_remote(loc * blocksize, col);
            SerialGemv(blocksize, AA, xx_remote, yy);
          }
        } else {
          if ((!async && !overlap) || (async && A_colind_at_j < num_local_rows)) {
            const auto loc                   = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
            const impl_scalar_type *const xx = &x(loc * blocksize, col);
            for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
              SerialDot(blocksize, lr, k, k0, colindsub_used, xx, yy);
          } else {
            const auto loc                          = A_colind_at_j - num_local_rows;
            const impl_scalar_type *const xx_remote = &x_remote(loc * blocksize, col);
            for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
              SerialDot(blocksize, lr, k, k0, colindsub_used, xx_remote, yy);
          }
        }
      }
      // move yy to y_packed
      if constexpr (overlap) {
        for (local_ordinal_type k = 0; k < blocksize; ++k)
          y_packed(pri, k, col)[v] += yy[k];
      } else {
        for (local_ordinal_type k = 0; k < blocksize; ++k)
          y_packed(pri, k, col)[v] = yy[k];
      }
    }
  }

  // GPU implementation for hasBlockCrsMatrix == true
  template <int B, bool async, bool overlap>
  KOKKOS_INLINE_FUNCTION void
  operator()(const GeneralTag<B, async, overlap, true> &, const member_type &member) const {
    const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);

    // constants
    const local_ordinal_type rowidx  = member.league_rank();
    const local_ordinal_type partidx = rowidx2part(rowidx);
    const local_ordinal_type pri     = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
    const local_ordinal_type v       = partidx % vector_length;

    const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
    const local_ordinal_type num_vectors    = y_packed_scalar.extent(2);
    const local_ordinal_type num_local_rows = lclrow.extent(0);

    // subview pattern
    using subview_1D_right_t  = decltype(Kokkos::subview(b, block_range, 0));
    using subview_1D_stride_t = decltype(Kokkos::subview(y_packed_scalar, 0, block_range, 0, 0));
    subview_1D_right_t bb(nullptr, blocksize);
    subview_1D_right_t xx(nullptr, blocksize);
    subview_1D_stride_t yy(nullptr, Kokkos::LayoutStride(blocksize, y_packed_scalar.stride_1()));
    auto A_block_cst = ConstUnmanaged<tpetra_block_access_view_type>(NULL, blocksize, blocksize);

    // Get shared allocation for a local copy of x, Ax, and A
    impl_scalar_type *local_Ax = reinterpret_cast<impl_scalar_type *>(member.team_scratch(0).get_shmem(blocksize * sizeof(impl_scalar_type)));
    impl_scalar_type *local_x  = reinterpret_cast<impl_scalar_type *>(member.thread_scratch(0).get_shmem(blocksize * sizeof(impl_scalar_type)));

    const local_ordinal_type lr  = lclrow(rowidx);
    const local_ordinal_type row = lr * blocksize;
    for (local_ordinal_type col = 0; col < num_vectors; ++col) {
      if (col)
        member.team_barrier();
      // y -= Rx
      // Initialize accumulation array
      Kokkos::parallel_for(Kokkos::TeamVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
        local_Ax[i] = 0;
      });
      member.team_barrier();

      int numEntries;
      if constexpr (!overlap) {
        numEntries = A_x_offsets.extent(2);
      } else {
        numEntries = A_x_offsets_remote.extent(2);
      }

      Kokkos::parallel_for(Kokkos::TeamThreadRange(member, 0, numEntries),
                           [&](const int k) {
                             int64_t A_offset = overlap ? A_x_offsets_remote(rowidx, 0, k) : A_x_offsets(rowidx, 0, k);
                             int64_t x_offset = overlap ? A_x_offsets_remote(rowidx, 1, k) : A_x_offsets(rowidx, 1, k);
                             if (A_offset != Kokkos::ArithTraits<int64_t>::min()) {
                               A_block_cst.assign_data(tpetra_values.data() + A_offset);
                               // Pull x into local memory
                               if constexpr (async) {
                                 size_type remote_cutoff = blocksize * num_local_rows;
                                 if (x_offset >= remote_cutoff)
                                   xx.assign_data(&x_remote(x_offset - remote_cutoff, col));
                                 else
                                   xx.assign_data(&x(x_offset, col));
                               } else {
                                 if constexpr (!overlap) {
                                   xx.assign_data(&x(x_offset, col));
                                 } else {
                                   xx.assign_data(&x_remote(x_offset, col));
                                 }
                               }

                               Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
                                 local_x[i] = xx(i);
                               });

                               // MatVec op Ax += A*x
                               Kokkos::parallel_for(
                                   Kokkos::ThreadVectorRange(member, blocksize),
                                   [&](const local_ordinal_type &k0) {
                                     impl_scalar_type val = 0;
                                     for (int k1 = 0; k1 < blocksize; k1++)
                                       val += A_block_cst(k0, k1) * local_x[k1];
                                     Kokkos::atomic_add(local_Ax + k0, val);
                                   });
                             }
                           });
      member.team_barrier();
      // Update y = b - local_Ax
      yy.assign_data(&y_packed_scalar(pri, 0, col, v));
      bb.assign_data(&b(row, col));
      Kokkos::parallel_for(Kokkos::TeamVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
        if (!overlap)
          yy(i) = bb(i) - local_Ax[i];
        else
          yy(i) -= local_Ax[i];
      });
    }
  }

  // GPU implementation for hasBlockCrsMatrix == false
  template <int B, bool async, bool overlap>
  KOKKOS_INLINE_FUNCTION void
  operator()(const GeneralTag<B, async, overlap, false> &, const member_type &member) const {
    const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);

    // constants
    const local_ordinal_type rowidx  = member.league_rank();
    const local_ordinal_type partidx = rowidx2part(rowidx);
    const local_ordinal_type pri     = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
    const local_ordinal_type v       = partidx % vector_length;

    const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
    const local_ordinal_type num_vectors    = y_packed_scalar.extent(2);
    const local_ordinal_type num_local_rows = lclrow.extent(0);

    // subview pattern
    using subview_1D_right_t  = decltype(Kokkos::subview(b, block_range, 0));
    using subview_1D_stride_t = decltype(Kokkos::subview(y_packed_scalar, 0, block_range, 0, 0));
    subview_1D_right_t bb(nullptr, blocksize);
    subview_1D_right_t xx(nullptr, blocksize);
    subview_1D_right_t xx_remote(nullptr, blocksize);
    subview_1D_stride_t yy(nullptr, Kokkos::LayoutStride(blocksize, y_packed_scalar.stride_1()));
    auto A_block_cst    = ConstUnmanaged<tpetra_block_access_view_type>(NULL, blocksize, blocksize);
    auto colindsub_used = overlap ? colindsub_remote : colindsub;
    auto rowptr_used    = overlap ? rowptr_remote : rowptr;

    const local_ordinal_type lr  = lclrow(rowidx);
    const local_ordinal_type row = lr * blocksize;
    for (local_ordinal_type col = 0; col < num_vectors; ++col) {
      yy.assign_data(&y_packed_scalar(pri, 0, col, v));
      if (!overlap) {
        // y := b
        bb.assign_data(&b(row, col));
        if (member.team_rank() == 0)
          VectorCopy(member, blocksize, bb, yy);
        member.team_barrier();
      }

      // y -= Rx
      const size_type A_k0 = A_block_rowptr[lr];
      Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr_used[lr], rowptr_used[lr + 1]),
                           [&](const local_ordinal_type &k) {
                             const size_type j                      = A_k0 + colindsub_used[k];
                             const local_ordinal_type A_colind_at_j = A_colind[j];
                             if ((async && A_colind_at_j < num_local_rows) || (!async && !overlap)) {
                               const auto loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
                               xx.assign_data(&x(loc * blocksize, col));
                               for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
                                 VectorDot(member, blocksize, lr, k, k0, colindsub_used, xx, yy);
                             } else {
                               const auto loc = A_colind_at_j - num_local_rows;
                               xx_remote.assign_data(&x_remote(loc * blocksize, col));
                               for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
                                 VectorDot(member, blocksize, lr, k, k0, colindsub_used, xx_remote, yy);
                             }
                           });
    }
  }

  // y = b - Rx; seq method
  template <typename MultiVectorLocalViewTypeY,
            typename MultiVectorLocalViewTypeB,
            typename MultiVectorLocalViewTypeX>
  void run(const MultiVectorLocalViewTypeY &y_,
           const MultiVectorLocalViewTypeB &b_,
           const MultiVectorLocalViewTypeX &x_) {
    IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
    IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<SeqTag>", ComputeResidual0, execution_space);

    y = y_;
    b = b_;
    x = x_;
    if constexpr (is_device<execution_space>::value) {
      const local_ordinal_type blocksize   = blocksize_requested;
      const local_ordinal_type team_size   = 8;
      const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size);
      const Kokkos::TeamPolicy<execution_space, SeqTag> policy(rowptr.extent(0) - 1, team_size, vector_size);
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<SeqTag>", policy, *this);
    } else {
      const Kokkos::RangePolicy<execution_space, SeqTag> policy(0, rowptr.extent(0) - 1);
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<SeqTag>", policy, *this);
    }
    IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
    IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
  }

  // y = b - R (x , x_remote)
  template <typename MultiVectorLocalViewTypeB,
            typename MultiVectorLocalViewTypeX,
            typename MultiVectorLocalViewTypeX_Remote>
  void run(const vector_type_3d_view &y_packed_,
           const MultiVectorLocalViewTypeB &b_,
           const MultiVectorLocalViewTypeX &x_,
           const MultiVectorLocalViewTypeX_Remote &x_remote_) {
    IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
    IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<AsyncTag>", ComputeResidual0, execution_space);

    b        = b_;
    x        = x_;
    x_remote = x_remote_;
    if constexpr (is_device<execution_space>::value) {
      y_packed_scalar = btdm_scalar_type_4d_view((btdm_scalar_type *)y_packed_.data(),
                                                 y_packed_.extent(0),
                                                 y_packed_.extent(1),
                                                 y_packed_.extent(2),
                                                 vector_length);
    } else {
      y_packed = y_packed_;
    }

    if constexpr (is_device<execution_space>::value) {
      const local_ordinal_type blocksize = blocksize_requested;
      // local_ordinal_type vl_power_of_two = 1;
      // for (;vl_power_of_two<=blocksize_requested;vl_power_of_two*=2);
      // vl_power_of_two *= (vl_power_of_two < blocksize_requested ? 2 : 1);
      // const local_ordinal_type vl = vl_power_of_two > vector_length ? vector_length : vl_power_of_two;
#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B)                                                                          \
  {                                                                                                                             \
    if (this->hasBlockCrsMatrix) {                                                                                              \
      const local_ordinal_type team_size   = 8;                                                                                 \
      const local_ordinal_type vector_size = 8;                                                                                 \
      const size_t shmem_team_size         = blocksize * sizeof(btdm_scalar_type);                                              \
      const size_t shmem_thread_size       = blocksize * sizeof(btdm_scalar_type);                                              \
      Kokkos::TeamPolicy<execution_space, AsyncTag<B, true>>                                                                    \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                                \
      policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size));                       \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<AsyncTag>",                                                        \
                           policy, *this);                                                                                      \
    } else {                                                                                                                    \
      const local_ordinal_type team_size   = 8;                                                                                 \
      const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size); \
      const Kokkos::TeamPolicy<execution_space, AsyncTag<B, false>>                                                             \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                                \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<AsyncTag>",                                                        \
                           policy, *this);                                                                                      \
    }                                                                                                                           \
  }                                                                                                                             \
  break
      switch (blocksize_requested) {
        case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
        case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
        case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
        case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
        case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
        case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
        case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
        case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
        case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
        default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
      }
#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
    } else {
#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B)                                                 \
  {                                                                                                    \
    if (this->hasBlockCrsMatrix) {                                                                     \
      const Kokkos::RangePolicy<execution_space, AsyncTag<B, true>> policy(0, rowidx2part.extent(0));  \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<AsyncTag>",                              \
                           policy, *this);                                                             \
    } else {                                                                                           \
      const Kokkos::RangePolicy<execution_space, AsyncTag<B, false>> policy(0, rowidx2part.extent(0)); \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<AsyncTag>",                              \
                           policy, *this);                                                             \
    }                                                                                                  \
  }                                                                                                    \
  break

      switch (blocksize_requested) {
        case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
        case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
        case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
        case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
        case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
        case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
        case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
        case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
        case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
        default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
      }
#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
    }
    IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
    IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
  }

  // y = b - R (y , y_remote)
  template <typename MultiVectorLocalViewTypeB,
            typename MultiVectorLocalViewTypeX,
            typename MultiVectorLocalViewTypeX_Remote>
  void run(const vector_type_3d_view &y_packed_,
           const MultiVectorLocalViewTypeB &b_,
           const MultiVectorLocalViewTypeX &x_,
           const MultiVectorLocalViewTypeX_Remote &x_remote_,
           const bool compute_owned) {
    IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
    IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<OverlapTag>", ComputeResidual0, execution_space);

    b        = b_;
    x        = x_;
    x_remote = x_remote_;
    if constexpr (is_device<execution_space>::value) {
      y_packed_scalar = btdm_scalar_type_4d_view((btdm_scalar_type *)y_packed_.data(),
                                                 y_packed_.extent(0),
                                                 y_packed_.extent(1),
                                                 y_packed_.extent(2),
                                                 vector_length);
    } else {
      y_packed = y_packed_;
    }

    if constexpr (is_device<execution_space>::value) {
      const local_ordinal_type blocksize = blocksize_requested;
      // local_ordinal_type vl_power_of_two = 1;
      // for (;vl_power_of_two<=blocksize_requested;vl_power_of_two*=2);
      // vl_power_of_two *= (vl_power_of_two < blocksize_requested ? 2 : 1);
      // const local_ordinal_type vl = vl_power_of_two > vector_length ? vector_length : vl_power_of_two;
#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B)                                                                        \
  if (this->hasBlockCrsMatrix) {                                                                                              \
    const local_ordinal_type team_size   = 8;                                                                                 \
    const local_ordinal_type vector_size = 8;                                                                                 \
    const size_t shmem_team_size         = blocksize * sizeof(btdm_scalar_type);                                              \
    const size_t shmem_thread_size       = blocksize * sizeof(btdm_scalar_type);                                              \
    if (compute_owned) {                                                                                                      \
      Kokkos::TeamPolicy<execution_space, OverlapTag<0, B, true>>                                                             \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                              \
      policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size));                     \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<0> >", policy, *this);                                \
    } else {                                                                                                                  \
      Kokkos::TeamPolicy<execution_space, OverlapTag<1, B, true>>                                                             \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                              \
      policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size));                     \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<1> >", policy, *this);                                \
    }                                                                                                                         \
  } else {                                                                                                                    \
    const local_ordinal_type team_size   = 8;                                                                                 \
    const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size); \
    if (compute_owned) {                                                                                                      \
      const Kokkos::TeamPolicy<execution_space, OverlapTag<0, B, false>>                                                      \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                              \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<0> >", policy, *this);                                \
    } else {                                                                                                                  \
      const Kokkos::TeamPolicy<execution_space, OverlapTag<1, B, false>>                                                      \
          policy(rowidx2part.extent(0), team_size, vector_size);                                                              \
      Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<1> >", policy, *this);                                \
    }                                                                                                                         \
  }                                                                                                                           \
  break
      switch (blocksize_requested) {
        case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
        case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
        case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
        case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
        case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
        case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
        case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
        case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
        case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
        default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
      }
#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
    } else {
#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B)                                          \
  if (this->hasBlockCrsMatrix) {                                                                \
    if (compute_owned) {                                                                        \
      const Kokkos::RangePolicy<execution_space, OverlapTag<0, B, true>>                        \
          policy(0, rowidx2part.extent(0));                                                     \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<0> >", policy, *this); \
    } else {                                                                                    \
      const Kokkos::RangePolicy<execution_space, OverlapTag<1, B, true>>                        \
          policy(0, rowidx2part.extent(0));                                                     \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<1> >", policy, *this); \
    }                                                                                           \
  } else {                                                                                      \
    if (compute_owned) {                                                                        \
      const Kokkos::RangePolicy<execution_space, OverlapTag<0, B, false>>                       \
          policy(0, rowidx2part.extent(0));                                                     \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<0> >", policy, *this); \
    } else {                                                                                    \
      const Kokkos::RangePolicy<execution_space, OverlapTag<1, B, false>>                       \
          policy(0, rowidx2part.extent(0));                                                     \
      Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<1> >", policy, *this); \
    }                                                                                           \
  }                                                                                             \
  break

      switch (blocksize_requested) {
        case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
        case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
        case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
        case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
        case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
        case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
        case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
        case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
        case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
        default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
      }
#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
    }
    IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
    IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
  }
};

}  // namespace BlockHelperDetails

}  // namespace Ifpack2

#endif