Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[torchlib] aten_lstm / aten_gru: runtime shape computation between layers makes multi-layer RNNs ~10% slower

Chiusa Adatta ai principianti
  • Qualcosa non va?
#3,096 0 commenti 1 reazione 2 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 3 giorni

@justinchuby ci sta già lavorando.

Dal 20/10/2026.

  • #3111 di @copilot-swe-agent — integrata

Valutazione

Difficoltà
2/5
Tempo stimato
1-3 ore
Idoneità per principianti
68/100
Tipo di issue
Refactoring
Chiarezza
Specificata chiaramente
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
machine-learning, performance

Direzione di ricerca

Inizia in onnxscript/function_libs/torch_lib/ops/core.py da aten_lstm e aten_gru, dove l'output di ciascun layer viene rimodellato come input del layer successivo. Aggiungi un helper che esegue lo squeeze dell'asse di direzione quando bidirectional è False e usa un reshape costante [0, 0, 2 * hidden_size] altrimenti. Esegui i test torch_lib per LSTM e GRU, inclusi i casi bidirezionali, per confermare che gli output non cambino e che i nodi Shape/Slice a runtime siano scomparsi.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Summary

aten_lstm and aten_gru (onnxscript/function_libs/torch_lib/ops/core.py, introduced in #2674) convert each layer's ONNX output
Y: [seq_length, num_directions, batch_size, hidden_size] into the next layer input
[seq_length, batch_size, num_directions * hidden_size] with:

Transpose(perm=[0, 2, 1, 3]) -> Shape -> 4x Slice -> Mul -> Reshape(-1) -> Concat -> Reshape

That is 10 nodes per layer, including a target shape computed at runtime. Two things about the shape are known when the graph is traced:

  • bidirectional is a Python bool. When it is False, num_directions == 1, so the whole sequence is just Squeeze(Y, axes=[1]), which is what the TorchScript exporter emits.
  • hidden_size is a static int (it is already passed as the LSTM/GRU hidden_size attribute). When bidirectional is True, the target shape can be the constant [0, 0, 2 * hidden_size], where 0 copies the input dimension.

With a dynamic batch axis, ONNX Runtime cannot constant-fold the shape computation, so every layer runs these extra tiny kernels.
In our case, it made the dynamo export slower than the TorchScript export of the same model (can be reproduced with dynamo=False).

Repro

Versions: torch 2.13.0+cpu, onnxscript main @ 9ff5a7c (same code in 0.7.1 / 0.7.2), onnxruntime 1.31.0, CPU, 1 thread.

import collections, sys, time
import numpy as np, onnx, onnxruntime as ort, torch


class StepLSTM(torch.nn.Module):
    def __init__(self, input_size=256, hidden_size=256, num_layers=4):
        super().__init__()
        self.lstm = torch.nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)

    def forward(self, x, h0, c0):
        return self.lstm(x, (h0, c0))


torch.manual_seed(0)
model = StepLSTM().eval()
batch = 5
inputs = (torch.randn(batch, 1, 256), torch.randn(4, batch, 256), torch.randn(4, batch, 256))
batch_dim = torch.export.Dim("batch")
path = sys.argv[1]
torch.onnx.export(
    model, inputs, path, dynamo=True,
    input_names=["x", "h0", "c0"],
    dynamic_shapes=({0: batch_dim}, {1: batch_dim}, {1: batch_dim}),
)

counts = collections.Counter(node.op_type for node in onnx.load(path).graph.node)
print(f"nodes={sum(counts.values())} {dict(sorted(counts.items()))}")

options = ort.SessionOptions()
options.intra_op_num_threads = 1
options.inter_op_num_threads = 1
session = ort.InferenceSession(path, options, providers=["CPUExecutionProvider"])
feeds = {name: tensor.numpy() for name, tensor in zip(["x", "h0", "c0"], inputs)}
for _ in range(500):
    session.run(None, feeds)
timings = []
for _ in range(5):
    start = time.perf_counter()
    for _ in range(2000):
        session.run(None, feeds)
    timings.append((time.perf_counter() - start) / 2000 * 1e6)
print(f"ort_latency={min(timings):.1f} us/step (best of 5)")
Results
nodes Reshape Shape Slice ORT latency (us/step)
main 104 8 4 56 991 - 1014
with the proposed change 68 0 0 40 908 - 911

Outputs are identical: max abs diff vs PyTorch is 8.3e-07 in both cases. That is about 9-10% faster for this model

Proposed fix

Factor the conversion into a helper used by both aten_lstm and aten_gru:

def _rnn_output_to_layer_output(Y: TFloat, bidirectional: bool, hidden_size: int) -> TFloat:
    """Y: [seq_length, num_directions, batch_size, hidden_size]
    Returns: [seq_length, batch_size, num_directions * hidden_size]
    """
    if not bidirectional:
        # num_directions == 1: dropping the axis is enough
        return op.Squeeze(Y, [1])
    # [seq_length, batch_size, num_directions, hidden_size]
    Y = op.Transpose(Y, perm=[0, 2, 1, 3])
    # 0 copies the input dimension (seq_length, batch_size): no runtime shape computation needed
    return op.Reshape(Y, op.Constant(value_ints=[0, 0, 2 * hidden_size]))

onnxscript-rnn-squeeze.patch

Image Image
Lingua principale
Python
Stelle
456
Fork
148
Merge medio
3g 13h
PR unite (30g)
21

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di microsoft/onnxscript

Tutte le issue di microsoft/onnxscript

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.