Commit 5957997
Migrate graph surgeries to onnx_ir + onnxscript rewriter (batch 1: infra + first surgeries) (microsoft#2550)
## Describe your changes
First batch of an incremental migration of `graph_surgeries.py` off the
protobuf / `OnnxDAG` approach onto the ONNX IR (`onnx_ir`) +
`onnxscript`
rewriter.
**Infrastructure**
- Add a `RewriteRuleSurgeon(Surgeon)` base class. Subclasses implement
`rules()`
returning an `onnxscript.rewriter.pattern.RewriteRuleSet`; the base
applies it
to the IR model via `call_ir`. This lets local subgraph pattern
replacements be
expressed declaratively, and the rewriter handles operand commutativity,
use-count bookkeeping, and dead-node cleanup for us.
**First surgeries ported**
- `ReciprocalMulToDiv`: `a * Reciprocal(x)` → `Div(a, x)`
(`commute=True` covers
both `Mul` operand orders; a shared `Reciprocal` is preserved
automatically).
- `ReplaceErfWithTanh`: `Erf(x)` → `Tanh(x * 605/503)`, emitting the
scale as an
initializer of the input's floating-point dtype; non-float inputs are
skipped.
This trims ~120 lines of manual proto walking. Subsequent batches will
port the
remaining pattern-based surgeries (`Gemm`↔`MatMul`+`Add`, QDQ passes,
RMSNorm
variants, decompositions, ...) and move the whole-graph surgeries
(rename/expose
I/O, `Non4D*`, dedup, `TieWordEmbeddings`, ...) to plain `onnx_ir`.
**Test change**
- `test_replace_erf_with_tanh` now reads the scale initializer via
`numpy_helper.to_array` instead of `.float_data`, so it is agnostic to
whether
the tensor is stored as `raw_data` or `float_data` (IR emits
`raw_data`).
No behavior change for users; the two ported surgeries produce
equivalent graphs.
## Checklist before requesting a review
- [x] Add unit tests for this change. (existing surgery tests cover
both; all pass)
- [x] Make sure all tests can pass. (`test_graph_surgeries.py`: 82
passed, 2 skipped)
- [ ] Update documents if necessary.
- [x] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? No — internal refactor, equivalent
output.
## (Optional) Issue link
---------
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>1 parent 1b34297 commit 5957997
2 files changed
Lines changed: 1017 additions & 1109 deletions
0 commit comments