From ce4020b26940a769c604358ee2b00f9c63f2d61e Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Fri, 16 Feb 2024 13:09:11 +0100
Subject: [PATCH] Massaging various files

---
 arbor/backends/event.hpp            |  4 +---
 arbor/backends/gpu/event_stream.hpp |  3 ++-
 arbor/backends/gpu/shared_state.cu  | 20 +++++++-------------
 arbor/fvm_lowered_cell_impl.hpp     |  8 +++++---
 4 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/arbor/backends/event.hpp b/arbor/backends/event.hpp
index dbc9f9f7cf..81df1b38ed 100644
--- a/arbor/backends/event.hpp
+++ b/arbor/backends/event.hpp
@@ -51,11 +51,9 @@ struct has_event_index<deliverable_event> : public std::true_type {};
 
 // Subset of event information required for mechanism delivery.
 struct deliverable_event_data {
-    // cell_local_size_type mech_id;    // same as target_handle::mech_id
     cell_local_size_type mech_index; // same as target_handle::mech_index
     float weight;
-    deliverable_event_data(cell_local_size_type id, cell_local_size_type idx, float w):
-        // mech_id(id),
+    deliverable_event_data(cell_local_size_type idx, float w):
         mech_index(idx),
         weight(w) {}
     ARB_SERDES_ENABLE(deliverable_event_data,
diff --git a/arbor/backends/gpu/event_stream.hpp b/arbor/backends/gpu/event_stream.hpp
index effac12015..4710014180 100644
--- a/arbor/backends/gpu/event_stream.hpp
+++ b/arbor/backends/gpu/event_stream.hpp
@@ -56,7 +56,8 @@ struct event_stream: public event_stream_base<Event> {
 
         // assign, copy to device (and potentially sort) the event data in parallel
         arb_assert(thread_pool_);
-        threading::parallel_for::apply(0, staged.size(), thread_pool_.get(),
+        arb_assert(base::ev_spans_.size() == staged.size() + 1);
+        threading::parallel_for::apply(0, base::ev_spans_.size() -1, thread_pool_.get(),
             [this, &staged](size_type i) {
                 const auto beg = base::ev_spans_[i];
                 const auto end = base::ev_spans_[i + 1];
diff --git a/arbor/backends/gpu/shared_state.cu b/arbor/backends/gpu/shared_state.cu
index 403067d70f..db22c03b0e 100644
--- a/arbor/backends/gpu/shared_state.cu
+++ b/arbor/backends/gpu/shared_state.cu
@@ -1,7 +1,5 @@
 // GPU kernels and wrappers for shared state methods.
 
-#include <cstdint>
-
 #include <backends/event.hpp>
 #include <backends/event_stream_state.hpp>
 
@@ -24,13 +22,11 @@ __global__ void add_scalar(unsigned n,
     }
 }
 
-__global__ void take_samples_impl(
-    const raw_probe_info* __restrict__ const begin_marked,
-    const raw_probe_info* __restrict__ const end_marked,
-    const arb_value_type time,
-    arb_value_type* __restrict__ const sample_time,
-    arb_value_type* __restrict__ const sample_value)
-{
+__global__ void take_samples_impl(const raw_probe_info* __restrict__ const begin_marked,
+                                  const raw_probe_info* __restrict__ const end_marked,
+                                  const arb_value_type time,
+                                  arb_value_type* __restrict__ const sample_time,
+                                  arb_value_type* __restrict__ const sample_value) {
     const unsigned i = threadIdx.x+blockIdx.x*blockDim.x;
     const unsigned nsamples = end_marked - begin_marked;
     if (i<nsamples) {
@@ -46,10 +42,8 @@ void add_scalar(std::size_t n, arb_value_type* data, arb_value_type v) {
     launch_1d(n, 128, kernel::add_scalar<arb_value_type>, n, data, v);
 }
 
-void take_samples_impl(
-    const event_stream_state<raw_probe_info>& s,
-    const arb_value_type& time, arb_value_type* sample_time, arb_value_type* sample_value)
-{
+void take_samples_impl(const event_stream_state<raw_probe_info>& s,
+                       const arb_value_type& time, arb_value_type* sample_time, arb_value_type* sample_value) {
     launch_1d(s.size(), 128, kernel::take_samples_impl, s.begin_marked, s.end_marked, time, sample_time, sample_value);
 }
 
diff --git a/arbor/fvm_lowered_cell_impl.hpp b/arbor/fvm_lowered_cell_impl.hpp
index f96e71b8c6..cc3867b6b4 100644
--- a/arbor/fvm_lowered_cell_impl.hpp
+++ b/arbor/fvm_lowered_cell_impl.hpp
@@ -578,12 +578,14 @@ fvm_lowered_cell_impl<Backend>::initialize(const std::vector<cell_gid_type>& gid
         }
     }
 
+    add_probes(gids, cells, rec, D, mechptr_by_name, mech_data, target_handles_, fvm_info.probe_map);
+
     // Create lookup structure for target ids.
     util::make_partition(target_handle_divisions_,
-        util::transform_view(gids, [&](cell_gid_type i) { return fvm_info.num_targets[i]; }));
-
-    add_probes(gids, cells, rec, D, mechptr_by_name, mech_data, target_handles_, fvm_info.probe_map);
+        util::transform_view(gids,
+                             [&](cell_gid_type i) { return fvm_info.num_targets[i]; }));
 
+    
     reset();
     return fvm_info;
 }