Folded view (#774)

* added the folded_view compilation option and propagated folded_views
* enable folded compilation for dynamically created workloads

---------

Co-authored-by: Marcus G K Williams <168222+mgkwill@users.noreply.github.com>

Author: ymeng-git

src/lava/magma/compiler/compiler.py

@@ -130,7 +130,9 @@ def compile(
         """
         # Group and sort all Processes connected to 'process' into a list of
         # ProcGroups.
-        proc_group_digraph = ProcGroupDiGraphs(process, run_cfg)
+        proc_group_digraph = ProcGroupDiGraphs(process,
+                                               run_cfg,
+                                               self._compile_config)
         proc_groups: ty.List[ProcGroup] = proc_group_digraph.get_proc_groups()
         # Get a flattened list of all AbstractProcesses
         process_list = list(itertools.chain.from_iterable(proc_groups))

src/lava/magma/compiler/compiler_graphs.py

@@ -163,6 +163,23 @@ def find_processes(proc: AbstractProcess,
     return seen_procs
 
 
+def annotate_folded_view(proc_list: ty.List[AbstractProcess],
+                         folded_procs: ty.List[str] = None):
+    """Annotate folded views and propagate them recursively
+    """
+    annotated : ty.Set(AbstractProcess) = set()
+    fv_inst_id : int = 0
+    for p in proc_list:
+        if p.__class__.__name__ in folded_procs:
+            p.folded_view = p.__class__
+            p.folded_view_inst_id = fv_inst_id
+            fv_inst_id += 1
+            annotated.add(p)
+
+    for p in annotated:
+        p.propagate_folded_views()
+
+
 class DiGraphBase(ntx.DiGraph):
     """Base class for directed graphs in the compiler.
 
@@ -600,26 +617,36 @@ class ProcGroupDiGraphs(AbstractProcGroupDiGraphs):
     sorted order of nodes.
     """
 
-    def __init__(self, proc: AbstractProcess, run_cfg: RunConfig):
+    def __init__(self, proc: AbstractProcess, run_cfg: RunConfig,
+                 compile_config: ty.Optional[ty.Dict[str, ty.Any]] = None):
 
         self._base_proc = proc  # Process on which compile/run was called
         self._run_cfg = run_cfg
+        self._compile_config = compile_config
         # 1. Find all Processes
         proc_list = find_processes(proc)
+
         # Check if any Process in proc_list is already compiled
         for p in proc_list:
             if p.is_compiled:
                 raise ex.ProcessAlreadyCompiled(p)
             p._is_compiled = True
+
         # Number of Processes before resolving HierarchicalProcesses
         self._num_procs_pre_pm_discovery = len(proc_list)
         # 2. Generate a ProcessGraph: This does not resolve
         # HierarchicalProcesses yet, because we have not discovered
         # ProcessModels
         self._raw_proc_digraph = ProcDiGraph(proc_list=proc_list)
         # 3. Find and select ProcessModels based on RunConfig:
+        elab_procs = []
         proc_procmodel_map = ProcGroupDiGraphs._map_proc_to_model(proc_list,
-                                                                  self._run_cfg)
+                                                                  self._run_cfg,
+                                                                  elab_procs)
+        if self._compile_config and 'folded_view' in self._compile_config:
+            folded_views = self._compile_config["folded_view"]
+            annotate_folded_view(elab_procs, folded_views)
+
         # Assign ProcessModels to Processes
         for p, pm in proc_procmodel_map.items():
             p._model_class = pm
@@ -939,7 +966,8 @@ def _propagate_var_ports(proc: AbstractProcess):
 
     @staticmethod
     def _expand_sub_proc_model(model_cls: ty.Type[AbstractSubProcessModel],
-                               proc: AbstractProcess, run_cfg: RunConfig):
+                               proc: AbstractProcess, run_cfg: RunConfig,
+                               elab_procs : ty.List[AbstractProcess]):
         """Expand a SubProcessModel by building it, extracting the
         sub-Processes contained within, and mapping the sub-Processes
         recursively to their ProcessModels.
@@ -969,11 +997,13 @@ def _expand_sub_proc_model(model_cls: ty.Type[AbstractSubProcessModel],
         proc.validate_var_aliases()
         # Recursively map sub processes to their ProcModel
         return ProcGroupDiGraphs._map_proc_to_model(list(sub_procs.values()),
-                                                    run_cfg)
+                                                    run_cfg,
+                                                    elab_procs)
 
     @staticmethod
     def _map_proc_to_model(procs: ty.List[AbstractProcess],
-                           run_cfg: RunConfig) -> ProcMap:
+                           run_cfg: RunConfig,
+                           elab_procs : ty.List[AbstractProcess]) -> ProcMap:
         """Associate each Process in a list of Processes to the corresponding
         ProcessModel as selected by run_cfg.
 
@@ -1006,17 +1036,20 @@ def _map_proc_to_model(procs: ty.List[AbstractProcess],
                 model_cls = ProcGroupDiGraphs._select_proc_models(proc,
                                                                   models_cls,
                                                                   run_cfg)
-            if issubclass(model_cls, AbstractSubProcessModel):
+            if model_cls and issubclass(model_cls, AbstractSubProcessModel):
                 # Recursively substitute SubProcModel by sub processes
                 sub_map = ProcGroupDiGraphs._expand_sub_proc_model(model_cls,
                                                                    proc,
-                                                                   run_cfg)
+                                                                   run_cfg,
+                                                                   elab_procs)
                 proc_map.update(sub_map)
                 proc._model_class = model_cls
-            else:
+            elif model_cls:
                 # Just map current Process to selected ProcessModel
                 proc_map[proc] = model_cls
 
+            elab_procs.append(proc)
+
         return proc_map
 
     def _collapse_isomodel_procs(self) -> ProcDiGraph:

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

@@ -211,6 +211,10 @@ def __init__(self, **proc_params) -> None:
         # Flag indicating whether process has been compiled already.
         self._is_compiled: bool = False
 
+        # folded view
+        self._folded_view : ty.Optional[AbstractProcess] = None
+        self._folded_view_inst_id : int = -1
+
         # Current runtime environment
         self._runtime: ty.Optional[Runtime] = None
 
@@ -253,6 +257,10 @@ def __init__(self, ..):
         self._init_proc_member_obj(attrs)
         self.vars.add_members(attrs)
 
+        attrs = self._find_attr_by_type(AbstractProcess)
+        self._init_proc_member_obj(attrs)
+        self.procs.add_members(attrs)
+
     def _find_attr_by_type(self, cls) -> ty.Dict:
         """Finds all class attributes of a certain class type."""
         attrs = dict()
@@ -303,6 +311,12 @@ def is_sub_proc_of(self, proc: AbstractProcess):
         else:
             return False
 
+    def propagate_folded_views(self):
+        for p in self.procs:
+            p._folded_view = self._folded_view
+            p._folded_view_inst_id = self._folded_view_inst_id
+            p.propagate_folded_views()
+
     def run(self,
             condition: AbstractRunCondition,
             run_cfg: ty.Optional[RunConfig] = None,
@@ -444,6 +458,23 @@ def is_compiled(self):
         """Returns True if process has been compiled."""
         return self._is_compiled
 
+    @property
+    def folded_view(self) -> ty.Type["AbstractProcess"]:
+        """ Return folded view process"""
+        return self._folded_view
+
+    @folded_view.setter
+    def folded_view(self, folded_view : ty.Optional[AbstractProcess] = None):
+        self._folded_view = folded_view
+
+    @property
+    def folded_view_inst_id(self) -> int:
+        return self._folded_view_inst_id
+
+    @folded_view_inst_id.setter
+    def folded_view_inst_id(self, inst_id : int = -1):
+        self._folded_view_inst_id = inst_id
+
 
 class ProcessParameters:
     """Wrapper around a dictionary that is used to pass parameters from a