Migrate OnnxDAG usage to onnx_ir + onnxscript rewriter and remove onnx_dag.py (microsoft#2558)

## Describe your changes

Replaces all `OnnxDAG` usage with `onnx_ir` (onnx-ir) and the
`onnxscript` rewriter, following the pattern from microsoft#2550, then deletes
the `onnx_dag.py` implementation.

**Production passes migrated to `onnx_ir`:**
- `vitis_ai/preprocess.py`, `static_llm.py`, `float16_conversion.py`,
`quantization.py`, `common.py`
- `mnb_to_qdq.py`
- `split.py` — `_run_for_config` rebuilt on ir; split graphs constructed
via per-split value maps, missing value-info resolved through
`SymbolicShapeInference`
- `compose.py` — `_get_composed_model` rebuilt on ir; merge-by-name with
`serde`-based node dedup and numpy shared-initializer equality
- `graph_surgeries.py` — `TieWordEmbeddings` moved from `ProtoSurgeon`
(proto/DAG) to `Surgeon` (`call_ir` on `ir.Model`); removed dead
`find_node` and the DAG-based base `add_reshape_node`
- `peephole_optimizer.py` — `fuse_reshape_operations` reimplemented on
`onnx_ir` (removing the `onnxruntime.transformers.onnx_model.OnnxModel`
dependency), replacing the previous stale OnnxDAG TODO

**onnx_ir idiom cleanups (from review feedback):**
- Use `ir_model.graphs()` (`common.py`), `ir_model.graph.all_nodes()`
(`mnb_to_qdq.py`), and `node.predecessors()` / `node.successors()`
(`split.py`) instead of hand-rolled helpers
- Use `ir.node` for node/attribute construction (`graph_surgeries.py`,
`mnb_to_qdq.py`, `vitis_ai/preprocess.py`)
- Use `ir.convenience.create_value_mapping` to build the value-info map
in `split.py`
- Use `value.dtype` directly in `compose.py`
- Moved the `onnx_ir` import to module top in `quantization.py`

**Correctness fixes for optional node names / opsets (from review
feedback):**
- `compose.py` — only dedup nodes by name when the name is non-empty, so
different unnamed nodes (`name=""`/`None`) no longer collide on the same
key and trigger false "Node mismatch" assertions
- `compose.py` — merge `opset_imports` across all component models (max
version per domain) instead of copying only from the first model, so
composing models with different domains/opset versions produces a valid
model
- `float16_conversion.py` — when `node_include_list` is set, raise a
clear error if the model contains unnamed nodes (which cannot be
included by name), preserving include-only semantics

**Tests:**
- `test_split_model.py`, `test_mnb_to_qdq.py` — migrated to ir idioms
- `test_graph_surgeries.py` — new read-only `GraphInspector` (ir-backed)
replaces `OnnxDAG` inspection sites; it reproduces OnnxDAG's
unique-naming of unnamed/duplicate nodes so name-based assertions stay
stable

**Cleanup:**
- Deleted `olive/passes/onnx/onnx_dag.py`
- No remaining `OnnxDAG` / `onnx_dag` references

Note: unlike OnnxDAG, ir does not auto-assign node names, so call sites
that keyed on names now guard against empty/`None` names (e.g.
`split.py`, `mnb_to_qdq.py`, `compose.py`, `float16_conversion.py`). The
quantized `TieWordEmbeddings` path is validated structurally only (no
ORT inference available offline).

## Checklist before requesting a review
- [x] Add unit tests for this change.
- [x] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [x] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Author: Copilot

olive/passes/onnx/common.py

@@ -22,7 +22,6 @@
 
 from olive.common.utils import StrEnumBase, hardlink_copy_file
 from olive.model import CompositeModelHandler, ONNXModelHandler
-from olive.passes.onnx.onnx_dag import OnnxDAG
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 from olive.resource_path import LocalFile, LocalFolder
 
@@ -539,12 +538,13 @@ def model_has_adapters_from_torchscript(
     :param model_path: The path to the model.
     :return: True if the model has adapters, False otherwise.
     """
-    dag = OnnxDAG(onnx.load(model_path, load_external_data=False))
-    if adapter_type == AdapterType.LOHA and is_loha_model(dag):
+    ir_model = ir.from_proto(onnx.load(model_path, load_external_data=False))
+    if adapter_type == AdapterType.LOHA and is_loha_model(ir_model):
         return True
     else:
-        for node_name in dag.get_node_names():
-            op_type = dag.get_node_op_type(node_name)
+        for node in ir_model.graph.all_nodes():
+            op_type = node.op_type
+            node_name = node.name or ""
             if (adapter_type == AdapterType.LORA and is_lora_node(op_type, node_name)) or (
                 adapter_type == AdapterType.DORA and is_dora_node(op_type, node_name)
             ):
@@ -564,10 +564,10 @@ def is_lora_node(op_type: str, node_name: str) -> bool:
     )
 
 
-def is_loha_model(dag: OnnxDAG) -> bool:
-    for graph in dag.graphs:
-        for initializer in graph.initializer:
-            if any(re.match(pattern, initializer.name) for pattern in LOHA_NAME_PATTERNS_TORCHSCRIPT):
+def is_loha_model(ir_model: ir.Model) -> bool:
+    for graph in ir_model.graphs():
+        for initializer_name in graph.initializers:
+            if any(re.match(pattern, initializer_name) for pattern in LOHA_NAME_PATTERNS_TORCHSCRIPT):
                 return True
     return False
 

olive/passes/onnx/compose.py

@@ -7,6 +7,9 @@
 from typing import Optional, Union
 
 import numpy as np
+import onnx_ir as ir
+from onnx_ir import serde
+from onnx_ir.passes.common import TopologicalSortPass
 
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import CompositeModelHandler, ONNXModelHandler
@@ -19,7 +22,6 @@
     model_proto_to_file,
     process_llm_pipeline,
 )
-from olive.passes.onnx.onnx_dag import OnnxDAG
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 logger = logging.getLogger(__name__)
@@ -109,91 +111,165 @@ def _get_composed_model(
         :param as_model_dir: Use model parent directory as output model_path.
         :return: Composed ONNX model.
         """
-        dags = [OnnxDAG.from_model_path(path) for path in onnx_model_paths]
+
+        def shape_list(value: ir.Value):
+            if value.shape is None:
+                return None
+            return [dim.value if isinstance(dim, ir.SymbolicDim) else dim for dim in value.shape]
+
+        ir_models = []
+        for path in onnx_model_paths:
+            ir_model = ir.load(path)
+            TopologicalSortPass()(ir_model)
+            ir_models.append(ir_model)
 
         seen_inputs = set()
         seen_outputs = set()
-        for dag in dags:
-            dag_inputs = set(dag.get_input_names())
-            dag_outputs = set(dag.get_output_names())
+        for ir_model in ir_models:
+            graph_inputs = {value.name for value in ir_model.graph.inputs}
+            graph_outputs = {value.name for value in ir_model.graph.outputs}
             # avoid circular connection, model_2 output cannot be model_1 input
-            assert dag_outputs.isdisjoint(seen_inputs), (
-                f"Output names {dag_outputs.intersection(seen_inputs)} are already used as input names."
+            assert graph_outputs.isdisjoint(seen_inputs), (
+                f"Output names {graph_outputs.intersection(seen_inputs)} are already used as input names."
             )
             # avoid reused output name
-            assert dag_outputs.isdisjoint(seen_outputs), (
-                f"Output names {dag_outputs.intersection(seen_outputs)} are already used as output names."
+            assert graph_outputs.isdisjoint(seen_outputs), (
+                f"Output names {graph_outputs.intersection(seen_outputs)} are already used as output names."
             )
 
             # update seen inputs and outputs
-            seen_inputs.update(dag_inputs)
-            seen_outputs.update(dag_outputs)
+            seen_inputs.update(graph_inputs)
+            seen_outputs.update(graph_outputs)
 
         # will only keep the unused outputs
         # inputs will be automatically taken care of during compose
         final_outputs = seen_outputs - seen_inputs
 
-        # compose
-        composed_dag = dags.pop(0)
-        while dags:
-            dag = dags.pop(0)
-
-            cd_input_names = set(composed_dag.get_input_names())
-            cd_output_names = set(composed_dag.get_output_names())
-            cd_initializer_names = set(composed_dag.get_initializer_names())
-            for input_name in dag.get_input_names():
-                if input_name in cd_input_names | cd_output_names:
-                    assert dag.get_io_shape(input_name) == composed_dag.get_io_shape(input_name), (
-                        f"Input shape mismatch: {input_name}"
-                    )
-                    assert dag.get_io_dtype(input_name) == composed_dag.get_io_dtype(input_name), (
-                        f"Input dtype mismatch: {input_name}"
-                    )
+        # compose by relinking values across models by name
+        composed_values: dict[str, ir.Value] = {}
+        composed_inputs: list[ir.Value] = []
+        composed_input_names: set[str] = set()
+        composed_initializers: dict[str, ir.Value] = {}
+        composed_nodes: list[ir.Node] = []
+        composed_node_names: dict[str, ir.Node] = {}
+        composed_output_names: list[str] = []
+        produced_names: set[str] = set()
+
+        def get_value(name: str) -> ir.Value:
+            value = composed_values.get(name)
+            if value is None:
+                value = ir.Value(name=name)
+                composed_values[name] = value
+            return value
+
+        for ir_model in ir_models:
+            graph = ir_model.graph
+
+            for inp in graph.inputs:
+                name = inp.name
+                if name in composed_input_names or name in produced_names:
+                    # already a graph input or an internal connection from a previous model
+                    existing = composed_values[name]
+                    assert shape_list(inp) == shape_list(existing), f"Input shape mismatch: {name}"
+                    assert inp.dtype == existing.dtype, f"Input dtype mismatch: {name}"
                     continue
 
-                # will add to graph 0 for now
-                # this will fail if the connection already exists in the graph = expected behavior.
-                composed_dag.add_input(dag.get_input_proto(input_name), 0)
-
-            for output_name in dag.get_output_names():
-                composed_dag.add_output(dag.get_output_proto(output_name), 0)
-
-            for init_name in dag.get_initializer_names():
-                if init_name in cd_initializer_names:
-                    (
-                        np.testing.assert_array_equal(
-                            dag.get_initializer_np_array(init_name), composed_dag.get_initializer_np_array(init_name)
-                        ),
-                        f"Initializer mismatch: {init_name}",
+                # will add to the composed graph inputs
+                value = get_value(name)
+                value.type = inp.type
+                value.shape = inp.shape
+                composed_inputs.append(value)
+                composed_input_names.add(name)
+
+            for init in graph.initializers.values():
+                name = init.name
+                if name in composed_initializers:
+                    np.testing.assert_array_equal(
+                        init.const_value.numpy(),
+                        composed_initializers[name].const_value.numpy(),
+                        err_msg=f"Initializer mismatch: {name}",
                     )
                     continue
 
-                composed_dag.add_initializer(dag.get_initializer_proto(init_name), 0)
-
-            for intermediate_name in dag.get_intermediate_names():
-                if value_info := dag.get_value_info_proto(intermediate_name):
-                    composed_dag.add_value_info(value_info, 0)
-
-            for node_name in dag.topological_sort():
-                node = dag.get_node_proto(node_name)
-                if composed_dag.has_node(node_name):
+                value = get_value(name)
+                value.const_value = init.const_value
+                value.type = init.type if init.type is not None else ir.TensorType(init.const_value.dtype)
+                value.shape = init.shape if init.shape is not None else ir.Shape(init.const_value.shape)
+                composed_initializers[name] = value
+
+            for node in graph:
+                name = node.name
+                # ONNX node names are optional, so only dedup nodes that carry a non-empty name.
+                # Unnamed nodes always get added to avoid different unnamed nodes colliding on the
+                # same "" / None key (which could trigger false "Node mismatch" assertions).
+                if name and name in composed_node_names:
                     # there might be some dq nodes for initializers that are common between models
                     # since split model keeps dq with the consumer op
-                    assert composed_dag.get_node_proto(node_name) == node, f"Node mismatch: {node_name}"
+                    assert (
+                        serde.serialize_node(composed_node_names[name]).SerializeToString()
+                        == serde.serialize_node(node).SerializeToString()
+                    ), f"Node mismatch: {name}"
                     continue
 
-                composed_dag.add_node(node, 0)
-
-        for output_name in composed_dag.get_output_names():
-            if output_name not in final_outputs:
-                composed_dag.remove_output(output_name)
-
-        # update the model graph
-        composed_dag.update()
+                new_inputs = [get_value(inp.name) if inp is not None else None for inp in node.inputs]
+                new_node = ir.Node(
+                    node.domain,
+                    node.op_type,
+                    inputs=new_inputs,
+                    attributes=list(node.attributes.values()),
+                    overload=node.overload,
+                    num_outputs=len(node.outputs),
+                    version=node.version,
+                    name=name,
+                )
+                for old_output, new_output in zip(node.outputs, new_node.outputs):
+                    out_name = old_output.name
+                    if not out_name:
+                        continue
+                    new_output.name = out_name
+                    new_output.type = old_output.type
+                    new_output.shape = old_output.shape
+                    composed_values[out_name] = new_output
+                    produced_names.add(out_name)
+                composed_nodes.append(new_node)
+                if name:
+                    composed_node_names[name] = new_node
+
+            for out in graph.outputs:
+                name = out.name
+                value = composed_values.get(name)
+                if value is None:
+                    value = get_value(name)
+                    value.type = out.type
+                    value.shape = out.shape
+                if name not in composed_output_names:
+                    composed_output_names.append(name)
+
+        # merge opset imports across all component models, taking the max version per domain
+        composed_opset_imports: dict[str, int] = {}
+        for ir_model in ir_models:
+            for domain, version in ir_model.graph.opset_imports.items():
+                if domain not in composed_opset_imports or version > composed_opset_imports[domain]:
+                    composed_opset_imports[domain] = version
+
+        composed_graph = ir.Graph(
+            inputs=composed_inputs,
+            outputs=[composed_values[name] for name in composed_output_names if name in final_outputs],
+            nodes=composed_nodes,
+            initializers=list(composed_initializers.values()),
+            name=ir_models[0].graph.name,
+            opset_imports=composed_opset_imports,
+        )
+        composed_model = ir.Model(
+            composed_graph,
+            ir_version=ir_models[0].ir_version,
+            producer_name="olive",
+        )
+        TopologicalSortPass()(composed_model)
 
         # save the composed model
         output_model_path = Path(output_model_path)
-        has_external_data = model_proto_to_file(composed_dag.model, output_model_path, **external_config)
+        has_external_data = model_proto_to_file(ir.to_proto(composed_model), output_model_path, **external_config)
 
         # copy over context binary files if any
         saved_cb_files = saved_cb_files if saved_cb_files is not None else {}

olive/passes/onnx/float16_conversion.py

@@ -5,12 +5,13 @@
 from pathlib import Path
 from typing import Union
 
+import onnx_ir as ir
+
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import ONNXModelHandler
 from olive.model.utils import resolve_onnx_path
 from olive.passes import Pass
 from olive.passes.onnx.common import get_external_data_config, model_proto_to_olive_model
-from olive.passes.onnx.onnx_dag import OnnxDAG
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
 
@@ -67,20 +68,27 @@ def _run_for_config(
         output_model_path = resolve_onnx_path(output_model_path, Path(model.model_path).name)
 
         loaded_model = model.load_model()
-        dag = OnnxDAG(loaded_model)
 
         op_block_list = config.op_block_list
-        if config.op_include_list:
-            if op_block_list is not None:
-                raise ValueError("op_include_list and op_block_list are mutually exclusive.")
-            op_block_list = [op_type for op_type in dag.get_node_op_types() if op_type not in config.op_include_list]
         node_block_list = config.node_block_list
-        if config.node_include_list:
-            if node_block_list is not None:
-                raise ValueError("node_include_list and node_block_list are mutually exclusive.")
-            node_block_list = [
-                node_name for node_name in dag.get_node_names() if node_name not in config.node_include_list
-            ]
+        if config.op_include_list or config.node_include_list:
+            ir_model = ir.from_proto(loaded_model)
+            all_nodes = list(ir_model.graph.all_nodes())
+            if config.op_include_list:
+                if op_block_list is not None:
+                    raise ValueError("op_include_list and op_block_list are mutually exclusive.")
+                op_block_list = list({node.op_type for node in all_nodes if node.op_type not in config.op_include_list})
+            if config.node_include_list:
+                if node_block_list is not None:
+                    raise ValueError("node_include_list and node_block_list are mutually exclusive.")
+                # node_include_list works by name, so every node must be named; otherwise unnamed
+                # nodes would silently be converted to float16 even though they can't be included.
+                if any(not node.name for node in all_nodes):
+                    raise ValueError(
+                        "node_include_list requires all nodes to be named, but the model contains unnamed nodes. "
+                        "Please ensure all nodes in the model have names or use node_block_list instead."
+                    )
+                node_block_list = [node.name for node in all_nodes if node.name not in config.node_include_list]
 
         # using the float16 converter from onnxruntime since it is regularly updated
         # and can handle large models (>2GB) as well as ort contrib ops