Make Pre Post Functions execution on board

src/lava/magma/compiler/compiler.py

@@ -62,7 +62,7 @@ class AbstractNcProcessModel:
 from lava.magma.core.process.process import AbstractProcess
 from lava.magma.core.resources import (CPU, LMT, Loihi1NeuroCore,
                                        Loihi2NeuroCore, NeuroCore)
-from lava.magma.core.run_configs import RunConfig
+from lava.magma.core.run_configs import RunConfig, AbstractLoihiHWRunCfg
 from lava.magma.core.sync.domain import SyncDomain
 from lava.magma.core.sync.protocols.async_protocol import AsyncProtocol
 from lava.magma.runtime.runtime import Runtime
@@ -144,12 +144,13 @@ def compile(
         (
             runtime_service_builders,
             proc_id_to_runtime_service_id_map,
-        ) = self._create_runtime_service_as_py_process_model(
+        ) = self._create_runtime_service_builder(
             node_to_sync_domain_dict,
             self.log,
             nc_builders,
             c_builders,
-            self._compile_config
+            run_cfg,
+            self._compile_config,
         )
         channel_builders = ChannelBuildersFactory().from_channel_map(
             channel_map, compile_config=self._compile_config
@@ -641,11 +642,12 @@ def _create_sync_domains(
         return list(sync_domains.values()), node_to_sync_domain_dict
 
     @staticmethod
-    def _create_runtime_service_as_py_process_model(
+    def _create_runtime_service_builder(
         node_to_sync_domain_dict: ty.Dict[Node, ty.Set[SyncDomain]],
         log: logging.getLoggerClass(),
         nc_builders: ty.Dict[AbstractProcess, NcProcessBuilder],
         c_builders: ty.Dict[AbstractProcess, CProcessBuilder],
+        run_cfg: RunConfig,
         compile_config: ty.Optional[ty.Dict[str, ty.Any]] = None
     ) -> ty.Tuple[
         ty.Dict[SyncDomain, RuntimeServiceBuilder], ty.Dict[int, int]
@@ -697,15 +699,22 @@ def _create_runtime_service_as_py_process_model(
                 log.debug("RuntimeService Class: " + str(rs_class.__name__))
                 model_ids: ty.List[int] = [p.id for p in sync_domain.processes]
 
+                rs_kwargs = {
+                    "c_builders":list(c_builders.values()),
+                    "nc_builders":list(nc_builders.values())
+                }
+                if isinstance(run_cfg, AbstractLoihiHWRunCfg):
+                    rs_kwargs["pre_run_fxs"] = run_cfg.pre_run_fxs
+                    rs_kwargs["post_run_fxs"] = run_cfg.post_run_fxs
+
                 rs_builder = RuntimeServiceBuilder(
                     rs_class,
                     sync_domain.protocol,
                     rs_id,
                     model_ids,
                     loihi_version,
                     log.level,
-                    c_builders=list(c_builders.values()),
-                    nc_builders=list(nc_builders.values()),
+                    **rs_kwargs
                 )
                 rs_builders[sync_domain] = rs_builder
                 for p in sync_domain.processes:

src/lava/magma/core/run_configs.py

@@ -330,7 +330,15 @@ def _order_according_to_resources(self, proc_models: ty.List[ty.Type[
         return list(range(len(proc_models)))
 
 
-class Loihi1SimCfg(AbstractLoihiRunCfg):
+class AbstractLoihiHWRunCfg(AbstractLoihiRunCfg):
+    pass
+
+
+class AbstractLoihiSimRunCfg(AbstractLoihiRunCfg):
+    pass
+
+
+class Loihi1SimCfg(AbstractLoihiSimRunCfg):
     """Run configuration selects appropriate ProcessModel -- either
     `SubProcessModel` for a hierarchical Process or else a `PyProcessModel`
     for a standard Process.
@@ -345,7 +353,7 @@ def _order_according_to_resources(self, proc_models: ty.List[ty.Type[
         return proc_models_ordered
 
 
-class Loihi1HwCfg(AbstractLoihiRunCfg):
+class Loihi1HwCfg(AbstractLoihiHWRunCfg):
     """
     A RunConfig for executing model on Loihi1 HW.
     For Loihi1 HW configurations, the preferred ProcModels are NcProcModels
@@ -354,6 +362,23 @@ class Loihi1HwCfg(AbstractLoihiRunCfg):
     a tag provided by the user. This RunConfig will default to a PyProcModel
     if no Loihi1-compatible ProcModel is being found.
     ."""
+    def __init__(self,
+                 custom_sync_domains: ty.Optional[ty.List[SyncDomain]] = None,
+                 select_tag: ty.Optional[str] = None,
+                 select_sub_proc_model: ty.Optional[bool] = False,
+                 exception_proc_model_map: ty.Optional[ty.Dict[
+                     ty.Type[AbstractProcess], ty.Type[
+                         AbstractProcessModel]]] = None,
+                 loglevel: int = logging.WARNING,
+                 pre_run_fxs: ty.List[ty.Callable] = [],
+                 post_run_fxs: ty.List[ty.Callable] = []):
+        super().__init__(custom_sync_domains,
+                         select_tag,
+                         select_sub_proc_model,
+                         exception_proc_model_map,
+                         loglevel)
+        self.pre_run_fxs: ty.List[ty.Callable] = pre_run_fxs
+        self.post_run_fxs: ty.List[ty.Callable] = post_run_fxs
 
     def _order_according_to_resources(self, proc_models: ty.List[ty.Type[
             AbstractProcessModel]]) -> ty.List[int]:
@@ -380,14 +405,13 @@ def _is_hw_supported(self, pm: ty.Type[AbstractProcessModel]) -> bool:
                  pm.required_resources)
                 and issubclass(pm, AbstractNcProcessModel))
 
-
 class Loihi2SimCfg(Loihi1SimCfg):
     """A RunConfig for simulating a Loihi 2 model CPU/GPU."""
 
     pass
 
 
-class Loihi2HwCfg(AbstractLoihiRunCfg):
+class Loihi2HwCfg(AbstractLoihiHWRunCfg):
     """
     A RunConfig for executing model on Loihi2 HW.
     For Loihi2 HW configurations, the preferred ProcModels are NcProcModels
@@ -396,6 +420,23 @@ class Loihi2HwCfg(AbstractLoihiRunCfg):
     a tag provided by the user. This RunConfig will default to a PyProcModel
     if no Loihi2-compatible ProcModel is being found.
     """
+    def __init__(self,
+                 custom_sync_domains: ty.Optional[ty.List[SyncDomain]] = None,
+                 select_tag: ty.Optional[str] = None,
+                 select_sub_proc_model: ty.Optional[bool] = False,
+                 exception_proc_model_map: ty.Optional[ty.Dict[
+                     ty.Type[AbstractProcess], ty.Type[
+                         AbstractProcessModel]]] = None,
+                 loglevel: int = logging.WARNING,
+                 pre_run_fxs: ty.List[ty.Callable] = [],
+                 post_run_fxs: ty.List[ty.Callable] = []):
+        super().__init__(custom_sync_domains,
+                         select_tag,
+                         select_sub_proc_model,
+                         exception_proc_model_map,
+                         loglevel)
+        self.pre_run_fxs: ty.List[ty.Callable] = pre_run_fxs
+        self.post_run_fxs: ty.List[ty.Callable] = post_run_fxs
 
     def _order_according_to_resources(self, proc_models: ty.List[ty.Type[
             AbstractProcessModel]]) -> ty.List[int]: