prototype implementing injector/extractor function, not wrapped in a Process

Author: gkarray

src/lava/magma/core/process/ports/ports.py

@@ -432,6 +432,17 @@ class OutPort(AbstractIOPort, AbstractSrcPort):
     sub processes.
     """
 
+    def __init__(self, shape: ty.Tuple[int, ...]):
+        super().__init__(shape)
+        self.external_pipe_flag = False
+        self.external_pipe_buffer_size = 64
+
+    def flag_external_pipe(self, buffer_size=None):
+        self.external_pipe_flag = True
+
+        if buffer_size is not None:
+            self.external_pipe_buffer_size = buffer_size
+
     def connect(
             self,
             ports: ty.Union["AbstractIOPort", ty.List["AbstractIOPort"]],
@@ -493,6 +504,15 @@ def __init__(
         super().__init__(shape)
         self._reduce_op = reduce_op
 
+        self.external_pipe_flag = False
+        self.external_pipe_buffer_size = 64
+
+    def flag_external_pipe(self, buffer_size=None):
+        self.external_pipe_flag = True
+
+        if buffer_size is not None:
+            self.external_pipe_buffer_size = buffer_size
+
     def connect(self,
                 ports: ty.Union["InPort", ty.List["InPort"]],
                 connection_configs: ty.Optional[ConnectionConfigs] = None):

src/lava/magma/runtime/runtime.py

@@ -10,6 +10,7 @@
 import typing as ty
 
 import numpy as np
+from lava.magma.core.model.py.ports import PyInPort, PyOutPort
 from scipy.sparse import csr_matrix
 from lava.magma.compiler.var_model import AbstractVarModel, LoihiSynapseVarModel
 from lava.magma.core.process.message_interface_enum import ActorType
@@ -26,7 +27,7 @@
 if ty.TYPE_CHECKING:
     from lava.magma.core.process.process import AbstractProcess
 from lava.magma.compiler.channels.pypychannel import CspRecvPort, CspSendPort, \
-    CspSelector
+    CspSelector, PyPyChannel
 from lava.magma.compiler.builders.channel_builder import (
     ChannelBuilderMp, RuntimeChannelBuilderMp, ServiceChannelBuilderMp,
     ChannelBuilderPyNc)
@@ -38,7 +39,7 @@
     ChannelType
 from lava.magma.compiler.executable import Executable
 from lava.magma.compiler.node import NodeConfig
-from lava.magma.core.process.ports.ports import create_port_id
+from lava.magma.core.process.ports.ports import create_port_id, InPort, OutPort
 from lava.magma.core.run_conditions import (AbstractRunCondition,
                                             RunContinuous, RunSteps)
 from lava.magma.compiler.channels.watchdog import WatchdogManagerInterface
@@ -308,6 +309,39 @@ def _build_processes(self):
                     proc._runtime = self
                     exception_q = Queue()
                     self.exception_q.append(exception_q)
+
+                    for name, py_port in proc_builder.py_ports.items():
+                        port = getattr(proc, name)
+
+                        if port.external_pipe_flag:
+                            if isinstance(port, InPort):
+                                pypychannel = PyPyChannel(
+                                    message_infrastructure=self._messaging_infrastructure,
+                                    src_name="src",
+                                    dst_name=name,
+                                    shape=py_port.shape,
+                                    dtype=py_port.d_type,
+                                    size=64)
+
+                                proc_builder.set_csp_ports([pypychannel.dst_port])
+
+                                port.external_pipe_csp_send_port = pypychannel.src_port
+                                port.external_pipe_csp_send_port.start()
+
+                            if isinstance(port, OutPort):
+                                pypychannel = PyPyChannel(
+                                    message_infrastructure=self._messaging_infrastructure,
+                                    src_name=name,
+                                    dst_name="dst",
+                                    shape=py_port.shape,
+                                    dtype=py_port.d_type,
+                                    size=64)
+
+                                proc_builder.set_csp_ports([pypychannel.src_port])
+
+                                port.external_pipe_csp_recv_port = pypychannel.dst_port
+                                port.external_pipe_csp_recv_port.start()
+
                     self._messaging_infrastructure.build_actor(target_fn,
                                                                proc_builder,
                                                                exception_q)

tests/lava/magma/runtime/test_external_pipe_io.py

@@ -0,0 +1,96 @@
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: LGPL 2.1 or later
+# See: https://spdx.org/licenses/
+import threading
+import unittest
+import typing as ty
+
+import numpy as np
+
+from lava.magma.core.decorator import implements, requires
+from lava.magma.core.model.py.model import PyLoihiProcessModel
+from lava.magma.core.model.py.ports import PyOutPort, PyInPort
+from lava.magma.core.model.py.type import LavaPyType
+from lava.magma.core.process.ports.ports import OutPort, InPort
+from lava.magma.core.process.process import AbstractProcess
+from lava.magma.core.resources import CPU
+from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
+from lava.magma.core.run_conditions import RunSteps
+from lava.magma.core.run_configs import Loihi2SimCfg
+
+
+class Process1(AbstractProcess):
+    def __init__(self, shape_1: ty.Tuple, **kwargs):
+        super().__init__(shape_1=shape_1, **kwargs)
+
+        self.in_1 = InPort(shape=shape_1)
+        self.out_1 = OutPort(shape=shape_1)
+
+
+@implements(proc=Process1, protocol=LoihiProtocol)
+@requires(CPU)
+class LoihiDenseSpkPyProcess1PM(PyLoihiProcessModel):
+    in_1: PyInPort = LavaPyType(PyInPort.VEC_DENSE, float)
+    out_1: PyOutPort = LavaPyType(PyOutPort.VEC_DENSE, float)
+
+    def __init__(self, proc_params):
+        super().__init__(proc_params)
+
+    def run_spk(self):
+        print(f"Receiving in Process...")
+        data_1 = self.in_1.recv()
+        print(f"Received {data_1} in Process...")
+
+        print(f"Sending {data_1} from Process...")
+        self.out_1.send(data_1)
+        print(f"Sent {data_1} from Process!")
+
+
+class TestExternalPipeIO(unittest.TestCase):
+    def test_run_steps_non_blocking(self):
+        data = [[1], [2], [3], [4], [5]]
+
+        relay = Process1(shape_1=(1,))
+        # Control buffer size with buffer_size arg, default is 64
+        relay.in_1.flag_external_pipe()
+        # Control buffer size with buffer_size arg, default is 64
+        relay.out_1.flag_external_pipe()
+
+        run_cfg = Loihi2SimCfg()
+        run_condition = RunSteps(num_steps=5, blocking=False)
+
+        def thread_inject_fn() -> None:
+            for send_data_single_item in data:
+                print(f"Sending {send_data_single_item} from thread_inject...")
+                # Use probe() before send() to know whether or not send() will
+                # block (i.e if the buffer of external_pipe_csp_send_port
+                # is full).
+                relay.in_1.external_pipe_csp_send_port.send(
+                    np.array(send_data_single_item))
+                print(f"Sent {send_data_single_item} from thread_inject!")
+
+        def thread_extract_fn() -> None:
+            for _ in range(len(data)):
+                print(f"Receiving in thread_extract...")
+                # Use probe() before recv() to know whether or not recv() will
+                # block (i.e if the buffer of external_pipe_csp_recv_port
+                # is empty).
+                received_data = relay.out_1.external_pipe_csp_recv_port.recv()
+                print(f"Received {received_data} in thread_extract!")
+
+        thread_inject = threading.Thread(target=thread_inject_fn,
+                                         daemon=True)
+        thread_extract = threading.Thread(target=thread_extract_fn,
+                                          daemon=True)
+
+        relay.run(condition=run_condition, run_cfg=run_cfg)
+
+        thread_inject.start()
+        thread_extract.start()
+
+        relay.wait()
+        relay.stop()
+
+
+if __name__ == "__main__":
+    unittest.main()