UNDF: UNDF-2026-000000514

pytorch-0003 — python_function.cpp tracer subgraph construction O(N²)

Location

torch/csrc/autograd/python_function.cpp, lines 1019–1030

Defective Code

// line 1019
for (it++; it != owning_block->nodes().end(); ++it) {   // O(N) outer loop over all nodes
    torch::jit::Node* node = *it;
    auto* clone_node =
        subgraph->insertNode(subgraph->createClone(node, value_map_func));
    for (size_t i = 0; i < node->outputs().size(); ++i) {   // O(K) outputs per node
        value_map[node->outputs()[i]] = clone_node->outputs()[i];
        auto trace_it = std::find(
            trace_outputs.begin(), trace_outputs.end(), node->outputs()[i]);  // O(T) scan
        if (trace_it != trace_outputs.end()) {
            subgraph->registerOutput(clone_node->outputs()[i]);
        }
    }
}

trace_outputs is a std::vector<Value*>. For each of N nodes with K outputs, std::find scans all T trace outputs: O(N × K × T).

Fixed Code

// Build a set once before the loop: O(T)
std::unordered_set<Value*> trace_outputs_set(
    trace_outputs.begin(), trace_outputs.end());

for (it++; it != owning_block->nodes().end(); ++it) {
    torch::jit::Node* node = *it;
    auto* clone_node =
        subgraph->insertNode(subgraph->createClone(node, value_map_func));
    for (size_t i = 0; i < node->outputs().size(); ++i) {
        value_map[node->outputs()[i]] = clone_node->outputs()[i];
        if (trace_outputs_set.count(node->outputs()[i])) {   // O(1)
            subgraph->registerOutput(clone_node->outputs()[i]);
        }
    }
}

Complexity Analysis

Path	Complexity
Slow (original)	O(N × K × T)
Fast (fixed)	O(T + N × K)

At N=200 nodes, K=4 outputs/node, T=50 trace outputs: 50× speedup.

Severity

MEDIUM — executed during JIT tracing of Python autograd functions. Triggered whenever torch.jit.trace or the tracing autograd path is used. T grows with the number of outputs being traced; N grows with the function's graph depth.

2 KiB Raw Blame History Unescape Escape