Refactor seeds buffer with runtime shape (#311)

* Define `SeedArray` container

* Refactor use of seed array

Author: sbaldu

include/CLUEstering/core/Clusterer.hpp

@@ -13,8 +13,9 @@
 #include "CLUEstering/data_structures/AssociationMap.hpp"
 #include "CLUEstering/data_structures/PointsHost.hpp"
 #include "CLUEstering/data_structures/PointsDevice.hpp"
-#include "CLUEstering/data_structures/internal/Tiles.hpp"
 #include "CLUEstering/data_structures/internal/Followers.hpp"
+#include "CLUEstering/data_structures/internal/SeedArray.hpp"
+#include "CLUEstering/data_structures/internal/Tiles.hpp"
 
 #include <array>
 #include <concepts>
@@ -48,18 +49,13 @@ namespace clue {
     std::array<uint8_t, Ndim> m_wrappedCoordinates;
 
     std::optional<TilesDevice> m_tiles;
-    std::optional<clue::device_buffer<Device, VecArray<int32_t, reserve>>> m_seeds;
+    std::optional<clue::internal::SeedArray<>> m_seeds;
     std::optional<FollowersDevice> m_followers;
 
     void setup(Queue& queue, const PointsHost& h_points, PointsDevice& dev_points) {
       detail::setup_tiles(queue, m_tiles, h_points, m_pointsPerTile, m_wrappedCoordinates);
       detail::setup_followers(queue, m_followers, h_points.size());
       clue::copyToDevice(queue, dev_points, h_points);
-
-      if (!m_seeds.has_value()) {
-        m_seeds = clue::make_device_buffer<VecArray<int32_t, reserve>>(queue);
-      }
-      alpaka::memset(queue, *m_seeds, 0x00);
     }
 
     template <concepts::convolutional_kernel Kernel>

include/CLUEstering/core/detail/Clusterer.hpp

@@ -7,10 +7,14 @@
 #include "CLUEstering/core/detail/ClusteringKernels.hpp"
 #include "CLUEstering/core/detail/ComputeTiles.hpp"
 #include "CLUEstering/core/detail/defines.hpp"
+#include "CLUEstering/core/detail/SetupFollowers.hpp"
+#include "CLUEstering/core/detail/SetupSeeds.hpp"
+#include "CLUEstering/core/detail/SetupTiles.hpp"
 #include "CLUEstering/data_structures/PointsHost.hpp"
 #include "CLUEstering/data_structures/PointsDevice.hpp"
-#include "CLUEstering/data_structures/internal/Tiles.hpp"
 #include "CLUEstering/data_structures/internal/Followers.hpp"
+#include "CLUEstering/data_structures/internal/SeedArray.hpp"
+#include "CLUEstering/data_structures/internal/Tiles.hpp"
 #include "CLUEstering/utils/get_clusters.hpp"
 
 #include <alpaka/mem/view/Traits.hpp>
@@ -132,10 +136,6 @@ namespace clue {
                                              std::size_t block_size) {
     detail::setup_tiles(queue, m_tiles, dev_points, m_pointsPerTile, m_wrappedCoordinates);
     detail::setup_followers(queue, m_followers, dev_points.size());
-    if (!m_seeds.has_value()) {
-      m_seeds = clue::make_device_buffer<VecArray<int32_t, reserve>>(queue);
-    }
-    alpaka::memset(queue, *m_seeds, 0x00);
     make_clusters_impl(dev_points, kernel, queue, block_size);
     alpaka::wait(queue);
   }
@@ -178,11 +178,13 @@ namespace clue {
 
     detail::computeLocalDensity<Acc>(
         queue, work_division, m_tiles->view(), dev_points.view(), kernel, m_dc, n_points);
+    auto seed_candidates = 0ul;
     detail::computeNearestHighers<Acc>(
-        queue, work_division, m_tiles->view(), dev_points.view(), m_dm, n_points);
+        queue, work_division, m_tiles->view(), dev_points.view(), m_dm, seed_candidates, n_points);
+    detail::setup_seeds(queue, m_seeds, seed_candidates);
     detail::findClusterSeeds<Acc>(queue,
                                   work_division,
-                                  m_seeds->data(),
+                                  m_seeds.value(),
                                   m_tiles->view(),
                                   dev_points.view(),
                                   m_seed_dc,
@@ -192,7 +194,7 @@ namespace clue {
     m_followers->template fill<Acc>(queue, dev_points);
 
     detail::assignPointsToClusters<Acc>(
-        queue, block_size, m_seeds->data(), m_followers->view(), dev_points.view());
+        queue, block_size, m_seeds.value(), m_followers->view(), dev_points.view());
 
     clue::copyToHost(queue, h_points, dev_points);
     h_points.mark_clustered();
@@ -213,11 +215,13 @@ namespace clue {
 
     detail::computeLocalDensity<Acc>(
         queue, work_division, m_tiles->view(), dev_points.view(), kernel, m_dc, n_points);
+    auto seed_candidates = 0ul;
     detail::computeNearestHighers<Acc>(
-        queue, work_division, m_tiles->view(), dev_points.view(), m_dm, n_points);
+        queue, work_division, m_tiles->view(), dev_points.view(), m_dm, seed_candidates, n_points);
+    detail::setup_seeds(queue, m_seeds, seed_candidates);
     detail::findClusterSeeds<Acc>(queue,
                                   work_division,
-                                  m_seeds->data(),
+                                  m_seeds.value(),
                                   m_tiles->view(),
                                   dev_points.view(),
                                   m_seed_dc,
@@ -227,7 +231,7 @@ namespace clue {
     m_followers->template fill<Acc>(queue, dev_points);
 
     detail::assignPointsToClusters<Acc>(
-        queue, block_size, m_seeds->data(), m_followers->view(), dev_points.view());
+        queue, block_size, m_seeds.value(), m_followers->view(), dev_points.view());
 
     alpaka::wait(queue);
     dev_points.mark_clustered();

include/CLUEstering/core/detail/ClusteringKernels.hpp

@@ -4,10 +4,11 @@
 #include "CLUEstering/core/ConvolutionalKernel.hpp"
 #include "CLUEstering/core/DistanceParameter.hpp"
 #include "CLUEstering/data_structures/PointsDevice.hpp"
+#include "CLUEstering/data_structures/internal/Followers.hpp"
+#include "CLUEstering/data_structures/internal/SearchBox.hpp"
+#include "CLUEstering/data_structures/internal/SeedArray.hpp"
 #include "CLUEstering/data_structures/internal/TilesView.hpp"
 #include "CLUEstering/data_structures/internal/VecArray.hpp"
-#include "CLUEstering/data_structures/internal/SearchBox.hpp"
-#include "CLUEstering/data_structures/internal/Followers.hpp"
 #include "CLUEstering/detail/make_array.hpp"
 #include "CLUEstering/internal/alpaka/work_division.hpp"
 #include "CLUEstering/internal/math/math.hpp"
@@ -154,6 +155,7 @@ namespace clue::detail {
                                   internal::TilesView<Ndim> dev_tiles,
                                   PointsView<Ndim> dev_points,
                                   DistanceParameter<Ndim> dm,
+                                  std::size_t* seed_candidates,
                                   int32_t n_points) const {
       for (auto i : alpaka::uniformElements(acc, n_points)) {
         float delta_i = std::numeric_limits<float>::max();
@@ -184,14 +186,17 @@ namespace clue::detail {
                                                        i);
 
         dev_points.nearest_higher[i] = nh_i;
+        if (nh_i == -1) {
+          alpaka::atomicAdd(acc, seed_candidates, 1ul);
+        }
       }
     }
   };
 
   struct KernelFindClusters {
     template <typename TAcc, std::size_t Ndim>
     ALPAKA_FN_ACC void operator()(const TAcc& acc,
-                                  VecArray<int32_t, reserve>* seeds,
+                                  clue::internal::SeedArrayView seeds,
                                   internal::TilesView<Ndim> tiles,
                                   PointsView<Ndim> dev_points,
                                   DistanceParameter<Ndim> seed_dc,
@@ -211,7 +216,7 @@ namespace clue::detail {
         if (is_seed) {
           dev_points.is_seed[i] = 1;
           dev_points.nearest_higher[i] = -1;
-          seeds->push_back(acc, i);
+          seeds.push_back(acc, i);
         } else {
           dev_points.is_seed[i] = 0;
         }
@@ -222,16 +227,15 @@ namespace clue::detail {
   struct KernelAssignClusters {
     template <typename TAcc, std::size_t Ndim>
     ALPAKA_FN_ACC void operator()(const TAcc& acc,
-                                  VecArray<int32_t, reserve>* seeds,
+                                  clue::internal::SeedArrayView seeds,
                                   clue::FollowersView followers,
                                   PointsView<Ndim> dev_points) const {
-      const auto& seeds_0 = *seeds;
-      const auto n_seeds = seeds_0.size();
+      const auto n_seeds = seeds.size();
       for (auto idx_cls : alpaka::uniformElements(acc, n_seeds)) {
         int local_stack[256] = {-1};
         int local_stack_size = 0;
 
-        int idx_this_seed = seeds_0[idx_cls];
+        int idx_this_seed = seeds[idx_cls];
         dev_points.cluster_index[idx_this_seed] = idx_cls;
         local_stack[local_stack_size] = idx_this_seed;
         ++local_stack_size;
@@ -279,33 +283,52 @@ namespace clue::detail {
                                     internal::TilesView<Ndim>& tiles,
                                     PointsView<Ndim>& dev_points,
                                     const DistanceParameter<Ndim>& dm,
+                                    std::size_t& seed_candidates,
                                     int32_t size) {
-    alpaka::exec<TAcc>(
-        queue, work_division, KernelCalculateNearestHigher{}, tiles, dev_points, dm, size);
+    auto d_seed_candidates = clue::make_device_buffer<std::size_t>(queue);
+    alpaka::memset(queue, d_seed_candidates, 0u);
+    alpaka::exec<TAcc>(queue,
+                       work_division,
+                       KernelCalculateNearestHigher{},
+                       tiles,
+                       dev_points,
+                       dm,
+                       d_seed_candidates.data(),
+                       size);
+    alpaka::memcpy(queue, clue::make_host_view(seed_candidates), d_seed_candidates);
+    alpaka::wait(queue);
   }
 
   template <concepts::accelerator TAcc, concepts::queue TQueue, std::size_t Ndim>
   inline void findClusterSeeds(TQueue& queue,
                                const WorkDiv& work_division,
-                               VecArray<int32_t, reserve>* seeds,
+                               clue::internal::SeedArray<>& seeds,
                                internal::TilesView<Ndim>& tiles,
                                PointsView<Ndim>& dev_points,
                                const DistanceParameter<Ndim>& seed_dc,
                                float rhoc,
                                int32_t size) {
-    alpaka::exec<TAcc>(
-        queue, work_division, KernelFindClusters{}, seeds, tiles, dev_points, seed_dc, rhoc, size);
+    alpaka::exec<TAcc>(queue,
+                       work_division,
+                       KernelFindClusters{},
+                       seeds.view(),
+                       tiles,
+                       dev_points,
+                       seed_dc,
+                       rhoc,
+                       size);
   }
 
   template <concepts::accelerator TAcc, concepts::queue TQueue, std::size_t Ndim>
   inline void assignPointsToClusters(TQueue& queue,
                                      std::size_t block_size,
-                                     VecArray<int32_t, reserve>* seeds,
+                                     clue::internal::SeedArray<>& seeds,
                                      clue::FollowersView followers,
                                      PointsView<Ndim> dev_points) {
-    const Idx grid_size = clue::divide_up_by(reserve, block_size);
+    const Idx grid_size = clue::divide_up_by(seeds.size(queue), block_size);
     const auto work_division = clue::make_workdiv<TAcc>(grid_size, block_size);
-    alpaka::exec<TAcc>(queue, work_division, KernelAssignClusters{}, seeds, followers, dev_points);
+    alpaka::exec<TAcc>(
+        queue, work_division, KernelAssignClusters{}, seeds.view(), followers, dev_points);
   }
 
 }  // namespace clue::detail

include/CLUEstering/core/detail/SetupSeeds.hpp

@@ -0,0 +1,27 @@
+
+#pragma once
+
+#include "CLUEstering/detail/concepts.hpp"
+#include "CLUEstering/data_structures/internal/SeedArray.hpp"
+#include <cstddef>
+#include <optional>
+
+namespace clue::detail {
+
+  template <concepts::queue TQueue,
+            concepts::device TDev = decltype(alpaka::getDev(std::declval<TQueue>()))>
+  inline void setup_seeds(TQueue& queue,
+                          std::optional<clue::internal::SeedArray<TDev>>& seeds,
+                          std::size_t seed_candidates) {
+    if (!seeds.has_value()) {
+      seeds = clue::internal::SeedArray<TDev>(queue, seed_candidates);
+    }
+    if (seeds->capacity() < seed_candidates) {
+      seeds = clue::internal::SeedArray<TDev>(queue, seed_candidates);
+    } else {
+      seeds->reset(queue);
+    }
+    alpaka::wait(queue);
+  }
+
+}  // namespace clue::detail

include/CLUEstering/data_structures/internal/SeedArray.hpp

@@ -0,0 +1,85 @@
+
+#pragma once
+
+#include "CLUEstering/detail/concepts.hpp"
+#include "CLUEstering/internal/alpaka/memory.hpp"
+#include <alpaka/alpaka.hpp>
+#include <cstddef>
+#include <cstdint>
+
+namespace clue::internal {
+
+  class SeedArrayView {
+    int32_t* m_data;
+    std::size_t* m_size;
+    std::size_t m_capacity;
+
+  public:
+    ALPAKA_FN_HOST_ACC constexpr SeedArrayView(int32_t* data,
+                                               std::size_t* size,
+                                               std::size_t capacity)
+        : m_data{data}, m_size{size}, m_capacity{capacity} {}
+
+    ALPAKA_FN_ACC constexpr auto& operator[](std::size_t index) { return m_data[index]; }
+    ALPAKA_FN_ACC constexpr const auto& operator[](std::size_t index) const {
+      return m_data[index];
+    }
+
+    ALPAKA_FN_ACC constexpr auto size() const {
+      // NOTE: not thread-safe
+      // Could restrict this to KernelAssignClusters, but maybe that's an overkill
+      return *m_size;
+    }
+
+    template <clue::concepts::accelerator TAcc>
+    ALPAKA_FN_ACC constexpr void push_back(const TAcc& acc, int32_t value) {
+      auto prev = alpaka::atomicAdd(acc, m_size, 1ul);
+      if (prev < m_capacity) {
+        m_data[prev] = value;
+      } else {
+        alpaka::atomicSub(acc, m_size, 1ul);
+      }
+    }
+  };
+
+  template <clue::concepts::device TDev = clue::Device>
+  class SeedArray {
+  private:
+    clue::device_buffer<TDev, int32_t[]> m_buffer;
+    clue::device_buffer<TDev, std::size_t> m_dsize;
+    std::optional<std::size_t> m_size;
+    std::size_t m_capacity;
+    SeedArrayView m_view;
+
+  public:
+    template <clue::concepts::queue TQueue>
+    SeedArray(TQueue& queue, std::size_t size)
+        : m_buffer{clue::make_device_buffer<int32_t[]>(queue, size)},
+          m_dsize{clue::make_device_buffer<std::size_t>(queue)},
+          m_size{std::nullopt},
+          m_capacity{size},
+          m_view{m_buffer.data(), m_dsize.data(), m_capacity} {}
+
+    ALPAKA_FN_HOST constexpr auto capacity() const { return m_capacity; }
+
+    template <clue::concepts::queue TQueue>
+    ALPAKA_FN_HOST auto size(TQueue& queue) {
+      if (!m_size.has_value()) {
+        m_size = std::make_optional<std::size_t>();
+        alpaka::memcpy(queue, clue::make_host_view(*m_size), m_dsize);
+        alpaka::wait(queue);
+      }
+      return *m_size;
+    }
+
+    template <clue::concepts::queue TQueue>
+    ALPAKA_FN_HOST auto reset(TQueue& queue) {
+      m_size = std::nullopt;
+      alpaka::memset(queue, m_dsize, 0u);
+    }
+
+    ALPAKA_FN_HOST const auto& view() const { return m_view; }
+    ALPAKA_FN_HOST auto& view() { return m_view; }
+  };
+
+}  // namespace clue::internal