onnx-graphsurgeon: Graph.toposort() hits RecursionError when a graph has a local function and an If/Loop node #4874

Description

Graph.toposort() in ONNX-GraphSurgeon fails with RecursionError when a graph has at least one local function and at least one node with a subgraph (If, Loop, Scan). The node doesn't need to call the function.

  File "onnx_graphsurgeon/ir/graph.py", line 594, in get_used_funcs
    inputs.update(get_used_funcs(subgraph.nodes))
  [Previous line repeated 990 more times]
RecursionError: maximum recursion depth exceeded

The problem is in get_used_funcs in onnx_graphsurgeon/ir/graph.py, which toposort() uses to order functions (default mode="full"):

def get_used_funcs(nodes):
    inputs = {}
    for subgraph in self.subgraphs():
        inputs.update(get_used_funcs(subgraph.nodes))
    ...

The loop uses self.subgraphs(), the subgraphs of the outer graph, not the subgraphs of the nodes that were passed in. self is the same on every recursive call, so each call loops over the same subgraphs again and the recursion never stops.

Even when the graph has no top-level subgraphs and it doesn't crash, the function never looks inside the subgraphs of the nodes it was given. If function A calls function B only from inside an If branch, toposort() misses the dependency and returns the functions in the order [A, B] instead of [B, A].

cleanup(), fold_constants() and export_onnx() work fine on the same graph. Only toposort() fails.

Environment

TensorRT Version: N/A, ONNX-GraphSurgeon only. Reproduced on main (16ed063, onnx-graphsurgeon 0.6.2) and on onnx-graphsurgeon 0.6.1 from PyPI. The code was added in the TensorRT 10.0 release.

NVIDIA GPU: N/A

NVIDIA Driver Version: N/A

CUDA Version: N/A

CUDNN Version: N/A

Operating System: Linux

Python Version (if applicable): 3.14 (onnx 1.23.2, numpy 2.5.4)

Tensorflow Version (if applicable): N/A

PyTorch Version (if applicable): N/A

Baremetal or Container (if so, version): Baremetal

Relevant Files

Model link: N/A, the script below builds the graph directly.

Steps To Reproduce

Commands or scripts:

import numpy as np
import onnx_graphsurgeon as gs

# Any local function
x = gs.Variable("x", dtype=np.float32, shape=(1,))
y = gs.Variable("y", dtype=np.float32, shape=(1,))
func = gs.Function("MyFunc", inputs=[x], outputs=[y], nodes=[gs.Node("Identity", inputs=[x], outputs=[y])], opset=18)

# Any node with a subgraph
def branch(name):
    out = gs.Variable(f"{name}_out", dtype=np.float32, shape=(1,))
    const = gs.Constant(f"{name}_const", np.ones((1,), dtype=np.float32))
    return gs.Graph(nodes=[gs.Node("Identity", inputs=[const], outputs=[out])], outputs=[out])

cond = gs.Variable("cond", dtype=np.bool_, shape=())
if_out = gs.Variable("if_out", dtype=np.float32, shape=(1,))
if_node = gs.Node("If", inputs=[cond], outputs=[if_out],
                  attrs={"then_branch": branch("then"), "else_branch": branch("else")})

graph = gs.Graph(nodes=[if_node], inputs=[cond], outputs=[if_out], functions=[func], opset=18)
graph.toposort()

Full traceback:

  File "onnx_graphsurgeon/ir/graph.py", line 636, in toposort
    func, level=get_hierarchy_level(func)
  File "onnx_graphsurgeon/ir/graph.py", line 611, in get_hierarchy_level
    for inp in get_inputs(node_or_func)
  File "onnx_graphsurgeon/ir/graph.py", line 604, in get_inputs
    inputs = get_used_funcs(node_or_func.nodes).values()
  File "onnx_graphsurgeon/ir/graph.py", line 594, in get_used_funcs
    inputs.update(get_used_funcs(subgraph.nodes))
  File "onnx_graphsurgeon/ir/graph.py", line 594, in get_used_funcs
    inputs.update(get_used_funcs(subgraph.nodes))
  [Previous line repeated 990 more times]
  File "onnx_graphsurgeon/ir/graph.py", line 593, in get_used_funcs
    for subgraph in self.subgraphs():
  File "onnx_graphsurgeon/ir/graph.py", line 381, in subgraphs
    for subgraph in node.subgraphs(recursive=recursive):
RecursionError: maximum recursion depth exceeded

Have you tried the latest release?: Yes. It reproduces on the latest onnx-graphsurgeon from PyPI (0.6.1) and on current main.

Can this model run on other frameworks?: N/A. This is a bug in graph manipulation, not in inference. onnx.checker.check_model accepts the exported graph.

Proposed fix

Recurse into the subgraphs of each node in nodes:

def get_used_funcs(nodes):
    inputs = {}
    for node in nodes:
        for subgraph in node.subgraphs():
            inputs.update(get_used_funcs(subgraph.nodes))
        func_id = (node.domain, node.op)
        if func_id in func_id_to_func:
            inputs[func_id] = func_id_to_func[func_id]
    return inputs

With this change, the repro above passes and the nested A/B case is sorted as [B, A]. I'd like to send a PR with this fix and regression tests in tests/ir/test_graph.py.