Add sparse stack

common/cuda_hip/CMakeLists.txt

@@ -46,6 +46,7 @@ set(CUDA_HIP_SOURCES
     solver/cb_gmres_kernels.cpp
     sketch/count_sketch_kernels.cpp
     sketch/gaussian_sketch_kernels.cpp
+    sketch/sparse_stack_kernels.cpp
     solver/idr_kernels.cpp
     solver/multigrid_kernels.cpp
     stop/criterion_kernels.cpp

common/cuda_hip/sketch/sparse_stack_kernels.cpp

@@ -0,0 +1,44 @@
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "core/sketch/sparse_stack_kernels.hpp"
+
+#include <ginkgo/core/base/exception_helpers.hpp>
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace sparse_stack {
+
+
+template <typename ValueType, typename IndexType>
+void generate(std::shared_ptr<const DefaultExecutor> exec,
+              size_type sketch_size, size_type input_size, size_type zeta,
+              array<IndexType>& hash_map, array<ValueType>& signs,
+              uint64 seed) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SPARSE_STACK_GENERATE);
+
+template <typename ValueType, typename IndexType>
+void apply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta,
+           const array<IndexType>& hash_map, const array<ValueType>& signs,
+           matrix::view::dense<const ValueType> b,
+           matrix::view::dense<ValueType> x) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_APPLY);
+
+template <typename ValueType, typename IndexType>
+void rapply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta,
+            const array<IndexType>& hash_map, const array<ValueType>& signs,
+            matrix::view::dense<const ValueType> b,
+            matrix::view::dense<ValueType> x) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_RAPPLY);
+
+
+}  // namespace sparse_stack
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

core/CMakeLists.txt

@@ -95,8 +95,6 @@ target_sources(
         matrix/scaled_permutation.cpp
         matrix/sellp.cpp
         matrix/sparsity_csr.cpp
-        sketch/count_sketch.cpp
-        sketch/gaussian_sketch.cpp
         multigrid/fixed_coarsening.cpp
         multigrid/pgm.cpp
         preconditioner/batch_jacobi.cpp
@@ -110,6 +108,9 @@ target_sources(
         reorder/mc64.cpp
         reorder/rcm.cpp
         reorder/scaled_reordered.cpp
+        sketch/count_sketch.cpp
+        sketch/gaussian_sketch.cpp
+        sketch/sparse_stack.cpp
         solver/batch_bicgstab.cpp
         solver/batch_cg.cpp
         solver/bicg.cpp

core/device_hooks/common_kernels.inc.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -57,6 +57,9 @@
 #include "core/preconditioner/jacobi_kernels.hpp"
 #include "core/preconditioner/sor_kernels.hpp"
 #include "core/reorder/rcm_kernels.hpp"
+#include "core/sketch/count_sketch_kernels.hpp"
+#include "core/sketch/gaussian_sketch_kernels.hpp"
+#include "core/sketch/sparse_stack_kernels.hpp"
 #include "core/solver/batch_bicgstab_kernels.hpp"
 #include "core/solver/batch_cg_kernels.hpp"
 #include "core/solver/bicg_kernels.hpp"
@@ -76,8 +79,6 @@
 #include "core/solver/multigrid_kernels.hpp"
 #include "core/solver/pipe_cg_kernels.hpp"
 #include "core/solver/upper_trs_kernels.hpp"
-#include "core/sketch/count_sketch_kernels.hpp"
-#include "core/sketch/gaussian_sketch_kernels.hpp"
 #include "core/stop/criterion_kernels.hpp"
 #include "core/stop/residual_norm_kernels.hpp"
 
@@ -1190,6 +1191,16 @@ GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COUNT_SKETCH_RAPPLY);
 
 }  // namespace count_sketch
 
+namespace sparse_stack {
+
+
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_GENERATE);
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_APPLY);
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK_RAPPLY);
+
+
+}  // namespace sparse_stack
+
 
 }  // namespace GKO_HOOK_MODULE
 }  // namespace kernels

core/sketch/count_sketch.cpp

@@ -50,6 +50,15 @@ CountSketch<ValueType, IndexType>::create(
 }
 
 
+template <typename ValueType, typename IndexType>
+std::unique_ptr<CountSketch<ValueType, IndexType>>
+CountSketch<ValueType, IndexType>::create(
+    std::shared_ptr<const Executor> exec)
+{
+    return std::unique_ptr<CountSketch>{new CountSketch(exec)};
+}
+
+
 template <typename ValueType, typename IndexType>
 void CountSketch<ValueType, IndexType>::apply_sketch_impl(
     const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const

core/sketch/gaussian_sketch.cpp

@@ -51,6 +51,15 @@ std::unique_ptr<GaussianSketch<ValueType>> GaussianSketch<ValueType>::create(
 }
 
 
+
+template <typename ValueType>
+std::unique_ptr<GaussianSketch<ValueType>> GaussianSketch<ValueType>::create(
+    std::shared_ptr<const Executor> exec)
+{
+    return std::unique_ptr<GaussianSketch>{new GaussianSketch(exec)};
+}
+
+
 template <typename ValueType>
 void GaussianSketch<ValueType>::apply_sketch_impl(
     const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const

core/sketch/sparse_stack.cpp

@@ -0,0 +1,88 @@
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/base/precision_dispatch.hpp>
+#include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/sketch/sparse_stack.hpp>
+
+#include "core/sketch/sparse_stack_kernels.hpp"
+
+
+namespace gko {
+namespace sketch {
+namespace sparse_stack {
+namespace {
+
+
+GKO_REGISTER_OPERATION(generate, sparse_stack::generate);
+GKO_REGISTER_OPERATION(apply, sparse_stack::apply);
+GKO_REGISTER_OPERATION(rapply, sparse_stack::rapply);
+
+
+}  // anonymous namespace
+}  // namespace sparse_stack
+
+
+template <typename ValueType, typename IndexType>
+SparseStack<ValueType, IndexType>::SparseStack(
+    std::shared_ptr<const Executor> exec, size_type sketch_size,
+    size_type input_size, size_type zeta, uint64 seed)
+    : EnableLinOp<SparseStack, SketchOperator<ValueType>>(
+          exec, dim<2>{sketch_size, input_size}),
+      zeta_{zeta},
+      hash_map_{exec, input_size * zeta},
+      signs_{exec, input_size * zeta},
+      seed_{seed}
+{
+    exec->run(sparse_stack::make_generate(sketch_size, input_size, zeta_,
+                                          hash_map_, signs_, seed_));
+}
+
+
+template <typename ValueType, typename IndexType>
+std::unique_ptr<SparseStack<ValueType, IndexType>>
+SparseStack<ValueType, IndexType>::create(std::shared_ptr<const Executor> exec,
+                                          size_type sketch_size,
+                                          size_type input_size, size_type zeta,
+                                          uint64 seed)
+{
+    return std::unique_ptr<SparseStack>{
+        new SparseStack(exec, sketch_size, input_size, zeta, seed)};
+}
+
+template <typename ValueType, typename IndexType>
+std::unique_ptr<SparseStack<ValueType, IndexType>>
+SparseStack<ValueType, IndexType>::create(std::shared_ptr<const Executor> exec)
+{
+    return std::unique_ptr<SparseStack>{new SparseStack(exec)};
+}
+
+
+template <typename ValueType, typename IndexType>
+void SparseStack<ValueType, IndexType>::apply_sketch_impl(
+    const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const
+{
+    this->get_executor()->run(sparse_stack::make_apply(
+        zeta_, hash_map_, signs_, b->get_const_device_view(),
+        x->get_device_view()));
+}
+
+
+template <typename ValueType, typename IndexType>
+void SparseStack<ValueType, IndexType>::rapply_sketch_impl(
+    const matrix::Dense<ValueType>* b, matrix::Dense<ValueType>* x) const
+{
+    this->get_executor()->run(sparse_stack::make_rapply(
+        zeta_, hash_map_, signs_, b->get_const_device_view(),
+        x->get_device_view()));
+}
+
+
+#define GKO_DECLARE_SPARSE_STACK(ValueType, IndexType) \
+    class SparseStack<ValueType, IndexType>
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SPARSE_STACK);
+
+
+}  // namespace sketch
+}  // namespace gko

core/sketch/sparse_stack_kernels.hpp

@@ -0,0 +1,61 @@
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_
+#define GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_
+
+
+#include <ginkgo/core/base/array.hpp>
+#include <ginkgo/core/base/types.hpp>
+#include <ginkgo/core/matrix/dense.hpp>
+
+#include "core/base/kernel_declaration.hpp"
+
+
+namespace gko {
+namespace kernels {
+
+
+#define GKO_DECLARE_SPARSE_STACK_GENERATE(ValueType, IndexType)                \
+    void generate(std::shared_ptr<const DefaultExecutor> exec,                 \
+                  size_type sketch_size, size_type input_size, size_type zeta, \
+                  array<IndexType>& hash_map, array<ValueType>& signs,         \
+                  gko::uint64 seed)
+
+#define GKO_DECLARE_SPARSE_STACK_APPLY(ValueType, IndexType)                \
+    void apply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta, \
+               const array<IndexType>& hash_map,                            \
+               const array<ValueType>& signs,                               \
+               matrix::view::dense<const ValueType> b,                      \
+               matrix::view::dense<ValueType> x)
+
+#define GKO_DECLARE_SPARSE_STACK_RAPPLY(ValueType, IndexType)                \
+    void rapply(std::shared_ptr<const DefaultExecutor> exec, size_type zeta, \
+                const array<IndexType>& hash_map,                            \
+                const array<ValueType>& signs,                               \
+                matrix::view::dense<const ValueType> b,                      \
+                matrix::view::dense<ValueType> x)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                         \
+    template <typename ValueType, typename IndexType>        \
+    GKO_DECLARE_SPARSE_STACK_GENERATE(ValueType, IndexType); \
+    template <typename ValueType, typename IndexType>        \
+    GKO_DECLARE_SPARSE_STACK_APPLY(ValueType, IndexType);    \
+    template <typename ValueType, typename IndexType>        \
+    GKO_DECLARE_SPARSE_STACK_RAPPLY(ValueType, IndexType)
+
+
+GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(sparse_stack,
+                                        GKO_DECLARE_ALL_AS_TEMPLATES);
+
+
+#undef GKO_DECLARE_ALL_AS_TEMPLATES
+
+
+}  // namespace kernels
+}  // namespace gko
+
+
+#endif  // GKO_CORE_SKETCH_SPARSE_STACK_KERNELS_HPP_

core/test/sketch/CMakeLists.txt

@@ -1,2 +1,3 @@
 ginkgo_create_test(gaussian_sketch)
 ginkgo_create_test(count_sketch)
+ginkgo_create_test(sparse_stack)