Optimise Memory Consumption during Setup

arbor/CMakeLists.txt

@@ -27,7 +27,7 @@ set(arbor_sources
     io/serialize_hex.cpp
     label_resolution.cpp
     lif_cell_group.cpp
-    mc_cell_group.cpp
+    cable_cell_group.cpp
     mechcat.cpp
     mechinfo.cpp
     memory/gpu_wrappers.cpp

arbor/benchmark_cell_group.cpp

@@ -15,6 +15,12 @@
 
 namespace arb {
 
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const benchmark_cell& c) {
+    src.add_cell();
+    src.add_label(c.source, {0, 1});
+    return 1;
+}
+
 benchmark_cell_group::benchmark_cell_group(const std::vector<cell_gid_type>& gids,
                                            const recipe& rec,
                                            cell_label_range& cg_sources,
@@ -29,14 +35,11 @@ benchmark_cell_group::benchmark_cell_group(const std::vector<cell_gid_type>& gid
 
     cells_.reserve(gids_.size());
     for (auto gid: gids_) {
-        cells_.push_back(util::any_cast<benchmark_cell>(rec.get_cell_description(gid)));
-    }
-
-    for (const auto& c: cells_) {
-        cg_sources.add_cell();
+        auto cell = util::any_cast<benchmark_cell>(rec.get_cell_description(gid));
+        cells_.push_back(cell);
         cg_targets.add_cell();
-        cg_sources.add_label(c.source, {0, 1});
-        cg_targets.add_label(c.target, {0, 1});
+        cg_targets.add_label(cell.target, {0, 1});
+        get_sources(cg_sources, cell);
     }
 
     benchmark_cell_group::reset();

arbor/benchmark_cell_group.hpp

@@ -40,5 +40,7 @@ class benchmark_cell_group: public cell_group {
     std::vector<cell_gid_type> gids_;
 };
 
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const benchmark_cell& c);
+
 } // namespace arb
 

arbor/mc_cell_group.cpp → arbor/cable_cell_group.cpp

@@ -16,7 +16,7 @@
 #include "event_binner.hpp"
 #include "fvm_lowered_cell.hpp"
 #include "label_resolution.hpp"
-#include "mc_cell_group.hpp"
+#include "cable_cell_group.hpp"
 #include "profile/profiler_macro.hpp"
 #include "sampler_map.hpp"
 #include "util/filter.hpp"
@@ -30,15 +30,25 @@ namespace arb {
 ARB_DEFINE_LEXICOGRAPHIC_ORDERING(arb::target_handle,(a.mech_id,a.mech_index,a.intdom_index),(b.mech_id,b.mech_index,b.intdom_index))
 ARB_DEFINE_LEXICOGRAPHIC_ORDERING(arb::deliverable_event,(a.time,a.handle,a.weight),(b.time,b.handle,b.weight))
 
-mc_cell_group::mc_cell_group(const std::vector<cell_gid_type>& gids,
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const cable_cell& c) {
+    src.add_cell();
+    cell_size_type count = 0;
+    for (const auto& [label, range]: c.detector_ranges()) {
+        src.add_label(label, range);
+        count += range.end - range.begin;
+    }
+    return count;
+}
+
+cable_cell_group::cable_cell_group(const std::vector<cell_gid_type>& gids,
                              const recipe& rec,
                              cell_label_range& cg_sources,
                              cell_label_range& cg_targets,
                              fvm_lowered_cell_ptr lowered):
     gids_(gids), lowered_(std::move(lowered))
 {
     // Default to no binning of events
-    mc_cell_group::set_binning_policy(binning_kind::none, 0);
+    cable_cell_group::set_binning_policy(binning_kind::none, 0);
 
     // Build lookup table for gid to local index.
     for (auto i: util::count_along(gids_)) {
@@ -70,7 +80,7 @@ mc_cell_group::mc_cell_group(const std::vector<cell_gid_type>& gids,
     spike_sources_.shrink_to_fit();
 }
 
-void mc_cell_group::reset() {
+void cable_cell_group::reset() {
     spikes_.clear();
 
     for (auto &entry: sampler_map_) {
@@ -84,7 +94,7 @@ void mc_cell_group::reset() {
     lowered_->reset();
 }
 
-void mc_cell_group::set_binning_policy(binning_kind policy, time_type bin_interval) {
+void cable_cell_group::set_binning_policy(binning_kind policy, time_type bin_interval) {
     binners_.clear();
     binners_.resize(gids_.size(), event_binner(policy, bin_interval));
 }
@@ -388,7 +398,7 @@ void run_samples(
     std::visit([&](auto& x) {run_samples(x, sc, raw_times, raw_samples, sample_records, scratch); }, sc.pdata_ptr->info);
 }
 
-void mc_cell_group::advance(epoch ep, time_type dt, const event_lane_subrange& event_lanes) {
+void cable_cell_group::advance(epoch ep, time_type dt, const event_lane_subrange& event_lanes) {
     time_type tstart = lowered_->time();
 
     // Bin and collate deliverable events from event lanes.
@@ -554,7 +564,7 @@ void mc_cell_group::advance(epoch ep, time_type dt, const event_lane_subrange& e
     }
 }
 
-void mc_cell_group::add_sampler(sampler_association_handle h, cell_member_predicate probeset_ids,
+void cable_cell_group::add_sampler(sampler_association_handle h, cell_member_predicate probeset_ids,
                                 schedule sched, sampler_function fn, sampling_policy policy)
 {
     std::lock_guard<std::mutex> guard(sampler_mex_);
@@ -568,17 +578,17 @@ void mc_cell_group::add_sampler(sampler_association_handle h, cell_member_predic
     }
 }
 
-void mc_cell_group::remove_sampler(sampler_association_handle h) {
+void cable_cell_group::remove_sampler(sampler_association_handle h) {
     std::lock_guard<std::mutex> guard(sampler_mex_);
     sampler_map_.erase(h);
 }
 
-void mc_cell_group::remove_all_samplers() {
+void cable_cell_group::remove_all_samplers() {
     std::lock_guard<std::mutex> guard(sampler_mex_);
     sampler_map_.clear();
 }
 
-std::vector<probe_metadata> mc_cell_group::get_probe_metadata(cell_member_type probeset_id) const {
+std::vector<probe_metadata> cable_cell_group::get_probe_metadata(cell_member_type probeset_id) const {
     // Probe associations are fixed after construction, so we do not need to grab the mutex.
 
     std::optional<probe_tag> maybe_tag = util::value_by_key(probe_map_.tag, probeset_id);

arbor/mc_cell_group.hpp → arbor/cable_cell_group.hpp

@@ -12,6 +12,7 @@
 #include <arbor/recipe.hpp>
 #include <arbor/sampling.hpp>
 #include <arbor/spike.hpp>
+#include <arbor/cable_cell.hpp>
 
 #include "backends/event.hpp"
 #include "cell_group.hpp"
@@ -24,11 +25,11 @@
 
 namespace arb {
 
-class ARB_ARBOR_API mc_cell_group: public cell_group {
+class ARB_ARBOR_API cable_cell_group: public cell_group {
 public:
-    mc_cell_group() = default;
+    cable_cell_group() = default;
 
-    mc_cell_group(const std::vector<cell_gid_type>& gids,
+    cable_cell_group(const std::vector<cell_gid_type>& gids,
                   const recipe& rec,
                   cell_label_range& cg_sources,
                   cell_label_range& cg_targets,
@@ -102,4 +103,6 @@ class ARB_ARBOR_API mc_cell_group: public cell_group {
     std::vector<std::size_t> target_handle_divisions_;
 };
 
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const cable_cell& c);
+
 } // namespace arb

arbor/cell_group_factory.cpp

@@ -9,7 +9,7 @@
 #include "execution_context.hpp"
 #include "fvm_lowered_cell.hpp"
 #include "lif_cell_group.hpp"
-#include "mc_cell_group.hpp"
+#include "cable_cell_group.hpp"
 #include "spike_source_cell_group.hpp"
 
 namespace arb {
@@ -27,7 +27,7 @@ ARB_ARBOR_API cell_group_factory cell_kind_implementation(
     switch (ck) {
     case cell_kind::cable:
         return [bk, ctx, seed](const gid_vector& gids, const recipe& rec, cell_label_range& cg_sources, cell_label_range& cg_targets) {
-            return make_cell_group<mc_cell_group>(gids, rec, cg_sources, cg_targets, make_fvm_lowered_cell(bk, ctx, seed));
+            return make_cell_group<cable_cell_group>(gids, rec, cg_sources, cg_targets, make_fvm_lowered_cell(bk, ctx, seed));
         };
 
     case cell_kind::spike_source:

arbor/communication/communicator.cpp

@@ -1,5 +1,6 @@
 #include <utility>
 #include <vector>
+#include <any>
 
 #include <arbor/assert.hpp>
 #include <arbor/common_types.hpp>
@@ -17,6 +18,10 @@
 #include "util/partition.hpp"
 #include "util/rangeutil.hpp"
 #include "util/span.hpp"
+#include "spike_source_cell_group.hpp"
+#include "cable_cell_group.hpp"
+#include "benchmark_cell_group.hpp"
+#include "lif_cell_group.hpp"
 
 #include "communication/communicator.hpp"
 
@@ -31,9 +36,31 @@ communicator::communicator(const recipe& rec,
                                                      distributed_{ctx.distributed},
                                                      thread_pool_{ctx.thread_pool} {}
 
-void communicator::update_connections(const connectivity& rec,
+// This is a bit nasty, as we basically reimplement things from all cell kinds ...
+cell_label_range get_sources(cell_gid_type gid, const recipe& rec) {
+    auto cell = rec.get_cell_description(gid);
+    auto kind = rec.get_cell_kind(gid);
+    cell_label_range result;
+    if (kind == cell_kind::lif) {
+        get_sources(result, util::any_cast<lif_cell>(cell));
+    }
+    else if (kind == cell_kind::spike_source) {
+        get_sources(result, util::any_cast<spike_source_cell>(cell));
+    }
+    else if (kind == cell_kind::benchmark) {
+        get_sources(result, util::any_cast<benchmark_cell>(cell));
+    }
+    else if (kind == cell_kind::cable) {
+        get_sources(result, util::any_cast<cable_cell>(cell));
+    }
+    else {
+        throw arbor_internal_error("Unknown cell kind");
+    }
+    return result;
+}    
+
+void communicator::update_connections(const recipe& rec,
                                       const domain_decomposition& dom_dec,
-                                      const label_resolution_map& source_resolution_map,
                                       const label_resolution_map& target_resolution_map) {
     // Forget all lingering information
     connections_.clear();
@@ -102,11 +129,20 @@ void communicator::update_connections(const connectivity& rec,
     auto offsets = connection_part_; // Copy, as we use this as the list of current target indices to write into
     auto src_domain = src_domains.begin();
     auto target_resolver = resolver(&target_resolution_map);
+    auto source_labels = std::unordered_map<cell_gid_type, std::unique_ptr<label_resolution_map>>{};
     for (const auto& cell: gid_infos) {
         auto index = cell.index_on_domain;
-        auto source_resolver = resolver(&source_resolution_map);
+        auto source_resolvers = std::unordered_map<cell_gid_type, resolver>{};
         for (const auto& c: cell.conns) {
-            auto src_lid = source_resolver.resolve(c.source);
+            auto sgid = c.source.gid;
+            if (!source_labels.count(sgid)) {
+                auto sources = cell_labels_and_gids{get_sources(sgid, rec), {sgid}};
+                source_labels.emplace(sgid, std::make_unique<label_resolution_map>(sources));
+            }
+            if (!source_resolvers.count(sgid)) {
+                source_resolvers.emplace(sgid, source_labels.at(sgid).get());
+            }
+            auto src_lid = source_resolvers.at(sgid).resolve(c.source);
             auto tgt_lid = target_resolver.resolve({cell.gid, c.dest});
             auto offset  = offsets[*src_domain]++;
             ++src_domain;

arbor/communication/communicator.hpp

@@ -68,9 +68,8 @@ class ARB_ARBOR_API communicator {
 
     void reset();
 
-    void update_connections(const connectivity& rec,
+    void update_connections(const recipe& rec,
                             const domain_decomposition& dom_dec,
-                            const label_resolution_map& source_resolution_map,
                             const label_resolution_map& target_resolution_map);
 
 private:

arbor/fvm_lowered_cell_impl.hpp

@@ -35,6 +35,7 @@
 #include "util/rangeutil.hpp"
 #include "util/strprintf.hpp"
 #include "util/transform.hpp"
+#include "cable_cell_group.hpp"
 
 namespace arb {
 
@@ -418,18 +419,12 @@ fvm_initialization_data fvm_lowered_cell_impl<Backend>::initialize(
         auto gid = gids[i];
         const auto& c = cells[i];
 
-        fvm_info.source_data.add_cell();
         fvm_info.target_data.add_cell();
         fvm_info.gap_junction_data.add_cell();
 
-        unsigned count = 0;
-        for (const auto& [label, range]: c.detector_ranges()) {
-            fvm_info.source_data.add_label(label, range);
-            count+=(range.end - range.begin);
-        }
-        fvm_info.num_sources[gid] = count;
+        fvm_info.num_sources[gid] = get_sources(fvm_info.source_data, c);
 
-        count = 0;
+        unsigned count = 0;
         for (const auto& [label, range]: c.synapse_ranges()) {
             fvm_info.target_data.add_label(label, range);
             count+=(range.end - range.begin);

arbor/include/arbor/simulation.hpp

@@ -45,7 +45,7 @@ class ARB_ARBOR_API simulation {
 
     static simulation_builder create(recipe const &);
 
-    void update(const connectivity& rec);
+    void update(const recipe& rec);
 
     void reset();
 

arbor/label_resolution.hpp

@@ -114,7 +114,6 @@ struct ARB_ARBOR_API resolver {
     resolver(const label_resolution_map* label_map): label_map_(label_map) {}
     cell_lid_type resolve(const cell_global_label_type& iden);
 
-private:
     using state_variant = std::variant<round_robin_state, round_robin_halt_state, assert_univalent_state>;
 
     state_variant construct_state(lid_selection_policy pol);

arbor/lif_cell_group.cpp

@@ -8,7 +8,13 @@
 #include "util/filter.hpp"
 #include "util/maputil.hpp"
 
-using namespace arb;
+namespace arb {
+
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const lif_cell& c) {
+    src.add_cell();
+    src.add_label(c.source, {0, 1});
+    return 1;
+}
 
 // Constructor containing gid of first cell in a group and a container of all cells.
 lif_cell_group::lif_cell_group(const std::vector<cell_gid_type>& gids,
@@ -246,3 +252,5 @@ std::vector<probe_metadata> lif_cell_group::get_probe_metadata(cell_member_type
         return {};
     }
 }
+
+} // namespace arb

arbor/lif_cell_group.hpp

@@ -78,4 +78,6 @@ class ARB_ARBOR_API lif_cell_group: public cell_group {
     std::unordered_map<cell_member_type, lif_probe_info> probes_;
 };
 
+cell_size_type ARB_ARBOR_API get_sources(cell_label_range& src, const lif_cell& c);
+
 } // namespace arb