Addressed comments

Signed-off-by: Keshav Vinayak Jha <[email protected]>

Author: keshavvinayak01

python/torch_mlir/extras/fx_importer.py

@@ -566,7 +566,9 @@ def __init__(
         self._hooks = hooks or FxImporterHooks()
         self.symbol_table = SymbolTable(self._m.operation)
         self._hooks.prepare_module(self._m.operation)
+        # Used specifically in HOPs to map module IDs to function names
         self._graph_module_to_func_name: Dict[int, str] = {}
+        # Handles collision of function names in the same module
         self._func_name_counter: int = 0
 
     def _config_check(self):
@@ -834,7 +836,7 @@ def import_program(
         # HOP operator graphs are stateless graphs with no mutation, and it is correct
         # to import them as stateless graphs.
         self._import_all_child_modules(
-            prog, func_name, import_symbolic_shape_expressions
+            prog.graph.owning_module, func_name, import_symbolic_shape_expressions
         )
 
         # Import all nodes and return.
@@ -853,49 +855,27 @@ def import_program(
 
     def _import_all_child_modules(
         self,
-        prog_or_module: Union[torch.export.ExportedProgram, GraphModule],
+        module: GraphModule,
         parent_name: str,
         import_symbolic_shape_expressions: bool = False,
     ):
-        """Recursively import all child modules that have graphs.
+        """Import all child modules by delegating to import_graph_module.
 
-        This simple approach imports all submodules recursively, which is sufficient
-        for HOP operations since they only reference existing submodules.
+        This is a thin wrapper that extracts the owning module and delegates to
+        import_graph_module for each child.
+
+        Note: This only imports children, not the parent module itself.
         """
-        # Get the owning module from either ExportedProgram or GraphModule
-        if isinstance(prog_or_module, GraphModule):
-            owning_module = prog_or_module
-        else:
-            owning_module = prog_or_module.graph.owning_module
 
-        for child_name, child_module in owning_module.named_children():
+        for child_name, child_module in module.named_children():
             if isinstance(child_module, GraphModule) and hasattr(child_module, "graph"):
-                # Check if we've already assigned a name to this module
-                module_id = id(child_module)
-                # Module already imported, skip it
-                if module_id in self._graph_module_to_func_name:
-                    continue
-                # Use the child_name directly - PyTorch already provides unique names
-                child_func_name = child_name
-                # Handle collision by adding counter suffix if name already exists
-                if child_func_name in self._graph_module_to_func_name.values():
-                    child_func_name = f"{child_name}_{self._func_name_counter}"
-                    self._func_name_counter += 1
-                # Store the mapping for future lookups
-                self._graph_module_to_func_name[module_id] = child_func_name
-                # Import the child as a stateless graph (private function)
-                self.import_stateless_graph(
-                    child_module.graph,
-                    func_name=child_func_name,
+                self.import_graph_module(
+                    child_module,
+                    func_name=child_name,
                     func_visibility="private",
                     import_symbolic_shape_expressions=import_symbolic_shape_expressions,
                 )
 
-                # Recursively import its children
-                self._import_all_child_modules(
-                    child_module, child_func_name, import_symbolic_shape_expressions
-                )
-
     def import_frozen_program(
         self,
         prog: torch.export.ExportedProgram,
@@ -993,13 +973,56 @@ def import_frozen_program(
             import_symbolic_shape_expressions=import_symbolic_shape_expressions,
         )
 
-    def import_graph_module(self, gm: GraphModule) -> Operation:
+    def import_graph_module(
+        self,
+        gm: GraphModule,
+        *,
+        func_name: str = "main",
+        func_visibility: Optional[str] = None,
+        import_symbolic_shape_expressions: bool = False,
+    ) -> Operation:
         """Low-level import of a GraphModule assuming that it has been functionalized.
 
+        This method recursively imports all child GraphModules first, then imports
+        the provided GraphModule itself. This ensures that any higher-order operations
+        that reference child modules will find them already imported.
+
         TODO: This mechanism is deprecated by the `import_program` entry-point and
         it should be removed when no longer required for backwards compatibility.
+
+        Note: This method should only be used for HOPs.
         """
-        return self.import_stateless_graph(gm.graph)
+        # Store the mapping for this module itself (HOPs will need to look this up)
+        module_id = id(gm)
+        if module_id not in self._graph_module_to_func_name:
+            # Ensure the func_name is unique
+            final_func_name = func_name
+            if func_name in self._graph_module_to_func_name.values():
+                final_func_name = f"{func_name}_{self._func_name_counter}"
+                self._func_name_counter += 1
+            self._graph_module_to_func_name[module_id] = final_func_name
+        else:
+            # Module already imported, use existing name
+            final_func_name = self._graph_module_to_func_name[module_id]
+
+        # First, recursively import all child modules
+        for child_name, child_module in gm.named_children():
+            if isinstance(child_module, GraphModule) and hasattr(child_module, "graph"):
+                # Recursively import this child (which will handle its own mapping)
+                self.import_graph_module(
+                    child_module,
+                    func_name=child_name,
+                    func_visibility="private",
+                    import_symbolic_shape_expressions=import_symbolic_shape_expressions,
+                )
+
+        # Then import this module's own graph
+        return self.import_stateless_graph(
+            gm.graph,
+            func_name=final_func_name,
+            func_visibility=func_visibility,
+            import_symbolic_shape_expressions=import_symbolic_shape_expressions,
+        )
 
     def import_stateless_graph(
         self,
@@ -1033,11 +1056,9 @@ def import_stateless_graph(
             entry_block,
         )
 
-        # Import child modules (for HOPs) before importing nodes
-        if hasattr(g, 'owning_module') and g.owning_module is not None:
-            self._import_all_child_modules(
-                g.owning_module, func_name, import_symbolic_shape_expressions
-            )
+        # Note: Child module importing is handled by import_graph_module, which is
+        # the recommended entry point. This method is deprecated and should only be
+        # used for stateless graphs that truly have no child modules.
 
         node_importer.import_nodes(
             g.nodes, import_symbolic_shape_expressions=import_symbolic_shape_expressions
@@ -1068,10 +1089,11 @@ def _graph_to_function_meta(self, g: Graph) -> Tuple[FunctionType, Location]:
                 # for single returns it might be a single Node.
                 output_arg = node.args[0]
                 # Handle both single Node and tuple/list of Nodes
-                if isinstance(output_arg, (list, tuple)):
-                    result_nodes = output_arg
-                else:
-                    result_nodes = [output_arg]
+                result_nodes = (
+                    output_arg
+                    if isinstance(output_arg, (list, tuple))
+                    else [output_arg]
+                )
 
                 for result_node in result_nodes:
                     if result_node is None:
@@ -1586,11 +1608,11 @@ def import_nodes(
                     # results.
                     output_arg = node.args[0]
                     # Handle both single Node and tuple/list of Nodes
-                    if isinstance(output_arg, (list, tuple)):
-                        result_nodes = output_arg
-                    else:
-                        result_nodes = [output_arg]
-
+                    result_nodes = (
+                        output_arg
+                        if isinstance(output_arg, (list, tuple))
+                        else [output_arg]
+                    )
                     operands = [self._import_argument(loc, arg) for arg in result_nodes]
                     func_dialect.ReturnOp(operands, loc=loc)
 
@@ -1722,8 +1744,8 @@ def _import_hop_while_loop(
         body_fn_module = getattr(root_module, body_fn_node.target, None)
 
         # Generate function names with module IDs for uniqueness
-        cond_fn_name = self.fx_importer._graph_module_to_func_name.get(id(cond_fn_module))
-        body_fn_name = self.fx_importer._graph_module_to_func_name.get(id(body_fn_module))
+        cond_fn_name = self.fx_importer._graph_module_to_func_name[id(cond_fn_module)]
+        body_fn_name = self.fx_importer._graph_module_to_func_name[id(body_fn_module)]
 
         # Import the carries (loop state variables)
         carry_values = []
@@ -1765,7 +1787,7 @@ def _import_hop_while_loop(
         with loc:
             max_iter = _make_constant_op(
                 "torch.constant.int",
-                self._cc.integer_attr(9223372036854775807, 64),
+                torch.iinfo(torch.int64).max,
                 self._cc.torch_int_type,
             )
 
@@ -2153,7 +2175,7 @@ def _import_scalar_as_tensor(self, loc: Location, arg: NodeArgument) -> Value:
 
     def _import_tuple_argument(
         self, loc: Location, arg: tuple, expected_jit_type
-    ) -> List[Value]:
+    ) -> list[Value]:
         """Import a tuple argument by importing each element separately."""
         # For tuples in while_loop carries, treat each element as a separate argument
         return [self._import_argument(loc, elem, expected_jit_type) for elem in arg]
@@ -2268,7 +2290,6 @@ def _import_getitem(self, loc: Location, node: torch.fx.Node):
             # NOTE: the length of the list must be knowable at compile time.
             if ref_node not in self._unpack_list_values:
                 node_result = self.resolve_node_value(ref_node, 0)
-                node_val = ref_node.meta.get("val")
 
                 if str(node_result.type) in TORCH_LIST_TYPES:
                     result_types = [