adds advanced scatter operation

MarcelKoch · MarcelKoch · commit d5da55e26377 · 2025-02-21T15:08:35.000+01:00
diff --git a/common/unified/matrix/row_scatterer.cpp b/common/unified/matrix/row_scatterer.cpp
@@ -40,6 +40,20 @@ GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE_2(
     GKO_DECLARE_ROW_SCATTER_SIMPLE_APPLY);
 
 
+template <typename ValueType, typename OutputType, typename IndexType>
+void advanced_row_scatter(std::shared_ptr<const DefaultExecutor> exec,
+                          const array<IndexType>* row_idxs,
+                          const matrix::Dense<ValueType>* alpha,
+                          const matrix::Dense<ValueType>* orig,
+                          const matrix::Dense<OutputType>* beta,
+                          matrix::Dense<OutputType>* target,
+                          bit_packed_span<bool, IndexType, uint32> mask,
+                          bool& invalid_access) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE_2(
+    GKO_DECLARE_ROW_SCATTER_ADVANCED_APPLY);
+
+
 }  // namespace row_scatter
 }  // namespace GKO_DEVICE_NAMESPACE
 }  // namespace kernels
diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -1122,9 +1122,10 @@ namespace row_scatter {
 
 
 GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE_2(GKO_DECLARE_ROW_SCATTER_SIMPLE_APPLY);
+GKO_STUB_MIXED_VALUE_AND_INDEX_TYPE_2(GKO_DECLARE_ROW_SCATTER_ADVANCED_APPLY);
 
 
-}
+}  // namespace row_scatter
 }  // namespace GKO_HOOK_MODULE
 }  // namespace kernels
 }  // namespace gko
diff --git a/core/matrix/row_scatterer.cpp b/core/matrix/row_scatterer.cpp
@@ -8,6 +8,7 @@
 #include <ginkgo/core/matrix/dense.hpp>
 
 #include "core/base/dispatch_helper.hpp"
+#include "core/components/bit_packed_storage.hpp"
 #include "core/matrix/row_scatterer_kernels.hpp"
 
 namespace gko {
@@ -16,9 +17,10 @@ namespace {
 
 
 GKO_REGISTER_OPERATION(row_scatter, row_scatter::row_scatter);
+GKO_REGISTER_OPERATION(advanced_row_scatter, row_scatter::advanced_row_scatter);
 
 
-}
+}  // namespace
 
 
 template <typename IndexType>
@@ -41,20 +43,20 @@ template <typename IndexType>
 RowScatterer<IndexType>::RowScatterer(std::shared_ptr<const Executor> exec,
                                       array<IndexType> idxs, size_type to_size)
     : EnableLinOp<RowScatterer<IndexType>>(exec, {to_size, idxs.get_size()}),
-      idxs_(exec, std::move(idxs))
+      idxs_(exec, std::move(idxs)),
+      mask_(exec,
+            bit_packed_span<bool, IndexType, uint32>::storage_size(to_size, 1))
 {}
 
 
 template <typename IndexType>
 void RowScatterer<IndexType>::apply_impl(const LinOp* b, LinOp* x) const
 {
     auto impl = [&](const auto* orig, auto* target) {
-        auto exec = orig->get_executor();
+        auto exec = this->get_executor();
         bool invalid_access = false;
 
-        exec->run(make_row_scatter(
-            make_temporary_clone(exec, &idxs_).get(), orig,
-            make_temporary_clone(exec, target).get(), invalid_access));
+        exec->run(make_row_scatter(&idxs_, orig, target, invalid_access));
 
         // TODO: find a uniform way to handle device-side errors
         if (invalid_access) {
@@ -79,19 +81,37 @@ template <typename IndexType>
 void RowScatterer<IndexType>::apply_impl(const LinOp* alpha, const LinOp* b,
                                          const LinOp* beta, LinOp* x) const
 {
-    auto x_copy = gko::clone(x);
+    auto impl = [&](const auto* orig, auto* target) {
+        auto exec = this->get_executor();
+        bool invalid_access = false;
+
+        auto dense_alpha = make_temporary_conversion<
+            typename std::decay_t<decltype(*orig)>::value_type>(alpha);
+        auto dense_beta = make_temporary_conversion<
+            typename std::decay_t<decltype(*target)>::value_type>(beta);
+
+        exec->run(make_advanced_row_scatter(
+            &idxs_, dense_alpha.get(), orig, dense_beta.get(), target,
+            bit_packed_span<bool, IndexType, uint32>(mask_.get_data(), 1,
+                                                     this->get_size()[0]),
+            invalid_access));
+
+        if (invalid_access) {
+            GKO_INVALID_STATE("Out-of-bounds scatter index detected.");
+        }
+    };
+
+    mask_.fill(uint32{});
+
     run<Dense,
 #if GINKGO_ENABLE_HALF
         gko::half, std::complex<gko::half>,
 #endif
         float, double, std::complex<float>, std::complex<double>>(
-        x, [&](auto* target) {
-            using dense_type = std::decay_t<decltype(*target)>;
-            as<dense_type>(x_copy)->fill(
-                gko::zero<typename dense_type::value_type>());
-            this->apply_impl(b, x_copy);
-            target->scale(beta);
-            target->add_scaled(alpha, x_copy);
+        b, [&](auto* orig) {
+            using value_type =
+                typename std::decay_t<decltype(*orig)>::value_type;
+            mixed_precision_dispatch_real_complex<value_type>(impl, orig, x);
         });
 }
 
diff --git a/core/matrix/row_scatterer_kernels.hpp b/core/matrix/row_scatterer_kernels.hpp
@@ -8,6 +8,7 @@
 #include <ginkgo/core/matrix/dense.hpp>
 
 #include "core/base/kernel_declaration.hpp"
+#include "core/components/bit_packed_storage.hpp"
 
 
 namespace gko {
@@ -19,10 +20,20 @@ namespace kernels {
                      const matrix::Dense<_vtype>* orig,              \
                      matrix::Dense<_otype>* target, bool& invalid_access)
 
+#define GKO_DECLARE_ROW_SCATTER_ADVANCED_APPLY(_vtype, _otype, _itype)        \
+    void advanced_row_scatter(                                                \
+        std::shared_ptr<const DefaultExecutor> exec,                          \
+        const array<_itype>* row_idxs, const matrix::Dense<_vtype>* alpha,    \
+        const matrix::Dense<_vtype>* orig, const matrix::Dense<_otype>* beta, \
+        matrix::Dense<_otype>* target,                                        \
+        bit_packed_span<bool, _itype, uint32> mask, bool& invalid_access)
 
-#define GKO_DECLARE_ALL_AS_TEMPLATES                                       \
-    template <typename ValueType, typename OutputType, typename IndexType> \
-    GKO_DECLARE_ROW_SCATTER_SIMPLE_APPLY(ValueType, OutputType, IndexType)
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                                        \
+    template <typename ValueType, typename OutputType, typename IndexType>  \
+    GKO_DECLARE_ROW_SCATTER_SIMPLE_APPLY(ValueType, OutputType, IndexType); \
+    template <typename ValueType, typename OutputType, typename IndexType>  \
+    GKO_DECLARE_ROW_SCATTER_ADVANCED_APPLY(ValueType, OutputType, IndexType)
 
 
 GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(row_scatter,
diff --git a/include/ginkgo/core/matrix/row_scatterer.hpp b/include/ginkgo/core/matrix/row_scatterer.hpp
@@ -7,9 +7,18 @@
 #include <ginkgo/core/base/lin_op.hpp>
 
 namespace gko {
+
+template <typename ValueType, typename IndexType, typename StorageType>
+class bit_packed_span;
+
+
 namespace matrix {
 
 
+/**
+ *
+ * @tparam IndexType type for defining the scatter-to indices
+ */
 template <typename IndexType = int32>
 class RowScatterer : public EnableLinOp<RowScatterer<IndexType>> {
     friend class EnablePolymorphicObject<RowScatterer<IndexType>, LinOp>;
@@ -32,6 +41,7 @@ class RowScatterer : public EnableLinOp<RowScatterer<IndexType>> {
                           array<IndexType> idxs, size_type to_size);
 
     array<IndexType> idxs_;
+    mutable array<uint32> mask_;
 };
 
 
diff --git a/reference/matrix/row_scatterer_kernels.cpp b/reference/matrix/row_scatterer_kernels.cpp
@@ -39,6 +39,40 @@ void row_scatter(std::shared_ptr<const ReferenceExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE_2(
     GKO_DECLARE_ROW_SCATTER_SIMPLE_APPLY);
 
+template <typename ValueType, typename OutputType, typename IndexType>
+void advanced_row_scatter(std::shared_ptr<const ReferenceExecutor> exec,
+                          const array<IndexType>* row_idxs,
+                          const matrix::Dense<ValueType>* alpha,
+                          const matrix::Dense<ValueType>* orig,
+                          const matrix::Dense<OutputType>* beta,
+                          matrix::Dense<OutputType>* target,
+                          bit_packed_span<bool, IndexType, uint32> mask,
+                          bool& invalid_access)
+{
+    using type = highest_precision<ValueType, OutputType>;
+    auto scalar_alpha = alpha->at(0, 0);
+    auto scalar_beta = beta->at(0, 0);
+    auto rows = row_idxs->get_const_data();
+    for (size_type i = 0; i < row_idxs->get_size(); ++i) {
+        if (rows[i] >= target->get_size()[0]) {
+            invalid_access = true;
+            return;
+        }
+
+        bool scaled = mask.get(rows[i]);
+        mask.set(rows[i], true);
+        for (size_type j = 0; j < orig->get_size()[1]; ++j) {
+            target->at(rows[i], j) =
+                static_cast<type>(scalar_alpha * orig->at(i, j)) +
+                (scaled ? type{1.0} : static_cast<type>(scalar_beta)) *
+                    static_cast<type>(target->at(rows[i], j));
+        }
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_MIXED_VALUE_AND_INDEX_TYPE_2(
+    GKO_DECLARE_ROW_SCATTER_ADVANCED_APPLY);
+
 
 }  // namespace row_scatter
 }  // namespace reference
diff --git a/reference/test/matrix/row_scatterer_kernels.cpp b/reference/test/matrix/row_scatterer_kernels.cpp
@@ -10,6 +10,7 @@
 
 #include <ginkgo/core/matrix/row_scatterer.hpp>
 
+#include "core/components/bit_packed_storage.hpp"
 #include "core/test/utils.hpp"
 
 
@@ -22,20 +23,25 @@ class RowScatter : public ::testing::Test {
     using DenseIn = gko::matrix::Dense<in_value_type>;
     using DenseOut = gko::matrix::Dense<out_value_type>;
     using Scatterer = gko::matrix::RowScatterer<index_type>;
-
+    using mask_type = gko::bit_packed_span<bool, index_type, gko::uint32>;
 
     std::shared_ptr<gko::ReferenceExecutor> exec =
         gko::ReferenceExecutor::create();
 
     std::unique_ptr<DenseIn> in =
         gko::initialize<DenseIn>(I<I<in_value_type>>{{1, 2}, {3, 4}}, exec);
-    std::unique_ptr<DenseIn> in_repeated =
-        gko::initialize<DenseIn>(I<I<in_value_type>>{{1, 2}, {3, 4}, {5, 6}}, exec);
+    std::unique_ptr<DenseIn> in_repeated = gko::initialize<DenseIn>(
+        I<I<in_value_type>>{{1, 2}, {3, 4}, {5, 6}}, exec);
     std::unique_ptr<DenseOut> out = gko::initialize<DenseOut>(
         I<I<out_value_type>>{{11, 22}, {33, 44}, {55, 66}, {77, 88}}, exec);
+    std::unique_ptr<DenseIn> alpha =
+        gko::initialize<DenseIn>({in_value_type{-1.0}}, this->exec);
+    std::unique_ptr<DenseOut> beta =
+        gko::initialize<DenseOut>({out_value_type{-2.0}}, this->exec);
 
     gko::array<index_type> idxs = {exec, {3, 1}};
     gko::array<index_type> idxs_repeated = {exec, {3, 3, 1}};
+    gko::array<gko::uint32> mask = {exec, mask_type::storage_size(4, 1)};
 
     std::unique_ptr<Scatterer> scatterer =
         Scatterer::create(exec, idxs, out->get_size()[0]);
@@ -50,7 +56,7 @@ TYPED_TEST_SUITE(RowScatter, gko::test::MixedPresisionValueIndexTypes,
 #endif
 
 
-TYPED_TEST(RowScatter, CanRowScatter)
+TYPED_TEST(RowScatter, CanSimpleRowScatter)
 {
     bool invalid_access = false;
 
@@ -72,8 +78,8 @@ TYPED_TEST(RowScatter, SimpleRowScatterIsAdditive)
     bool invalid_access = false;
 
     gko::kernels::reference::row_scatter::row_scatter(
-        this->exec, &this->idxs_repeated, this->in_repeated.get(), this->out.get(),
-        invalid_access);
+        this->exec, &this->idxs_repeated, this->in_repeated.get(),
+        this->out.get(), invalid_access);
 
     ASSERT_FALSE(invalid_access);
     auto expected = gko::initialize<typename TestFixture::DenseOut>(
@@ -84,6 +90,48 @@ TYPED_TEST(RowScatter, SimpleRowScatterIsAdditive)
 }
 
 
+TYPED_TEST(RowScatter, CanAdvancedRowScatter)
+{
+    using mask_type = typename TestFixture::mask_type;
+    bool invalid_access = false;
+    this->mask.fill(0);
+
+    gko::kernels::reference::row_scatter::advanced_row_scatter(
+        this->exec, &this->idxs, this->alpha.get(), this->in.get(),
+        this->beta.get(), this->out.get(),
+        mask_type(this->mask.get_data(), 1, this->out->get_size()[0]),
+        invalid_access);
+
+    ASSERT_FALSE(invalid_access);
+    auto expected = gko::initialize<typename TestFixture::DenseOut>(
+        I<I<typename TestFixture::out_value_type>>{
+            {11, 22}, {-69, -92}, {55, 66}, {-155, -178}},
+        this->exec);
+    GKO_ASSERT_MTX_NEAR(this->out, expected, 0.0);
+}
+
+
+TYPED_TEST(RowScatter, AdvancedRowScatterIsAdditive)
+{
+    using mask_type = typename TestFixture::mask_type;
+    bool invalid_access = false;
+    this->mask.fill(0);
+
+    gko::kernels::reference::row_scatter::advanced_row_scatter(
+        this->exec, &this->idxs_repeated, this->alpha.get(),
+        this->in_repeated.get(), this->beta.get(), this->out.get(),
+        mask_type(this->mask.get_data(), 1, this->out->get_size()[0]),
+        invalid_access);
+
+    ASSERT_FALSE(invalid_access);
+    auto expected = gko::initialize<typename TestFixture::DenseOut>(
+        I<I<typename TestFixture::out_value_type>>{
+            {11, 22}, {-71, -94}, {55, 66}, {-158, -182}},
+        this->exec);
+    GKO_ASSERT_MTX_NEAR(this->out, expected, 0.0);
+}
+
+
 TYPED_TEST(RowScatter, CanDetectInvalidAccess)
 {
     bool invalid_access = false;
@@ -106,3 +154,16 @@ TYPED_TEST(RowScatter, CanRowScatterSimpleApply)
         this->exec);
     GKO_ASSERT_MTX_NEAR(this->out, expected, 0.0);
 }
+
+
+TYPED_TEST(RowScatter, CanRowScatterAdvancedApply)
+{
+    this->scatterer->apply(this->alpha, this->in.get(), this->beta,
+                           this->out.get());
+
+    auto expected = gko::initialize<typename TestFixture::DenseOut>(
+        I<I<typename TestFixture::out_value_type>>{
+            {11, 22}, {-69, -92}, {55, 66}, {-155, -178}},
+        this->exec);
+    GKO_ASSERT_MTX_NEAR(this->out, expected, 0.0);
+}
