Skip to content

Commit 83b22df

Browse files
cflintocflinto
authored
Merge pull request #2014 from ginkgo-project/feat/rand-develop-poc
Add sparse stack
2 parents 7340d04 + 9d2c752 commit 83b22df

18 files changed

Lines changed: 786 additions & 41 deletions

‎common/cuda_hip/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ set(CUDA_HIP_SOURCES
4646
solver/cb_gmres_kernels.cpp
4747
sketch/count_sketch_kernels.cpp
4848
sketch/gaussian_sketch_kernels.cpp
49+
sketch/sparse_stack_kernels.cpp
4950
solver/idr_kernels.cpp
5051
solver/multigrid_kernels.cpp
5152
stop/criterion_kernels.cpp

‎common/cuda_hip/sketch/sparse_stack_kernels.cpp‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
2+
//
3+
// SPDX-License-Identifier: BSD-3-Clause
4+
5+
#include "core/sketch/sparse_stack_kernels.hpp"
6+
7+
#include <ginkgo/core/base/exception_helpers.hpp>
8+
9+
namespace gko {
10+
namespace kernels {
11+
namespace GKO_DEVICE_NAMESPACE {
12+
namespace sparse_stack {
13+
14+
15+
template <typename ValueType, typename IndexType>
16+
void generate(std::shared_ptr<const DefaultExecutor> exec,
17+
size_type sketch_size, size_type input_size, size_type zeta,
18+
array<IndexType>& hash_map, array<ValueType>& signs,
19+
uint64 seed) GKO_NOT_IMPLEMENTED;
20+
21+
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
22+
GKO_DECLARE_SPARSE_STACK_GENERATE);
23+
24+
template <typename ValueType, typename IndexType>
25+
void apply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta,
26+
const array<IndexType>& hash_map, const array<ValueType>& signs,
27+
matrix::view::dense<const ValueType> b,
28+
matrix::view::dense<ValueType> x) GKO_NOT_IMPLEMENTED;
29+
30+
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_APPLY);
31+
32+
template <typename ValueType, typename IndexType>
33+
void rapply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta,
34+
const array<IndexType>& hash_map, const array<ValueType>& signs,
35+
matrix::view::dense<const ValueType> b,
36+
matrix::view::dense<ValueType> x) GKO_NOT_IMPLEMENTED;
37+
38+
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_RAPPLY);
39+
40+
41+
} // namespace sparse_stack
42+
} // namespace GKO_DEVICE_NAMESPACE
43+
} // namespace kernels
44+
} // namespace gko

‎core/CMakeLists.txt‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,6 @@ target_sources(
9595
matrix/scaled_permutation.cpp
9696
matrix/sellp.cpp
9797
matrix/sparsity_csr.cpp
98-
sketch/count_sketch.cpp
99-
sketch/gaussian_sketch.cpp
10098
multigrid/fixed_coarsening.cpp
10199
multigrid/pgm.cpp
102100
preconditioner/batch_jacobi.cpp
@@ -110,6 +108,9 @@ target_sources(
110108
reorder/mc64.cpp
111109
reorder/rcm.cpp
112110
reorder/scaled_reordered.cpp
111+
sketch/count_sketch.cpp
112+
sketch/gaussian_sketch.cpp
113+
sketch/sparse_stack.cpp
113114
solver/batch_bicgstab.cpp
114115
solver/batch_cg.cpp
115116
solver/bicg.cpp

‎core/device_hooks/common_kernels.inc.cpp‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
1+
// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
22
//
33
// SPDX-License-Identifier: BSD-3-Clause
44

@@ -57,6 +57,9 @@
5757
#include "core/preconditioner/jacobi_kernels.hpp"
5858
#include "core/preconditioner/sor_kernels.hpp"
5959
#include "core/reorder/rcm_kernels.hpp"
60+
#include "core/sketch/count_sketch_kernels.hpp"
61+
#include "core/sketch/gaussian_sketch_kernels.hpp"
62+
#include "core/sketch/sparse_stack_kernels.hpp"
6063
#include "core/solver/batch_bicgstab_kernels.hpp"
6164
#include "core/solver/batch_cg_kernels.hpp"
6265
#include "core/solver/bicg_kernels.hpp"
@@ -76,8 +79,6 @@
7679
#include "core/solver/multigrid_kernels.hpp"
7780
#include "core/solver/pipe_cg_kernels.hpp"
7881
#include "core/solver/upper_trs_kernels.hpp"
79-
#include "core/sketch/count_sketch_kernels.hpp"
80-
#include "core/sketch/gaussian_sketch_kernels.hpp"
8182
#include "core/stop/criterion_kernels.hpp"
8283
#include "core/stop/residual_norm_kernels.hpp"
8384

@@ -1190,6 +1191,16 @@ GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COUNT_SKETCH_RAPPLY);
11901191

11911192
} // namespace count_sketch
11921193

1194+
namespace sparse_stack {
1195+
1196+
1197+
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_GENERATE);
1198+
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_APPLY);
1199+
GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_RAPPLY);
1200+
1201+
1202+
} // namespace sparse_stack
1203+
11931204

11941205
} // namespace GKO_HOOK_MODULE
11951206
} // namespace kernels

‎core/sketch/count_sketch.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,15 @@ CountSketch<ValueType, IndexType>::create(
5050
}
5151

5252

53+
template <typename ValueType, typename IndexType>
54+
std::unique_ptr<CountSketch<ValueType, IndexType>>
55+
CountSketch<ValueType, IndexType>::create(
56+
std::shared_ptr<const Executor> exec)
57+
{
58+
return std::unique_ptr<CountSketch>{new CountSketch(exec)};
59+
}
60+
61+
5362
template <typename ValueType, typename IndexType>
5463
void CountSketch<ValueType, IndexType>::apply_sketch_impl(
5564
const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const

‎core/sketch/gaussian_sketch.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,15 @@ std::unique_ptr<GaussianSketch<ValueType>> GaussianSketch<ValueType>::create(
5151
}
5252

5353

54+
55+
template <typename ValueType>
56+
std::unique_ptr<GaussianSketch<ValueType>> GaussianSketch<ValueType>::create(
57+
std::shared_ptr<const Executor> exec)
58+
{
59+
return std::unique_ptr<GaussianSketch>{new GaussianSketch(exec)};
60+
}
61+
62+
5463
template <typename ValueType>
5564
void GaussianSketch<ValueType>::apply_sketch_impl(
5665
const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const

‎core/sketch/sparse_stack.cpp‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
2+
//
3+
// SPDX-License-Identifier: BSD-3-Clause
4+
5+
#include <ginkgo/core/base/precision_dispatch.hpp>
6+
#include <ginkgo/core/matrix/dense.hpp>
7+
#include <ginkgo/core/sketch/sparse_stack.hpp>
8+
9+
#include "core/sketch/sparse_stack_kernels.hpp"
10+
11+
12+
namespace gko {
13+
namespace sketch {
14+
namespace sparse_stack {
15+
namespace {
16+
17+
18+
GKO_REGISTER_OPERATION(generate, sparse_stack::generate);
19+
GKO_REGISTER_OPERATION(apply, sparse_stack::apply);
20+
GKO_REGISTER_OPERATION(rapply, sparse_stack::rapply);
21+
22+
23+
} // anonymous namespace
24+
} // namespace sparse_stack
25+
26+
27+
template <typename ValueType, typename IndexType>
28+
SparseStack<ValueType, IndexType>::SparseStack(
29+
std::shared_ptr<const Executor> exec, size_type sketch_size,
30+
size_type input_size, size_type zeta, uint64 seed)
31+
: EnableLinOp<SparseStack, SketchOperator<ValueType>>(
32+
exec, dim<2>{sketch_size, input_size}),
33+
zeta_{zeta},
34+
hash_map_{exec, input_size * zeta},
35+
signs_{exec, input_size * zeta},
36+
seed_{seed}
37+
{
38+
exec->run(sparse_stack::make_generate(sketch_size, input_size, zeta_,
39+
hash_map_, signs_, seed_));
40+
}
41+
42+
43+
template <typename ValueType, typename IndexType>
44+
std::unique_ptr<SparseStack<ValueType, IndexType>>
45+
SparseStack<ValueType, IndexType>::create(std::shared_ptr<const Executor> exec,
46+
size_type sketch_size,
47+
size_type input_size, size_type zeta,
48+
uint64 seed)
49+
{
50+
return std::unique_ptr<SparseStack>{
51+
new SparseStack(exec, sketch_size, input_size, zeta, seed)};
52+
}
53+
54+
template <typename ValueType, typename IndexType>
55+
std::unique_ptr<SparseStack<ValueType, IndexType>>
56+
SparseStack<ValueType, IndexType>::create(std::shared_ptr<const Executor> exec)
57+
{
58+
return std::unique_ptr<SparseStack>{new SparseStack(exec)};
59+
}
60+
61+
62+
template <typename ValueType, typename IndexType>
63+
void SparseStack<ValueType, IndexType>::apply_sketch_impl(
64+
const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const
65+
{
66+
this->get_executor()->run(sparse_stack::make_apply(
67+
zeta_, hash_map_, signs_, b->get_const_device_view(),
68+
x->get_device_view()));
69+
}
70+
71+
72+
template <typename ValueType, typename IndexType>
73+
void SparseStack<ValueType, IndexType>::rapply_sketch_impl(
74+
const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const
75+
{
76+
this->get_executor()->run(sparse_stack::make_rapply(
77+
zeta_, hash_map_, signs_, b->get_const_device_view(),
78+
x->get_device_view()));
79+
}
80+
81+
82+
#define GKO_DECLARE_SPARSE_STACK(ValueType, IndexType) \
83+
class SparseStack<ValueType, IndexType>
84+
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK);
85+
86+
87+
} // namespace sketch
88+
} // namespace gko

‎core/sketch/sparse_stack_kernels.hpp‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
2+
//
3+
// SPDX-License-Identifier: BSD-3-Clause
4+
5+
#ifndef GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_
6+
#define GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_
7+
8+
9+
#include <ginkgo/core/base/array.hpp>
10+
#include <ginkgo/core/base/types.hpp>
11+
#include <ginkgo/core/matrix/dense.hpp>
12+
13+
#include "core/base/kernel_declaration.hpp"
14+
15+
16+
namespace gko {
17+
namespace kernels {
18+
19+
20+
#define GKO_DECLARE_SPARSE_STACK_GENERATE(ValueType, IndexType) \
21+
void generate(std::shared_ptr<const DefaultExecutor> exec, \
22+
size_type sketch_size, size_type input_size, size_type zeta, \
23+
array<IndexType>& hash_map, array<ValueType>& signs, \
24+
gko::uint64 seed)
25+
26+
#define GKO_DECLARE_SPARSE_STACK_APPLY(ValueType, IndexType) \
27+
void apply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta, \
28+
const array<IndexType>& hash_map, \
29+
const array<ValueType>& signs, \
30+
matrix::view::dense<const ValueType> b, \
31+
matrix::view::dense<ValueType> x)
32+
33+
#define GKO_DECLARE_SPARSE_STACK_RAPPLY(ValueType, IndexType) \
34+
void rapply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta, \
35+
const array<IndexType>& hash_map, \
36+
const array<ValueType>& signs, \
37+
matrix::view::dense<const ValueType> b, \
38+
matrix::view::dense<ValueType> x)
39+
40+
41+
#define GKO_DECLARE_ALL_AS_TEMPLATES \
42+
template <typename ValueType, typename IndexType> \
43+
GKO_DECLARE_SPARSE_STACK_GENERATE(ValueType, IndexType); \
44+
template <typename ValueType, typename IndexType> \
45+
GKO_DECLARE_SPARSE_STACK_APPLY(ValueType, IndexType); \
46+
template <typename ValueType, typename IndexType> \
47+
GKO_DECLARE_SPARSE_STACK_RAPPLY(ValueType, IndexType)
48+
49+
50+
GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(sparse_stack,
51+
GKO_DECLARE_ALL_AS_TEMPLATES);
52+
53+
54+
#undef GKO_DECLARE_ALL_AS_TEMPLATES
55+
56+
57+
} // namespace kernels
58+
} // namespace gko
59+
60+
61+
#endif // GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_

‎core/test/sketch/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
ginkgo_create_test(gaussian_sketch)
22
ginkgo_create_test(count_sketch)
3+
ginkgo_create_test(sparse_stack)

0 commit comments

Comments
 (0)