Merge pull request #67 from a-vartenkov/mpback_stdp

STDP implementation in the multhreaded backend.

Author: artiomn

examples/mnist-learn/main.cpp

@@ -44,7 +44,7 @@ int main(int argc, char **argv)
 
     // Defines path to backend, on which to run a network.
     std::filesystem::path path_to_backend =
-        std::filesystem::path(argv[0]).parent_path() / "knp-cpu-single-threaded-backend";
+        std::filesystem::path(argv[0]).parent_path() / "knp-cpu-multi-threaded-backend";
 
     // Read data from corresponding files.
     auto spike_frames = read_spike_frames(argv[1]);

knp/backends/cpu/cpu-library/include/knp/backends/cpu-library/impl/additive_stdp_impl.h

@@ -28,6 +28,7 @@
 
 #include <spdlog/spdlog.h>
 
+#include <algorithm>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -152,11 +153,12 @@ constexpr bool is_additive_stdp_synapse()
 
 
 template <class DeltaLikeSynapse>
-void register_additive_stdp_spikes(
+void register_additive_stdp_spikes_part(
     knp::core::Projection<knp::synapse_traits::STDP<knp::synapse_traits::STDPAdditiveRule, DeltaLikeSynapse>>
         &projection,
-    std::vector<SpikeMessage> &all_messages)
+    std::vector<SpikeMessage> &all_messages, uint64_t part_start, uint64_t part_end)
 {
+    if (part_start != 0) return;  // Not much sense to parallelize this by projection, so it's calculated just once.
     SPDLOG_DEBUG("Calculating additive STDP delta synapse projection...");
 
     using ProjectionType = typename std::decay_t<decltype(projection)>;
@@ -209,13 +211,15 @@ void register_additive_stdp_spikes(
 
 
 template <class DeltaLikeSynapse>
-void update_projection_weights_additive_stdp(
+void update_projection_weights_additive_stdp_part(
     knp::core::Projection<knp::synapse_traits::STDP<knp::synapse_traits::STDPAdditiveRule, DeltaLikeSynapse>>
-        &projection)
+        &projection,
+    uint64_t part_start, uint64_t part_end)
 {
     // Update projection parameters.
-    for (auto &proj : projection)
+    for (uint64_t i = part_start; i < std::min(projection.size(), part_end); ++i)
     {
+        auto &proj = projection[i];
         SPDLOG_TRACE("Applying STDP rule...");
         auto &rule = std::get<knp::core::synapse_data>(proj).rule_;
         const auto period = rule.tau_plus_ + rule.tau_minus_;
@@ -238,19 +242,32 @@ template <class DeltaLikeSynapse>
 struct WeightUpdateSTDP<synapse_traits::STDP<synapse_traits::STDPAdditiveRule, DeltaLikeSynapse>>
 {
     using Synapse = synapse_traits::STDP<synapse_traits::STDPAdditiveRule, DeltaLikeSynapse>;
-    static void init_projection(
-        knp::core::Projection<Synapse> &projection, std::vector<SpikeMessage> &all_messages, knp::core::Step step)
+
+    static void init_projection_part(
+        knp::core::Projection<Synapse> &projection, std::vector<SpikeMessage> &all_messages, knp::core::Step step,
+        uint64_t part_start, uint64_t part_end)
     {
-        register_additive_stdp_spikes(projection, all_messages);
+        register_additive_stdp_spikes_part(projection, all_messages, part_start, part_end);
     }
 
     static void init_synapse(const knp::synapse_traits::synapse_parameters<Synapse> &projection, knp::core::Step step)
     {
     }
 
+    static void modify_weights_part(knp::core::Projection<Synapse> &projection, uint64_t part_start, uint64_t part_end)
+    {
+        update_projection_weights_additive_stdp_part(projection, part_start, part_end);
+    }
+
+    static void init_projection(
+        knp::core::Projection<Synapse> &projection, std::vector<SpikeMessage> &all_messages, knp::core::Step step)
+    {
+        init_projection_part(projection, all_messages, step, 0, projection.size());
+    }
+
     static void modify_weights(knp::core::Projection<Synapse> &projection)
     {
-        update_projection_weights_additive_stdp(projection);
+        modify_weights_part(projection, 0, projection.size());
     }
 };
 

knp/backends/cpu/cpu-library/include/knp/backends/cpu-library/impl/base_stdp_impl.h

@@ -4,45 +4,56 @@
  * @kaspersky_support A. Vartenkov
  * @date 04.11.2023
  * @license Apache 2.0
- * @copyright © 2024 AO Kaspersky Lab
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
+ * @copyright © 2024 AO Kaspersky Lab
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
  * limitations under the License.
  */
 #pragma once
 #include <knp/core/messaging/messaging.h>
 #include <knp/core/projection.h>
 
+#include <unordered_map>
 #include <vector>
 
 /**
  * @brief Namespace for CPU backends.
  */
 namespace knp::backends::cpu
 {
-
 template <class DeltaLikeSynapse>
 struct WeightUpdateSTDP
 {
-    static void init_projection(
+    static void init_projection_part(
         const knp::core::Projection<DeltaLikeSynapse> &projection,
-        const std::vector<core::messaging::SpikeMessage> &messages, uint64_t step)
+        const std::vector<core::messaging::SpikeMessage> &messages, uint64_t part_begin, uint64_t part_end,
+        uint64_t step)
     {
     }
 
-    static void init_synapse(const knp::synapse_traits::synapse_parameters<DeltaLikeSynapse> &projection, uint64_t step)
+    static void init_synapse(knp::synapse_traits::synapse_parameters<DeltaLikeSynapse> &params, uint64_t step) {}
+
+    static void modify_weights_part(
+        const knp::core::Projection<DeltaLikeSynapse> &projection, uint64_t part_begin, uint64_t part_end)
     {
     }
 
+    static void init_projection(
+        const knp::core::Projection<DeltaLikeSynapse> &projection,
+        const std::vector<core::messaging::SpikeMessage> &messages, uint64_t step)
+    {
+    }
+
+
     static void modify_weights(const knp::core::Projection<DeltaLikeSynapse> &projection) {}
 };
 

knp/backends/cpu/cpu-library/include/knp/backends/cpu-library/impl/blifat_population_impl.h

@@ -314,6 +314,17 @@ void calculate_neurons_post_input_state_part(
 }
 
 
+/**
+ * Anything that works after the population has been calculated goes here. Mostly for STDP and other learning.
+ */
+template <class BlifatLikeNeuron, class ProjectionContainer>
+void finalize_population(
+    knp::core::Population<BlifatLikeNeuron> &population, const knp::core::messaging::SpikeMessage &message,
+    ProjectionContainer &projections, knp::core::Step step)
+{
+}
+
+
 /**
  * @brief Process BLIFAT neuron population and return spiked neuron indexes.
  * @tparam BlifatLikeNeuron type of neuron which inference can be calculated the same as BLIFAT.

knp/backends/cpu/cpu-library/include/knp/backends/cpu-library/impl/delta_synapse_projection_impl.h

@@ -24,6 +24,7 @@
 #include <knp/core/message_bus.h>
 #include <knp/core/projection.h>
 #include <knp/synapse-traits/delta.h>
+#include <knp/synapse-traits/stdp_synaptic_resource_rule.h>
 
 #include <spdlog/spdlog.h>
 
@@ -132,15 +133,15 @@ template <class DeltaLikeSynapse>
 void calculate_projection_part_impl(
     knp::core::Projection<DeltaLikeSynapse> &projection,
     const std::unordered_map<knp::core::Step, size_t> &message_in_data, MessageQueue &future_messages, uint64_t step_n,
-    size_t part_start, size_t part_size, std::mutex &mutex)
+    uint64_t part_start, uint64_t part_size, std::mutex &mutex)
 {
     size_t part_end = std::min(part_start + part_size, projection.size());
     std::vector<std::pair<uint64_t, knp::core::messaging::SynapticImpact>> container;
+    WeightUpdateStdpMp<DeltaLikeSynapse>::init_projection_part(projection, message_in_data, step_n);
     for (size_t synapse_index = part_start; synapse_index < part_end; ++synapse_index)
     {
         auto &synapse = projection[synapse_index];
         // update_step(synapse.params_, step_n);
-        // TODO: Move update logic here too.
         auto iter = message_in_data.find(std::get<core::source_neuron_id>(synapse));
         if (iter == message_in_data.end())
         {
@@ -150,6 +151,7 @@ void calculate_projection_part_impl(
         // Add new impact.
         // The message is sent on step N - 1, received on step N.
         uint64_t key = std::get<core::synapse_data>(synapse).delay_ + step_n - 1;
+        WeightUpdateStdpMp<DeltaLikeSynapse>::init_synapse(std::get<core::synapse_data>(synapse), step_n);
 
         knp::core::messaging::SynapticImpact impact{
             synapse_index, std::get<core::synapse_data>(synapse).weight_ * iter->second,
@@ -183,6 +185,7 @@ void calculate_projection_part_impl(
             future_messages.insert(std::make_pair(value.first, message_out));
         }
     }
+    WeightUpdateStdpMp<DeltaLikeSynapse>::modify_weights_part(projection);
 }
 
 

knp/backends/cpu/cpu-library/include/knp/backends/cpu-library/impl/synaptic_resource_stdp_impl.h

@@ -32,7 +32,9 @@
 #include <algorithm>
 #include <limits>
 #include <numeric>
+#include <unordered_map>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include <boost/mp11.hpp>
@@ -381,6 +383,40 @@ struct WeightUpdateSTDP<synapse_traits::STDP<synapse_traits::STDPSynapticResourc
 };
 
 
+template <class DeltaLikeSynapse>
+struct WeightUpdateStdpMp
+{
+    using Synapse = DeltaLikeSynapse;
+    static void init_projection_part(
+        const knp::core::Projection<Synapse> &projection,
+        const std::unordered_map<knp::core::Step, size_t> &message_data, uint64_t step)
+    {
+    }
+
+    static void init_synapse(knp::synapse_traits::synapse_parameters<Synapse> &params, uint64_t step) {}
+
+    static void modify_weights_part(const knp::core::Projection<Synapse> &projection) {}
+};
+
+
+template <class DeltaLikeSynapse>
+struct WeightUpdateStdpMp<synapse_traits::STDP<synapse_traits::STDPSynapticResourceRule, DeltaLikeSynapse>>
+{
+    using Synapse = synapse_traits::STDP<synapse_traits::STDPSynapticResourceRule, DeltaLikeSynapse>;
+    static void init_projection_part(
+        const knp::core::Projection<Synapse> &projection,
+        const std::unordered_map<knp::core::Step, size_t> &message_data, uint64_t step)
+    {
+    }
+
+    static void init_synapse(knp::synapse_traits::synapse_parameters<Synapse> &params, uint64_t step)
+    {
+        params.rule_.last_spike_step_ = step;
+    }
+
+    static void modify_weights_part(const knp::core::Projection<Synapse> &projection) {}
+};
+
 template <class NeuronType, class SynapseType>
 void do_STDP_resource_plasticity(
     knp::core::Population<knp::neuron_traits::SynapticResourceSTDPNeuron<NeuronType>> &population,
@@ -402,4 +438,6 @@ void do_STDP_resource_plasticity(
     // 3. Renormalize resources if needed.
     knp::backends::cpu::renormalize_resource(working_projections, population, step);
 }
+
+
 }  // namespace knp::backends::cpu

knp/backends/cpu/cpu-multi-threaded-backend/CMakeLists.txt

@@ -48,6 +48,7 @@ knp_add_library("${PROJECT_NAME}"
     BOTH
         impl/backend.cpp
         impl/get_network.cpp
+        impl/template_specs.cpp
         ${${PROJECT_NAME}_headers}
     ALIAS KNP::Backends::CPUMultiThreaded
     LINK_PRIVATE

knp/backends/cpu/cpu-multi-threaded-backend/impl/backend.cpp

@@ -21,6 +21,9 @@
 
 #include <knp/backends/cpu-library/blifat_population.h>
 #include <knp/backends/cpu-library/delta_synapse_projection.h>
+#include <knp/backends/cpu-library/impl/altai_lif_population_impl.h>
+#include <knp/backends/cpu-library/impl/blifat_population_impl.h>
+#include <knp/backends/cpu-library/impl/synaptic_resource_stdp_impl.h>
 #include <knp/backends/cpu-library/init.h>
 #include <knp/backends/cpu-multi-threaded/backend.h>
 #include <knp/backends/thread_pool/thread_pool.h>
@@ -33,6 +36,7 @@
 
 #include <functional>
 #include <optional>
+#include <unordered_map>
 #include <vector>
 
 #include <boost/mp11.hpp>
@@ -153,10 +157,28 @@ std::vector<knp::core::messaging::SpikeMessage> MultiThreadedCPUBackend::calcula
                         std::ref(pop), std::ref(message), neuron_index, population_part_size_, std::ref(ep_mutex_));
                 },
                 population);
+        }
+    }
+    calc_pool_->join();
+    for (size_t pop_id = 0; pop_id < populations_.size(); ++pop_id)
+    {
+        auto &message = spike_container[pop_id];
+        std::visit(
+            [this, &message](auto &pop)
+            {
+                using T = std::decay_t<decltype(pop)>;
+                auto call_finalize = [](T &pop_ref, knp::core::messaging::SpikeMessage &message_ref,
+                                        ProjectionContainer &proj_ref, knp::core::Step step)
+                {
+                    knp::backends::cpu::finalize_population<typename T::PopulationNeuronType, ProjectionContainer>(
+                        pop_ref, message_ref, proj_ref, step);
+                };
+                calc_pool_->post(call_finalize, std::ref(pop), std::ref(message), std::ref(projections_), get_step());
+            },
+            populations_[pop_id]);
 #if defined(_MSC_VER)
 #    pragma warning(pop)
 #endif
-        }
     }
     calc_pool_->join();
     return spike_container;
@@ -394,5 +416,4 @@ MultiThreadedCPUBackend::ProjectionConstIterator MultiThreadedCPUBackend::end_pr
 
 
 BOOST_DLL_ALIAS(knp::backends::multi_threaded_cpu::MultiThreadedCPUBackend::create, create_knp_backend)
-
 }  // namespace knp::backends::multi_threaded_cpu