Skip to content

Commit

Permalink
wip: review updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@pratikvn
pratikvn committed Apr 2, 2024
1 parent cb730a3 commit 233c13c
Show file tree
Hide file tree
Showing 13 changed files with 136 additions and 161 deletions.
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ repos:
name: reuse-annotate
entry: reuse annotate --license BSD-3-Clause --copyright 'The Ginkgo authors' --style c --merge-copyright
language: python
additional_dependencies: [reuse]
additional_dependencies: [reuse==2.1.0]
types_or: [c, c++, cuda, inc]
exclude: |
(?x)^(
Expand Down
3 changes: 1 addition & 2 deletions core/device_hooks/common_kernels.inc.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -768,8 +768,7 @@ namespace batch_jacobi {

GKO_STUB_INDEX_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_CUMULATIVE_BLOCK_STORAGE);
GKO_STUB_INDEX_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK);
GKO_STUB_INDEX_TYPE(GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP);
GKO_STUB_VALUE_AND_INT32_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_EXTRACT_PATTERN_KERNEL);
GKO_STUB_VALUE_AND_INT32_TYPE(GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_KERNEL);
Expand Down
45 changes: 39 additions & 6 deletions core/preconditioner/batch_jacobi.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,49 @@ GKO_REGISTER_OPERATION(extract_common_blocks_pattern,
batch_jacobi::extract_common_blocks_pattern);
GKO_REGISTER_OPERATION(compute_block_jacobi,
batch_jacobi::compute_block_jacobi);
GKO_REGISTER_OPERATION(find_row_is_part_of_which_block,
batch_jacobi::find_row_is_part_of_which_block);
GKO_REGISTER_OPERATION(find_row_block_map, batch_jacobi::find_row_block_map);
GKO_REGISTER_OPERATION(compute_cumulative_block_storage,
batch_jacobi::compute_cumulative_block_storage);


} // namespace jacobi


template <typename ValueType, typename IndexType>
Jacobi<ValueType, IndexType>::Jacobi(std::shared_ptr<const Executor> exec)
: EnableBatchLinOp<Jacobi>(exec),
num_blocks_{},
blocks_(exec),
row_block_map_info_(exec),
blocks_cumulative_storage_(exec),
blocks_storage_scheme_{batched_jacobi_blocks_storage_scheme<index_type>()}
{
parameters_.block_pointers.set_executor(this->get_executor());
}


template <typename ValueType, typename IndexType>
Jacobi<ValueType, IndexType>::Jacobi(
const Factory* factory, std::shared_ptr<const BatchLinOp> system_matrix)
: EnableBatchLinOp<Jacobi>(factory->get_executor(),
gko::transpose(system_matrix->get_size())),
parameters_{factory->get_parameters()},
num_blocks_{parameters_.block_pointers.get_size() > 0
? parameters_.block_pointers.get_size() - 1
: 0},
blocks_(factory->get_executor()),
row_block_map_info_(factory->get_executor(),
system_matrix->get_common_size()[0]),
blocks_cumulative_storage_(factory->get_executor(), num_blocks_ + 1),
blocks_storage_scheme_{batched_jacobi_blocks_storage_scheme<index_type>()}

{
parameters_.block_pointers.set_executor(this->get_executor());
GKO_ASSERT_BATCH_HAS_SQUARE_DIMENSIONS(system_matrix);
this->generate_precond(system_matrix.get());
}


template <typename ValueType, typename IndexType>
void Jacobi<ValueType, IndexType>::detect_blocks(
const size_type num_batch,
Expand Down Expand Up @@ -101,9 +135,9 @@ void Jacobi<ValueType, IndexType>::generate_precond(

blocks_.resize_and_reset(this->compute_storage_space(num_batch));

exec->run(jacobi::make_find_row_is_part_of_which_block(
exec->run(jacobi::make_find_row_block_map(
num_blocks_, parameters_.block_pointers.get_const_data(),
row_part_of_which_block_info_.get_data()));
row_block_map_info_.get_data()));

// Note: Row-major order offers advantage in terms of
// performance in both preconditioner generation and application for both
Expand All @@ -125,8 +159,7 @@ void Jacobi<ValueType, IndexType>::generate_precond(
first_sys_csr.get(), num_blocks_, blocks_storage_scheme_,
blocks_cumulative_storage_.get_const_data(),
parameters_.block_pointers.get_const_data(),
row_part_of_which_block_info_.get_const_data(),
blocks_pattern.get_data()));
row_block_map_info_.get_const_data(), blocks_pattern.get_data()));

exec->run(jacobi::make_compute_block_jacobi(
sys_csr.get(), parameters_.max_block_size, num_blocks_,
Expand Down
64 changes: 31 additions & 33 deletions core/preconditioner/batch_jacobi_kernels.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,11 @@ namespace gko {
namespace kernels {


#define GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK( \
IndexType) \
void find_row_is_part_of_which_block( \
std::shared_ptr<const DefaultExecutor> exec, \
const size_type num_blocks, const IndexType* block_pointers, \
IndexType* row_part_of_which_block_info)
#define GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP(IndexType) \
void find_row_block_map(std::shared_ptr<const DefaultExecutor> exec, \
const size_type num_blocks, \
const IndexType* block_pointers, \
IndexType* row_block_map_info)

#define GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_CUMULATIVE_BLOCK_STORAGE( \
IndexType) \
Expand All @@ -36,17 +35,16 @@ namespace kernels {
const size_type num_blocks, const IndexType* block_pointers, \
IndexType* blocks_cumulative_storage)

#define GKO_DECLARE_BATCH_BLOCK_JACOBI_EXTRACT_PATTERN_KERNEL(ValueType, \
IndexType) \
void extract_common_blocks_pattern( \
std::shared_ptr<const DefaultExecutor> exec, \
const matrix::Csr<ValueType, IndexType>* first_sys_csr, \
const size_type num_blocks, \
const batch::preconditioner::batched_jacobi_blocks_storage_scheme< \
IndexType>& storage_scheme, \
const IndexType* cumulative_block_storage, \
const IndexType* block_pointers, \
const IndexType* row_part_of_which_block_info, \
#define GKO_DECLARE_BATCH_BLOCK_JACOBI_EXTRACT_PATTERN_KERNEL(ValueType, \
IndexType) \
void extract_common_blocks_pattern( \
std::shared_ptr<const DefaultExecutor> exec, \
const matrix::Csr<ValueType, IndexType>* first_sys_csr, \
const size_type num_blocks, \
const batch::preconditioner::batched_jacobi_blocks_storage_scheme< \
IndexType>& storage_scheme, \
const IndexType* cumulative_block_storage, \
const IndexType* block_pointers, const IndexType* row_block_map_info, \
IndexType* blocks_pattern)


Expand Down Expand Up @@ -80,7 +78,7 @@ namespace kernels {
batched_jacobi_blocks_storage_scheme<IndexType>& storage_scheme, \
const IndexType* cumulative_block_storage, \
const ValueType* blocks_array, const IndexType* block_ptrs, \
const IndexType* row_part_of_which_block_info, \
const IndexType* row_block_map_info, \
const batch::MultiVector<ValueType>* r, \
batch::MultiVector<ValueType>* z)

Expand All @@ -93,24 +91,24 @@ namespace kernels {
batched_jacobi_blocks_storage_scheme<IndexType>& storage_scheme, \
const IndexType* cumulative_block_storage, \
const ValueType* blocks_array, const IndexType* block_ptrs, \
const IndexType* row_part_of_which_block_info, \
const IndexType* row_block_map_info, \
const batch::MultiVector<ValueType>* r, \
batch::MultiVector<ValueType>* z)

#define GKO_DECLARE_ALL_AS_TEMPLATES \
template <typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_CUMULATIVE_BLOCK_STORAGE( \
IndexType); \
template <typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK(IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_EXTRACT_PATTERN_KERNEL(ValueType, \
IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_JACOBI_ELL_APPLY_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
#define GKO_DECLARE_ALL_AS_TEMPLATES \
template <typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_CUMULATIVE_BLOCK_STORAGE( \
IndexType); \
template <typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP(IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_EXTRACT_PATTERN_KERNEL(ValueType, \
IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_BLOCK_JACOBI_COMPUTE_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_JACOBI_ELL_APPLY_KERNEL(ValueType, IndexType); \
template <typename ValueType, typename IndexType> \
GKO_DECLARE_BATCH_JACOBI_APPLY_KERNEL(ValueType, IndexType)


Expand Down
10 changes: 5 additions & 5 deletions cuda/preconditioner/batch_jacobi_kernels.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -47,7 +47,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -66,13 +66,13 @@ GKO_INSTANTIATE_FOR_INT32_TYPE(


template <typename IndexType>
void find_row_is_part_of_which_block(
void find_row_block_map(
std::shared_ptr<const DefaultExecutor> exec, const size_type num_blocks,
const IndexType* const block_pointers,
IndexType* const row_part_of_which_block_info) GKO_NOT_IMPLEMENTED;
IndexType* const row_block_map_info) GKO_NOT_IMPLEMENTED;

GKO_INSTANTIATE_FOR_INT32_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK);
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP);


template <typename ValueType, typename IndexType>
Expand Down
10 changes: 5 additions & 5 deletions dpcpp/preconditioner/batch_jacobi_kernels.dp.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -47,7 +47,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -66,13 +66,13 @@ GKO_INSTANTIATE_FOR_INT32_TYPE(


template <typename IndexType>
void find_row_is_part_of_which_block(
void find_row_block_map(
std::shared_ptr<const DefaultExecutor> exec, const size_type num_blocks,
const IndexType* const block_pointers,
IndexType* const row_part_of_which_block_info) GKO_NOT_IMPLEMENTED;
IndexType* const row_block_map_info) GKO_NOT_IMPLEMENTED;

GKO_INSTANTIATE_FOR_INT32_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK);
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP);


template <typename ValueType, typename IndexType>
Expand Down
10 changes: 5 additions & 5 deletions hip/preconditioner/batch_jacobi_kernels.hip.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -47,7 +47,7 @@ void batch_jacobi_apply(
IndexType>& storage_scheme,
const IndexType* const cumulative_block_storage,
const ValueType* const blocks_array, const IndexType* const block_ptrs,
const IndexType* const row_part_of_which_block_info,
const IndexType* const row_block_map_info,
const batch::MultiVector<ValueType>* const r,
batch::MultiVector<ValueType>* const z) GKO_NOT_IMPLEMENTED;

Expand All @@ -66,13 +66,13 @@ GKO_INSTANTIATE_FOR_INT32_TYPE(


template <typename IndexType>
void find_row_is_part_of_which_block(
void find_row_block_map(
std::shared_ptr<const DefaultExecutor> exec, const size_type num_blocks,
const IndexType* const block_pointers,
IndexType* const row_part_of_which_block_info) GKO_NOT_IMPLEMENTED;
IndexType* const row_block_map_info) GKO_NOT_IMPLEMENTED;

GKO_INSTANTIATE_FOR_INT32_TYPE(
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_IS_PART_OF_WHICH_BLOCK);
GKO_DECLARE_BATCH_BLOCK_JACOBI_FIND_ROW_BLOCK_MAP);


template <typename ValueType, typename IndexType>
Expand Down
Loading

0 comments on commit 233c13c

Please sign in to comment.