Prune unused ops

facebookincubator · Jul 12, 2023 · eeab8f7 · eeab8f7
1 parent b1495b9
commit eeab8f7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/aitemplate/compiler/transform/transform_merge_view_ops.py b/python/aitemplate/compiler/transform/transform_merge_view_ops.py
@@ -91,5 +91,9 @@ def merge_view_ops(sorted_graph: List[Tensor], workdir: str = None) -> List[Tens
             break
 
     if changed:
+        # Prune tensors that may have become unused after view op merging
+        for t in sorted_graph:
+            if len(t._attrs["dst_ops"]) == 0 and not t._attrs["is_output"]:
+                transform_utils.remove_tensor_from_sorted_graph(t)
         return transform_utils.sanitize_sorted_graph(toposort(sorted_graph))
     return sorted_graph