Checkpoint/Resume

CMakeLists.txt

@@ -136,9 +136,20 @@ if(ARB_GPU STREQUAL "cuda")
     set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER})
     enable_language(CUDA)
     find_package(CUDAToolkit)
-    if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
-        set(CMAKE_CUDA_ARCHITECTURES 60 70 80)
+    if(${CUDAToolkit_VERSION_MAJOR} GREATER_EQUAL 12)
+        if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+            # Pascal, Volta, Ampere, Hopper
+            set(CMAKE_CUDA_ARCHITECTURES 60 70 80 90)
+        endif()
+    elseif(${CUDAToolkit_VERSION_MAJOR} GREATER_EQUAL 11)
+        if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+            # Pascal, Volta, Ampere
+            set(CMAKE_CUDA_ARCHITECTURES 60 70 80)
+        endif()
+    else()
+        message(FATAL_ERROR "Need at least CUDA 11, got ${CUDAToolkit_VERSION_MAJOR}")
     endif()
+
     # We _still_ need this otherwise CUDA symbols will not be exported
     # from libarbor.a leading to linker errors when link external clients.
     # Unit tests are NOT external enough. Re-review this somewhere in the
@@ -181,7 +192,7 @@ include("CheckCompilerXLC")
 include("CompilerOptions")
 add_compile_options("$<$<COMPILE_LANGUAGE:CXX>:${CXXOPT_WALL}>")
 set(CMAKE_CXX_STANDARD 17)
-set(CMAKE_CUDA_STANDARD 14)
+set(CMAKE_CUDA_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
 
@@ -241,19 +252,18 @@ install(TARGETS arborio-public-deps EXPORT arborio-targets)
 
 install(PROGRAMS scripts/arbor-build-catalogue DESTINATION ${CMAKE_INSTALL_BINDIR})
 install(FILES mechanisms/BuildModules.cmake DESTINATION ${ARB_INSTALL_DATADIR})
-# External libraries in `ext` sub-directory: json, tinyopt and randon123.
-# Creates interface libraries `ext-json`, `ext-tinyopt` and `ext-random123`
+
+# External libraries in `ext` sub-directory: json, tinyopt and random123.
+# Creates interface libraries `ext-tinyopt` and `ext-random123`
 
 cmake_dependent_option(ARB_USE_BUNDLED_FMT "Use bundled FMT lib." ON "ARB_USE_BUNDLED_LIBS" OFF)
 cmake_dependent_option(ARB_USE_BUNDLED_PUGIXML "Use bundled XML lib." ON "ARB_USE_BUNDLED_LIBS" OFF)
 cmake_dependent_option(ARB_USE_BUNDLED_GTEST "Use bundled GoogleTest." ON "ARB_USE_BUNDLED_LIBS" OFF)
 
 cmake_dependent_option(ARB_USE_BUNDLED_JSON "Use bundled Niels Lohmann's json library." ON "ARB_USE_BUNDLED_LIBS" OFF)
 if(NOT ARB_USE_BUNDLED_JSON)
-  find_package(nlohmann_json)
-  set(json_library_name nlohmann_json::nlohmann_json)
-else()
-  unset(nlohmann_json_DIR CACHE)
+  find_package(nlohmann_json 3.11.2 CONFIG REQUIRED)
+  message(STATUS "Using external JSON = ${nlohmann_json_VERSION}")
 endif()
 
 cmake_dependent_option(ARB_USE_BUNDLED_RANDOM123 "Use bundled Random123 lib." ON "ARB_USE_BUNDLED_LIBS" OFF)

arbor/backends/event.hpp

@@ -2,6 +2,7 @@
 
 #include <arbor/common_types.hpp>
 #include <arbor/fvm_types.hpp>
+#include <arbor/serdes.hpp>
 #include <arbor/mechanism_abi.h>
 #include <arbor/generic_event.hpp>
 
@@ -17,10 +18,22 @@ struct target_handle {
     cell_local_size_type mech_index; // instance of the mechanism
 
     target_handle() = default;
+
     target_handle(cell_local_size_type mech_id, cell_local_size_type mech_index):
         mech_id(mech_id), mech_index(mech_index) {}
+
+    ARB_SERDES_ENABLE(target_handle, mech_id, mech_index);
 };
 
+}
+
+template<typename K>
+void serialize(arb::serializer &ser, const K &k, const arb::target_handle&);
+template<typename K>
+void deserialize(arb::serializer &ser, const K &k, arb::target_handle&);
+
+namespace arb {
+
 struct deliverable_event {
     time_type time = 0;
     float weight = 0;
@@ -29,11 +42,21 @@ struct deliverable_event {
     deliverable_event() = default;
     deliverable_event(time_type time, target_handle handle, float weight):
         time(time), weight(weight), handle(handle) {}
+
+    ARB_SERDES_ENABLE(deliverable_event, time, weight, handle);
 };
 
 template<>
 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;
+    ARB_SERDES_ENABLE(deliverable_event_data, mech_id, mech_index, weight);
+};
+
 // Stream index accessor function for multi_event_stream:
 inline cell_local_size_type event_index(const arb_deliverable_event_data& ed) {
     return ed.mech_index;

arbor/backends/event_stream_base.hpp

@@ -58,7 +58,6 @@ class event_stream_base {
         ev_spans_.clear();
         index_ = 0;
     }
-
 };
 
 } // namespace arb

arbor/backends/gpu/event_stream.hpp

@@ -9,12 +9,17 @@
 #include "util/rangeutil.hpp"
 #include "util/transform.hpp"
 #include "threading/threading.hpp"
+#include <arbor/mechanism_abi.h>
+
+ARB_SERDES_ENABLE_EXT(arb_deliverable_event_data, mech_index, weight);
 
 namespace arb {
 namespace gpu {
 
 template <typename Event>
-class event_stream : public event_stream_base<Event, typename memory::device_vector<::arb::event_data_type<Event>>::view_type> {
+class event_stream :
+        public event_stream_base<Event,
+                                 typename memory::device_vector<::arb::event_data_type<Event>>::view_type> {
 public:
     using base = event_stream_base<Event, typename memory::device_vector<::arb::event_data_type<Event>>::view_type>;
     using size_type = typename base::size_type;
@@ -70,12 +75,14 @@ class event_stream : public event_stream_base<Event, typename memory::device_vec
                 // host span
                 auto host_span = memory::make_view(base::ev_data_)(offset, offset + size);
                 // make event data and copy
-                std::copy_n(util::transform_view(staged[i], [](const auto& x) {
-                    return event_data(x);}).begin(), size, host_span.begin());
+                std::copy_n(util::transform_view(staged[i],
+                                                 [](const auto& x) { return event_data(x); }).begin(),
+                            size,
+                            host_span.begin());
                 // sort if necessary
                 if constexpr (has_event_index<Event>::value) {
-                    util::stable_sort_by(host_span, [](const event_data_type& ed) {
-                        return event_index(ed); });
+                    util::stable_sort_by(host_span,
+                                         [](const event_data_type& ed) { return event_index(ed); });
                 }
                 // copy to device
                 memory::copy_async(host_span, base::ev_spans_[i]);
@@ -84,6 +91,49 @@ class event_stream : public event_stream_base<Event, typename memory::device_vec
         arb_assert(num_events == base::ev_data_.size());
     }
 
+    friend void serialize(serializer& ser, const std::string& k, const event_stream<Event>& t) {
+        ser.begin_write_map(::arb::to_serdes_key(k));
+        ARB_SERDES_WRITE(ev_data_);
+        ser.begin_write_map("ev_spans_");
+        auto base_ptr = t.device_ev_data_.data();
+        for (size_t ix = 0; ix < t.ev_spans_.size(); ++ix) {
+            ser.begin_write_map(std::to_string(ix));
+            const auto& span = t.ev_spans_[ix];
+            ser.write("offset", static_cast<unsigned long long>(span.begin() - base_ptr));
+            ser.write("size", static_cast<unsigned long long>(span.size()));
+            ser.end_write_map();
+        }
+        ser.end_write_map();
+        ARB_SERDES_WRITE(index_);
+        ARB_SERDES_WRITE(device_ev_data_);
+        ARB_SERDES_WRITE(offsets_);
+        ser.end_write_map();
+    }
+
+    friend void deserialize(serializer& ser, const std::string& k, event_stream<Event>& t) {
+        ser.begin_read_map(::arb::to_serdes_key(k));
+        ARB_SERDES_READ(ev_data_);
+        ser.begin_read_map("ev_spans_");
+        for (size_t ix = 0; ser.next_key(); ++ix) {
+            ser.begin_read_map(std::to_string(ix));
+            unsigned long long offset = 0, size = 0;
+            ser.read("offset", offset);
+            ser.read("size", size);
+            typename base::span_type span{t.ev_data_.data() + offset, size};
+            if (ix < t.ev_spans_.size()) {
+                t.ev_spans_[ix] = span;
+            } else {
+                t.ev_spans_.emplace_back(span);
+            }
+            ser.end_read_map();
+        }
+        ser.end_read_map();
+        ARB_SERDES_READ(index_);
+        ARB_SERDES_READ(device_ev_data_);
+        ARB_SERDES_READ(offsets_);
+        ser.end_read_map();
+    }
+
 private:
     template<typename D>
     static void resize(D& d, std::size_t size) {

arbor/backends/gpu/rand.hpp

@@ -4,6 +4,7 @@
 #include <vector>
 
 #include <arbor/mechanism.hpp>
+#include <arbor/serdes.hpp>
 
 #include <util/pimpl.hpp>
 #include <backends/rand_fwd.hpp>
@@ -19,6 +20,8 @@ class random_numbers {
 
     void update(mechanism& m);
 
+    ARB_SERDES_ENABLE(random_numbers, data_, random_number_update_counter_);
+
   private:
     // random number device storage
     array data_;

arbor/backends/gpu/shared_state.cpp

@@ -22,6 +22,8 @@
 #include "util/range.hpp"
 #include "util/strprintf.hpp"
 
+#include <iostream>
+
 using arb::memory::make_const_view;
 
 namespace arb {
@@ -175,11 +177,11 @@ shared_state::shared_state(task_system_handle tp,
                            const std::vector<arb_value_type>& diam,
                            const std::vector<arb_value_type>& area,
                            const std::vector<arb_index_type>& src_to_spike_,
-                           const fvm_detector_info& detector,
+                           const fvm_detector_info& detector_info,
                            unsigned, // align parameter ignored
                            arb_seed_type cbprng_seed_):
     thread_pool(tp),
-    n_detector(detector.count),
+    n_detector(detector_info.count),
     n_cv(n_cv_),
     cv_to_cell(make_const_view(cv_to_cell_vec)),
     voltage(n_cv_),
@@ -193,11 +195,7 @@ shared_state::shared_state(task_system_handle tp,
     src_to_spike(make_const_view(src_to_spike_)),
     cbprng_seed(cbprng_seed_),
     sample_events(thread_pool),
-    watcher{n_cv_,
-            src_to_spike.data(),
-            detector.cv,
-            detector.threshold,
-            detector.ctx}
+    watcher{n_cv_, src_to_spike.data(), detector_info}
 {
     memory::fill(time_since_spike, -1.0);
     add_scalar(temperature_degC.size(), temperature_degC.data(), -273.15);
@@ -242,8 +240,8 @@ void shared_state::instantiate(mechanism& m,
     m.ppack_.time_since_spike = time_since_spike.data();
     m.ppack_.n_detectors      = n_detector;
 
-    if (storage.find(id) != storage.end()) throw arb::arbor_internal_error("Duplicate mech id in shared state");
-    auto& store = storage.emplace(id, thread_pool).first->second;
+    if (storage.count(id)) throw arb::arbor_internal_error("Duplicate mech id in shared state");
+    auto& store = storage.emplace(id, mech_storage{thread_pool}).first->second;
 
     // Allocate view pointers
     store.state_vars_ = std::vector<arb_value_type*>(m.mech_.n_state_vars);

arbor/backends/gpu/shared_state.hpp

@@ -50,7 +50,7 @@ struct ARB_ARBOR_API ion_state {
     array Xi_;          // (mM) internal concentration
     array Xd_;          // (mM) diffusive concentration
     array Xo_;          // (mM) external concentration
-    array gX_;             // (kS/m²) per-species conductivity
+    array gX_;          // (kS/m²) per-species conductivity
 
     array init_Xi_;     // (mM) area-weighted initial internal concentration
     array init_Xo_;     // (mM) area-weighted initial external concentration
@@ -64,9 +64,7 @@ struct ARB_ARBOR_API ion_state {
 
     ion_state() = default;
 
-    ion_state(const fvm_ion_config& ion_data,
-              unsigned align,
-              solver_ptr ptr);
+    ion_state(const fvm_ion_config& ion_data, unsigned align, solver_ptr ptr);
 
     // Set ion concentrations to weighted proportion of default concentrations.
     void init_concentration();
@@ -116,24 +114,23 @@ struct ARB_ARBOR_API istim_state {
     istim_state() = default;
 };
 
-struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_state> {
-    struct mech_storage {
-        mech_storage() = default;
-        mech_storage(task_system_handle tp) : deliverable_events_(tp) {}
-
-        array data_;
-        iarray indices_;
-        std::vector<arb_value_type>  globals_;
-        std::vector<arb_value_type*> parameters_;
-        std::vector<arb_value_type*> state_vars_;
-        std::vector<arb_ion_state>   ion_states_;
-        memory::device_vector<arb_value_type*> parameters_d_;
-        memory::device_vector<arb_value_type*> state_vars_d_;
-        memory::device_vector<arb_ion_state>   ion_states_d_;
-        random_numbers random_numbers_;
-        deliverable_event_stream deliverable_events_;
-    };
+struct mech_storage {
+    mech_storage() = default;
+    mech_storage(task_system_handle tp) : deliverable_events_(tp) {}
+    array data_;
+    iarray indices_;
+    std::vector<arb_value_type>  globals_;
+    std::vector<arb_value_type*> parameters_;
+    std::vector<arb_value_type*> state_vars_;
+    std::vector<arb_ion_state>   ion_states_;
+    memory::device_vector<arb_value_type*> parameters_d_;
+    memory::device_vector<arb_value_type*> state_vars_d_;
+    memory::device_vector<arb_ion_state>   ion_states_d_;
+    random_numbers random_numbers_;
+    deliverable_event_stream deliverable_events_;
+};
 
+struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_state> {
     task_system_handle thread_pool;
 
     using cable_solver = arb::gpu::matrix_state_fine<arb_value_type, arb_index_type>;
@@ -183,7 +180,7 @@ struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_st
                  const std::vector<arb_index_type>& cv_to_cell_vec,
                  const fvm_cv_discretization& D,
                  const std::vector<arb_index_type>& src_to_spike,
-                 const fvm_detector_info& detector,
+                 const fvm_detector_info& detector_info,
                  const std::unordered_map<std::string, fvm_ion_config>& ions,
                  const fvm_stimulus_config& stims,
                  unsigned align,
@@ -197,7 +194,7 @@ struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_st
                        D.diam_um,
                        D.cv_area,
                        src_to_spike,
-                       detector,
+                       detector_info,
                        align,
                        cbprng_seed_}
     {
@@ -216,7 +213,7 @@ struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_st
                  const std::vector<arb_value_type>& diam,
                  const std::vector<arb_value_type>& area,
                  const std::vector<arb_index_type>& src_to_spike,
-                 const fvm_detector_info& detector,
+                 const fvm_detector_info& detector_info,
                  unsigned, // align parameter ignored
                  arb_seed_type cbprng_seed_ = 0u);
 
@@ -251,4 +248,21 @@ struct ARB_ARBOR_API shared_state: shared_state_base<shared_state, array, ion_st
 ARB_ARBOR_API std::ostream& operator<<(std::ostream& o, shared_state& s);
 
 } // namespace gpu
+
+ARB_SERDES_ENABLE_EXT(gpu::ion_state, Xd_, gX_);
+ARB_SERDES_ENABLE_EXT(gpu::mech_storage,
+                      data_,
+                      // NOTE(serdes) ion_states_, this is just a bunch of pointers
+                      random_numbers_,
+                      deliverable_events_);
+ARB_SERDES_ENABLE_EXT(gpu::shared_state,
+                      cbprng_seed,
+                      ion_data,
+                      storage,
+                      voltage,
+                      current_density,
+                      conductivity,
+                      time_since_spike,
+                      time, time_to,
+                      dt);
 } // namespace arb

arbor/backends/gpu/stimulus.cu

@@ -1,11 +1,11 @@
 #include <cmath>
 
-#include <arbor/fvm_types.hpp>
 #include <arbor/gpu/gpu_api.hpp>
 #include <arbor/gpu/gpu_common.hpp>
 #include <arbor/gpu/math_cu.hpp>
 
-#include "backends/gpu/stimulus.hpp"
+#include "stimulus.hpp"
+
 
 namespace arb {
 namespace gpu {