Refactoring py/ports

src/lava/magma/compiler/builder.py

@@ -25,10 +25,7 @@
 from lava.magma.compiler.utils import VarInitializer, PortInitializer, \
     VarPortInitializer
 from lava.magma.core.model.py.ports import (
-    AbstractPyPort,
-    PyInPort,
-    PyOutPort,
-    PyRefPort, PyVarPort,
+    PyInPort, PyOutPort, PyRefPort, PyVarPort, AbstractIOPyPort
 )
 from lava.magma.compiler.channels.interfaces import AbstractCspPort, Channel, \
     ChannelType
@@ -363,7 +360,7 @@ def build(self):
         for name, p in self.py_ports.items():
             # Build PyPort
             lt = self._get_lava_type(name)
-            port_cls = ty.cast(ty.Type[AbstractPyPort], lt.cls)
+            port_cls = ty.cast(ty.Type[AbstractIOPyPort], lt.cls)
             csp_ports = []
             if name in self.csp_ports:
                 csp_ports = self.csp_ports[name]

src/lava/magma/compiler/compiler.py

@@ -34,7 +34,7 @@
 from lava.magma.core.model.py.ports import RefVarTypeMapping
 from lava.magma.core.model.sub.model import AbstractSubProcessModel
 from lava.magma.core.process.ports.ports import AbstractPort, VarPort, \
-    ImplicitVarPort
+    ImplicitVarPort, RefPort
 from lava.magma.core.process.process import AbstractProcess
 from lava.magma.core.resources import CPU, NeuroCore
 from lava.magma.core.run_configs import RunConfig
@@ -217,21 +217,10 @@ def _propagate_var_ports(proc: AbstractProcess):
         for vp in proc.var_ports:
             v = vp.var.aliased_var
             if v is not None:
-                sub_proc = v.process
                 # Create an implicit Var port in the sub process
-                new_vp = ImplicitVarPort(v)
-                # Propagate name and parent process of Var to VarPort
-                new_vp.name = "_" + v.name + "_implicit_port"
-                new_vp.process = sub_proc
-                # VarPort name could shadow existing attribute
-                if hasattr(sub_proc, new_vp.name):
-                    raise AssertionError(
-                        "Name of implicit VarPort might conflict"
-                        " with existing attribute.")
-                setattr(sub_proc, new_vp.name, new_vp)
-                sub_proc.var_ports.add_members({new_vp.name: new_vp})
+                imp_vp = RefPort.create_implicit_var_port(v)
                 # Connect the VarPort to the new VarPort
-                vp.connect(new_vp)
+                vp.connect(imp_vp)
 
     def _expand_sub_proc_model(self,
                                model_cls: ty.Type[AbstractSubProcessModel],

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

@@ -9,6 +9,7 @@
 from lava.magma.compiler.channels.interfaces import AbstractCspPort
 from lava.magma.compiler.channels.pypychannel import CspSendPort, CspRecvPort
 from lava.magma.core.model.interfaces import AbstractPortImplementation
+from lava.magma.core.model.model import AbstractProcessModel
 from lava.magma.runtime.mgmt_token_enums import enum_to_np, enum_equal
 
 
@@ -20,7 +21,22 @@ def csp_ports(self) -> ty.List[AbstractCspPort]:
         pass
 
 
-class PyInPort(AbstractPyPort):
+class AbstractIOPyPort(AbstractPyPort):
+    def __init__(self,
+                 csp_ports: ty.List[AbstractCspPort],
+                 process_model: AbstractProcessModel,
+                 shape: ty.Tuple[int, ...] = tuple(),
+                 d_type: type = int):
+        self._csp_ports = csp_ports
+        super().__init__(process_model, shape, d_type)
+
+    @property
+    def csp_ports(self) -> ty.List[AbstractCspPort]:
+        """Returns all csp ports of the port."""
+        return self._csp_ports
+
+
+class PyInPort(AbstractIOPyPort):
     """Python implementation of InPort used within AbstractPyProcessModel.
     If buffer is empty, recv() will be blocking.
     """
@@ -30,15 +46,6 @@ class PyInPort(AbstractPyPort):
     SCALAR_DENSE: ty.Type["PyInPortScalarDense"] = None
     SCALAR_SPARSE: ty.Type["PyInPortScalarSparse"] = None
 
-    def __init__(self, csp_recv_ports: ty.List[CspRecvPort], *args):
-        self._csp_recv_ports = csp_recv_ports
-        super().__init__(*args)
-
-    @property
-    def csp_ports(self) -> ty.List[AbstractCspPort]:
-        """Returns all csp ports of the port."""
-        return self._csp_recv_ports
-
     @abstractmethod
     def recv(self):
         pass
@@ -57,7 +64,7 @@ def probe(self) -> bool:
         # Returns True only when probe returns True for all _csp_recv_ports.
         return ft.reduce(
             lambda acc, csp_port: acc and csp_port.probe(),
-            self._csp_recv_ports,
+            self.csp_ports,
             True,
         )
 
@@ -66,14 +73,14 @@ class PyInPortVectorDense(PyInPort):
     def recv(self) -> np.ndarray:
         return ft.reduce(
             lambda acc, csp_port: acc + csp_port.recv(),
-            self._csp_recv_ports,
+            self.csp_ports,
             np.zeros(self._shape, self._d_type),
         )
 
     def peek(self) -> np.ndarray:
         return ft.reduce(
             lambda acc, csp_port: acc + csp_port.peek(),
-            self._csp_recv_ports,
+            self.csp_ports,
             np.zeros(self._shape, self._d_type),
         )
 
@@ -108,23 +115,14 @@ def peek(self) -> ty.Tuple[int, int]:
 PyInPort.SCALAR_SPARSE = PyInPortScalarSparse
 
 
-class PyOutPort(AbstractPyPort):
+class PyOutPort(AbstractIOPyPort):
     """Python implementation of OutPort used within AbstractPyProcessModels."""
 
     VEC_DENSE: ty.Type["PyOutPortVectorDense"] = None
     VEC_SPARSE: ty.Type["PyOutPortVectorSparse"] = None
     SCALAR_DENSE: ty.Type["PyOutPortScalarDense"] = None
     SCALAR_SPARSE: ty.Type["PyOutPortScalarSparse"] = None
 
-    def __init__(self, csp_send_ports: ty.List[CspSendPort], *args):
-        self._csp_send_ports = csp_send_ports
-        super().__init__(*args)
-
-    @property
-    def csp_ports(self) -> ty.List[AbstractCspPort]:
-        """Returns all csp ports of the port."""
-        return self._csp_send_ports
-
     @abstractmethod
     def send(self, data: ty.Union[np.ndarray, int]):
         pass
@@ -136,7 +134,7 @@ def flush(self):
 class PyOutPortVectorDense(PyOutPort):
     def send(self, data: np.ndarray):
         """Sends data only if port is not dangling."""
-        for csp_port in self._csp_send_ports:
+        for csp_port in self.csp_ports:
             csp_port.send(data)
 
 
@@ -176,10 +174,13 @@ class PyRefPort(AbstractPyPort):
 
     def __init__(self,
                  csp_send_port: ty.Optional[CspSendPort],
-                 csp_recv_port: ty.Optional[CspRecvPort], *args):
+                 csp_recv_port: ty.Optional[CspRecvPort],
+                 process_model: AbstractProcessModel,
+                 shape: ty.Tuple[int, ...] = tuple(),
+                 d_type: type = int):
         self._csp_recv_port = csp_recv_port
         self._csp_send_port = csp_send_port
-        super().__init__(*args)
+        super().__init__(process_model, shape, d_type)
 
     @property
     def csp_ports(self) -> ty.List[AbstractCspPort]:
@@ -270,11 +271,14 @@ class PyVarPort(AbstractPyPort):
     def __init__(self,
                  var_name: str,
                  csp_send_port: ty.Optional[CspSendPort],
-                 csp_recv_port: ty.Optional[CspRecvPort], *args):
+                 csp_recv_port: ty.Optional[CspRecvPort],
+                 process_model: AbstractProcessModel,
+                 shape: ty.Tuple[int, ...] = tuple(),
+                 d_type: type = int):
         self._csp_recv_port = csp_recv_port
         self._csp_send_port = csp_send_port
         self.var_name = var_name
-        super().__init__(*args)
+        super().__init__(process_model, shape, d_type)
 
     @property
     def csp_ports(self) -> ty.List[AbstractCspPort]:

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

@@ -46,11 +46,11 @@ def __init__(self, shape: ty.Tuple):
         self.out_connections: ty.List[AbstractPort] = []
 
     def _validate_ports(
-        self,
-        ports: ty.List["AbstractPort"],
-        port_type: ty.Type["AbstractPort"],
-        assert_same_shape: bool = True,
-        assert_same_type: bool = False,
+            self,
+            ports: ty.List["AbstractPort"],
+            port_type: ty.Type["AbstractPort"],
+            assert_same_shape: bool = True,
+            assert_same_type: bool = False,
     ):
         """Checks that each port in 'ports' is of type 'port_type' and that
         shapes of each port is identical to this port's shape."""
@@ -87,11 +87,11 @@ def _add_outputs(self, outputs: ty.List["AbstractPort"]):
         self.out_connections += outputs
 
     def _connect_forward(
-        self,
-        ports: ty.List["AbstractPort"],
-        port_type: ty.Type["AbstractPort"],
-        assert_same_shape: bool = True,
-        assert_same_type: bool = True,
+            self,
+            ports: ty.List["AbstractPort"],
+            port_type: ty.Type["AbstractPort"],
+            assert_same_shape: bool = True,
+            assert_same_type: bool = True,
     ):
         """Creates a forward connection from this AbstractPort to other
         ports by adding other ports to this AbstractPort's out_connection and
@@ -107,11 +107,11 @@ def _connect_forward(
             p._add_inputs([self])
 
     def _connect_backward(
-        self,
-        ports: ty.List["AbstractPort"],
-        port_type: ty.Type["AbstractPort"],
-        assert_same_shape: bool = True,
-        assert_same_type: bool = True,
+            self,
+            ports: ty.List["AbstractPort"],
+            port_type: ty.Type["AbstractPort"],
+            assert_same_shape: bool = True,
+            assert_same_type: bool = True,
     ):
         """Creates a backward connection from other ports to this
         AbstractPort by adding other ports to this AbstractPort's
@@ -181,9 +181,9 @@ def flatten(self) -> "ReshapePort":
         return self.reshape((self.size,))
 
     def concat_with(
-        self,
-        ports: ty.Union["AbstractPort", ty.List["AbstractPort"]],
-        axis: int,
+            self,
+            ports: ty.Union["AbstractPort", ty.List["AbstractPort"]],
+            axis: int,
     ) -> "ConcatPort":
         """Concatenates this port with other ports in given order along given
         axis by deriving and returning a new virtual ConcatPort. This implies
@@ -254,7 +254,7 @@ class OutPort(AbstractIOPort, AbstractSrcPort):
     """
 
     def connect(
-        self, ports: ty.Union["AbstractIOPort", ty.List["AbstractIOPort"]]
+            self, ports: ty.Union["AbstractIOPort", ty.List["AbstractIOPort"]]
     ):
         """Connects this OutPort to other InPort(s) of another process
         or to OutPort(s) of its parent process.
@@ -287,9 +287,9 @@ class InPort(AbstractIOPort, AbstractDstPort):
     """
 
     def __init__(
-        self,
-        shape: ty.Tuple,
-        reduce_op: ty.Optional[ty.Type[AbstractReduceOp]] = None,
+            self,
+            shape: ty.Tuple,
+            reduce_op: ty.Optional[ty.Type[AbstractReduceOp]] = None,
     ):
         super().__init__(shape)
         self._reduce_op = reduce_op
@@ -305,7 +305,7 @@ def connect(self, ports: ty.Union["InPort", ty.List["InPort"]]):
         self._connect_forward(to_list(ports), InPort)
 
     def connect_from(
-        self, ports: ty.Union["AbstractIOPort", ty.List["AbstractIOPort"]]
+            self, ports: ty.Union["AbstractIOPort", ty.List["AbstractIOPort"]]
     ):
         """Connects other OutPort(s) to this InPort or connects other
         InPort(s) of parent process to this InPort.
@@ -339,7 +339,7 @@ class RefPort(AbstractRVPort, AbstractSrcPort):
     RefPort to a Var via the connect_var(..) method."""
 
     def connect(
-        self, ports: ty.Union["AbstractRVPort", ty.List["AbstractRVPort"]]
+            self, ports: ty.Union["AbstractRVPort", ty.List["AbstractRVPort"]]
     ):
         """Connects this RefPort to other VarPort(s) of another process
         or to RefPort(s) of its parent process.
@@ -440,19 +440,7 @@ def connect_var(self, variables: ty.Union[Var, ty.List[Var]]):
             if var_shape != v.shape:
                 raise AssertionError("All 'vars' must have same shape.")
             # Create a VarPort to wrap Var
-            vp = ImplicitVarPort(v)
-            # Propagate name and parent process of Var to VarPort
-            vp.name = "_" + v.name + "_implicit_port"
-            if v.process is not None:
-                # Only assign when parent process is already assigned
-                vp.process = v.process
-                # VarPort name could shadow existing attribute
-                if hasattr(v.process, vp.name):
-                    raise AssertionError(
-                        "Name of implicit VarPort might conflict"
-                        " with existing attribute.")
-                setattr(v.process, vp.name, vp)
-                v.process.var_ports.add_members({vp.name: vp})
+            vp = self.create_implicit_var_port(v)
             var_ports.append(vp)
         # Connect RefPort to VarPorts that wrap Vars
         self.connect(var_ports)
@@ -461,6 +449,26 @@ def get_dst_vars(self) -> ty.List[Var]:
         """Returns destination Vars this RefPort is connected to."""
         return [ty.cast(VarPort, p).var for p in self.get_dst_ports()]
 
+    @staticmethod
+    def create_implicit_var_port(var: Var) -> "ImplicitVarPort":
+        """Creates and returns an ImplicitVarPort for the given Var."""
+        # Create a VarPort to wrap Var
+        vp = ImplicitVarPort(var)
+        # Propagate name and parent process of Var to VarPort
+        vp.name = "_" + var.name + "_implicit_port"
+        if var.process is not None:
+            # Only assign when parent process is already assigned
+            vp.process = var.process
+            # VarPort name could shadow existing attribute
+            if hasattr(var.process, vp.name):
+                raise AssertionError(
+                    "Name of implicit VarPort might conflict"
+                    " with existing attribute.")
+            setattr(var.process, vp.name, vp)
+            var.process.var_ports.add_members({vp.name: vp})
+
+        return vp
+
 
 # TODO: (PP) enable connecting multiple VarPorts/RefPorts to a VarPort
 class VarPort(AbstractRVPort, AbstractDstPort):
@@ -521,7 +529,7 @@ def connect(self, ports: ty.Union["VarPort", ty.List["VarPort"]]):
         self._connect_forward(to_list(ports), VarPort)
 
     def connect_from(
-        self, ports: ty.Union["AbstractRVPort", ty.List["AbstractRVPort"]]
+            self, ports: ty.Union["AbstractRVPort", ty.List["AbstractRVPort"]]
     ):
         """Connects other RefPort(s) to this VarPort or connects other
         VarPort(s) of parent process to this VarPort.

tests/lava/magma/core/model/test_py_model.py

@@ -432,14 +432,14 @@ def test_build_with_dangling_ports(self):
         # Validate that the Process with no OutPorts indeed has no output
         # CspPort
         self.assertIsInstance(
-            pm_with_no_out_ports.in_port._csp_recv_ports[0], FakeCspPort)
-        self.assertEqual(pm_with_no_out_ports.out_port._csp_send_ports, [])
+            pm_with_no_out_ports.in_port.csp_ports[0], FakeCspPort)
+        self.assertEqual(pm_with_no_out_ports.out_port.csp_ports, [])
 
         # Validate that the Process with no InPorts indeed has no input
         # CspPort
-        self.assertEqual(pm_with_no_in_ports.in_port._csp_recv_ports, [])
+        self.assertEqual(pm_with_no_in_ports.in_port.csp_ports, [])
         self.assertIsInstance(
-            pm_with_no_in_ports.out_port._csp_send_ports[0], FakeCspPort)
+            pm_with_no_in_ports.out_port.csp_ports[0], FakeCspPort)
 
     def test_set_ref_var_ports(self):
         """Check RefPorts and VarPorts can be set."""