Skip to content

Twostage gs streams #2158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 10 commits into
base: develop
Choose a base branch
from
Open

Twostage gs streams #2158

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
1ecbc6c
Stream support for Gauss-Seidel: Symbolic, Numeric, Apply (Twostage)
e10harvey Sep 26, 2023
7408367
sparse/unit_test: Test twostage with streams
e10harvey Sep 26, 2023
f4e3df5
Update omega_tsgs_classic type
e10harvey Sep 26, 2023
6f510b3
Move sptrsv to streams
e10harvey Sep 27, 2023
7505956
Add lightweighthint
e10harvey Sep 27, 2023
35b8474
update symbolic check_count
e10harvey Oct 31, 2023
d11fd83
BLAS - Axpby: only print debug info at high debug level
lucbv Mar 28, 2024
923ebf4
Sparse - twostage GS: implementing Brian's review
lucbv Mar 28, 2024
e3e1122
Applying clang-format
lucbv Mar 28, 2024
d52f809
BLAS - Axpby: fixing issue with output level
lucbv Mar 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion blas/impl/KokkosBlas1_axpby_unification_attempt_traits.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -768,7 +768,7 @@ struct AxpbyUnificationAttemptTraits {
// ********************************************************************
// Routine to print information on input variables and internal variables
// ********************************************************************
#if (KOKKOSKERNELS_DEBUG_LEVEL > 0)
#if (KOKKOSKERNELS_DEBUG_LEVEL > 1)
static void printInformation(std::ostream& os, std::string const& headerMsg) {
os << headerMsg << ": AV = "
<< typeid(AV).name()
Expand Down
12 changes: 12 additions & 0 deletions blas/unit_test/Test_Blas1_axpby_unification.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,12 @@
#include <KokkosBlas1_axpby.hpp>
#include <KokkosKernels_TestUtils.hpp>

#if (KOKKOSKERNELS_DEBUG_LEVEL > 0)
#define KOKKOSKERNELS_DEBUG_LEVEL_OLD KOKKOSKERNELS_DEBUG_LEVEL
#undef KOKKOSKERNELS_DEBUG_LEVEL
#define KOKKOSKERNELS_DEBUG_LEVEL 2
#endif

static constexpr int numVecsAxpbyTest = 15;

namespace Test {
Expand Down Expand Up @@ -2537,6 +2543,12 @@ void impl_test_axpby_mv_unification(int const N, int const K) {

} // namespace Test

#if (KOKKOSKERNELS_DEBUG_LEVEL > 1)
#undef KOKKOSKERNELS_DEBUG_LEVEL
#define KOKKOSKERNELS_DEBUG_LEVEL KOKKOSKERNELS_DEBUG_LEVEL_OLD
#undef KOKKOSKERNELS_DEBUG_LEVEL_OLD
#endif

template <class tScalarA, class tScalarX, class tScalarB, class tScalarY,
class Device>
int test_axpby_unification() {
Expand Down
8 changes: 4 additions & 4 deletions sparse/impl/KokkosSparse_sptrsv_symbolic_impl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,8 +285,8 @@ void lower_tri_symbolic(ExecSpaceIn& space, TriSolveHandle& thandle,
check_count = 0; // reset
Kokkos::parallel_reduce(
"check_count device",
Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>(
0, dnodes_per_level.extent(0)),
Kokkos::RangePolicy<ExecutionSpace>(space, 0,
dnodes_per_level.extent(0)),
KOKKOS_LAMBDA(const long i, long& update) {
update += dnodes_per_level(i);
},
Expand Down Expand Up @@ -740,8 +740,8 @@ void upper_tri_symbolic(ExecutionSpace& space, TriSolveHandle& thandle,
check_count = 0; // reset
Kokkos::parallel_reduce(
"check_count device",
Kokkos::RangePolicy<Kokkos::DefaultExecutionSpace>(
0, dnodes_per_level.extent(0)),
Kokkos::RangePolicy<ExecutionSpace>(space, 0,
dnodes_per_level.extent(0)),
KOKKOS_LAMBDA(const long i, long& update) {
update += dnodes_per_level(i);
},
Expand Down
226 changes: 130 additions & 96 deletions sparse/impl/KokkosSparse_twostage_gauss_seidel_impl.hpp
View file
Open in desktop

Large diffs are not rendered by default.

28 changes: 18 additions & 10 deletions sparse/src/KokkosKernels_Handle.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -609,15 +609,20 @@ class KokkosKernelsHandle {
* @param handle_exec_space The execution space instance to execute kernels on.
* @param num_streams The number of streams to allocate memory for.
* @param gs_algorithm Specifies which algorithm to use:
*
* KokkosSpace::GS_DEFAULT PointGaussSeidel
* KokkosSpace::GS_PERMUTED ??
* KokkosSpace::GS_TEAM ??
* KokkosSpace::GS_CLUSTER ??
* KokkosSpace::GS_TWOSTAGE ??
*
* KokkosSpace::GS_DEFAULT PointGaussSeidel or BlockGaussSeidel, depending on matrix type.
* KokkosSpace::GS_PERMUTED Reorders rows/cols into colors to improve locality. Uses RangePolicy over rows.
* KokkosSpace::GS_TEAM Uses TeamPolicy over batches of rows with ThreadVector within rows.
* KokkosSpace::GS_CLUSTER Uses independent clusters of nodes in the graph. Within a cluster, x is updated sequentially.
* For more information, see: https://arxiv.org/pdf/2204.02934.pdf.
* KokkosSpace::GS_TWOSTAGE Uses spmv to parallelize inner sweeps of x.
* For more information, see: https://arxiv.org/pdf/2104.01196.pdf.
* @param coloring_algorithm Specifies which coloring algorithm to color the graph with:
*
* KokkosGraph::COLORING_DEFAULT ??
*
* KokkosGraph::COLORING_DEFAULT Depends on execution space:
* COLORING_SERIAL on Kokkos::Serial;
* COLORING_EB on GPUs;
* COLORING_VBBIT on Kokkos::Sycl or elsewhere.
* KokkosGraph::COLORING_SERIAL Serial Greedy Coloring
* KokkosGraph::COLORING_VB Vertex Based Coloring
* KokkosGraph::COLORING_VBBIT Vertex Based Coloring with bit array
Expand Down Expand Up @@ -753,8 +758,11 @@ class KokkosKernelsHandle {
* KokkosSparse::NUM_CLUSTERING_ALGORITHMS ??
* @param hint_verts_per_cluster Hint how many verticies to use per cluster
* @param coloring_algorithm Specifies which coloring algorithm to color the graph with:
*
* KokkosGraph::COLORING_DEFAULT ??
*
* KokkosGraph::COLORING_DEFAULT Depends on execution space:
* COLORING_SERIAL on Kokkos::Serial;
* COLORING_EB on GPUs;
* COLORING_VBBIT on Kokkos::Sycl or elsewhere.
* KokkosGraph::COLORING_SERIAL Serial Greedy Coloring
* KokkosGraph::COLORING_VB Vertex Based Coloring
* KokkosGraph::COLORING_VBBIT Vertex Based Coloring with bit array
Expand Down
96 changes: 48 additions & 48 deletions sparse/src/KokkosSparse_gauss_seidel.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,13 @@ namespace Experimental {
/// @brief Gauss-Seidel preconditioner setup (first phase, based on sparsity
/// pattern only)
///
/// @tparam ExecutionSpace This kernels execution space type.
/// @tparam ExecutionSpace This kernels execution space type
/// @tparam KernelHandle A specialization of
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @param space The execution space instance this kernel will be run on.
/// @param handle KernelHandle instance
/// @param space The execution space instance this kernel will be run on
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -112,7 +112,7 @@ void gauss_seidel_symbolic(const ExecutionSpace &space, KernelHandle *handle,
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @param handle KernelHandle instance
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -141,7 +141,7 @@ void gauss_seidel_symbolic(KernelHandle *handle,
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @param handle KernelHandle instance
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param block_size The number of degrees of freedom per block
Expand Down Expand Up @@ -173,15 +173,15 @@ void block_gauss_seidel_symbolic(
/// @brief Gauss-Seidel preconditioner setup (second phase, based on matrix's
/// numeric values)
///
/// @tparam ExecutionSpace This kernels execution space type.
/// @tparam ExecutionSpace This kernels execution space type
/// @tparam format The matrix storage format, CRS or BSR
/// @tparam KernelHandle A specialization of
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @param space The execution space instance this kernel will be run on.
/// @param handle KernelHandle instance
/// @param space The execution space instance this kernel will be run on
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -279,8 +279,8 @@ void gauss_seidel_numeric(const ExecutionSpace &space, KernelHandle *handle,
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type. The user-provided
/// inverse diagonal must share this type.
/// @param handle KernelHandle instance
/// inverse diagonal must share this type
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -313,16 +313,16 @@ void gauss_seidel_numeric(KernelHandle *handle,
/// numeric values). This version accepts the matrix's inverse diagonal from the
/// user.
///
/// @tparam ExecutionSpace This kernels execution space type.
/// @tparam ExecutionSpace This kernels execution space type
/// @tparam format The matrix storage format, CRS or BSR
/// @tparam KernelHandle A specialization of
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type. The user-provided
/// inverse diagonal must share this type.
/// @param space The execution space instance this kernel will be run on.
/// @param handle KernelHandle instance
/// inverse diagonal must share this type
/// @param space The execution space instance this kernel will be run on
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -427,8 +427,8 @@ void gauss_seidel_numeric(const ExecutionSpace &space, KernelHandle *handle,
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type. The user-provided
/// inverse diagonal must share this type.
/// @param handle KernelHandle instance
/// inverse diagonal must share this type
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -468,7 +468,7 @@ void gauss_seidel_numeric(KernelHandle *handle,
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @param handle handle A KokkosKernelsHandle instance
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param block_size The number of degrees of freedom per block
Expand Down Expand Up @@ -504,20 +504,20 @@ void block_gauss_seidel_numeric(
/// @brief Apply symmetric (forward + backward) Gauss-Seidel preconditioner to
/// system AX=Y
///
/// @tparam ExecutionSpace This kernels execution space type.
/// @tparam ExecutionSpace This kernels execution space type
/// @tparam format The matrix storage format, CRS or BSR
/// @tparam KernelHandle A specialization of
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// rank-1 or rank-2 View
/// @param space The execution space instance this kernel will be run
/// on. NOTE: Currently only used for GS_DEFAULT.
/// @param handle handle A KokkosKernelsHandle instance
/// on. NOTE: Currently only used for GS_DEFAULT and GS_TWOSTAGE
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -680,8 +680,8 @@ void symmetric_gauss_seidel_apply(
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle handle A KokkosKernelsHandle instance
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -727,10 +727,10 @@ void symmetric_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle handle A KokkosKernelsHandle instance.
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param block_size The number of degrees of freedom per block
Expand Down Expand Up @@ -793,12 +793,12 @@ void symmetric_block_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// rank-1 or rank-2 View
/// @param space The execution space instance this kernel will be run
/// on. NOTE: Currently only used for GS_DEFAULT.
/// @param handle KernelHandle instance
/// on. NOTE: Currently only used for GS_DEFAULT and GS_TWOSTAGE
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -960,10 +960,10 @@ void forward_sweep_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle KernelHandle instance
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -1008,10 +1008,10 @@ void forward_sweep_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle KernelHandle instance
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param block_size The number of degrees of freedom per block
Expand Down Expand Up @@ -1067,20 +1067,20 @@ void forward_sweep_block_gauss_seidel_apply(
///
/// @brief Apply backward Gauss-Seidel preconditioner to system AX=Y
///
/// @tparam ExecutionSpace This kernels execution space type.
/// @tparam ExecutionSpace This kernels execution space type
/// @tparam format The matrix storage format, CRS or BSR
/// @tparam KernelHandle A specialization of
/// KokkosKernels::Experimental::KokkosKernelsHandle
/// @tparam lno_row_view_t_ The matrix's rowmap type
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// rank-1 or rank-2 View
/// @param space The execution space instance this kernel will be run
/// on. NOTE: Currently only used for GS_DEFAULT.
/// @param handle KernelHandle instance
/// on. NOTE: Currently only used for GS_DEFAULT and GS_TWOSTAGE
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -1242,10 +1242,10 @@ void backward_sweep_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle KernelHandle instance
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param row_map The matrix's rowmap
Expand Down Expand Up @@ -1290,10 +1290,10 @@ void backward_sweep_gauss_seidel_apply(
/// @tparam lno_nnz_view_t_ The matrix's entries type
/// @tparam scalar_nnz_view_t_ The matrix's values type
/// @tparam x_scalar_view_t The type of the X (left-hand side, unknown) vector.
/// May be rank-1 or rank-2 View.
/// May be rank-1 or rank-2 View
/// @tparam y_scalar_view_t The type of the Y (right-hand side) vector. May be
/// rank-1 or rank-2 View.
/// @param handle KernelHandle instance
/// rank-1 or rank-2 View
/// @param handle A KokkosKernelsHandle instance
/// @param num_rows Number of rows in the matrix
/// @param num_cols Number of columns in the matrix
/// @param block_size The number of degrees of freedom per block
Expand Down
15 changes: 12 additions & 3 deletions sparse/src/KokkosSparse_gauss_seidel_handle.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -767,9 +767,18 @@ class TwoStageGaussSeidelHandle

void initVectors(int nrows_, int nrhs_) {
if (this->nrows != nrows_ || this->nrhs != nrhs_) {
this->localR = vector_view_t("temp", nrows_, nrhs_);
this->localT = vector_view_t("temp", nrows_, nrhs_);
this->localZ = vector_view_t("temp", nrows_, nrhs_);
vector_view_t r(
Kokkos::view_alloc(this->execution_space, std::string("temp")),
nrows_, nrhs_);
this->localR = r;
vector_view_t t(
Kokkos::view_alloc(this->execution_space, std::string("temp")),
nrows_, nrhs_);
this->localT = t;
vector_view_t z(
Kokkos::view_alloc(this->execution_space, std::string("temp")),
nrows_, nrhs_);
this->localZ = z;
this->nrows = nrows_;
this->nrhs = nrhs_;
}
Expand Down
Loading