Shmem Ports & Channel refactor development

src/lava/magma/compiler/builders/channel_builder.py

@@ -10,10 +10,14 @@
     AbstractProcessModel
 from lava.magma.compiler.builders. \
     runtimeservice_builder import RuntimeServiceBuilder
-from lava.magma.compiler.channels.interfaces import (
+from message_infrastructure import (
     Channel,
-    ChannelType,
+    ChannelTransferType,
+    get_channel_factory,
+    ShmemChannel,
+    SharedMemory
 )
+from lava.magma.compiler.channels.interfaces import ChannelType
 from lava.magma.compiler.utils import PortInitializer
 from message_infrastructure \
     .message_infrastructure_interface import (MessageInfrastructureInterface)
@@ -29,7 +33,7 @@ class ChannelBuilderMp(AbstractChannelBuilder):
     and multi processing backbone.
     """
 
-    channel_type: ChannelType
+    channel_type: ChannelTransferType
     src_process: "AbstractProcess"
     dst_process: "AbstractProcess"
     src_port_initializer: PortInitializer
@@ -54,9 +58,15 @@ def build(
         Exception
             Can't build channel of type specified
         """
-        channel_class = messaging_infrastructure.channel_class(
-            channel_type=self.channel_type
-        )
+        channel_factory = get_channel_factory()
+        shm = messaging_infrastructure.smm.alloc_mem(
+            self.src_port_initializer.size)
+        return channel_factory.get_channel(ChannelTransferType.SHMEMCHANNEL,
+                                           shm,
+                                           self.src_port_initializer.d_type,
+                                           self.src_port_initializer.size,
+                                           self.src_port_initializer.shape,
+                                           self.src_port_initializer.name)
         return channel_class(
             messaging_infrastructure,
             self.src_port_initializer.name,
@@ -73,7 +83,7 @@ class ServiceChannelBuilderMp(AbstractChannelBuilder):
     as messaging and multi processing backbone.
     """
 
-    channel_type: ChannelType
+    channel_type: ChannelTransferType
     src_process: ty.Union[RuntimeServiceBuilder,
                           ty.Type["AbstractProcessModel"]]
     dst_process: ty.Union[RuntimeServiceBuilder,
@@ -99,9 +109,15 @@ def build(
         Exception
             Can't build channel of type specified
         """
-        channel_class = messaging_infrastructure.channel_class(
-            channel_type=self.channel_type
-        )
+        channel_factory = get_channel_factory()
+        shm = messaging_infrastructure.smm.alloc_mem(
+            self.port_initializer.size)
+        return channel_factory.get_channel(ChannelTransferType.SHMEMCHANNEL,
+                                           shm,
+                                           self.port_initializer.d_type,
+                                           self.port_initializer.size,
+                                           self.port_initializer.shape,
+                                           self.port_initializer.name)
 
         channel_name: str = self.port_initializer.name
         return channel_class(
@@ -120,7 +136,7 @@ class RuntimeChannelBuilderMp(AbstractChannelBuilder):
     used as messaging and multi processing backbone.
     """
 
-    channel_type: ChannelType
+    channel_type: ChannelTransferType
     src_process: ty.Union[RuntimeServiceBuilder, ty.Type["Runtime"]]
     dst_process: ty.Union[RuntimeServiceBuilder, ty.Type["Runtime"]]
     port_initializer: PortInitializer
@@ -144,9 +160,15 @@ def build(
         Exception
             Can't build channel of type specified
         """
-        channel_class = messaging_infrastructure.channel_class(
-            channel_type=self.channel_type
-        )
+        channel_factory = get_channel_factory()
+        shm = messaging_infrastructure.smm.alloc_mem(
+            self.port_initializer.size)
+        return channel_factory.get_channel(ChannelTransferType.SHMEMCHANNEL,
+                                           shm,
+                                           self.port_initializer.d_type,
+                                           self.port_initializer.size,
+                                           self.port_initializer.shape,
+                                           self.port_initializer.name)
 
         channel_name: str = self.port_initializer.name
         return channel_class(
@@ -165,7 +187,7 @@ class ChannelBuilderNx(AbstractChannelBuilder):
     infrastructure.
     """
 
-    channel_type: ChannelType
+    channel_type: ChannelTransferType
     src_process: "AbstractProcess"
     dst_process: "AbstractProcess"
     src_port_initializer: PortInitializer

src/lava/magma/compiler/builders/py_builder.py

@@ -7,15 +7,18 @@
 import numpy as np
 from lava.magma.compiler.builders.interfaces import AbstractProcessBuilder
 
-from lava.magma.compiler.channels.interfaces import AbstractCspPort
-from lava.magma.compiler.channels.pypychannel import CspRecvPort, CspSendPort
+from message_infrastructure import (
+    AbstractTransferPort,
+    RecvPort,
+    SendPort
+)
 from lava.magma.compiler.utils import (PortInitializer, VarInitializer,
                                        VarPortInitializer)
 from lava.magma.core.model.py.model import AbstractPyProcessModel
-from lava.magma.core.model.py.ports import (AbstractPyIOPort,
-                                            IdentityTransformer, PyInPort,
-                                            PyOutPort, PyRefPort, PyVarPort,
-                                            VirtualPortTransformer)
+from message_infrastructure.ports import (AbstractPyIOPort,
+                                          IdentityTransformer, PyInPort,
+                                          PyOutPort, PyRefPort, PyVarPort,
+                                          VirtualPortTransformer)
 from lava.magma.core.model.py.type import LavaPyType
 
 
@@ -60,10 +63,11 @@ def __init__(
         self.py_ports: ty.Dict[str, PortInitializer] = {}
         self.ref_ports: ty.Dict[str, PortInitializer] = {}
         self.var_ports: ty.Dict[str, VarPortInitializer] = {}
-        self.csp_ports: ty.Dict[str, ty.List[AbstractCspPort]] = {}
-        self._csp_port_map: ty.Dict[str, ty.Dict[str, AbstractCspPort]] = {}
-        self.csp_rs_send_port: ty.Dict[str, CspSendPort] = {}
-        self.csp_rs_recv_port: ty.Dict[str, CspRecvPort] = {}
+        self.csp_ports: ty.Dict[str, ty.List[AbstractTransferPort]] = {}
+        self._csp_port_map: ty.Dict[str,
+                                    ty.Dict[str, AbstractTransferPort]] = {}
+        self.csp_rs_send_port: ty.Dict[str, SendPort] = {}
+        self.csp_rs_recv_port: ty.Dict[str, RecvPort] = {}
         self.proc_params = proc_params
 
     def check_all_vars_and_ports_set(self):
@@ -165,13 +169,13 @@ def set_var_ports(self, var_ports: ty.List[VarPortInitializer]):
         self._check_not_assigned_yet(self.var_ports, new_ports.keys(), "ports")
         self.var_ports.update(new_ports)
 
-    def set_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
+    def set_csp_ports(self, csp_ports: ty.List[AbstractTransferPort]):
         """Appends the given list of CspPorts to the ProcessModel. Used by the
         runtime to configure csp ports during initialization (_build_channels).
 
         Parameters
         ----------
-        csp_ports : ty.List[AbstractCspPort]
+        csp_ports : ty.List[AbstractTransferPort]
 
 
         Raises
@@ -181,7 +185,7 @@ def set_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
         """
         new_ports = {}
         for p in csp_ports:
-            new_ports.setdefault(p.name, []).extend(
+            new_ports.setdefault(p.name(), []).extend(
                 p if isinstance(p, list) else [p]
             )
 
@@ -197,7 +201,8 @@ def set_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
             else:
                 self.csp_ports[port_name] = new_ports[port_name]
 
-    def add_csp_port_mapping(self, py_port_id: str, csp_port: AbstractCspPort):
+    def add_csp_port_mapping(self, py_port_id: str,
+                             csp_port: AbstractTransferPort):
         """Appends a mapping from a PyPort ID to a CSP port. This is used
         to associate a CSP port in a PyPort with transformation functions
         that implement the behavior of virtual ports.
@@ -212,10 +217,10 @@ def add_csp_port_mapping(self, py_port_id: str, csp_port: AbstractCspPort):
         """
         # Add or update the mapping
         self._csp_port_map.setdefault(
-            csp_port.name, {}
+            csp_port.name(), {}
         ).update({py_port_id: csp_port})
 
-    def set_rs_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
+    def set_rs_csp_ports(self, csp_ports: ty.List[AbstractTransferPort]):
         """Set RS CSP Ports
 
         Parameters
@@ -224,10 +229,10 @@ def set_rs_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
 
         """
         for port in csp_ports:
-            if isinstance(port, CspSendPort):
-                self.csp_rs_send_port.update({port.name: port})
-            if isinstance(port, CspRecvPort):
-                self.csp_rs_recv_port.update({port.name: port})
+            if isinstance(port, SendPort):
+                self.csp_rs_send_port.update({port.name(): port})
+            if isinstance(port, RecvPort):
+                self.csp_rs_recv_port.update({port.name(): port})
 
     def _get_lava_type(self, name: str) -> LavaPyType:
         return getattr(self.proc_model, name)
@@ -302,9 +307,9 @@ def build(self):
             if name in self.csp_ports:
                 csp_ports = self.csp_ports[name]
                 csp_recv = csp_ports[0] if isinstance(
-                    csp_ports[0], CspRecvPort) else csp_ports[1]
+                    csp_ports[0], RecvPort) else csp_ports[1]
                 csp_send = csp_ports[0] if isinstance(
-                    csp_ports[0], CspSendPort) else csp_ports[1]
+                    csp_ports[0], SendPort) else csp_ports[1]
 
             transformer = VirtualPortTransformer(
                 self._csp_port_map[name],
@@ -329,9 +334,9 @@ def build(self):
             if name in self.csp_ports:
                 csp_ports = self.csp_ports[name]
                 csp_recv = csp_ports[0] if isinstance(
-                    csp_ports[0], CspRecvPort) else csp_ports[1]
+                    csp_ports[0], RecvPort) else csp_ports[1]
                 csp_send = csp_ports[0] if isinstance(
-                    csp_ports[0], CspSendPort) else csp_ports[1]
+                    csp_ports[0], SendPort) else csp_ports[1]
 
             transformer = VirtualPortTransformer(
                 self._csp_port_map[name],

src/lava/magma/compiler/builders/runtimeservice_builder.py

@@ -5,8 +5,12 @@
 import logging
 import typing as ty
 
-from lava.magma.compiler.channels.interfaces import AbstractCspPort
-from lava.magma.compiler.channels.pypychannel import CspRecvPort, CspSendPort
+from message_infrastructure import (
+    AbstractTransferPort,
+    RecvPort,
+    SendPort
+)
+
 from lava.magma.core.sync.protocol import AbstractSyncProtocol
 from lava.magma.runtime.runtime_services.enums import LoihiVersion
 from lava.magma.runtime.runtime_services.runtime_service import \
@@ -49,17 +53,17 @@ def __init__(
         self.log.setLevel(loglevel)
         self._runtime_service_id = runtime_service_id
         self._model_ids: ty.List[int] = model_ids
-        self.csp_send_port: ty.Dict[str, CspSendPort] = {}
-        self.csp_recv_port: ty.Dict[str, CspRecvPort] = {}
-        self.csp_proc_send_port: ty.Dict[str, CspSendPort] = {}
-        self.csp_proc_recv_port: ty.Dict[str, CspRecvPort] = {}
+        self.csp_send_port: ty.Dict[str, SendPort] = {}
+        self.csp_recv_port: ty.Dict[str, RecvPort] = {}
+        self.csp_proc_send_port: ty.Dict[str, SendPort] = {}
+        self.csp_proc_recv_port: ty.Dict[str, RecvPort] = {}
         self.loihi_version: ty.Type[LoihiVersion] = loihi_version
 
     @property
     def runtime_service_id(self):
         return self._runtime_service_id
 
-    def set_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
+    def set_csp_ports(self, csp_ports: ty.List[AbstractTransferPort]):
         """Set CSP Ports
 
         Parameters
@@ -68,12 +72,12 @@ def set_csp_ports(self, csp_ports: ty.List[AbstractCspPort]):
 
         """
         for port in csp_ports:
-            if isinstance(port, CspSendPort):
-                self.csp_send_port.update({port.name: port})
-            if isinstance(port, CspRecvPort):
-                self.csp_recv_port.update({port.name: port})
+            if isinstance(port, SendPort):
+                self.csp_send_port.update({port.name(): port})
+            if isinstance(port, RecvPort):
+                self.csp_recv_port.update({port.name(): port})
 
-    def set_csp_proc_ports(self, csp_ports: ty.List[AbstractCspPort]):
+    def set_csp_proc_ports(self, csp_ports: ty.List[AbstractTransferPort]):
         """Set CSP Process Ports
 
         Parameters
@@ -82,10 +86,10 @@ def set_csp_proc_ports(self, csp_ports: ty.List[AbstractCspPort]):
 
         """
         for port in csp_ports:
-            if isinstance(port, CspSendPort):
-                self.csp_proc_send_port.update({port.name: port})
-            if isinstance(port, CspRecvPort):
-                self.csp_proc_recv_port.update({port.name: port})
+            if isinstance(port, SendPort):
+                self.csp_proc_send_port.update({port.name(): port})
+            if isinstance(port, RecvPort):
+                self.csp_proc_recv_port.update({port.name(): port})
 
     def build(self) -> AbstractRuntimeService:
         """Build the runtime service
@@ -115,21 +119,21 @@ def build(self) -> AbstractRuntimeService:
 
         if not nxsdk_rts:
             for port in self.csp_proc_send_port.values():
-                if "service_to_process" in port.name:
+                if "service_to_process" in port.name():
                     rs.service_to_process.append(port)
 
             for port in self.csp_proc_recv_port.values():
-                if "process_to_service" in port.name:
+                if "process_to_service" in port.name():
                     rs.process_to_service.append(port)
 
             self.log.debug("Setup 'RuntimeService <--> Rrocess; ports")
 
         for port in self.csp_send_port.values():
-            if "service_to_runtime" in port.name:
+            if "service_to_runtime" in port.name():
                 rs.service_to_runtime = port
 
         for port in self.csp_recv_port.values():
-            if "runtime_to_service" in port.name:
+            if "runtime_to_service" in port.name():
                 rs.runtime_to_service = port
 
         self.log.debug("Setup 'Runtime <--> RuntimeService' ports")

src/lava/magma/compiler/channels/pypychannel.py

@@ -19,6 +19,7 @@
     from message_infrastructure.message_infrastructure_interface \
         import (
             MessageInfrastructureInterface)
+"""Depricated"""
 
 
 @dataclass

src/lava/magma/compiler/compiler.py

@@ -68,6 +68,8 @@ class AbstractNcProcessModel:
 from lava.magma.runtime.runtime import Runtime
 from lava.magma.runtime.runtime_services.enums import LoihiVersion
 
+from message_infrastructure import ChannelTransferType
+
 
 class Compiler:
     """Lava processes Compiler, called from any process in a process network.
@@ -728,7 +730,7 @@ def _create_sync_channel_builders(
         sync_channel_builders: ty.List[AbstractChannelBuilder] = []
         for sync_domain in rsb:
             runtime_to_service = RuntimeChannelBuilderMp(
-                ChannelType.PyPy,
+                ChannelTransferType.SHMEMCHANNEL,
                 Runtime,
                 rsb[sync_domain],
                 self._create_mgmt_port_initializer(
@@ -738,7 +740,7 @@ def _create_sync_channel_builders(
             sync_channel_builders.append(runtime_to_service)
 
             service_to_runtime = RuntimeChannelBuilderMp(
-                ChannelType.PyPy,
+                ChannelTransferType.SHMEMCHANNEL,
                 rsb[sync_domain],
                 Runtime,
                 self._create_mgmt_port_initializer(
@@ -750,7 +752,7 @@ def _create_sync_channel_builders(
             for process in sync_domain.processes:
                 if issubclass(process.model_class, AbstractPyProcessModel):
                     service_to_process = ServiceChannelBuilderMp(
-                        ChannelType.PyPy,
+                        ChannelTransferType.SHMEMCHANNEL,
                         rsb[sync_domain],
                         process,
                         self._create_mgmt_port_initializer(
@@ -760,7 +762,7 @@ def _create_sync_channel_builders(
                     sync_channel_builders.append(service_to_process)
 
                     process_to_service = ServiceChannelBuilderMp(
-                        ChannelType.PyPy,
+                        ChannelTransferType.SHMEMCHANNEL,
                         process,
                         rsb[sync_domain],
                         self._create_mgmt_port_initializer(