Tpetra_Details_MatrixApplyHelper.hpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_MATRIX_APPLY_HELPER_HPP
#define TPETRA_MATRIX_APPLY_HELPER_HPP

#include "Kokkos_Core.hpp"
#include "KokkosSparse_spmv_handle.hpp"
#include "Tpetra_Details_IntRowPtrHelper.hpp"

namespace Tpetra {
namespace Details {

template<typename LocalMatrix, typename IntLocalMatrix, typename MultiVectorLocalView>
struct MatrixApplyHelper : public IntRowPtrHelper<LocalMatrix, IntLocalMatrix>
{
  // using Rowptrs = typename LocalMatrix::row_map_type;
  // using IntRowptrs = typename IntLocalMatrix::row_map_type;
  using XVectorType = typename MultiVectorLocalView::const_type;
  using YVectorType = MultiVectorLocalView;
  using SPMVHandle = KokkosSparse::SPMVHandle<typename LocalMatrix::device_type, LocalMatrix, XVectorType, YVectorType>;
  using SPMVHandleInt = KokkosSparse::SPMVHandle<typename LocalMatrix::device_type, IntLocalMatrix, XVectorType, YVectorType>;

  MatrixApplyHelper(size_t nnz, const typename LocalMatrix::row_map_type& rowptrs, KokkosSparse::SPMVAlgorithm algo = KokkosSparse::SPMV_DEFAULT)
    : IntRowPtrHelper<LocalMatrix, IntLocalMatrix>(nnz, rowptrs), handle_int(algo) {}
  

  // SPMVHandles are lazily initialized by actual spmv calls.
  // We declare both here because we don't know until runtime (based on nnz) which should be used.
  SPMVHandle handle;
  SPMVHandleInt handle_int;
};
}
}

#endif // TPETRA_MATRIX_APPLY_HELPER_HPP