Skip to content

Replace shape-manipulating subgraphs with ComputeShape operator#1187

Draft
robertknight wants to merge 2 commits intomainfrom
infer-shapes-values
Draft

Replace shape-manipulating subgraphs with ComputeShape operator#1187
robertknight wants to merge 2 commits intomainfrom
infer-shapes-values

Conversation

@robertknight
Copy link
Owner

Use symbolic tensor values produced by shape inference to replace subgraphs in models with a ComputeShape operator.

In the Llama 3 3B model, this reduces the total number of operators run on each inference pass from 1740 to 1192 (-31%). This reduces interpreter overhead, although the impact is usually small because shape-manipulating subgraphs are
operating on small values. More importantly, it removes more potentially fusion-blocking value consumers from the graph. This generalizes the idea of ComputeShapeFusion.

Fixes #270

TODO:

  • Tests
  • Rename ComputeShape to something more general?

@robertknight robertknight force-pushed the infer-shapes-values branch 2 times, most recently from fdeceaa to cd8e3be Compare January 11, 2026 16:00
Use symbolic tensor values produced by shape inference to replace
subgraphs in models with a `ComputeShape` operator.

In the Llama 3 model, this reduces the total number of operators run on each
inference pass from 1740 to 1192 (-31%). This reduces interpreter overhead,
although the impact is usually small because shape-manipulating subgraphs are
operating on small values. More importantly, it removes more potentially
fusion-blocking value consumers from the graph. This generalizes the idea of
`ComputeShapeFusion`.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Memoize or precompute subgraphs that depend only on input shapes

1 participant