From 8f04e7015ab74d29bfb8e310a103e9f6061f4070 Mon Sep 17 00:00:00 2001
From: Lukas Rothenberger <lukas.rothenberger@tu-darmstadt.de>
Date: Wed, 22 Nov 2023 14:44:37 +0100
Subject: [PATCH] feat(PET): parallel computation of function metadata

---
 discopop_explorer/PETGraphX.py      | 48 ++++++++++++++++-------------
 discopop_explorer/parallel_utils.py | 35 +++++++++++++++++++++
 2 files changed, 62 insertions(+), 21 deletions(-)
 create mode 100644 discopop_explorer/parallel_utils.py

diff --git a/discopop_explorer/PETGraphX.py b/discopop_explorer/PETGraphX.py
index d4a14fc7d..3885b1484 100644
--- a/discopop_explorer/PETGraphX.py
+++ b/discopop_explorer/PETGraphX.py
@@ -22,6 +22,7 @@
 from .parser import LoopData, readlineToCUIdMap, writelineToCUIdMap, DependenceItem
 from .variable import Variable
 
+global_pet = None
 
 # unused
 # node_props = [
@@ -362,6 +363,7 @@ def get_exit_cu_ids(self, pet: PETGraphX) -> Set[NodeID]:
         return exit_cu_ids
 
     def calculate_reachability_pairs(self, pet: PETGraphX):
+        reachability_pairs: Dict[NodeID, Set[NodeID]] = dict()
         # create graph copy and remove all but successor edges
         copied_graph = pet.g.copy()
 
@@ -376,10 +378,10 @@ def calculate_reachability_pairs(self, pet: PETGraphX):
 
         # calculate dfs successors for children CUs
         for node_id in cast(List[NodeID], self.children_cu_ids):
-            self.reachability_pairs[node_id] = {node_id}
+            reachability_pairs[node_id] = {node_id}
             successors = [t for s, t in nx.dfs_tree(copied_graph, node_id).edges()]
-            self.reachability_pairs[node_id].update(successors)
-        pass
+            reachability_pairs[node_id].update(successors)
+        return reachability_pairs
 
     def get_immediate_post_dominators(self, pet: PETGraphX) -> Dict[NodeID, NodeID]:
         if self.immediate_post_dominators_present:
@@ -724,24 +726,28 @@ def calculateFunctionMetadata(self) -> None:
         # store id of parent function in each node
         # and store in each function node a list of all children ids
         func_nodes = self.all_nodes(FunctionNode)
-        print("Calculating metadata for functions: ")
-        with alive_bar(len(func_nodes)) as progress_bar:
-            for func_node in func_nodes:
-                stack: List[Node] = self.direct_children(func_node)
-                func_node.children_cu_ids = [node.id for node in stack]
-
-                while stack:
-                    child = stack.pop()
-                    child.parent_function_id = func_node.id
-                    children = self.direct_children(child)
-                    func_node.children_cu_ids.extend([node.id for node in children])
-                    stack.extend(children)
-
-                func_node.calculate_reachability_pairs(self)
-                progress_bar()
-        print("\tDone.")
-
-        print("Metadata calculation done.")
+        print("Calculating local metadata results for functions...")
+        import tqdm  # type: ignore
+        from multiprocessing import Pool
+        from .parallel_utils import pet_function_metadata_initialize_worker, pet_function_metadata_parse_func
+
+        param_list = func_nodes
+        with Pool(initializer=pet_function_metadata_initialize_worker, initargs=(self,)) as pool:
+            tmp_result = list(
+                tqdm.tqdm(pool.imap_unordered(pet_function_metadata_parse_func, param_list), total=len(param_list))
+            )
+        # calculate global result
+        print("Calculating global result...")
+        global_reachability_dict: Dict[NodeID, Set[NodeID]] = dict()
+        for local_result in tmp_result:
+            parsed_function_id, local_reachability_dict, local_children_ids = local_result
+            # set reachability values to function nodes
+            cast(FunctionNode, self.node_at(parsed_function_id)).reachability_pairs = local_reachability_dict
+            # set parent function for visited nodes
+            for child_id in local_children_ids:
+                self.node_at(child_id).parent_function_id = parsed_function_id
+
+        print("\tMetadata calculation done.")
 
         # cleanup dependencies (remove dependencies, if it is overwritten by a more specific Intra-iteration dependency
         print("Cleaning duplicated dependencies...")
diff --git a/discopop_explorer/parallel_utils.py b/discopop_explorer/parallel_utils.py
new file mode 100644
index 000000000..0948bd6f9
--- /dev/null
+++ b/discopop_explorer/parallel_utils.py
@@ -0,0 +1,35 @@
+# This file is part of the DiscoPoP software (http://www.discopop.tu-darmstadt.de)
+#
+# Copyright (c) 2020, Technische Universitaet Darmstadt, Germany
+#
+# This software may be modified and distributed under the terms of
+# the 3-Clause BSD License.  See the LICENSE file in the package base
+# directory for details.
+
+from .PETGraphX import Node, NodeID
+from typing import List, Set
+
+global_pet = None
+
+
+def pet_function_metadata_initialize_worker(pet):
+    global global_pet
+    global_pet = pet
+
+
+def pet_function_metadata_parse_func(param_tuple):
+    func_node = param_tuple
+    stack: List[Node] = global_pet.direct_children(func_node)
+    func_node.children_cu_ids = [node.id for node in stack]
+    local_children: Set[NodeID] = set()
+
+    while stack:
+        child = stack.pop()
+        local_children.add(child.id)
+        children = global_pet.direct_children(child)
+        func_node.children_cu_ids.extend([node.id for node in children])
+        stack.extend(children)
+
+    local_reachability_dict = func_node.calculate_reachability_pairs(global_pet)
+    local_result = func_node.id, local_reachability_dict, local_children
+    return local_result