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();
@@ -94,6 +121,17 @@ void communicator::update_connections(const connectivity& rec,
         }
     }
 
+    struct label_map {
+        std::unique_ptr<label_resolution_map> map;
+        resolver res;
+
+        label_map(cell_gid_type gid, const recipe& rec):
+            map{std::make_unique<label_resolution_map>(cell_labels_and_gids{get_sources(gid, rec), {gid}})},
+            res{map.get()}
+                {}
+        cell_lid_type resolve(const cell_global_label_type& lbl) { return res.resolve(lbl); }
+    };
+
     // Construct the connections.
     // The loop above gave the information required to construct in place
     // the connections as partitioned by the domain of their source gid.
@@ -104,9 +142,11 @@ void communicator::update_connections(const connectivity& rec,
     auto target_resolver = resolver(&target_resolution_map);
     for (const auto& cell: gid_infos) {
         auto index = cell.index_on_domain;
-        auto source_resolver = resolver(&source_resolution_map);
+        auto sources = std::unordered_map<cell_gid_type, label_map>{};
         for (const auto& c: cell.conns) {
-            auto src_lid = source_resolver.resolve(c.source);
+            auto sgid = c.source.gid;
+            if (!sources.count(sgid)) sources.emplace(sgid, label_map{sgid, rec});
+            auto src_lid = sources.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/profile/profiler.hpp

@@ -13,6 +13,7 @@ namespace arb {
 
 namespace profile {
 
+
 // type used for region identifiers
 using region_id_type = std::size_t;
 
@@ -27,6 +28,10 @@ struct profile {
     // the accumulated time spent in each region.
     std::vector<double> times;
 
+    // Maximum Resident Set Size up to this region's end.
+    std::vector<long> max_rss;
+    std::vector<long> fst_rss;
+
     // the number of threads for which profiling information was recorded.
     std::size_t num_threads;
 

arbor/include/arbor/recipe.hpp

@@ -61,41 +61,8 @@ struct gap_junction_connection {
         peer(std::move(peer)), local(std::move(local)), weight(g) {}
 };
 
-struct ARB_ARBOR_API has_gap_junctions {
-    virtual std::vector<gap_junction_connection> gap_junctions_on(cell_gid_type) const {
-        return {};
-    }
-    virtual ~has_gap_junctions() {}
-};
-
-struct ARB_ARBOR_API has_synapses {
-    virtual std::vector<cell_connection> connections_on(cell_gid_type) const {
-        return {};
-    }
-    virtual ~has_synapses() {}
-};
-
-struct ARB_ARBOR_API has_probes {
-    virtual std::vector<probe_info> get_probes(cell_gid_type gid) const {
-        return {};
-    }
-    virtual ~has_probes() {}
-};
-
-struct ARB_ARBOR_API has_generators {
-    virtual std::vector<event_generator> event_generators(cell_gid_type) const {
-        return {};
-    }
-    virtual ~has_generators() {}
-};
-
-// Toppings allow updating a simulation
-struct ARB_ARBOR_API connectivity: public has_synapses, has_generators {
-    virtual ~connectivity() {}
-};
-
 // Recipes allow building a simulation by lazy queries
-struct ARB_ARBOR_API recipe: public has_gap_junctions, has_probes, connectivity {
+struct ARB_ARBOR_API recipe {
     // number of cells to build
     virtual cell_size_type num_cells() const = 0;
     // Cell description type will be specific to cell kind of cell with given gid.
@@ -104,6 +71,14 @@ struct ARB_ARBOR_API recipe: public has_gap_junctions, has_probes, connectivity
     virtual cell_kind get_cell_kind(cell_gid_type) const = 0;
     // Global property type will be specific to given cell kind.
     virtual std::any get_global_properties(cell_kind) const { return std::any{}; };
+    // list of gap junctions incoming to this gid
+    virtual std::vector<gap_junction_connection> gap_junctions_on(cell_gid_type) const { return {}; }
+    // list of synapses incoming to this gid
+    virtual std::vector<cell_connection> connections_on(cell_gid_type) const { return {}; }
+    // list of probes on this gid
+    virtual std::vector<probe_info> get_probes(cell_gid_type gid) const { return {}; }
+    // list of generators on this gid
+    virtual std::vector<event_generator> event_generators(cell_gid_type) const { return {}; }
 
     virtual ~recipe() {}
 };