Address feedback: optimize tensor fingerprinting and traverse subgraphs

Author: AbhishekHerbertSamuel

src/onnx_ir/passes/common/deduplicate_initializers.py

@@ -1,49 +1,48 @@
 from onnx_ir import ir
 from onnx_ir.passes.base import GraphTransformPass
+from onnx_ir.traversal import iterate_graph
 
 
 class DeduplicateInitializersPass(GraphTransformPass):
     """
-    This pass removes duplicate initializer tensors from the graph.
+    Graph transformation pass to remove duplicate initializer tensors.
 
-    It identifies duplicates based on a content-based fingerprint consisting of:
-    - Tensor byte content (`tobytes()`)
-    - Data type (`dtype`)
+    Identifies duplicates based on:
+    - Data type
     - Shape
+    - Byte content (used only if dtype and shape match)
 
-    All duplicates are replaced with the first (canonical) occurrence, and node
-    inputs referring to redundant initializers are updated accordingly.
+    Updates all node inputs (including subgraphs) to refer to the canonical tensor.
     """
 
     def apply(self, graph: ir.Graph) -> ir.Graph:
-        seen = {}      # Maps (tobytes, dtype, shape) -> canonical initializer name
-        name_map = {}  # Maps duplicate initializer name -> canonical name
+        seen = {}      # (dtype, shape) → {tobytes: name}
+        name_map = {}  # Duplicate name → canonical name
 
-        # Iterate over all initializers in the graph
+        # Iterate through initializers and group by dtype and shape first
         for initializer in list(graph.initializers.values()):
-            key = (
-                initializer.const_value.tobytes(),              # Content fingerprint
-                initializer.const_value.dtype,                  # Data type
-                tuple(initializer.const_value.shape),           # Shape tuple
-            )
-
-            if key in seen:
-                # Found a duplicate: store the name mapping and remove it from graph
-                canonical_name = seen[key]
+            dtype = initializer.const_value.dtype
+            shape = tuple(initializer.const_value.shape)
+            content = initializer.const_value.tobytes()
+
+            if (dtype, shape) not in seen:
+                seen[(dtype, shape)] = {}
+
+            group = seen[(dtype, shape)]
+            if content in group:
+                # Duplicate found
+                canonical_name = group[content]
                 name_map[initializer.name] = canonical_name
                 graph.initializers.pop(initializer.name)
             else:
-                # First time seeing this tensor → keep it
-                seen[key] = initializer.name
+                group[content] = initializer.name
 
-        # Update node inputs to use the canonical initializer names
-        for node in graph:
+        # Update all node inputs (including subgraphs)
+        for node in iterate_graph(graph):
             for i, input_value in enumerate(node.inputs):
                 if input_value is not None and input_value.name in name_map:
-                    # Replace input with the deduplicated initializer
-                    new_name = name_map[input_value.name]
-                    replacement = graph.initializers[new_name]
+                    canonical = name_map[input_value.name]
+                    replacement = graph.initializers[canonical]
                     node.replace_input_with(i, replacement)
 
         return graph
-