Dependency graph construction bug fix in progress

et_converter/pytorch2chakra_converter.py

@@ -146,10 +146,10 @@ def convert(self) -> None:
         self.split_cpu_nodes_with_gpu_child()
 
         for pytorch_nid, pytorch_node in self.pytorch_nodes.items():
-            if pytorch_node.is_cpu_op():
+            if (pytorch_node.get_op_type() == PyTorchNodeType.CPU_OP)\
+            or (pytorch_node.get_op_type() == PyTorchNodeType.LABEL):
                 self.update_input_tensor_map(pytorch_node.id, pytorch_node.inputs)
                 self.update_output_tensor_map(pytorch_node.id, pytorch_node.outputs)
-
                 if pytorch_node.child_gpu:
                     pytorch_gpu_node = pytorch_node.child_gpu
                     self.update_input_tensor_map(pytorch_gpu_node.id, pytorch_gpu_node.inputs)
@@ -344,6 +344,7 @@ def split_cpu_nodes_with_gpu_child(self) -> None:
                 cpu_node.id = new_cpu_node_id
                 for child_node in cpu_node.children:
                     child_node.parent = cpu_node.id
+                    child_node.add_ctrl_dep(cpu_node)
                 updated_pytorch_nodes[new_cpu_node_id] = cpu_node
             else:
                 gpu_node = cpu_node.child_gpu
@@ -354,9 +355,9 @@ def split_cpu_nodes_with_gpu_child(self) -> None:
 
                 cpu_node_first, cpu_node_second, updated_gpu_node =\
                         self._split_cpu_node(cpu_node, gpu_node)
-                updated_pytorch_nodes[cpu_node_first.id] = cpu_node_first
-                updated_pytorch_nodes[cpu_node_second.id] = cpu_node_second
-                updated_pytorch_nodes[updated_gpu_node.id] = updated_gpu_node
+                updated_pytorch_nodes[cpu_node_first.id] = copy.deepcopy(cpu_node_first)
+                updated_pytorch_nodes[cpu_node_second.id] = copy.deepcopy(cpu_node_second)
+                updated_pytorch_nodes[updated_gpu_node.id] = copy.deepcopy(updated_gpu_node)
 
         self.pytorch_nodes = updated_pytorch_nodes
 
@@ -382,9 +383,10 @@ def _split_cpu_node(
         cpu_node_first.id = self.id_assigner.assign_unique_id(cpu_node.id)
         cpu_node_first.ts = cpu_node.ts
         cpu_node_first.dur = gpu_node.ts - cpu_node.ts
-        cpu_node_first.set_child_gpu = gpu_node
-        for child_node in cpu_node_first.children:
+        cpu_node_first.set_child_gpu(gpu_node)
+        for child_node in cpu_node.children:
             child_node.parent = cpu_node_first.id
+            child_node.add_ctrl_dep(cpu_node_first)
         if cpu_node_first.ts >= gpu_node.ts or cpu_node_first.dur <= 0:
             err_msg = f"Invalid timestamps for the first split of CPU node {cpu_node.id}"
             self.logger.error(err_msg)
@@ -399,8 +401,9 @@ def _split_cpu_node(
         cpu_node_second.dur = cpu_node.dur - (gpu_node.ts - cpu_node.ts)
         cpu_node_second.set_child_gpu(None)
         cpu_node_second.add_data_dep(cpu_node_first)
-        for child_node in cpu_node_second.children:
+        for child_node in cpu_node.children:
             child_node.parent = cpu_node_second.id
+            child_node.add_ctrl_dep(cpu_node_second)
         if cpu_node_second.ts <= cpu_node_first.ts or cpu_node_second.dur <= 0:
             err_msg = f"Invalid timestamps for the first split of CPU node {cpu_node.id}"
             self.logger.error(err_msg)
@@ -445,7 +448,8 @@ def convert_to_chakra_node(self, pytorch_node: PyTorchNode) -> ChakraNode:
         chakra_node.id = pytorch_node.id
         chakra_node.name = pytorch_node.name
         chakra_node.type = self.get_chakra_node_type_from_pytorch_node(pytorch_node)
-        chakra_node.ctrl_deps.append(pytorch_node.parent)
+        for ctrl_dep in pytorch_node.ctrl_deps:
+            chakra_node.ctrl_deps.append(ctrl_dep.id)
         chakra_node.duration_micros = pytorch_node.dur if pytorch_node.has_dur() else 0
         chakra_node.inputs.values = str(pytorch_node.inputs)
         chakra_node.inputs.shapes = str(pytorch_node.input_shapes)