Skip to content
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

Jump to bottom

Split batched solver compilation #1629

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

Split batched solver compilation #1629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
feaeabd
[batch] provide default index type for matrix device types
MarcelKoch Jun 20, 2024
f0497f2
[batch] handle constness of index type same as value type
MarcelKoch Jun 20, 2024
bc63683
[batch] add macro to instantiate batched solver
MarcelKoch Jun 20, 2024
b30e6f6
[batch] split bicgstab compilation (hip)
MarcelKoch Sep 17, 2024
b673a45
[batch] split bicgstab compilation (cuda)
MarcelKoch Sep 17, 2024
1255d08
[batch] split cg compilation (hip)
MarcelKoch Sep 17, 2024
045ad1c
[batch] split cg compilation (cuda)
MarcelKoch Sep 17, 2024
85b80df
[batch] review updates:
MarcelKoch Oct 1, 2024
6cdfded
[batch] add launch bounds and fix register check
MarcelKoch Oct 2, 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
21 changes: 15 additions & 6 deletions common/cuda_hip/solver/batch_bicgstab_kernels.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
#ifndef GKO_COMMON_CUDA_HIP_SOLVER_BATCH_BICGSTAB_KERNELS_HPP_
#define GKO_COMMON_CUDA_HIP_SOLVER_BATCH_BICGSTAB_KERNELS_HPP_

#include "core/solver/batch_bicgstab_kernels.hpp"

#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/types.hpp>
Expand All @@ -25,6 +27,11 @@
namespace gko {
namespace kernels {
namespace GKO_DEVICE_NAMESPACE {


constexpr int max_bicgstab_threads = 1024;


namespace batch_single_kernels {


Expand Down Expand Up @@ -168,12 +175,14 @@ __device__ __forceinline__ void update_x_middle(
template <typename StopType, int n_shared, bool prec_shared_bool,
typename PrecType, typename LogType, typename BatchMatrixType,
typename ValueType>
__global__ void apply_kernel(
const gko::kernels::batch_bicgstab::storage_config sconf,
const int max_iter, const gko::remove_complex<ValueType> tol,
LogType logger, PrecType prec_shared, const BatchMatrixType mat,
const ValueType* const __restrict__ b, ValueType* const __restrict__ x,
ValueType* const __restrict__ workspace = nullptr)
__global__ void __launch_bounds__(max_bicgstab_threads)
apply_kernel(const gko::kernels::batch_bicgstab::storage_config sconf,
const int max_iter, const gko::remove_complex<ValueType> tol,
LogType logger, PrecType prec_shared,
const BatchMatrixType mat,
const ValueType* const __restrict__ b,
ValueType* const __restrict__ x,
ValueType* const __restrict__ workspace = nullptr)
{
using real_type = typename gko::remove_complex<ValueType>;
const auto num_batch_items = mat.num_batch_items;
Expand Down
23 changes: 15 additions & 8 deletions common/cuda_hip/solver/batch_cg_kernels.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
#define GKO_COMMON_CUDA_HIP_SOLVER_BATCH_CG_KERNELS_HPP_


#include "core/solver/batch_cg_kernels.hpp"

#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/types.hpp>
Expand All @@ -27,6 +29,11 @@
namespace gko {
namespace kernels {
namespace GKO_DEVICE_NAMESPACE {


constexpr int max_cg_threads = 1024;


namespace batch_single_kernels {


Expand Down Expand Up @@ -113,14 +120,14 @@ __device__ __forceinline__ void update_x_and_r(
template <typename StopType, const int n_shared, const bool prec_shared_bool,
typename PrecType, typename LogType, typename BatchMatrixType,
typename ValueType>
__global__ void apply_kernel(const gko::kernels::batch_cg::storage_config sconf,
const int max_iter,
const gko::remove_complex<ValueType> tol,
LogType logger, PrecType prec_shared,
const BatchMatrixType mat,
const ValueType* const __restrict__ b,
ValueType* const __restrict__ x,
ValueType* const __restrict__ workspace = nullptr)
__global__ void __launch_bounds__(max_cg_threads)
apply_kernel(const gko::kernels::batch_cg::storage_config sconf,
const int max_iter, const gko::remove_complex<ValueType> tol,
LogType logger, PrecType prec_shared,
const BatchMatrixType mat,
const ValueType* const __restrict__ b,
ValueType* const __restrict__ x,
ValueType* const __restrict__ workspace = nullptr)
{
using real_type = typename gko::remove_complex<ValueType>;
const auto num_batch_items = mat.num_batch_items;
Expand Down
20 changes: 10 additions & 10 deletions core/matrix/batch_struct.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,14 +22,14 @@ namespace csr {
/**
* Encapsulates one matrix from a batch of csr matrices.
*/
template <typename ValueType, typename IndexType>
template <typename ValueType, typename IndexType = const int32>
struct batch_item {
using value_type = ValueType;
using index_type = IndexType;

ValueType* values;
const index_type* col_idxs;
const index_type* row_ptrs;
index_type* col_idxs;
index_type* row_ptrs;
index_type num_rows;
index_type num_cols;
index_type num_nnz_per_item;
Expand All @@ -44,15 +44,15 @@ struct batch_item {
/**
* A 'simple' structure to store a global uniform batch of csr matrices.
*/
template <typename ValueType, typename IndexType>
template <typename ValueType, typename IndexType = const int32>
struct uniform_batch {
using value_type = ValueType;
using index_type = IndexType;
using entry_type = batch_item<value_type, index_type>;

ValueType* values;
const index_type* col_idxs;
const index_type* row_ptrs;
index_type* col_idxs;
index_type* row_ptrs;
size_type num_batch_items;
index_type num_rows;
index_type num_cols;
Expand Down Expand Up @@ -119,13 +119,13 @@ namespace ell {
/**
* Encapsulates one matrix from a batch of ell matrices.
*/
template <typename ValueType, typename IndexType>
template <typename ValueType, typename IndexType = const int32>
struct batch_item {
using value_type = ValueType;
using index_type = IndexType;

ValueType* values;
const index_type* col_idxs;
index_type* col_idxs;
index_type stride;
index_type num_rows;
index_type num_cols;
Expand All @@ -141,14 +141,14 @@ struct batch_item {
/**
* A 'simple' structure to store a global uniform batch of ell matrices.
*/
template <typename ValueType, typename IndexType>
template <typename ValueType, typename IndexType = const int32>
struct uniform_batch {
using value_type = ValueType;
using index_type = IndexType;
using entry_type = batch_item<value_type, index_type>;

ValueType* values;
const index_type* col_idxs;
index_type* col_idxs;
size_type num_batch_items;
index_type stride;
index_type num_rows;
Expand Down
35 changes: 35 additions & 0 deletions core/solver/batch_dispatch.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,41 @@ enum class log_type { simple_convergence_completion };
} // namespace log


#define GKO_BATCH_INSTANTIATE_STOP(macro, ...) \
macro(__VA_ARGS__, \
::gko::batch::solver::device::batch_stop::SimpleAbsResidual); \
template macro( \
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the template here (and in the other macros below) could be removed, if the value/index type instantiation macros would accept variable number or arguments.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That doesn't work until C++20. A macro with (arg, ...) requires two arguments before c++20.

__VA_ARGS__, \
::gko::batch::solver::device::batch_stop::SimpleRelResidual)

#define GKO_BATCH_INSTANTIATE_PRECONDITIONER(macro, ...) \
GKO_BATCH_INSTANTIATE_STOP( \
macro, __VA_ARGS__, \
::gko::batch::solver::device::batch_preconditioner::Identity); \
template GKO_BATCH_INSTANTIATE_STOP( \
macro, __VA_ARGS__, \
::gko::batch::solver::device::batch_preconditioner::ScalarJacobi); \
template GKO_BATCH_INSTANTIATE_STOP( \
macro, __VA_ARGS__, \
::gko::batch::solver::device::batch_preconditioner::BlockJacobi)

#define GKO_BATCH_INSTANTIATE_LOGGER(macro, ...) \
GKO_BATCH_INSTANTIATE_PRECONDITIONER( \
macro, __VA_ARGS__, \
::gko::batch::solver::device::batch_log::SimpleFinalLogger)

#define GKO_BATCH_INSTANTIATE_MATRIX(macro, ...) \
GKO_BATCH_INSTANTIATE_LOGGER(macro, __VA_ARGS__, \
batch::matrix::ell::uniform_batch); \
template GKO_BATCH_INSTANTIATE_LOGGER( \
macro, __VA_ARGS__, batch::matrix::dense::uniform_batch); \
template GKO_BATCH_INSTANTIATE_LOGGER(macro, __VA_ARGS__, \
batch::matrix::csr::uniform_batch)

#define GKO_BATCH_INSTANTIATE(macro, ...) \
GKO_BATCH_INSTANTIATE_MATRIX(macro, __VA_ARGS__)


/**
* Handles dispatching to the correct instantiation of a batched solver
* depending on runtime parameters.
Expand Down
4 changes: 4 additions & 0 deletions cuda/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ add_library(ginkgo_cuda $<TARGET_OBJECTS:ginkgo_cuda_device> "")
include(${PROJECT_SOURCE_DIR}/cmake/template_instantiation.cmake)
add_instantiation_files(${PROJECT_SOURCE_DIR}/common/cuda_hip matrix/csr_kernels.instantiate.cpp CSR_INSTANTIATE)
add_instantiation_files(${PROJECT_SOURCE_DIR}/common/cuda_hip matrix/fbcsr_kernels.instantiate.cpp FBCSR_INSTANTIATE)
add_instantiation_files(. solver/batch_bicgstab_launch.instantiate.cu BATCH_BICGSTAB_INSTANTIATE)
add_instantiation_files(. solver/batch_cg_launch.instantiate.cu BATCH_CG_INSTANTIATE)
# we don't split up the dense kernels into distinct compilations
list(APPEND GKO_UNIFIED_COMMON_SOURCES ${PROJECT_SOURCE_DIR}/common/unified/matrix/dense_kernels.instantiate.cpp)
target_sources(ginkgo_cuda
Expand All @@ -21,7 +23,9 @@ target_sources(ginkgo_cuda
matrix/fft_kernels.cu
preconditioner/batch_jacobi_kernels.cu
solver/batch_bicgstab_kernels.cu
${BATCH_BICGSTAB_INSTANTIATE}
solver/batch_cg_kernels.cu
${BATCH_CG_INSTANTIATE}
solver/lower_trs_kernels.cu
solver/upper_trs_kernels.cu
${GKO_UNIFIED_COMMON_SOURCES}
Expand Down
Loading
Loading