UNDF: UNDF-2026-000000512

pytorch-0001 — graph_fuser.cpp fuseChunkByReusingExistingFusedChunk O(N²)

Location

torch/csrc/jit/passes/graph_fuser.cpp, lines 513–527

Defective Code

void fuseChunkByReusingExistingFusedChunk(
    Node* group,
    Node* chunk,
    Node* existingFusedChunk) {
  if (chunk->outputs().size() != existingFusedChunk->outputs().size()) {
    return;
  }
  auto& subgraph = getSubgraph(group);
  for (size_t i = 0; i < chunk->outputs().size(); ++i) {  // O(C) outer loop
    // Find the input to the FusionGroup (group)
    auto* replacement_val = existingFusedChunk->outputs().at(i);
    auto* val = chunk->outputs().at(i);
    auto it = std::find(group->inputs().begin(), group->inputs().end(), val);  // O(I) scan
    auto input_index = it - group->inputs().begin();
    // Rewrite the graph to use replacement_val
    auto group_input = subgraph.inputs().at(input_index);
    group_input->replaceAllUsesWith(replacement_val);
    group->removeInput(input_index);
    subgraph.eraseInput(input_index);
  }
  chunk->destroy();
}

The outer loop iterates over chunk->outputs() (size C = chunk count). For each iteration, std::find performs a linear scan of group->inputs() (size I = fusion group inputs). Total: O(C × I).

In a large model with many fusion groups and chunk splits, C and I both grow with graph depth. fuseChunkByReusingExistingFusedChunk is called from canFuseChunk which is itself called in a scan loop over all nodes, making the aggregate complexity O(N × C × I).

Fixed Code

void fuseChunkByReusingExistingFusedChunk(
    Node* group,
    Node* chunk,
    Node* existingFusedChunk) {
  if (chunk->outputs().size() != existingFusedChunk->outputs().size()) {
    return;
  }
  auto& subgraph = getSubgraph(group);

  // Build a reverse index: Value* -> input position, O(I) once
  std::unordered_map<Value*, size_t> input_index_map;
  const auto inputs = group->inputs();
  for (size_t k = 0; k < inputs.size(); ++k) {
    input_index_map[inputs[k]] = k;
  }

  // Process in reverse order so that removeInput(index) doesn't shift earlier indices
  for (int i = static_cast<int>(chunk->outputs().size()) - 1; i >= 0; --i) {
    auto* replacement_val = existingFusedChunk->outputs().at(i);
    auto* val = chunk->outputs().at(i);
    auto map_it = input_index_map.find(val);    // O(1)
    if (map_it == input_index_map.end()) continue;
    size_t input_index = map_it->second;
    auto group_input = subgraph.inputs().at(input_index);
    group_input->replaceAllUsesWith(replacement_val);
    group->removeInput(input_index);
    subgraph.eraseInput(input_index);
  }
  chunk->destroy();
}

Complexity Analysis

Path	Complexity
Slow (original)	O(C × I) per call, O(N × C × I) aggregate
Fast (fixed)	O(I + C) per call, O(N × (I + C)) aggregate

At C=32 chunks and I=128 group inputs: 32× speedup per call.

Severity

HIGH — executed in the kernel fusion pass over every fusion group during model compilation. Large transformer models with torch.compile() fuse hundreds of groups with dozens of chunk splits.

3.1 KiB Raw Blame History Unescape Escape