Simplified recursive import_stateless_graph call

Author: keshavvinayak01

python/torch_mlir/extras/fx_importer.py

@@ -840,72 +840,44 @@ def import_program(
         # Even though import_stateless_graph is deprecated as an entrypoint mechanism,
         # HOP operator graphs are stateless graphs with no mutation, and it is correct
         # to import them as stateless graphs.
-        self._import_child_graph_modules(
+        self._import_all_child_modules(
             prog,
             func_name,
             import_symbolic_shape_expressions
         )
 
         return func_op
 
-    def _import_child_graph_modules(
+    def _import_all_child_modules(
         self,
         prog: torch.export.ExportedProgram,
         parent_name: str,
         import_symbolic_shape_expressions: bool = False
     ):
-        """Scan graph for get_attr nodes and import referenced submodules as stateless graphs.
+        """Recursively import all child modules that have graphs.
 
-        This is used to import HOP subgraphs. When we encounter a get_attr node, it references
-        a submodule that needs to be imported as a private function that can be called
-        by the HOP operation.
+        This simple approach imports all submodules recursively, which is sufficient
+        for HOP operations since they only reference existing submodules.
         """
-        # Collect all get_attr nodes by recursively searching through all graphs
-        get_attr_nodes = []
-        self._collect_get_attr_nodes(prog.graph, get_attr_nodes, set())
-        # Import each referenced module as a stateless graph
-        for node in get_attr_nodes:
-            attr_name = node.target
-
-            # Find the module that contains this attribute
-            # The node's graph should have an owning_module that contains the submodule
-            if hasattr(node.graph, 'owning_module'):
-                parent_module = node.graph.owning_module
-
-                # Check if the parent module has the referenced attribute
-                if hasattr(parent_module, attr_name):
-                    child_module = getattr(parent_module, attr_name)
-
-                    # Only import if it's a GraphModule with a graph
-                    if isinstance(child_module, GraphModule) and hasattr(child_module, 'graph'):
-                        # Generate function name: parent_childname_id for uniqueness
-                        child_func_name = f"{parent_name}_{attr_name}_{id(child_module)}"
-
-                        # Import the child as a stateless graph (private function)
-                        self.import_stateless_graph(
-                            child_module.graph,
-                            func_name=child_func_name,
-                            func_visibility="private",
-                            import_symbolic_shape_expressions=import_symbolic_shape_expressions,
-                        )
-
-    def _collect_get_attr_nodes(self, graph: Graph, get_attr_nodes: List[Node], visited: Set[int] = None):
-        """Recursively collect all get_attr nodes from a graph and its subgraphs."""
-        if visited is None:
-            visited = set()
-
-        graph_id = id(graph)
-        if graph_id in visited:
-            return  # Already processed this graph
-        visited.add(graph_id)
-
-        for node in graph.nodes:
-            if node.op == "get_attr":
-                get_attr_nodes.append(node)
+        for child_name, child_module in prog.graph.owning_module.named_children():
+            if isinstance(child_module, GraphModule) and hasattr(child_module, 'graph'):
+                # Generate function name: parent_childname
+                child_func_name = f"{parent_name}_{child_name}_{id(child_module)}"
+
+                # Import the child as a stateless graph (private function)
+                self.import_stateless_graph(
+                    child_module.graph,
+                    func_name=child_func_name,
+                    func_visibility="private",
+                    import_symbolic_shape_expressions=import_symbolic_shape_expressions,
+                )
 
-            # If this node has a subgraph (like HOP operations), search it too
-            if hasattr(node, 'graph') and node.graph is not None:
-                self._collect_get_attr_nodes(node.graph, get_attr_nodes, visited)
+                # Recursively import its children
+                self._import_all_child_modules(
+                    child_module,
+                    child_func_name,
+                    import_symbolic_shape_expressions
+                )
 
     def import_frozen_program(
         self,