Skip to content

Commit 5957997

Browse files
justinchubyCopilot
andauthored
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

File tree

0 commit comments

Comments
 (0)