From 9197c790432d1029a6aa6fd903e2d607bef769b8 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:03 +0000 Subject: [PATCH 1/9] [DLA_transforms_3] [3/17] Standalone DLA graph cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 839 ++++++++++++ .../dla_transforms/test_dla_graph_cleanup.py | 1176 +++++++++++++++++ 2 files changed, 2015 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py new file mode 100644 index 00000000000..e619d6eeb4d --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -0,0 +1,839 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Combined graph-cleanup pass for DLA ONNX models. + +Applies five cleanup steps in sequence: + +1. **Canonicalize non-4D model inputs** — Each 1-D/2-D/3-D graph input gets a single + canonical ``Unsqueeze`` inserted at the front of the graph, padding to 4-D. + Existing ``Unsqueeze`` consumers of that input are either deduped (same output shape) + or replaced by a ``Reshape`` (different output shape). All other consumers are + redirected to the canonical Unsqueeze output. + +2. **Replace intermediary Squeeze / Unsqueeze with 4-D Reshape** — Any ``Squeeze`` or + ``Unsqueeze`` whose output is **not** a graph output is replaced by a ``Reshape`` whose + output shape is the original output shape padded to 4-D by prepending 1-s. + If the padded input and output shapes are identical the node is removed entirely + and consumers are wired to the input directly. + +3. **Collapse Reshape chains** — Consecutive ``Reshape`` nodes whose output is consumed + only by the next ``Reshape`` are collapsed. A chain where first-input-shape equals + last-output-shape is removed entirely; otherwise it is reduced to a single ``Reshape``. + +4. **Fold constants** — Uses ``onnx-graphsurgeon`` to evaluate constant subgraphs at + compile time, skipping ``DequantizeLinear`` and ``Identity`` nodes. Silently skipped + if ``onnx-graphsurgeon`` is unavailable. + +5. **Remove unused initializers** — Initializers that are not referenced as a node input + are dropped. +""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + _node_inputs, + add_unique_initializers, + add_value_info, + batch_replace_nodes, + get_node_attr_i, + infer_shapes, + insert_nodes_at_position, + pad4d, +) + +# ── Internal utilities ──────────────────────────────────────────────────────── + + +# ── Step 1: Canonicalize non-4D model inputs ───────────────────────────────── + + +def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: + """Insert one canonical Unsqueeze per non-4D graph input; dedup existing Unsqueeze consumers.""" + graph = model.graph + unique_init: set[str] = set() + consumer_map = cache.consumer_map + nodes_to_insert: list[tuple] = [] # (new_node, reference_node) + nodes_to_remove: list = [] + cnt = 0 + + for graph_input in list(graph.input): + tt = graph_input.type.tensor_type + if not tt.HasField("shape"): + continue + rank = sum(1 for _ in tt.shape.dim) + if rank == 4: + continue + if rank not in (1, 2, 3): + logger.debug( + "_canonicalize_inputs: skipping input %r with rank %d.", graph_input.name, rank + ) + continue + + # Read concrete input shape (may be None for symbolic dims) + input_shape: list[int | None] = [ + d.dim_value if d.HasField("dim_value") else None for d in tt.shape.dim + ] + + n_new = 4 - rank + axes_arr = np.arange(n_new, dtype=np.int64) + canonical_4d = [1] * n_new + list(input_shape) + + axes_name = f"{graph_input.name}_unsqueeze_axes" + out_name = f"{graph_input.name}_4d" + add_unique_initializers( + graph, unique_init, [numpy_helper.from_array(axes_arr, name=axes_name)] + ) + + unsq = helper.make_node( + "Unsqueeze", + inputs=[graph_input.name, axes_name], + outputs=[out_name], + name=f"{graph_input.name}_unsqueeze", + ) + add_value_info(graph, out_name, tt.elem_type, canonical_4d) + + consumers = consumer_map.get(graph_input.name, []) + consumers = [ + (c_node, c_idx) + for c_node, c_idx in consumers + if len(c_node.input) > 0 + and graph_input.name in c_node.input + and not ( + (c_node.op_type == "Squeeze" and "Gather" in c_node.name) + or c_node.op_type in ("GroupQueryAttention", "Gather") + ) + ] + if not consumers: + continue + + # Find the first consumer node in graph order for insertion point + consumer_set = {id(n) for n, _ in consumers} + ref_node = next((n for n in graph.node if id(n) in consumer_set), None) + nodes_to_insert.append((unsq, ref_node)) + + for c_node, c_idx in consumers: + if c_node.op_type == "Unsqueeze": + old_out_shape = cache.get_shape(c_node.output[0]) + if old_out_shape is not None: + old_4d = pad4d(old_out_shape) + if old_4d == canonical_4d: + if ( + c_node.output[0] in cache.graph_output_names + and old_out_shape != canonical_4d + ): + continue + if c_node.output[0] in cache.graph_output_names: + nodes_to_insert.append( + ( + helper.make_node( + "Identity", + [out_name], + list(c_node.output), + name=cache.unique_name(c_node.name + "_identity"), + ), + c_node, + ) + ) + # Same 4-D shape → redirect downstream consumers, remove old Unsqueeze + for dn, di in consumer_map.get(c_node.output[0], []): + dn.input[di] = out_name + else: + # Different 4-D shape → Reshape from canonical output + rs_name = f"{c_node.name}_reshape_shape" + add_unique_initializers( + graph, + unique_init, + [ + numpy_helper.from_array( + np.array(old_4d, dtype=np.int64), name=rs_name + ) + ], + ) + rs = helper.make_node( + "Reshape", + inputs=[out_name, rs_name], + outputs=list(c_node.output), + name=f"{c_node.name}_reshape", + ) + cache.update_shape(c_node.output[0], old_4d) + nodes_to_insert.append((rs, c_node)) + else: + # Unknown shape — just update the input to canonical unsqueeze output + continue + nodes_to_remove.append(c_node) + else: + c_node.input[c_idx] = out_name + cnt += 1 + logger.debug( + "_canonicalize_inputs: processed graph input %r (rank %d).", graph_input.name, rank + ) + + for new_node, ref in nodes_to_insert: + insert_nodes_at_position(graph, [new_node], ref) + for node in nodes_to_remove: + if node in graph.node: + graph.node.remove(node) + logger.debug("canonicalize_inputs: processed %d non-4D graph input(s).", cnt) + return model + + +# ── Step 2: Replace intermediary Squeeze / Unsqueeze with 4-D Reshape ──────── + + +def _compute_sq_unsq_out_shape( + node: onnx.NodeProto, graph: onnx.GraphProto, cache: GraphCache +) -> list[int] | None: + """Compute the output shape of a Squeeze/Unsqueeze from its axes and input shape. + + Used as a fallback when shape inference did not populate ``value_info``. + Returns ``None`` if the shape cannot be determined. + """ + in_shape = cache.get_shape(node.input[0]) + if in_shape is None: + return None + + # Resolve axes from input[1]: initializer or inline Constant node + axes: list[int] | None = None + if len(node.input) > 1 and node.input[1]: + for init in graph.initializer: + if init.name == node.input[1]: + axes = numpy_helper.to_array(init).tolist() + break + if axes is None: + for n in graph.node: + if n.op_type == "Constant" and n.output and n.output[0] == node.input[1]: + for attr in n.attribute: + if attr.name == "value": + axes = numpy_helper.to_array(attr.t).tolist() + break + break + + if node.op_type == "Squeeze": + if axes is None: + return [d for d in in_shape if d != 1] + out = list(in_shape) + for ax in sorted([a if a >= 0 else len(out) + a for a in axes], reverse=True): + if 0 <= ax < len(out): + out.pop(ax) + return out + else: # Unsqueeze + if axes is None: + return None + out = list(in_shape) + for ax in sorted([a if a >= 0 else len(out) + a + 1 for a in axes]): + out.insert(ax, 1) + return out + + +_GATHER_OPS = frozenset({"Gather", "GatherElements"}) +_GQA_OPS = frozenset({"GroupQueryAttention"}) + + +def _should_skip_squeeze_unsqueeze( + node: onnx.NodeProto, + cache: GraphCache, +) -> bool: + """Return True if this Squeeze/Unsqueeze should be left in place. + + Skip conditions: + * ``Unsqueeze`` whose data input is a graph input — canonical 4D-promotion node. + * Output feeds ``Gather`` / ``GatherElements`` ``input[1]`` (indices must stay INT32). + * Output feeds any ``GroupQueryAttention`` input (GQA manages its own shapes). + * ``Unsqueeze`` whose data input is a GQA output AND the GQA input was 3D while + the GQA output is 4D — this Unsqueeze was placed to adapt the 3D→4D boundary + and must be preserved. + + The check looks one level through an intervening ``Cast`` node since + ``Squeeze → Cast → Gather/GQA`` also requires the original shape to be preserved. + """ + if node.op_type == "Unsqueeze" and node.input and node.input[0] in cache.graph_input_names: + return True + + # Keep Unsqueeze whose input is a GQA output, the GQA output is 3D, + # and the Unsqueeze output is 4D — this is the canonical 3D→4D bridge + # after GroupQueryAttention and must be preserved. + if node.op_type == "Unsqueeze" and node.input and node.output: + prod = cache.get_producer(node.input[0]) + if prod is not None and prod.op_type in _GQA_OPS: + gqa_out_shape = cache.get_shape(node.input[0]) + unsq_out_shape = cache.get_shape(node.output[0]) + if ( + gqa_out_shape is not None + and len(gqa_out_shape) == 3 + and unsq_out_shape is not None + and len(unsq_out_shape) == 4 + ): + return True + + for c_node, c_idx in cache.get_consumers(node.output[0]): + if c_node.op_type in _GATHER_OPS and c_idx == 1 and node.op_type == "Squeeze": + return True + if c_node.op_type in _GQA_OPS: + return True + # Look one level deeper through Cast + if c_node.op_type == "Cast" and c_node.output: + for gc_node, gc_idx in cache.get_consumers(c_node.output[0]): + if gc_node.op_type in _GATHER_OPS and gc_idx == 1 and node.op_type == "Squeeze": + return True + if gc_node.op_type in _GQA_OPS: + return True + return False + + +def _replace_intermediary_squeeze_unsqueeze( + model: onnx.ModelProto, cache: GraphCache +) -> onnx.ModelProto: + """Replace Squeeze/Unsqueeze nodes with 4-D Reshape nodes. + + Skip conditions: + * Input is a graph input (Unsqueeze) — canonical 4D-promotion node, must be preserved. + * Output feeds Gather/GatherElements ``input[1]`` (indices must remain INT32). + * Output feeds GroupQueryAttention (GQA manages its own dimensionality). + * Output of GroupQueryAttention feeds to Unsqueeze node with 3D input. + """ + graph = model.graph + unique_init: set[str] = set() + graph_output_names = cache.graph_output_names + + nodes_to_remove: list = [] + nodes_to_insert: list[tuple] = [] + cnt_replace = 0 + + for node in list(graph.node): + if node.op_type not in ("Squeeze", "Unsqueeze"): + continue + if not node.input or not node.output: + continue + if _should_skip_squeeze_unsqueeze(node, cache): + logger.debug( + "%s %r: skipped (output feeds Gather indices, GQA, or follows GQA output).", + node.op_type, + node.name, + ) + continue + out_shape = cache.get_shape(node.output[0]) or _compute_sq_unsq_out_shape( + node, graph, cache + ) + if out_shape is None: + logger.debug("%s %r: output shape unknown, skipping.", node.op_type, node.name) + continue + + is_graph_output = node.output[0] in graph_output_names + + # Graph output: use declared shape (preserves name binding); else pad to 4D. + reshape_target = out_shape if is_graph_output else pad4d(out_shape) + + rs_name = f"{node.name}_4d_shape" + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(np.array(reshape_target, dtype=np.int64), name=rs_name)], + ) + rs = helper.make_node( + "Reshape", + inputs=[node.input[0], rs_name], + outputs=list(node.output), + name=f"{node.name}_to_reshape", + ) + nodes_to_insert.append((rs, node)) + nodes_to_remove.append(node) + cache.update_shape(node.output[0], reshape_target) + cnt_replace += 1 + logger.debug("%s %r: replaced by Reshape → %s.", node.op_type, node.name, reshape_target) + + for new_node, ref in nodes_to_insert: + insert_nodes_at_position(graph, [new_node], ref) + for node in nodes_to_remove: + if node in graph.node: + graph.node.remove(node) + logger.debug("replace_intermediary_sq_unsq: replaced %d node(s) with Reshape.", cnt_replace) + return model + + +# ── Step 3: Collapse Reshape chains ────────────────────────────────────────── + + +def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: + """Remove noop Reshapes then collapse consecutive Reshape chains. + + Pass 1 — Remove every noop Reshape (``in_shape == out_shape``) completely: + consumers are wired directly to the Reshape input; at graph-output boundaries + an ``Identity`` node is inserted to keep the output name alive. + + Pass 2 — Collapse consecutive Reshape chains (each member's output consumed + only by the next Reshape) into a single Reshape; chains that are noops as + a whole are removed the same way. + """ + graph = model.graph + graph_output_names = cache.graph_output_names + noop_count = 0 + chain_count = 0 + chain_nodes_count = 0 + # ── Pass 1: Remove standalone noop Reshapes ─────────────────────────────── + node_idx = {id(n): i for i, n in enumerate(graph.node)} + pass1_batch: dict[int, tuple[list, list]] = {} + + for node in cache.nodes_by_op("Reshape"): + in_shape = cache.get_shape(node.input[0]) + out_shape = cache.get_shape(node.output[0]) + if in_shape is None or out_shape is None: + continue + if list(in_shape) != list(out_shape): + continue + # Noop reshape: rewire downstream node-consumers to bypass this node + inp = node.input[0] + out_tensor = node.output[0] + for c_node, c_idx in cache.get_consumers(out_tensor): + if c_node.input[c_idx] == out_tensor: + c_node.input[c_idx] = inp + if out_tensor in graph_output_names: + id_node = helper.make_node( + "Identity", + [inp], + [out_tensor], + name=f"cleanup_noop_id_{out_tensor}", + ) + pass1_batch[node_idx[id(node)]] = ([node], [id_node]) + else: + pass1_batch[node_idx[id(node)]] = ([node], []) + noop_count += 1 + + batch_replace_nodes(graph, pass1_batch) + + # Rebuild cache/index after Pass 1 mutations so Pass 2 sees the updated + # topology, consumers, and shape metadata. + cache.rebuild() + node_idx = {id(n): i for i, n in enumerate(graph.node)} + pass2_batch: dict[int, tuple[list, list]] = {} + + # ── Pass 2: Collapse Reshape chains (Solution 2 — output-to-chain map) ──── + # Build a map: current_tail_output → chain_list. + # For each Reshape, check if its input tensor is an existing chain tail. + # If yes (and that Reshape is the sole consumer of that tensor), extend the + # chain. If no, start a new chain keyed by this Reshape's output. + # This correctly handles non-Reshape nodes interspersed between chain + # members (e.g. r1→r2→r5→r3 where r5 is not a Reshape: r1,r2,r3 still + # form a chain as long as each Reshape is the sole consumer of the previous + # Reshape's output). + all_reshapes = [n for n in graph.node if n.op_type == "Reshape"] + + # output_name → chain (list of Reshape nodes whose last output == output_name) + tail_to_chain: dict[str, list] = {} + + for node in all_reshapes: + inp_tensor = node.input[0] if node.input else None + if inp_tensor and inp_tensor in tail_to_chain: + # Extend existing chain only if this node is the sole consumer + consumers = cache.get_consumers(inp_tensor) + if len(consumers) == 1 and inp_tensor not in cache.graph_output_names: + chain = tail_to_chain.pop(inp_tensor) + chain.append(node) + tail_to_chain[node.output[0]] = chain + continue + # Start a new singleton chain + tail_to_chain[node.output[0]] = [node] + + # Flush all chains with length > 1 + for ch in tail_to_chain.values(): + if len(ch) <= 1: + continue + in_shape = cache.get_shape(ch[0].input[0]) + out_shape = cache.get_shape(ch[-1].output[0]) + if in_shape is None or out_shape is None: + continue + inp = ch[0].input[0] + out_tensor = ch[-1].output[0] + if list(in_shape) == list(out_shape): + # Entire chain is a noop + for c_node, c_idx in cache.get_consumers(out_tensor): + if c_node.input[c_idx] == out_tensor: + c_node.input[c_idx] = inp + if out_tensor in graph_output_names: + id_node = helper.make_node( + "Identity", + [inp], + [out_tensor], + name=f"cleanup_chain_id_{out_tensor}", + ) + pass2_batch[node_idx[id(ch[0])]] = (ch[:], [id_node]) + else: + pass2_batch[node_idx[id(ch[0])]] = (ch[:], []) + else: + rs_shape = numpy_helper.from_array( + np.array(out_shape, dtype=np.int64), + name=ch[-1].name + "_chain_shape", + ) + graph.initializer.append(rs_shape) + rs = helper.make_node( + "Reshape", + inputs=[inp, rs_shape.name], + outputs=[out_tensor], + name=ch[-1].name + "_chain", + ) + pass2_batch[node_idx[id(ch[0])]] = (ch[:], [rs]) + chain_count += 1 + chain_nodes_count += len(ch) + + batch_replace_nodes(graph, pass2_batch) + + logger.debug( + "remove_reshape_chains: removed %d noop Reshape(s); collapsed %d chain(s) (%d nodes) to single Reshape.", + noop_count, + chain_count, + chain_nodes_count, + ) + return model + + +# ── Step 4: Fold constants ──────────────────────────────────────────────────── + + +def _fold_constants(model: onnx.ModelProto) -> onnx.ModelProto: + """Fold constant-valued subgraphs using onnx-graphsurgeon. + + ``DequantizeLinear`` and ``Identity`` nodes are excluded from folding. + Returns the (possibly new) model. If onnx-graphsurgeon is unavailable the + original model is returned unchanged. + """ + try: + import onnx_graphsurgeon as gs + except ImportError: + logger.warning("onnx-graphsurgeon not installed; skipping constant folding.") + return model + + def _exclude(node) -> bool: + return node.op in ("DequantizeLinear", "Identity") + + prev_ir = onnx.IR_VERSION + onnx.IR_VERSION = model.ir_version + try: + g = gs.import_onnx(model) + g.fold_constants(should_exclude_node=_exclude).cleanup().toposort() + model2 = gs.export_onnx(g) + finally: + onnx.IR_VERSION = prev_ir + + folded = len(model.graph.node) - len(model2.graph.node) + if folded > 0: + logger.debug("fold_constants: folded %d node(s).", folded) + model2 = infer_shapes(model2) + return model2 + + +# ── Step 5: Replace graph-output Reshape with Squeeze where possible ───────── + + +def _replace_graph_output_reshape_with_squeeze( + model: onnx.ModelProto, cache: GraphCache +) -> onnx.ModelProto: + """Replace a Reshape that feeds a graph output with an equivalent Squeeze. + + A Reshape ``[*in_shape] → [*out_shape]`` at a graph output can be replaced by + ``Squeeze(axes)`` if and only if: + + * Both shapes are fully known (no symbolic dims). + * The output shape is a strict sub-sequence of the input shape obtained by + removing only size-1 dimensions (i.e. the total element count is unchanged + and no dimension values change — only unit dims are dropped). + + The squeeze axes are the indices in ``in_shape`` of the 1-dims that are not + present in ``out_shape``. + + Example:: + + [1, 3, 4, 1] → Reshape → [3, 4] (graph output) + becomes + [1, 3, 4, 1] → Squeeze(axes=[0, 3]) → [3, 4] + """ + graph = model.graph + unique_init: set[str] = set() + graph_output_names = cache.graph_output_names + + go_node_idx = {id(n): i for i, n in enumerate(graph.node)} + go_batch: dict[int, tuple[list, list]] = {} + cnt = 0 + + for node in list(graph.node): + if node.op_type != "Reshape": + continue + if not node.output or node.output[0] not in graph_output_names: + continue + + in_name = node.input[0] + out_name = node.output[0] + + in_shape = cache.get_shape(in_name) + out_shape = cache.get_shape(out_name) + + if in_shape is None or out_shape is None: + continue # shapes unknown — can't reason + if None in in_shape or None in out_shape: + continue # symbolic dims — skip + + # Verify element count is unchanged + in_elems = 1 + for d in in_shape: + in_elems *= d + out_elems = 1 + for d in out_shape: + out_elems *= d + if in_elems != out_elems: + continue + + # Check that out_shape is a subsequence of in_shape with only 1-dims removed. + # Match from RIGHT-TO-LEFT so that the rightmost (non-unit) dims are anchored + # first. The leftmost 1-dims that cannot be matched become the squeeze axes, + # guaranteeing we always squeeze from the lowest axis indices (left to right). + # + # Example: [1,1,1,512] → [1,512] + # Right-to-left: 512 matches, then 1 matches, leaving leading [1,1] → axes=[0,1] + # (The old left-to-right approach would yield axes=[1,2] instead.) + squeeze_axes: list[int] = [] + j = len(out_shape) - 1 # pointer into out_shape, going backwards + matched = True + for i in range(len(in_shape) - 1, -1, -1): + d = in_shape[i] + if j >= 0 and d == out_shape[j]: + j -= 1 # matched this dim from the right + elif d == 1: + squeeze_axes.append(i) # unit dim to drop + else: + matched = False + break + pfx = node.name or f"go_sq_{cnt}" + + if not matched or j != -1: + # Cannot express as a pure Squeeze of the existing input. + # If the output is non-4D, pad the Reshape target to 4D with a + # leading 1-dim and add a Squeeze(axis=[0]) to strip it back. + # + # Example: + # [5,6,7,8] → Reshape([30,14,4]) → graph_output[30,14,4] + # becomes: + # [5,6,7,8] → Reshape([1,30,14,4]) → Squeeze([0]) → graph_output[30,14,4] + if len(out_shape) >= 4: + continue # already 4D or higher — leave as-is + # Build a 4D shape by prepending 1 + shape_4d = [1, *out_shape] + rs_shape_name = f"{pfx}_4d_shape" + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(np.array(shape_4d, dtype=np.int64), name=rs_shape_name)], + ) + # Intermediate 4D tensor name + rs_4d_out = f"{pfx}_4d_out" + # Update the Reshape node to target the 4D shape + node.input[1] = rs_shape_name + node.output[0] = rs_4d_out + # Squeeze to remove the leading 1-dim + sq_axes_name = f"{pfx}_sq0_axes" + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(np.array([0], dtype=np.int64), name=sq_axes_name)], + ) + sq = helper.make_node( + "Squeeze", + inputs=[rs_4d_out, sq_axes_name], + outputs=[out_name], + name=f"{pfx}_squeeze_leading", + ) + go_batch[go_node_idx[id(node)]] = ([node], [node, sq]) + cnt += 1 + logger.debug( + "Reshape %r → Reshape(%s) + Squeeze([0]) for non-4D graph output %r.", + node.name, + shape_4d, + out_name, + ) + continue + + squeeze_axes.sort() # ascending axis order required by ONNX Squeeze + + # Build Squeeze node + axes_name = f"{pfx}_squeeze_axes" + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(np.array(squeeze_axes, dtype=np.int64), name=axes_name)], + ) + sq = helper.make_node( + "Squeeze", + inputs=[in_name, axes_name], + outputs=[out_name], + name=f"{pfx}_squeeze", + ) + go_batch[go_node_idx[id(node)]] = ([node], [sq]) + cnt += 1 + logger.debug( + "Reshape %r → Squeeze(axes=%s) for graph output %r.", + node.name, + squeeze_axes, + out_name, + ) + + batch_replace_nodes(graph, go_batch) + + logger.debug("replace_graph_output_reshape_with_squeeze: replaced %d Reshape(s).", cnt) + return model + + +# ── Step 7: Remove unused initializers ─────────────────────────────────────── + + +def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: + """Drop initializers that are not referenced by any node input.""" + graph = model.graph + used: set[str] = set() + for node in graph.node: + used.update(_node_inputs(node)) + # Graph outputs are kept alive by convention even if no node reads them + for o in graph.output: + used.add(o.name) + + kept = [init for init in graph.initializer if init.name in used] + removed = len(graph.initializer) - len(kept) + graph.ClearField("initializer") + graph.initializer.extend(kept) + logger.debug("remove_unused_initializers: removed %d unused initializer(s).", removed) + return model + + +# ── Step 7: Remove consecutive Cast chains ────────────────────────────────── + + +def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: + """Collapse consecutive Cast nodes into a single Cast with the last node's target dtype. + + A chain ``Cast(to=A) → Cast(to=B) → Cast(to=C)`` is replaced by a single + ``Cast(to=C)`` from the first Cast's input to the last Cast's output. + Chains of length 1 are left untouched. + """ + graph = model.graph + node_idx = {id(n): i for i, n in enumerate(graph.node)} + cnt = 0 + + # Build chains: tail_output → [cast_node, ...] + all_casts = [n for n in graph.node if n.op_type == "Cast"] + tail_to_chain: dict[str, list] = {} + + for node in all_casts: + inp_tensor = node.input[0] if node.input else None + if inp_tensor and inp_tensor in tail_to_chain: + consumers = cache.get_consumers(inp_tensor) + if len(consumers) == 1 and inp_tensor not in cache.graph_output_names: + chain = tail_to_chain.pop(inp_tensor) + chain.append(node) + tail_to_chain[node.output[0]] = chain + continue + tail_to_chain[node.output[0]] = [node] + + # Collapse chains with length > 1 + replacements: dict[int, tuple[list, list]] = {} + for ch in tail_to_chain.values(): + if len(ch) <= 1: + continue + first_input = ch[0].input[0] + source_type = cache.get_dtype(first_input) + # Removing a lossy intermediate cast changes rounding/truncation. + lossless = { + onnx.TensorProto.FLOAT16: {onnx.TensorProto.FLOAT, onnx.TensorProto.DOUBLE}, + onnx.TensorProto.FLOAT: {onnx.TensorProto.DOUBLE}, + onnx.TensorProto.INT8: { + onnx.TensorProto.INT16, + onnx.TensorProto.INT32, + onnx.TensorProto.INT64, + onnx.TensorProto.FLOAT, + onnx.TensorProto.DOUBLE, + }, + onnx.TensorProto.INT16: { + onnx.TensorProto.INT32, + onnx.TensorProto.INT64, + onnx.TensorProto.FLOAT, + onnx.TensorProto.DOUBLE, + }, + onnx.TensorProto.INT32: {onnx.TensorProto.INT64, onnx.TensorProto.DOUBLE}, + } + if any( + get_node_attr_i(node, "to", 1) not in {source_type, *lossless.get(source_type, set())} + for node in ch[:-1] + ): + continue + last_output = ch[-1].output[0] + last_to = get_node_attr_i(ch[-1], "to", 1) # ONNX Cast default: FLOAT + + new_cast = helper.make_node( + "Cast", + inputs=[first_input], + outputs=[last_output], + to=last_to, + name=f"{ch[-1].name}_chain", + ) + idx = node_idx[id(ch[0])] + replacements[idx] = (ch[:], [new_cast]) + cnt += len(ch) + + batch_replace_nodes(graph, replacements) + logger.debug( + "remove_consecutive_casts: collapsed %d Cast node(s) in %d chain(s).", + cnt, + len(replacements), + ) + return model + + +# ── Master transform ────────────────────────────────────────────────────────── + + +def _apply_graph_cleanup(model: onnx.ModelProto) -> onnx.ModelProto: + """Apply all five cleanup steps in sequence.""" + model = infer_shapes(model) + cache = GraphCache(model.graph) + + # Step 1 — intermediary Squeeze/Unsqueeze → 4-D Reshape + model = _replace_intermediary_squeeze_unsqueeze(model, cache) + cache.rebuild() + + # Step 2 — non-4D inputs → canonical Unsqueeze + model = _canonicalize_inputs(model, cache) + model = infer_shapes(model) + cache.rebuild() + + # Step 3 — collapse Reshape chains + model = _remove_reshape_chains(model, cache) + cache.rebuild() + + # Step 4 — replace graph-output Reshape with Squeeze where possible + model = _replace_graph_output_reshape_with_squeeze(model, cache) + + # Step 5 — fold constants (graphsurgeon creates new model — cache invalidated) + model = _fold_constants(model) + + # Step 6 — collapse consecutive Cast chains + cache = GraphCache(model.graph) + model = _remove_consecutive_casts(model, cache) + + # Step 7 — remove unused initializers + model = _remove_unused_initializers(model) + + return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py new file mode 100644 index 00000000000..6383799c015 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -0,0 +1,1176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_graph_cleanup`. + +"""Structural and ORT parity tests for ``dla_graph_cleanup``. + +The combined transform applies five steps: + 1. Canonicalize non-4D inputs (adds Unsqueeze, deduplicates existing ones) + 2. Replace intermediary Squeeze/Unsqueeze with 4-D Reshape + 3. Collapse Reshape chains (noop → remove; multi → single Reshape) + 4. Fold constants via onnx-graphsurgeon + 5. Remove unused initializers + +Steps 1 and 2 intentionally change tensor shapes, so ORT parity only works +for end-to-end models that happen to keep graph I/O shapes the same. +Tests for those steps verify structural changes (node types, counts). +Steps 3-5 preserve graph I/O shapes, so ORT parity is verified for them. +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms import dla_graph_cleanup as _mod + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "graph_cleanup_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_graph_cleanup +_fold_consts = _mod._fold_constants +_remove_unused_inits = _mod._remove_unused_initializers + +from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache + + +def _canonicalize_inputs(model): + return _mod._canonicalize_inputs(model, GraphCache(model.graph)) + + +def _replace_sq_unsq(model): + return _mod._replace_intermediary_squeeze_unsqueeze(model, GraphCache(model.graph)) + + +def _remove_reshapes(model): + return _mod._remove_reshape_chains(model, GraphCache(model.graph)) + + +def _reshape_to_squeeze(model): + return _mod._replace_graph_output_reshape_with_squeeze(model, GraphCache(model.graph)) + + +def _remove_cast_chains(model): + return _mod._remove_consecutive_casts(model, GraphCache(model.graph)) + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model: onnx.ModelProto) -> None: + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + inferred = SymbolicShapeInference.infer_shapes(model) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + except Exception: + pass + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + import onnxruntime as ort # type: ignore[import-untyped] + + _clamp_ir(model) + sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +_RNG = np.random.default_rng(2024) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 1 — Canonicalize non-4D inputs +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_canonicalize_3d_input_unsqueeze_added(): + """3-D graph input X [2,3,4] → Unsqueeze to 4-D [1,2,3,4] is inserted. + + Downstream Relu receives the 4-D tensor. + """ + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "g", + [_vi("X", TensorProto.FLOAT, [2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [2, 3, 4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _canonicalize_inputs(tr) + + unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] + assert len(unsq_nodes) == 1, "one Unsqueeze must be inserted for the 3-D input" + unsq = unsq_nodes[0] + assert unsq.input[0] == "X", "Unsqueeze reads the graph input" + + # Relu must now read the Unsqueeze output, not X directly + relu_node = next(n for n in tr.graph.node if n.op_type == "Relu") + assert relu_node.input[0] == unsq.output[0], "Relu must be redirected to Unsqueeze output" + + +def test_canonicalize_2d_input_unsqueeze_added(): + """2-D graph input X [4,5] → Unsqueeze at axes [0,1] to 4-D [1,1,4,5].""" + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "g", + [_vi("X", TensorProto.FLOAT, [4, 5])], + [_vi("Y", TensorProto.FLOAT, [4, 5])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _canonicalize_inputs(tr) + + unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] + assert len(unsq_nodes) == 1 + axes_init = next(i for i in tr.graph.initializer if i.name.endswith("_unsqueeze_axes")) + axes = numpy_helper.to_array(axes_init).tolist() + assert axes == [0, 1], f"axes should be [0,1] for 2-D input, got {axes}" + + +def test_canonicalize_3d_input_duplicate_unsqueezes_deduped(): + """3-D input consumed by two Unsqueeze nodes → only one canonical Unsqueeze remains. + + Both existing Unsqueezes (same axes) are replaced by the single canonical one. + """ + axes_init = numpy_helper.from_array(np.array([0], dtype=np.int64), name="ax") + unsq_a = helper.make_node("Unsqueeze", ["X", "ax"], ["Xa"], name="unsq_a") + unsq_b = helper.make_node("Unsqueeze", ["X", "ax"], ["Xb"], name="unsq_b") + relu_a = helper.make_node("Relu", ["Xa"], ["Ya"], name="relu_a") + relu_b = helper.make_node("Relu", ["Xb"], ["Yb"], name="relu_b") + graph = helper.make_graph( + [unsq_a, unsq_b, relu_a, relu_b], + "g", + [_vi("X", TensorProto.FLOAT, [2, 3, 4])], + [_vi("Ya", TensorProto.FLOAT, [1, 2, 3, 4]), _vi("Yb", TensorProto.FLOAT, [1, 2, 3, 4])], + initializer=[axes_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _canonicalize_inputs(tr) + + unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] + # Only the single canonical Unsqueeze should remain; old ones are removed/replaced + assert len(unsq_nodes) == 1, ( + f"exactly one Unsqueeze should remain after dedup, got {len(unsq_nodes)}" + ) + + +def test_canonicalize_4d_input_unchanged(): + """4-D graph input → step 1 is a noop.""" + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "g", + [_vi("X", TensorProto.FLOAT, [1, 2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 3, 4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_nodes = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _canonicalize_inputs(tr) + + assert len(tr.graph.node) == orig_nodes, "4-D input: no nodes should be added" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 2 — Replace intermediary Squeeze / Unsqueeze with 4-D Reshape +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_intermediary_squeeze_replaced_with_reshape(): + """Squeeze [2,1,6,8] → [2,6,8] (axis=1) is replaced by Reshape → [1,2,6,8]. + + pad([2,1,6,8]) = [2,1,6,8], pad([2,6,8]) = [1,2,6,8] — shapes differ → Reshape inserted. + The squeezed axis must be 1 (a unit dimension) for ORT symbolic shape inference to accept it. + """ + sq_axes = numpy_helper.from_array(np.array([1], dtype=np.int64), name="sq_ax") + relu = helper.make_node("Relu", ["X"], ["X_relu"], name="relu") + squeeze = helper.make_node("Squeeze", ["X_relu", "sq_ax"], ["Sq_out"], name="sq") + relu2 = helper.make_node("Relu", ["Sq_out"], ["Y"], name="relu2") + graph = helper.make_graph( + [relu, squeeze, relu2], + "g", + [_vi("X", TensorProto.FLOAT, [2, 1, 6, 8])], + [_vi("Y", TensorProto.FLOAT, [2, 6, 8])], + initializer=[sq_axes], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + assert not any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze must be removed" + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert reshape_nodes, "a Reshape must be inserted" + + # Shape init name follows pattern {node.name}_4d_shape → "sq_4d_shape" + shape_init = next((i for i in tr.graph.initializer if i.name == "sq_4d_shape"), None) + assert shape_init is not None, "Reshape shape initializer must exist" + shape = numpy_helper.to_array(shape_init).tolist() + assert shape == [1, 2, 6, 8], f"expected [1,2,6,8], got {shape}" + + +def test_intermediary_unsqueeze_replaced_with_reshape(): + """Unsqueeze [2,4,6] → [2,4,1,6] (axis=2) is replaced by Reshape → [2,4,1,6]. + + pad([2,4,6]) = [1,2,4,6], pad([2,4,1,6]) = [2,4,1,6] — shapes differ → Reshape inserted. + """ + unsq_axes = numpy_helper.from_array(np.array([2], dtype=np.int64), name="unsq_ax") + relu = helper.make_node("Relu", ["X"], ["X_relu"], name="relu") + unsqueeze = helper.make_node("Unsqueeze", ["X_relu", "unsq_ax"], ["Uq_out"], name="uq") + relu2 = helper.make_node("Relu", ["Uq_out"], ["Y"], name="relu2") + graph = helper.make_graph( + [relu, unsqueeze, relu2], + "g", + [_vi("X", TensorProto.FLOAT, [2, 4, 6])], + [_vi("Y", TensorProto.FLOAT, [2, 4, 1, 6])], + initializer=[unsq_axes], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + assert not any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze must be removed" + assert any(n.op_type == "Reshape" for n in tr.graph.node), "Reshape must be inserted" + # Shape init name: {node.name}_4d_shape → "uq_4d_shape" + shape_init = next((i for i in tr.graph.initializer if i.name == "uq_4d_shape"), None) + assert shape_init is not None, "Reshape shape initializer must exist" + shape = numpy_helper.to_array(shape_init).tolist() + assert shape == [2, 4, 1, 6], f"expected [2,4,1,6], got {shape}" + + +def test_intermediary_squeeze_at_graph_output_replaced(): + """Squeeze whose output IS a graph output is now replaced like any other intermediary Squeeze. + + Squeeze [1,2,1,5] → [1,2,5] (axis=2): pad([1,2,1,5])=[1,2,1,5] ≠ pad([1,2,5])=[1,1,2,5] + so a Reshape is inserted and the Squeeze is removed. + """ + sq_axes = numpy_helper.from_array(np.array([2], dtype=np.int64), name="sq_ax2") + relu = helper.make_node("Relu", ["X"], ["X_relu"], name="relu") + squeeze = helper.make_node("Squeeze", ["X_relu", "sq_ax2"], ["Y"], name="sq") + graph = helper.make_graph( + [relu, squeeze], + "g", + [_vi("X", TensorProto.FLOAT, [1, 2, 1, 5])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 5])], + initializer=[sq_axes], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + assert not any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze at graph output must now be replaced" + ) + assert any(n.op_type == "Reshape" for n in tr.graph.node), "Reshape must be inserted" + + +def test_intermediary_squeeze_identity_at_graph_output_uses_reshape(): + """Squeeze identity in 4-D AND graph output → Reshape with graph output shape. + + Squeeze [1,1,5] → [1,5] (axis=0): padded shapes are identical so normally the node + would be removed. Because Y is a graph output, a Reshape is inserted instead and its + target shape must equal the declared graph output shape [1,5] — not the 4D padded shape. + """ + sq_axes = numpy_helper.from_array(np.array([0], dtype=np.int64), name="sq_ax_id") + relu = helper.make_node("Relu", ["X"], ["X_relu"], name="relu") + squeeze = helper.make_node("Squeeze", ["X_relu", "sq_ax_id"], ["Y"], name="sq_id") + graph = helper.make_graph( + [relu, squeeze], + "g", + [_vi("X", TensorProto.FLOAT, [1, 1, 5])], + [_vi("Y", TensorProto.FLOAT, [1, 5])], + initializer=[sq_axes], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + assert not any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be replaced even for identity case at graph output" + ) + assert any(n.op_type == "Reshape" for n in tr.graph.node), ( + "Reshape must be inserted to preserve the graph output name" + ) + + # Reshape target shape must equal the graph output shape [1,5], not 4D padded [1,1,1,5] + rs_shape_init = next((i for i in tr.graph.initializer if i.name == "sq_id_4d_shape"), None) + assert rs_shape_init is not None, "Reshape shape initializer must exist" + assert numpy_helper.to_array(rs_shape_init).tolist() == [1, 5], ( + "Reshape target must equal graph output shape [1,5]" + ) + + assert any(o.name == "Y" for o in tr.graph.output), "Y must remain a graph output" + + +def test_unsqueeze_on_graph_input_skipped(): + """Unsqueeze whose data input IS a graph input must not be replaced. + + This Unsqueeze is the canonical 4D-promotion node placed by convert_ops_to_4d + or canonicalize_inputs. Replacing it with a Reshape would break the graph + because the input tensor shape stays 3D but the Reshape would expect a 4D input. + """ + unsq_axes = numpy_helper.from_array(np.array([0], dtype=np.int64), name="gi_unsq_ax") + # Unsqueeze directly on graph input X [2,4,6] → [1,2,4,6] + unsqueeze = helper.make_node("Unsqueeze", ["X", "gi_unsq_ax"], ["X_4d"], name="gi_unsq") + relu = helper.make_node("Relu", ["X_4d"], ["Y"], name="relu") + graph = helper.make_graph( + [unsqueeze, relu], + "g", + [_vi("X", TensorProto.FLOAT, [2, 4, 6])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 4, 6])], + initializer=[unsq_axes], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_node_count = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + # Unsqueeze on graph input must be preserved exactly as-is + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + "Unsqueeze on graph input must be preserved" + ) + assert not any(n.op_type == "Reshape" for n in tr.graph.node), ( + "No Reshape must be inserted for graph-input Unsqueeze" + ) + assert len(tr.graph.node) == orig_node_count, "Node count must be unchanged" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 3 — Collapse Reshape chains +# ══════════════════════════════════════════════════════════════════════════════ + + +def _make_reshape_graph(shapes: list[list[int]]) -> onnx.ModelProto: + """Build a graph: X → Reshape1 → Reshape2 → … → Y, each using a fixed shape.""" + assert len(shapes) >= 2 + inits = [] + nodes = [] + prev_name = "X" + for i, shape in enumerate(shapes[1:], 1): + sh_init = numpy_helper.from_array(np.array(shape, dtype=np.int64), name=f"sh{i}") + inits.append(sh_init) + out_name = "Y" if i == len(shapes) - 1 else f"T{i}" + nodes.append(helper.make_node("Reshape", [prev_name, f"sh{i}"], [out_name], name=f"rs{i}")) + prev_name = out_name + graph = helper.make_graph( + nodes, + "g", + [_vi("X", TensorProto.FLOAT, shapes[0])], + [_vi("Y", TensorProto.FLOAT, shapes[-1])], + initializer=inits, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + return model + + +def test_reshape_chain_noop_removed(): + """X[1,2,3,4] → Reshape[1,2,3,4] → Reshape[1,2,3,4] → Y: both removed, ORT parity.""" + model = _make_reshape_graph([[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]) + x = _RNG.standard_normal([1, 2, 3, 4]).astype(np.float32) + + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_reshapes(tr) + + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(reshape_nodes) == 0, "identity-reshape chain must be fully removed" + + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_reshape_chain_collapsed_to_one(): + """X[6] → Reshape[2,3] → Reshape[1,1,2,3] → Y: collapsed to single Reshape, ORT parity.""" + model = _make_reshape_graph([[6], [2, 3], [1, 1, 2, 3]]) + x = _RNG.standard_normal([6]).astype(np.float32) + + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_reshapes(tr) + + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(reshape_nodes) == 1, "two-step chain must collapse to one Reshape" + + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_single_reshape_unchanged(): + """A single Reshape is not touched.""" + model = _make_reshape_graph([[6], [2, 3]]) + orig_count = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_reshapes(tr) + + assert len(tr.graph.node) == orig_count, "single Reshape must not be changed" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 4 — Fold constants +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_fold_constants_const_add_folded(): + """Constant + Add → output is a constant: the Add is folded away. + + After folding, the graph should have fewer nodes and ORT output must match. + """ + a_init = numpy_helper.from_array(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="A") + b_init = numpy_helper.from_array(np.array([4.0, 5.0, 6.0], dtype=np.float32), name="B") + add = helper.make_node("Add", ["A", "B"], ["Y"], name="add") + graph = helper.make_graph( + [add], + "g", + [], + [_vi("Y", TensorProto.FLOAT, [3])], + initializer=[a_init, b_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {}) + + tr = _fold_consts(copy.deepcopy(model)) + _prepare_model(tr) + got = _run_ort(tr, {}) + + # After folding the Add node should be gone (its result is a new initializer) + assert not any(n.op_type == "Add" for n in tr.graph.node), "Add must be folded away" + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_fold_constants_dynamic_input_not_folded(): + """Add with a dynamic input is not folded.""" + a_init = numpy_helper.from_array(np.array([1.0, 2.0, 3.0], dtype=np.float32), name="A") + add = helper.make_node("Add", ["X", "A"], ["Y"], name="add") + graph = helper.make_graph( + [add], + "g", + [_vi("X", TensorProto.FLOAT, [3])], + [_vi("Y", TensorProto.FLOAT, [3])], + initializer=[a_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_count = len(model.graph.node) + + tr = _fold_consts(copy.deepcopy(model)) + + assert len(tr.graph.node) == orig_count, "dynamic Add must not be folded" + + x = _RNG.standard_normal([3]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 5 — Remove unused initializers +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_remove_unused_initializers(): + """Initializer not referenced by any node is dropped; ORT parity preserved.""" + used = numpy_helper.from_array(np.array([1.0, 2.0], dtype=np.float32), name="used") + orphan = numpy_helper.from_array(np.array([99.0], dtype=np.float32), name="orphan") + relu = helper.make_node("Relu", ["used"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "g", + [], + [_vi("Y", TensorProto.FLOAT, [2])], + initializer=[used, orphan], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {}) + + tr = copy.deepcopy(model) + tr = _remove_unused_inits(tr) + + init_names = {i.name for i in tr.graph.initializer} + assert "used" in init_names, "'used' initializer must be kept" + assert "orphan" not in init_names, "'orphan' initializer must be removed" + + _prepare_model(tr) + got = _run_ort(tr, {}) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_remove_unused_initializers_all_used_unchanged(): + """When all initializers are in use, none are removed.""" + a = numpy_helper.from_array(np.array([1.0, 2.0], dtype=np.float32), name="A") + b = numpy_helper.from_array(np.array([3.0, 4.0], dtype=np.float32), name="B") + add = helper.make_node("Add", ["A", "B"], ["Y"], name="add") + graph = helper.make_graph( + [add], + "g", + [], + [_vi("Y", TensorProto.FLOAT, [2])], + initializer=[a, b], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_count = len(model.graph.initializer) + + tr = copy.deepcopy(model) + tr = _remove_unused_inits(tr) + + assert len(tr.graph.initializer) == orig_count, "no initializers should be removed" + + +# ══════════════════════════════════════════════════════════════════════════════ +# Combined: noop graph → all steps are noops, ORT parity +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_combined_noop_graph(): + """4-D input, single Relu, no Reshape chains, no constants → combined transform is a noop.""" + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "g", + [_vi("X", TensorProto.FLOAT, [1, 2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 3, 4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_count = len(model.graph.node) + orig_inits = len(model.graph.initializer) + + x = _RNG.standard_normal([1, 2, 3, 4]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert len(tr.graph.node) == orig_count, "noop: node count unchanged" + assert len(tr.graph.initializer) == orig_inits, "noop: initializer count unchanged" + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Combined: Reshape chain + orphan init + constant fold +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_combined_reshape_and_unused_init_cleaned(): + """Combined: Reshape chain (noop) removed + orphan init removed, ORT parity.""" + x_data = _RNG.standard_normal([1, 2, 3, 4]).astype(np.float32) + orphan = numpy_helper.from_array(np.array([0.0], dtype=np.float32), name="orphan2") + sh1 = numpy_helper.from_array(np.array([1, 2, 3, 4], dtype=np.int64), name="sh1") + sh2 = numpy_helper.from_array(np.array([1, 2, 3, 4], dtype=np.int64), name="sh2") + + rs1 = helper.make_node("Reshape", ["X", "sh1"], ["T1"], name="rs1") + rs2 = helper.make_node("Reshape", ["T1", "sh2"], ["Y"], name="rs2") + graph = helper.make_graph( + [rs1, rs2], + "g", + [_vi("X", TensorProto.FLOAT, [1, 2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 3, 4])], + initializer=[orphan, sh1, sh2], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {"X": x_data}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x_data}) + + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(reshape_nodes) == 0, "identity-Reshape chain must be fully removed" + assert "orphan2" not in {i.name for i in tr.graph.initializer}, "orphan init must be removed" + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_interleaved_reshape_chains_collapsed(): + """Two Reshape chains interleaved in node order are both collapsed correctly. + + Graph topology: + X1 [1,2,4,4] --r1--> [1,4,2,4] --r2--> [2,4,4,1] --\\ + Add --> Y [2,4,4,1] + X2 [1,2,4,4] --r4--> [1,4,2,4] --r5--> [2,4,4,1] --/ + + Node sequence in the graph: [r1, r4, r2, r5, add] + (r4 is interleaved between r1 and r2 in the node list) + + The old sequential algorithm misses both chains: + - Sees r1, then r4 (not consuming T1) → flushes singleton [r1], starts [r4] + - Then r2 (not consuming T4) → flushes singleton [r4], starts [r2] + - Then r5 (not consuming T2) → flushes singleton [r2], starts [r5] + - End: flushes singleton [r5] + → All 4 Reshape nodes remain unchanged. + + The new output-to-chain map algorithm finds both chains: + - r1 → {T1: [r1]}; r4 → {T1:[r1], T4:[r4]} + - r2 input=T1 in map → {T4:[r4], T2:[r1,r2]} + - r5 input=T4 in map → {T2:[r1,r2], T5:[r4,r5]} + - Flush [r1,r2]: in=[1,2,4,4] != out=[2,4,4,1] → collapse to 1 Reshape + - Flush [r4,r5]: same → collapse to 1 Reshape + → 2 Reshape nodes remain (one per chain), ORT parity. + """ + # All shapes are 4D; chain input [1,2,4,4] != chain output [2,4,4,1] + # so each 2-node chain collapses to a single Reshape (not removed as noop). + sh_a = numpy_helper.from_array(np.array([1, 4, 2, 4], dtype=np.int64), name="sh_a") + sh_b = numpy_helper.from_array(np.array([2, 4, 4, 1], dtype=np.int64), name="sh_b") + sh_c = numpy_helper.from_array(np.array([1, 4, 2, 4], dtype=np.int64), name="sh_c") + sh_d = numpy_helper.from_array(np.array([2, 4, 4, 1], dtype=np.int64), name="sh_d") + + r1 = helper.make_node("Reshape", ["X1", "sh_a"], ["T1"], name="r1") + r2 = helper.make_node("Reshape", ["T1", "sh_b"], ["T2"], name="r2") + r4 = helper.make_node("Reshape", ["X2", "sh_c"], ["T4"], name="r4") + r5 = helper.make_node("Reshape", ["T4", "sh_d"], ["T5"], name="r5") + add = helper.make_node("Add", ["T2", "T5"], ["Y"], name="add1") + + # Interleaved: r4 sits between r1 and r2 in the node list + graph = helper.make_graph( + [r1, r4, r2, r5, add], + "interleaved_chains", + [ + _vi("X1", TensorProto.FLOAT, [1, 2, 4, 4]), + _vi("X2", TensorProto.FLOAT, [1, 2, 4, 4]), + ], + [_vi("Y", TensorProto.FLOAT, [2, 4, 4, 1])], + initializer=[sh_a, sh_b, sh_c, sh_d], + value_info=[ + _vi("T1", TensorProto.FLOAT, [1, 4, 2, 4]), + _vi("T2", TensorProto.FLOAT, [2, 4, 4, 1]), + _vi("T4", TensorProto.FLOAT, [1, 4, 2, 4]), + _vi("T5", TensorProto.FLOAT, [2, 4, 4, 1]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x1 = _RNG.standard_normal([1, 2, 4, 4]).astype(np.float32) + x2 = _RNG.standard_normal([1, 2, 4, 4]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X1": x1, "X2": x2}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_reshapes(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X1": x1, "X2": x2}) + + # Each 2-node chain collapses to 1 Reshape → 2 Reshape nodes total (down from 4) + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(reshape_nodes) == 2, ( + f"Each chain collapses to 1 Reshape (2 total), got {[n.name for n in reshape_nodes]}" + ) + # Verify ORT parity + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-6) + + +# ══════════════════════════════════════════════════════════════════════════════ +# Step 4 — Replace graph-output Reshape with Squeeze +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_graph_output_reshape_replaced_with_squeeze(): + """[1,3,4,1] → Reshape → [3,4] at graph output → replaced by Squeeze(axes=[0,3]); ORT parity.""" + sh_init = numpy_helper.from_array(np.array([3, 4], dtype=np.int64), name="rs_shape") + reshape = helper.make_node("Reshape", ["X", "rs_shape"], ["Y"], name="rs") + graph = helper.make_graph( + [reshape], + "g", + [_vi("X", TensorProto.FLOAT, [1, 3, 4, 1])], + [_vi("Y", TensorProto.FLOAT, [3, 4])], + initializer=[sh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal([1, 3, 4, 1]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _reshape_to_squeeze(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not any(n.op_type == "Reshape" for n in tr.graph.node), "Reshape must be replaced" + sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] + assert len(sq_nodes) == 1, "One Squeeze must be inserted" + + # Verify axes initializer contains [0, 3] + axes_init = next((i for i in tr.graph.initializer if "squeeze_axes" in i.name), None) + assert axes_init is not None, "Squeeze axes initializer must exist" + assert numpy_helper.to_array(axes_init).tolist() == [0, 3], ( + f"Expected axes [0,3], got {numpy_helper.to_array(axes_init).tolist()}" + ) + + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_graph_output_reshape_squeeze_axes_left_to_right(): + """[1,1,1,512] → Reshape → [1,512]: squeeze axes must be [0,1] (leftmost 1s), not [1,2]. + + Verifies that when multiple 1-dims could be squeezed, the algorithm always + removes from the lowest (leftmost) axis indices first. + """ + sh_init = numpy_helper.from_array(np.array([1, 512], dtype=np.int64), name="rs_shape_ltr") + reshape = helper.make_node("Reshape", ["X", "rs_shape_ltr"], ["Y"], name="rs_ltr") + graph = helper.make_graph( + [reshape], + "g", + [_vi("X", TensorProto.FLOAT, [1, 1, 1, 512])], + [_vi("Y", TensorProto.FLOAT, [1, 512])], + initializer=[sh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal([1, 1, 1, 512]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _reshape_to_squeeze(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not any(n.op_type == "Reshape" for n in tr.graph.node), "Reshape must be replaced" + sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] + assert len(sq_nodes) == 1, "One Squeeze must be inserted" + + axes_init = next((i for i in tr.graph.initializer if "squeeze_axes" in i.name), None) + assert axes_init is not None + axes = numpy_helper.to_array(axes_init).tolist() + assert axes == [0, 1], f"Expected axes [0,1] (leftmost 1s), got {axes}" + + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_graph_output_reshape_not_squeezable_padded_to_4d(): + """[2,3,4] → Reshape → [6,4] at graph output: not a pure squeeze, output is 2D (<4D). + + New behaviour: pad Reshape target to [1,6,4], then Squeeze([0]) → [6,4]. + This ensures the Reshape produces a 4D tensor before being squeezed back. + """ + sh_init = numpy_helper.from_array(np.array([6, 4], dtype=np.int64), name="rs_shape2") + reshape = helper.make_node("Reshape", ["X", "rs_shape2"], ["Y"], name="rs2") + graph = helper.make_graph( + [reshape], + "g", + [_vi("X", TensorProto.FLOAT, [2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [6, 4])], + initializer=[sh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal([2, 3, 4]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _reshape_to_squeeze(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + # Reshape must still exist (now targeting [1,6,4]) + rs_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(rs_nodes) == 1, "Reshape must remain (with updated shape)" + # Squeeze([0]) must be inserted after Reshape + sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] + assert len(sq_nodes) == 1, "Squeeze([0]) must be inserted" + # Verify Reshape now targets a 4D shape [1,6,4] + rs_shape_init = next((i for i in tr.graph.initializer if i.name == rs_nodes[0].input[1]), None) + assert rs_shape_init is not None + assert numpy_helper.to_array(rs_shape_init).tolist() == [1, 6, 4], ( + f"Reshape shape must be [1,6,4], got {numpy_helper.to_array(rs_shape_init).tolist()}" + ) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_graph_output_reshape_4d_output_kept(): + """[5,6,7,8] → Reshape → [30,14,4] at graph output: not a pure squeeze, BUT padded to [1,30,14,4] + Squeeze([0]).""" + sh_init = numpy_helper.from_array(np.array([30, 14, 4], dtype=np.int64), name="rs_shape_3d") + reshape = helper.make_node("Reshape", ["X", "rs_shape_3d"], ["Y"], name="rs_3d") + graph = helper.make_graph( + [reshape], + "g", + [_vi("X", TensorProto.FLOAT, [5, 6, 7, 8])], + [_vi("Y", TensorProto.FLOAT, [30, 14, 4])], + initializer=[sh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal([5, 6, 7, 8]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _reshape_to_squeeze(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + # Reshape must target [1,30,14,4] + rs_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + assert len(rs_nodes) == 1 + rs_shape_init = next((i for i in tr.graph.initializer if i.name == rs_nodes[0].input[1]), None) + assert rs_shape_init is not None + assert numpy_helper.to_array(rs_shape_init).tolist() == [1, 30, 14, 4], ( + f"Reshape shape must be [1,30,14,4], got {numpy_helper.to_array(rs_shape_init).tolist()}" + ) + # Squeeze([0]) strips the leading 1 back to [30,14,4] + sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] + assert len(sq_nodes) == 1 + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +def test_graph_output_reshape_not_at_graph_output_unchanged(): + """Reshape that is NOT a graph output must not be touched by this step.""" + sh_init = numpy_helper.from_array(np.array([3, 4], dtype=np.int64), name="rs_shape3") + reshape = helper.make_node("Reshape", ["X", "rs_shape3"], ["T"], name="rs3") + relu = helper.make_node("Relu", ["T"], ["Y"], name="relu") + graph = helper.make_graph( + [reshape, relu], + "g", + [_vi("X", TensorProto.FLOAT, [1, 3, 4, 1])], + [_vi("Y", TensorProto.FLOAT, [3, 4])], + initializer=[sh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _reshape_to_squeeze(tr) + + # Reshape feeds Relu (not a graph output) → must stay as Reshape + assert any(n.op_type == "Reshape" for n in tr.graph.node), ( + "Reshape not at graph output must be unchanged" + ) + assert not any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "No Squeeze must be inserted for non-output Reshape" + ) + + +# ══════════════════════════════════════════════════════════════════════════════ +# GroupQueryAttention — Unsqueeze after GQA output must be preserved +# ══════════════════════════════════════════════════════════════════════════════ + + +def test_unsqueeze_after_gqa_3d_to_4d_skipped(): + """Unsqueeze whose input is a GQA 3D output and whose output is 4D must not be replaced. + + Pattern: X [b,s,h] → GQA → gqa_out [b,s,h] (3D) + gqa_out → Unsqueeze(axes=[0]) → gqa_4d [1,b,s,h] (4D) ← must be preserved + gqa_4d → Relu → Y [1,b,s,h] + + The Unsqueeze is the canonical 3D→4D bridge after GQA and must not be replaced + by a Reshape. + """ + b, s, h = 2, 4, 8 + unsq_axes = numpy_helper.from_array(np.array([0], dtype=np.int64), name="gqa_unsq_ax") + + # Stub GQA node with 3D input/output (com.microsoft domain, no real attributes needed) + gqa = helper.make_node( + "GroupQueryAttention", + inputs=["X", "", ""], + outputs=["gqa_out"], + name="gqa_node", + domain="com.microsoft", + num_heads=2, + kv_num_heads=2, + ) + unsqueeze = helper.make_node( + "Unsqueeze", ["gqa_out", "gqa_unsq_ax"], ["gqa_4d"], name="gqa_unsq" + ) + relu = helper.make_node("Relu", ["gqa_4d"], ["Y"], name="relu") + + graph = helper.make_graph( + [gqa, unsqueeze, relu], + "g", + [_vi("X", TensorProto.FLOAT, [b, s, h])], + [_vi("Y", TensorProto.FLOAT, [1, b, s, h])], + value_info=[ + _vi("gqa_out", TensorProto.FLOAT, [b, s, h]), # 3D GQA output + _vi("gqa_4d", TensorProto.FLOAT, [1, b, s, h]), # 4D Unsqueeze output + ], + initializer=[unsq_axes], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", _OPSET_VERSION), + helper.make_opsetid("com.microsoft", 1), + ], + ) + _clamp_ir(model) + orig_node_count = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + # Unsqueeze after GQA (3D→4D) must be preserved unchanged + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + "Unsqueeze after GQA 3D output must be preserved" + ) + assert not any(n.op_type == "Reshape" for n in tr.graph.node), ( + "No Reshape must replace the GQA→Unsqueeze bridge" + ) + assert len(tr.graph.node) == orig_node_count, ( + "Node count must be unchanged — no nodes added or removed" + ) + + +def test_unsqueeze_after_gqa_4d_output_not_skipped(): + """Unsqueeze after a GQA node whose output is already 4D is NOT special — treated normally. + + If the GQA output is 4D (not 3D), the skip condition does not apply. + """ + b, s, h = 1, 4, 8 + unsq_axes = numpy_helper.from_array(np.array([1], dtype=np.int64), name="gqa4d_unsq_ax") + + gqa = helper.make_node( + "GroupQueryAttention", + inputs=["X", "", ""], + outputs=["gqa_out"], + name="gqa_node2", + domain="com.microsoft", + num_heads=2, + kv_num_heads=2, + ) + unsqueeze = helper.make_node( + "Unsqueeze", ["gqa_out", "gqa4d_unsq_ax"], ["gqa_5d"], name="gqa_unsq2" + ) + relu = helper.make_node("Relu", ["gqa_5d"], ["Y"], name="relu") + + graph = helper.make_graph( + [gqa, unsqueeze, relu], + "g", + [_vi("X", TensorProto.FLOAT, [b, s, h])], + [_vi("Y", TensorProto.FLOAT, [b, 1, s, h])], + value_info=[ + # GQA output is 4D here (not 3D) — skip condition does NOT apply + _vi("gqa_out", TensorProto.FLOAT, [b, 1, s, h]), + _vi("gqa_5d", TensorProto.FLOAT, [b, 1, 1, s, h]), + ], + initializer=[unsq_axes], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", _OPSET_VERSION), + helper.make_opsetid("com.microsoft", 1), + ], + ) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _replace_sq_unsq(tr) + + # GQA output is 4D → Unsqueeze is NOT skipped by the GQA rule; + # it should be replaced (output has 5D — beyond 4D, so the generic handler skips too, + # but the important thing is the GQA-specific skip did NOT fire). + # Verify no GQA-specific skip was applied by confirming the Unsqueeze was processed + # (either replaced or removed, not left as-is due to the GQA rule). + # Here output is 5D so the transform skips it anyway — just assert no crash. + + +# ── Tests: _remove_consecutive_casts ───────────────────────────────────────── + + +def test_cast_chain_collapsed_to_single(): + """Cast(INT8→INT64) → Cast(INT64→INT32) collapsed to Cast(INT8→INT32).""" + cast1 = helper.make_node("Cast", ["X"], ["c1"], name="cast1", to=TensorProto.INT64) + cast2 = helper.make_node("Cast", ["c1"], ["Y"], name="cast2", to=TensorProto.INT32) + graph = helper.make_graph( + [cast1, cast2], + "g", + [_vi("X", TensorProto.INT8, [2, 4])], + [_vi("Y", TensorProto.INT32, [2, 4])], + value_info=[_vi("c1", TensorProto.INT64, [2, 4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_cast_chains(tr) + + cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] + assert len(cast_nodes) == 1, f"Expected 1 Cast, got {len(cast_nodes)}" + # Final Cast should target INT32 + to_val = next(a.i for a in cast_nodes[0].attribute if a.name == "to") + assert to_val == TensorProto.INT32 + # Input should be X, output should be Y + assert cast_nodes[0].input[0] == "X" + assert cast_nodes[0].output[0] == "Y" + + +def test_cast_chain_three_nodes(): + """Three consecutive Casts collapsed to one.""" + cast1 = helper.make_node("Cast", ["X"], ["c1"], name="cast1", to=TensorProto.INT64) + cast2 = helper.make_node("Cast", ["c1"], ["c2"], name="cast2", to=TensorProto.FLOAT) + cast3 = helper.make_node("Cast", ["c2"], ["Y"], name="cast3", to=TensorProto.INT32) + graph = helper.make_graph( + [cast1, cast2, cast3], + "g", + [_vi("X", TensorProto.INT8, [3])], + [_vi("Y", TensorProto.INT32, [3])], + value_info=[ + _vi("c1", TensorProto.INT64, [3]), + _vi("c2", TensorProto.FLOAT, [3]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_cast_chains(tr) + + cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] + assert len(cast_nodes) == 1 + to_val = next(a.i for a in cast_nodes[0].attribute if a.name == "to") + assert to_val == TensorProto.INT32 + + +def test_single_cast_unchanged(): + """A lone Cast node is not modified.""" + cast1 = helper.make_node("Cast", ["X"], ["Y"], name="cast1", to=TensorProto.INT32) + graph = helper.make_graph( + [cast1], + "g", + [_vi("X", TensorProto.FLOAT, [4])], + [_vi("Y", TensorProto.INT32, [4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_cast_chains(tr) + + cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] + assert len(cast_nodes) == 1 + assert cast_nodes[0].input[0] == "X" + assert cast_nodes[0].output[0] == "Y" + + +def test_cast_chain_broken_by_multi_consumer(): + """Cast chain is broken when intermediate tensor has multiple consumers.""" + cast1 = helper.make_node("Cast", ["X"], ["c1"], name="cast1", to=TensorProto.INT64) + cast2 = helper.make_node("Cast", ["c1"], ["Y"], name="cast2", to=TensorProto.INT32) + relu = helper.make_node("Relu", ["c1"], ["Z"], name="relu") + graph = helper.make_graph( + [cast1, cast2, relu], + "g", + [_vi("X", TensorProto.FLOAT, [2, 4])], + [_vi("Y", TensorProto.INT32, [2, 4]), _vi("Z", TensorProto.INT64, [2, 4])], + value_info=[_vi("c1", TensorProto.INT64, [2, 4])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _remove_cast_chains(tr) + + # Both Casts should remain (chain broken by multi-consumer c1) + cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] + assert len(cast_nodes) == 2 From 92ae732349c7928020d42779381be52b0bd25ca7 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 11:41:19 +0000 Subject: [PATCH 2/9] Preserve rank-sensitive paths in DLA graph cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 242 ++++++++---- .../dla_transforms/test_dla_graph_cleanup.py | 362 +++++++++++++----- 2 files changed, 439 insertions(+), 165 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index e619d6eeb4d..18a04ca8400 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -13,32 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Combined graph-cleanup pass for DLA ONNX models. +"""Graph cleanup for DLA ONNX models. -Applies five cleanup steps in sequence: - -1. **Canonicalize non-4D model inputs** — Each 1-D/2-D/3-D graph input gets a single - canonical ``Unsqueeze`` inserted at the front of the graph, padding to 4-D. - Existing ``Unsqueeze`` consumers of that input are either deduped (same output shape) - or replaced by a ``Reshape`` (different output shape). All other consumers are - redirected to the canonical Unsqueeze output. - -2. **Replace intermediary Squeeze / Unsqueeze with 4-D Reshape** — Any ``Squeeze`` or - ``Unsqueeze`` whose output is **not** a graph output is replaced by a ``Reshape`` whose - output shape is the original output shape padded to 4-D by prepending 1-s. - If the padded input and output shapes are identical the node is removed entirely - and consumers are wired to the input directly. - -3. **Collapse Reshape chains** — Consecutive ``Reshape`` nodes whose output is consumed - only by the next ``Reshape`` are collapsed. A chain where first-input-shape equals - last-output-shape is removed entirely; otherwise it is reduced to a single ``Reshape``. - -4. **Fold constants** — Uses ``onnx-graphsurgeon`` to evaluate constant subgraphs at - compile time, skipping ``DequantizeLinear`` and ``Identity`` nodes. Silently skipped - if ``onnx-graphsurgeon`` is unavailable. - -5. **Remove unused initializers** — Initializers that are not referenced as a node input - are dropped. +Promote eligible input and Squeeze/Unsqueeze paths to 4D while retaining original +ranks at rank-sensitive consumers. Collapse reshape chains, adapt graph-output +reshapes, fold constants, collapse lossless cast chains, and remove unused +initializers. Constant folding requires the optional onnx-graphsurgeon package. """ from __future__ import annotations @@ -63,16 +43,108 @@ # ── Internal utilities ──────────────────────────────────────────────────────── +_RANK_AGNOSTIC_OPS = frozenset( + { + "Abs", + "Acos", + "Acosh", + "Add", + "And", + "Asin", + "Asinh", + "Atan", + "Atanh", + "Cast", + "Ceil", + "Clip", + "Cos", + "Cosh", + "Div", + "Elu", + "Equal", + "Erf", + "Exp", + "Floor", + "Greater", + "GreaterOrEqual", + "HardSigmoid", + "HardSwish", + "Identity", + "IsInf", + "IsNaN", + "LeakyRelu", + "Less", + "LessOrEqual", + "Log", + "Max", + "Mean", + "Min", + "Mod", + "Mul", + "Neg", + "Not", + "Or", + "PRelu", + "Pow", + "Reciprocal", + "Relu", + "Round", + "Selu", + "Sigmoid", + "Sign", + "Sin", + "Sinh", + "Softplus", + "Softsign", + "Sqrt", + "Sub", + "Sum", + "Tan", + "Tanh", + "ThresholdedRelu", + "Where", + "Xor", + } +) + + +def _rank_sensitive_consumers( + graph: onnx.GraphProto, cache: GraphCache +) -> tuple[set[str], set[tuple[int, int]]]: + """Find edges where leading unit dimensions can change downstream semantics. + + Propagate restrictions through elementwise paths until an already-4D output + prevents an input promotion from changing downstream ranks. + """ + sensitive_values: set[str] = set() + sensitive_edges: set[tuple[int, int]] = set() + for node in reversed(graph.node): + rank_agnostic = node.domain in ("", "ai.onnx") and node.op_type in _RANK_AGNOSTIC_OPS + if rank_agnostic: + rank_agnostic = not any( + output in sensitive_values and len(cache.get_shape(output) or []) < 4 + for output in node.output + ) + if not rank_agnostic: + sensitive_values.update(_node_inputs(node)) + for index, name in enumerate(node.input): + if name and (not rank_agnostic or (node.op_type == "Clip" and index > 0)): + sensitive_values.add(name) + sensitive_edges.add((id(node), index)) + return sensitive_values, sensitive_edges + + # ── Step 1: Canonicalize non-4D model inputs ───────────────────────────────── def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: - """Insert one canonical Unsqueeze per non-4D graph input; dedup existing Unsqueeze consumers.""" + """Promote eligible input consumers to 4D without changing rank-sensitive paths.""" graph = model.graph unique_init: set[str] = set() consumer_map = cache.consumer_map nodes_to_insert: list[tuple] = [] # (new_node, reference_node) nodes_to_remove: list = [] + promotion_safety = None cnt = 0 for graph_input in list(graph.input): @@ -88,6 +160,10 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode ) continue + if promotion_safety is None: + promotion_safety = _rank_sensitive_consumers(graph, cache) + sensitive_values, sensitive_edges = promotion_safety + # Read concrete input shape (may be None for symbolic dims) input_shape: list[int | None] = [ d.dim_value if d.HasField("dim_value") else None for d in tt.shape.dim @@ -97,8 +173,29 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode axes_arr = np.arange(n_new, dtype=np.int64) canonical_4d = [1] * n_new + list(input_shape) - axes_name = f"{graph_input.name}_unsqueeze_axes" - out_name = f"{graph_input.name}_4d" + consumers = [] + for c_node, c_idx in consumer_map.get(graph_input.name, []): + if (c_node.op_type == "Squeeze" and "Gather" in c_node.name) or c_node.op_type in ( + "GroupQueryAttention", + "Gather", + ): + continue + if c_node.op_type == "Unsqueeze": + old_shape = cache.get_shape(c_node.output[0]) + if old_shape is None or ( + len(old_shape) < 4 + and c_node.output[0] not in cache.graph_output_names + and c_node.output[0] in sensitive_values + ): + continue + elif (id(c_node), c_idx) in sensitive_edges: + continue + consumers.append((c_node, c_idx)) + if not consumers: + continue + + axes_name = cache.unique_name(f"{graph_input.name}_unsqueeze_axes") + out_name = cache.unique_name(f"{graph_input.name}_4d") add_unique_initializers( graph, unique_init, [numpy_helper.from_array(axes_arr, name=axes_name)] ) @@ -107,24 +204,10 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode "Unsqueeze", inputs=[graph_input.name, axes_name], outputs=[out_name], - name=f"{graph_input.name}_unsqueeze", + name=cache.unique_name(f"{graph_input.name}_unsqueeze"), ) add_value_info(graph, out_name, tt.elem_type, canonical_4d) - consumers = consumer_map.get(graph_input.name, []) - consumers = [ - (c_node, c_idx) - for c_node, c_idx in consumers - if len(c_node.input) > 0 - and graph_input.name in c_node.input - and not ( - (c_node.op_type == "Squeeze" and "Gather" in c_node.name) - or c_node.op_type in ("GroupQueryAttention", "Gather") - ) - ] - if not consumers: - continue - # Find the first consumer node in graph order for insertion point consumer_set = {id(n) for n, _ in consumers} ref_node = next((n for n in graph.node if id(n) in consumer_set), None) @@ -158,13 +241,17 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode dn.input[di] = out_name else: # Different 4-D shape → Reshape from canonical output - rs_name = f"{c_node.name}_reshape_shape" + is_graph_output = c_node.output[0] in cache.graph_output_names + reshape_target = old_out_shape if is_graph_output else old_4d + rs_name = cache.unique_name( + f"{c_node.name or c_node.op_type}_reshape_shape" + ) add_unique_initializers( graph, unique_init, [ numpy_helper.from_array( - np.array(old_4d, dtype=np.int64), name=rs_name + np.array(reshape_target, dtype=np.int64), name=rs_name ) ], ) @@ -172,9 +259,10 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode "Reshape", inputs=[out_name, rs_name], outputs=list(c_node.output), - name=f"{c_node.name}_reshape", + name=cache.unique_name(f"{c_node.name or c_node.op_type}_reshape"), ) - cache.update_shape(c_node.output[0], old_4d) + if not is_graph_output: + cache.update_shape(c_node.output[0], old_4d) nodes_to_insert.append((rs, c_node)) else: # Unknown shape — just update the input to canonical unsqueeze output @@ -227,6 +315,9 @@ def _compute_sq_unsq_out_shape( break break + if node.op_type == "Squeeze" and axes is None: + axes = next((list(attr.ints) for attr in node.attribute if attr.name == "axes"), None) + if node.op_type == "Squeeze": if axes is None: return [d for d in in_shape if d != 1] @@ -258,8 +349,8 @@ def _should_skip_squeeze_unsqueeze( * ``Unsqueeze`` whose data input is a graph input — canonical 4D-promotion node. * Output feeds ``Gather`` / ``GatherElements`` ``input[1]`` (indices must stay INT32). * Output feeds any ``GroupQueryAttention`` input (GQA manages its own shapes). - * ``Unsqueeze`` whose data input is a GQA output AND the GQA input was 3D while - the GQA output is 4D — this Unsqueeze was placed to adapt the 3D→4D boundary + * ``Unsqueeze`` whose data input is a 3D GQA output and whose own output is 4D + — this Unsqueeze was placed to adapt the 3D→4D boundary and must be preserved. The check looks one level through an intervening ``Cast`` node since @@ -309,10 +400,12 @@ def _replace_intermediary_squeeze_unsqueeze( * Output feeds Gather/GatherElements ``input[1]`` (indices must remain INT32). * Output feeds GroupQueryAttention (GQA manages its own dimensionality). * Output of GroupQueryAttention feeds to Unsqueeze node with 3D input. + * Promoting the output would change a downstream rank-sensitive operation. """ graph = model.graph unique_init: set[str] = set() graph_output_names = cache.graph_output_names + sensitive_values: set[str] | None = None nodes_to_remove: list = [] nodes_to_insert: list[tuple] = [] @@ -338,11 +431,16 @@ def _replace_intermediary_squeeze_unsqueeze( continue is_graph_output = node.output[0] in graph_output_names + if not is_graph_output and len(out_shape) < 4: + if sensitive_values is None: + sensitive_values, _ = _rank_sensitive_consumers(graph, cache) + if node.output[0] in sensitive_values: + continue # Graph output: use declared shape (preserves name binding); else pad to 4D. reshape_target = out_shape if is_graph_output else pad4d(out_shape) - rs_name = f"{node.name}_4d_shape" + rs_name = cache.unique_name(f"{node.name or node.op_type}_4d_shape") add_unique_initializers( graph, unique_init, @@ -352,7 +450,7 @@ def _replace_intermediary_squeeze_unsqueeze( "Reshape", inputs=[node.input[0], rs_name], outputs=list(node.output), - name=f"{node.name}_to_reshape", + name=cache.unique_name(f"{node.name or node.op_type}_to_reshape"), ) nodes_to_insert.append((rs, node)) nodes_to_remove.append(node) @@ -410,7 +508,7 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo "Identity", [inp], [out_tensor], - name=f"cleanup_noop_id_{out_tensor}", + name=cache.unique_name(f"cleanup_noop_id_{out_tensor}"), ) pass1_batch[node_idx[id(node)]] = ([node], [id_node]) else: @@ -472,7 +570,7 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo "Identity", [inp], [out_tensor], - name=f"cleanup_chain_id_{out_tensor}", + name=cache.unique_name(f"cleanup_chain_id_{out_tensor}"), ) pass2_batch[node_idx[id(ch[0])]] = (ch[:], [id_node]) else: @@ -480,14 +578,14 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo else: rs_shape = numpy_helper.from_array( np.array(out_shape, dtype=np.int64), - name=ch[-1].name + "_chain_shape", + name=cache.unique_name(f"{ch[-1].name or 'reshape'}_chain_shape"), ) graph.initializer.append(rs_shape) rs = helper.make_node( "Reshape", inputs=[inp, rs_shape.name], outputs=[out_tensor], - name=ch[-1].name + "_chain", + name=cache.unique_name(f"{ch[-1].name or 'reshape'}_chain"), ) pass2_batch[node_idx[id(ch[0])]] = (ch[:], [rs]) chain_count += 1 @@ -521,6 +619,7 @@ def _fold_constants(model: onnx.ModelProto) -> onnx.ModelProto: return model def _exclude(node) -> bool: + """Keep quantization boundaries and identity bindings out of constant folding.""" return node.op in ("DequantizeLinear", "Identity") prev_ir = onnx.IR_VERSION @@ -623,8 +722,7 @@ def _replace_graph_output_reshape_with_squeeze( if not matched or j != -1: # Cannot express as a pure Squeeze of the existing input. - # If the output is non-4D, pad the Reshape target to 4D with a - # leading 1-dim and add a Squeeze(axis=[0]) to strip it back. + # Pad the Reshape target to 4D and restore the declared output rank. # # Example: # [5,6,7,8] → Reshape([30,14,4]) → graph_output[30,14,4] @@ -632,38 +730,42 @@ def _replace_graph_output_reshape_with_squeeze( # [5,6,7,8] → Reshape([1,30,14,4]) → Squeeze([0]) → graph_output[30,14,4] if len(out_shape) >= 4: continue # already 4D or higher — leave as-is - # Build a 4D shape by prepending 1 - shape_4d = [1, *out_shape] - rs_shape_name = f"{pfx}_4d_shape" + shape_4d = pad4d(out_shape) + squeeze_axes = list(range(4 - len(out_shape))) + rs_shape_name = cache.unique_name(f"{pfx}_4d_shape") add_unique_initializers( graph, unique_init, [numpy_helper.from_array(np.array(shape_4d, dtype=np.int64), name=rs_shape_name)], ) # Intermediate 4D tensor name - rs_4d_out = f"{pfx}_4d_out" + rs_4d_out = cache.unique_name(f"{pfx}_4d_out") # Update the Reshape node to target the 4D shape node.input[1] = rs_shape_name node.output[0] = rs_4d_out - # Squeeze to remove the leading 1-dim - sq_axes_name = f"{pfx}_sq0_axes" + sq_axes_name = cache.unique_name(f"{pfx}_sq0_axes") add_unique_initializers( graph, unique_init, - [numpy_helper.from_array(np.array([0], dtype=np.int64), name=sq_axes_name)], + [ + numpy_helper.from_array( + np.array(squeeze_axes, dtype=np.int64), name=sq_axes_name + ) + ], ) sq = helper.make_node( "Squeeze", inputs=[rs_4d_out, sq_axes_name], outputs=[out_name], - name=f"{pfx}_squeeze_leading", + name=cache.unique_name(f"{pfx}_squeeze_leading"), ) go_batch[go_node_idx[id(node)]] = ([node], [node, sq]) cnt += 1 logger.debug( - "Reshape %r → Reshape(%s) + Squeeze([0]) for non-4D graph output %r.", + "Reshape %r → Reshape(%s) + Squeeze(%s) for non-4D graph output %r.", node.name, shape_4d, + squeeze_axes, out_name, ) continue @@ -671,7 +773,7 @@ def _replace_graph_output_reshape_with_squeeze( squeeze_axes.sort() # ascending axis order required by ONNX Squeeze # Build Squeeze node - axes_name = f"{pfx}_squeeze_axes" + axes_name = cache.unique_name(f"{pfx}_squeeze_axes") add_unique_initializers( graph, unique_init, @@ -681,7 +783,7 @@ def _replace_graph_output_reshape_with_squeeze( "Squeeze", inputs=[in_name, axes_name], outputs=[out_name], - name=f"{pfx}_squeeze", + name=cache.unique_name(f"{pfx}_squeeze"), ) go_batch[go_node_idx[id(node)]] = ([node], [sq]) cnt += 1 @@ -787,7 +889,7 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx inputs=[first_input], outputs=[last_output], to=last_to, - name=f"{ch[-1].name}_chain", + name=cache.unique_name(f"{ch[-1].name or 'cast'}_chain"), ) idx = node_idx[id(ch[0])] replacements[idx] = (ch[:], [new_cast]) @@ -806,7 +908,7 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx def _apply_graph_cleanup(model: onnx.ModelProto) -> onnx.ModelProto: - """Apply all five cleanup steps in sequence.""" + """Apply the seven DLA cleanup stages while retaining rank-sensitive paths.""" model = infer_shapes(model) cache = GraphCache(model.graph) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index 6383799c015..ee15e261358 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -44,9 +44,13 @@ from onnx import TensorProto, helper, numpy_helper from modelopt.onnx.graph_surgery.dla_transforms import dla_graph_cleanup as _mod +from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") +import onnxruntime as ort +from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + _REPO_ROOT = Path(__file__).resolve().parents[5] _SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "graph_cleanup_test_models" _OPSET_VERSION = 17 @@ -57,26 +61,29 @@ _fold_consts = _mod._fold_constants _remove_unused_inits = _mod._remove_unused_initializers -from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache - def _canonicalize_inputs(model): + """Canonicalize model inputs using a fresh graph cache.""" return _mod._canonicalize_inputs(model, GraphCache(model.graph)) def _replace_sq_unsq(model): + """Replace eligible Squeeze and Unsqueeze nodes using current shapes.""" return _mod._replace_intermediary_squeeze_unsqueeze(model, GraphCache(model.graph)) def _remove_reshapes(model): + """Collapse reshape chains using a fresh graph cache.""" return _mod._remove_reshape_chains(model, GraphCache(model.graph)) def _reshape_to_squeeze(model): + """Adapt graph-output reshapes without changing their declared outputs.""" return _mod._replace_graph_output_reshape_with_squeeze(model, GraphCache(model.graph)) def _remove_cast_chains(model): + """Collapse eligible cast chains using current graph metadata.""" return _mod._remove_consecutive_casts(model, GraphCache(model.graph)) @@ -84,30 +91,25 @@ def _remove_cast_chains(model): def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Build a typed graph value with the declared shape.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir(model: onnx.ModelProto) -> None: + """Keep fixtures within ORT's supported IR range.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _prepare_model(model: onnx.ModelProto) -> None: - try: - from onnxruntime.tools.symbolic_shape_infer import ( - SymbolicShapeInference, # type: ignore[import-untyped] - ) - - inferred = SymbolicShapeInference.infer_shapes(model) - if inferred is not None: - del model.graph.value_info[:] - model.graph.value_info.extend(inferred.graph.value_info) - except Exception: - pass + """Populate intermediate shapes and surface unexpected inference failures.""" + inferred = SymbolicShapeInference.infer_shapes(model) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-untyped] - + """Run all graph outputs with ORT's CPU execution provider.""" _clamp_ir(model) sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} @@ -122,23 +124,23 @@ def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: def test_canonicalize_3d_input_unsqueeze_added(): - """3-D graph input X [2,3,4] → Unsqueeze to 4-D [1,2,3,4] is inserted. - - Downstream Relu receives the 4-D tensor. - """ + """Promote the safe branch of a shared input while preserving a Transpose branch.""" relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + transpose = helper.make_node("Transpose", ["X"], ["Z"], perm=[2, 1, 0]) graph = helper.make_graph( - [relu], + [relu, transpose], "g", [_vi("X", TensorProto.FLOAT, [2, 3, 4])], - [_vi("Y", TensorProto.FLOAT, [2, 3, 4])], + [_vi("Y", TensorProto.FLOAT, [2, 3, 4]), _vi("Z", TensorProto.FLOAT, [4, 3, 2])], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + feeds = {"X": _RNG.standard_normal([2, 3, 4]).astype(np.float32)} + expected = _run_ort(model, feeds) tr = copy.deepcopy(model) - _prepare_model(tr) - tr = _canonicalize_inputs(tr) + tr = _apply(tr) + onnx.checker.check_model(tr, full_check=True) unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] assert len(unsq_nodes) == 1, "one Unsqueeze must be inserted for the 3-D input" @@ -148,19 +150,29 @@ def test_canonicalize_3d_input_unsqueeze_added(): # Relu must now read the Unsqueeze output, not X directly relu_node = next(n for n in tr.graph.node if n.op_type == "Relu") assert relu_node.input[0] == unsq.output[0], "Relu must be redirected to Unsqueeze output" + transpose_node = next(n for n in tr.graph.node if n.op_type == "Transpose") + assert transpose_node.input[0] == "X" + actual = _run_ort(tr, feeds) + np.testing.assert_array_equal(actual["Y"], expected["Y"].reshape(1, 2, 3, 4)) + np.testing.assert_array_equal(actual["Z"], expected["Z"]) def test_canonicalize_2d_input_unsqueeze_added(): - """2-D graph input X [4,5] → Unsqueeze at axes [0,1] to 4-D [1,1,4,5].""" - relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + """Promote a 2D input when broadcasting already fixes downstream rank at four.""" + add = helper.make_node("Add", ["X", "W"], ["sum"], name="add") + transpose = helper.make_node("Transpose", ["sum"], ["Y"], perm=[0, 1, 3, 2]) graph = helper.make_graph( - [relu], + [add, transpose], "g", [_vi("X", TensorProto.FLOAT, [4, 5])], - [_vi("Y", TensorProto.FLOAT, [4, 5])], + [_vi("Y", TensorProto.FLOAT, [1, 1, 5, 4])], + [numpy_helper.from_array(np.zeros([1, 1, 4, 5], np.float32), "W")], + value_info=[_vi("sum", TensorProto.FLOAT, [1, 1, 4, 5])], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + feeds = {"X": _RNG.standard_normal([4, 5]).astype(np.float32)} + expected = _run_ort(model, feeds) tr = copy.deepcopy(model) _prepare_model(tr) @@ -171,6 +183,10 @@ def test_canonicalize_2d_input_unsqueeze_added(): axes_init = next(i for i in tr.graph.initializer if i.name.endswith("_unsqueeze_axes")) axes = numpy_helper.to_array(axes_init).tolist() assert axes == [0, 1], f"axes should be [0,1] for 2-D input, got {axes}" + add_node = next(node for node in tr.graph.node if node.op_type == "Add") + assert add_node.input[0] == unsq_nodes[0].output[0] + onnx.checker.check_model(tr, full_check=True) + np.testing.assert_array_equal(_run_ort(tr, feeds)["Y"], expected["Y"]) def test_canonicalize_3d_input_duplicate_unsqueezes_deduped(): @@ -229,6 +245,189 @@ def test_canonicalize_4d_input_unchanged(): # ══════════════════════════════════════════════════════════════════════════════ +@pytest.mark.parametrize("stage", ["input", "squeeze", "unsqueeze"]) +@pytest.mark.parametrize("consumer", ["Transpose", "Softmax", "MatMul", "transpose_via_relu"]) +def test_rank_promotion_preserves_rank_sensitive_consumers(stage, consumer): + """Preserve rank-sensitive behavior, including consumers behind an elementwise node.""" + input_shape = [1, 2, 3] if stage == "squeeze" else [2, 3] + source_shape = [1, 2, 3] if stage == "unsqueeze" else [2, 3] + nodes, initializers, values = [], [], [] + source = "X" + if stage != "input": + initializers.append(numpy_helper.from_array(np.array([0], np.int64), "axes")) + if stage == "unsqueeze": + nodes.append(helper.make_node("Identity", ["X"], ["middle"])) + values.append(_vi("middle", TensorProto.FLOAT, input_shape)) + source = "middle" + nodes.append(helper.make_node(stage.title(), [source, "axes"], ["adapted"])) + source = "adapted" + values.append(_vi(source, TensorProto.FLOAT, source_shape)) + if consumer == "transpose_via_relu": + nodes.append(helper.make_node("Relu", [source], ["activated"])) + source = "activated" + values.append(_vi(source, TensorProto.FLOAT, source_shape)) + consumer = "Transpose" + if consumer == "Transpose": + nodes.append( + helper.make_node( + consumer, [source], ["Y"], perm=list(reversed(range(len(source_shape)))) + ) + ) + output_shape = list(reversed(source_shape)) + elif consumer == "Softmax": + nodes.append(helper.make_node(consumer, [source], ["Y"], axis=len(source_shape) - 2)) + output_shape = source_shape + else: + initializers.append( + numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "W") + ) + nodes.append(helper.make_node(consumer, [source, "W"], ["Y"])) + output_shape = [*source_shape[:-1], 4] + graph = helper.make_graph( + nodes, + "rank_sensitive", + [_vi("X", TensorProto.FLOAT, input_shape)], + [_vi("Y", TensorProto.FLOAT, output_shape)], + initializers, + value_info=values, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + feeds = {"X": np.arange(np.prod(input_shape), dtype=np.float32).reshape(input_shape)} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + transformed = copy.deepcopy(model) + transformed = ( + _canonicalize_inputs(transformed) if stage == "input" else _replace_sq_unsq(transformed) + ) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected["Y"]) + assert transformed.SerializeToString() == before + + +@pytest.mark.parametrize("stage", ["canonical", "squeeze"]) +@pytest.mark.parametrize("names", ["empty", "duplicate", "existing"]) +def test_generated_reshape_targets_do_not_collide(stage, names): + """Keep distinct reshape targets for unnamed nodes and existing tensor names.""" + prefix = "uq" if stage == "canonical" else "sq" + first_name = "" if names == "empty" else prefix + second_name = prefix + "_other" if names == "existing" else first_name + initializers = [] + if stage == "canonical": + initializers.extend( + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "axes1"), + numpy_helper.from_array(np.array([1, 3], np.int64), "axes2"), + ] + ) + nodes = [ + helper.make_node("Unsqueeze", ["X", "axes1"], ["Y1"], name=first_name), + helper.make_node("Unsqueeze", ["X", "axes2"], ["Y2"], name=second_name), + ] + inputs = [_vi("X", TensorProto.FLOAT, [2, 3])] + outputs = [ + _vi("Y1", TensorProto.FLOAT, [2, 3, 1, 1]), + _vi("Y2", TensorProto.FLOAT, [2, 1, 3, 1]), + ] + feeds = {"X": np.arange(6, dtype=np.float32).reshape(2, 3)} + suffix = "_reshape_shape" + else: + initializers.append(numpy_helper.from_array(np.array([0], np.int64), "axes")) + nodes = [ + helper.make_node("Squeeze", ["X1", "axes"], ["Y1"], name=first_name), + helper.make_node("Squeeze", ["X2", "axes"], ["Y2"], name=second_name), + ] + inputs = [_vi("X1", TensorProto.FLOAT, [1, 2, 3]), _vi("X2", TensorProto.FLOAT, [1, 3, 2])] + outputs = [_vi("Y1", TensorProto.FLOAT, [2, 3]), _vi("Y2", TensorProto.FLOAT, [3, 2])] + feeds = { + "X1": np.arange(6, dtype=np.float32).reshape(1, 2, 3), + "X2": np.arange(6, dtype=np.float32).reshape(1, 3, 2), + } + suffix = "_4d_shape" + if names == "existing": + initializers.append(numpy_helper.from_array(np.array([99], np.int64), prefix + suffix)) + model = helper.make_model( + helper.make_graph(nodes, "target_names", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + reference = model + if names == "duplicate": + # ORT rejects duplicate labels; normalize only the reference fixture. + reference = copy.deepcopy(model) + for index, node in enumerate(reference.graph.node): + node.name = f"reference_{index}" + expected = _run_ort(reference, feeds) + transformed = copy.deepcopy(model) + transformed = ( + _canonicalize_inputs(transformed) if stage == "canonical" else _replace_sq_unsq(transformed) + ) + onnx.checker.check_model(transformed, full_check=True) + actual = _run_ort(transformed, feeds) + for name, value in expected.items(): + np.testing.assert_array_equal(actual[name], value) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_canonicalize_preserves_graph_output_shape(axis): + """Preserve a non-4D graph-output Unsqueeze's declared shape and values.""" + target = [2, 3] + target.insert(axis, 1) + axes = numpy_helper.from_array(np.array([axis], np.int64), "axes") + model = helper.make_model( + helper.make_graph( + [helper.make_node("Unsqueeze", ["X", "axes"], ["Y"], name="uq")], + "output_boundary", + [_vi("X", TensorProto.FLOAT, [2, 3])], + [_vi("Y", TensorProto.FLOAT, target)], + [axes], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir(model) + feeds = {"X": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _run_ort(model, feeds) + transformed = _canonicalize_inputs(copy.deepcopy(model)) + assert transformed.graph.output[0].type == model.graph.output[0].type + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected["Y"]) + + +@pytest.mark.parametrize("opset", [12, 17]) +@pytest.mark.parametrize("axes", [[1], [-3]]) +def test_squeeze_fallback_uses_selected_axes(opset, axes): + """Infer only selected squeeze axes when intermediate shape metadata is absent.""" + initializers = [] + if opset < 13: + squeeze = helper.make_node("Squeeze", ["X"], ["middle"], axes=axes, name="sq") + else: + initializers.append(numpy_helper.from_array(np.array(axes, np.int64), "axes")) + squeeze = helper.make_node("Squeeze", ["X", "axes"], ["middle"], name="sq") + model = helper.make_model( + helper.make_graph( + [squeeze, helper.make_node("Identity", ["middle"], ["Y"])], + "axes_fallback", + [_vi("X", TensorProto.FLOAT, [2, 1, 1, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 1, 3])], + initializers, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + _clamp_ir(model) + feeds = {"X": np.arange(6, dtype=np.float32).reshape(2, 1, 1, 3)} + expected = _run_ort(model, feeds) + transformed = _replace_sq_unsq(copy.deepcopy(model)) + reshape = next(node for node in transformed.graph.node if node.op_type == "Reshape") + target = next(init for init in transformed.graph.initializer if init.name == reshape.input[1]) + assert numpy_helper.to_array(target).tolist() == [1, 2, 1, 3] + transformed = _mod.infer_shapes(transformed) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal( + _run_ort(transformed, feeds)["Y"], expected["Y"].reshape(1, 2, 1, 3) + ) + + def test_intermediary_squeeze_replaced_with_reshape(): """Squeeze [2,1,6,8] → [2,6,8] (axis=1) is replaced by Reshape → [1,2,6,8]. @@ -675,7 +874,8 @@ def test_combined_reshape_and_unused_init_cleaned(): np.testing.assert_array_equal(ref["Y"], got["Y"]) -def test_interleaved_reshape_chains_collapsed(): +@pytest.mark.parametrize("names", ["named", "empty", "duplicate", "existing"]) +def test_interleaved_reshape_chains_collapsed(names): """Two Reshape chains interleaved in node order are both collapsed correctly. Graph topology: @@ -713,6 +913,9 @@ def test_interleaved_reshape_chains_collapsed(): r4 = helper.make_node("Reshape", ["X2", "sh_c"], ["T4"], name="r4") r5 = helper.make_node("Reshape", ["T4", "sh_d"], ["T5"], name="r5") add = helper.make_node("Add", ["T2", "T5"], ["Y"], name="add1") + if names in ("empty", "duplicate"): + for node in (r1, r2, r4, r5): + node.name = "" if names == "empty" else "reshape" # Interleaved: r4 sits between r1 and r2 in the node list graph = helper.make_graph( @@ -734,14 +937,26 @@ def test_interleaved_reshape_chains_collapsed(): model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + if names == "existing": + model.graph.initializer.append( + numpy_helper.from_array(np.array([99], np.int64), "r2_chain_shape") + ) + x1 = _RNG.standard_normal([1, 2, 4, 4]).astype(np.float32) x2 = _RNG.standard_normal([1, 2, 4, 4]).astype(np.float32) _prepare_model(model) - ref = _run_ort(model, {"X1": x1, "X2": x2}) + onnx.checker.check_model(model, full_check=True) + reference = model + if names == "duplicate": + reference = copy.deepcopy(model) + for index, node in enumerate(reference.graph.node): + node.name = f"reference_{index}" + ref = _run_ort(reference, {"X1": x1, "X2": x2}) tr = copy.deepcopy(model) _prepare_model(tr) tr = _remove_reshapes(tr) + onnx.checker.check_model(tr, full_check=True) _prepare_model(tr) got = _run_ort(tr, {"X1": x1, "X2": x2}) @@ -837,19 +1052,20 @@ def test_graph_output_reshape_squeeze_axes_left_to_right(): np.testing.assert_array_equal(ref["Y"], got["Y"]) -def test_graph_output_reshape_not_squeezable_padded_to_4d(): +@pytest.mark.parametrize("output_shape", [[6, 4], [24]]) +def test_graph_output_reshape_not_squeezable_padded_to_4d(output_shape): """[2,3,4] → Reshape → [6,4] at graph output: not a pure squeeze, output is 2D (<4D). - New behaviour: pad Reshape target to [1,6,4], then Squeeze([0]) → [6,4]. + Pad the Reshape target to rank four and squeeze every prepended unit dimension. This ensures the Reshape produces a 4D tensor before being squeezed back. """ - sh_init = numpy_helper.from_array(np.array([6, 4], dtype=np.int64), name="rs_shape2") + sh_init = numpy_helper.from_array(np.array(output_shape, dtype=np.int64), name="rs_shape2") reshape = helper.make_node("Reshape", ["X", "rs_shape2"], ["Y"], name="rs2") graph = helper.make_graph( [reshape], "g", [_vi("X", TensorProto.FLOAT, [2, 3, 4])], - [_vi("Y", TensorProto.FLOAT, [6, 4])], + [_vi("Y", TensorProto.FLOAT, output_shape)], initializer=[sh_init], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) @@ -865,18 +1081,19 @@ def test_graph_output_reshape_not_squeezable_padded_to_4d(): _prepare_model(tr) got = _run_ort(tr, {"X": x}) - # Reshape must still exist (now targeting [1,6,4]) + # The boundary restores the original output while the intermediate stays 4D. rs_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] assert len(rs_nodes) == 1, "Reshape must remain (with updated shape)" - # Squeeze([0]) must be inserted after Reshape sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] - assert len(sq_nodes) == 1, "Squeeze([0]) must be inserted" - # Verify Reshape now targets a 4D shape [1,6,4] + assert len(sq_nodes) == 1 rs_shape_init = next((i for i in tr.graph.initializer if i.name == rs_nodes[0].input[1]), None) assert rs_shape_init is not None - assert numpy_helper.to_array(rs_shape_init).tolist() == [1, 6, 4], ( - f"Reshape shape must be [1,6,4], got {numpy_helper.to_array(rs_shape_init).tolist()}" + assert ( + numpy_helper.to_array(rs_shape_init).tolist() + == [1] * (4 - len(output_shape)) + output_shape ) + axes = next(i for i in tr.graph.initializer if i.name == sq_nodes[0].input[1]) + assert numpy_helper.to_array(axes).tolist() == list(range(4 - len(output_shape))) np.testing.assert_array_equal(ref["Y"], got["Y"]) @@ -1016,65 +1233,11 @@ def test_unsqueeze_after_gqa_3d_to_4d_skipped(): ) -def test_unsqueeze_after_gqa_4d_output_not_skipped(): - """Unsqueeze after a GQA node whose output is already 4D is NOT special — treated normally. - - If the GQA output is 4D (not 3D), the skip condition does not apply. - """ - b, s, h = 1, 4, 8 - unsq_axes = numpy_helper.from_array(np.array([1], dtype=np.int64), name="gqa4d_unsq_ax") - - gqa = helper.make_node( - "GroupQueryAttention", - inputs=["X", "", ""], - outputs=["gqa_out"], - name="gqa_node2", - domain="com.microsoft", - num_heads=2, - kv_num_heads=2, - ) - unsqueeze = helper.make_node( - "Unsqueeze", ["gqa_out", "gqa4d_unsq_ax"], ["gqa_5d"], name="gqa_unsq2" - ) - relu = helper.make_node("Relu", ["gqa_5d"], ["Y"], name="relu") - - graph = helper.make_graph( - [gqa, unsqueeze, relu], - "g", - [_vi("X", TensorProto.FLOAT, [b, s, h])], - [_vi("Y", TensorProto.FLOAT, [b, 1, s, h])], - value_info=[ - # GQA output is 4D here (not 3D) — skip condition does NOT apply - _vi("gqa_out", TensorProto.FLOAT, [b, 1, s, h]), - _vi("gqa_5d", TensorProto.FLOAT, [b, 1, 1, s, h]), - ], - initializer=[unsq_axes], - ) - model = helper.make_model( - graph, - opset_imports=[ - helper.make_opsetid("", _OPSET_VERSION), - helper.make_opsetid("com.microsoft", 1), - ], - ) - _clamp_ir(model) - - tr = copy.deepcopy(model) - _prepare_model(tr) - tr = _replace_sq_unsq(tr) - - # GQA output is 4D → Unsqueeze is NOT skipped by the GQA rule; - # it should be replaced (output has 5D — beyond 4D, so the generic handler skips too, - # but the important thing is the GQA-specific skip did NOT fire). - # Verify no GQA-specific skip was applied by confirming the Unsqueeze was processed - # (either replaced or removed, not left as-is due to the GQA rule). - # Here output is 5D so the transform skips it anyway — just assert no crash. - - # ── Tests: _remove_consecutive_casts ───────────────────────────────────────── -def test_cast_chain_collapsed_to_single(): +@pytest.mark.parametrize("colliding_name", [False, True]) +def test_cast_chain_collapsed_to_single(colliding_name): """Cast(INT8→INT64) → Cast(INT64→INT32) collapsed to Cast(INT8→INT32).""" cast1 = helper.make_node("Cast", ["X"], ["c1"], name="cast1", to=TensorProto.INT64) cast2 = helper.make_node("Cast", ["c1"], ["Y"], name="cast2", to=TensorProto.INT32) @@ -1087,6 +1250,11 @@ def test_cast_chain_collapsed_to_single(): ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + if colliding_name: + model.graph.node.append(helper.make_node("Identity", ["X"], ["Z"], name="cast2_chain")) + model.graph.output.append(_vi("Z", TensorProto.INT8, [2, 4])) + feeds = {"X": np.arange(8, dtype=np.int8).reshape(2, 4)} + expected = _run_ort(model, feeds) tr = copy.deepcopy(model) _prepare_model(tr) @@ -1100,6 +1268,10 @@ def test_cast_chain_collapsed_to_single(): # Input should be X, output should be Y assert cast_nodes[0].input[0] == "X" assert cast_nodes[0].output[0] == "Y" + onnx.checker.check_model(tr, full_check=True) + actual = _run_ort(tr, feeds) + for name, value in expected.items(): + np.testing.assert_array_equal(actual[name], value) def test_cast_chain_three_nodes(): From 408267d17c13942f41c46e76358624d33e0ef291 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 12:07:58 +0000 Subject: [PATCH 3/9] Preserve axes and Cast semantics in DLA cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 13 +- .../dla_transforms/test_dla_graph_cleanup.py | 146 +++++++++++++++--- 2 files changed, 130 insertions(+), 29 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index 18a04ca8400..b4d028e41d2 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -180,7 +180,7 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode "Gather", ): continue - if c_node.op_type == "Unsqueeze": + if c_node.op_type == "Unsqueeze" and c_idx == 0: old_shape = cache.get_shape(c_node.output[0]) if old_shape is None or ( len(old_shape) < 4 @@ -330,7 +330,8 @@ def _compute_sq_unsq_out_shape( if axes is None: return None out = list(in_shape) - for ax in sorted([a if a >= 0 else len(out) + a + 1 for a in axes]): + out_rank = len(in_shape) + len(axes) + for ax in sorted(a if a >= 0 else out_rank + a for a in axes): out.insert(ax, 1) return out @@ -825,7 +826,7 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: - """Collapse consecutive Cast nodes into a single Cast with the last node's target dtype. + """Collapse standard-domain Cast chains while preserving the final Cast's attributes. A chain ``Cast(to=A) → Cast(to=B) → Cast(to=C)`` is replaced by a single ``Cast(to=C)`` from the first Cast's input to the last Cast's output. @@ -836,7 +837,7 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx cnt = 0 # Build chains: tail_output → [cast_node, ...] - all_casts = [n for n in graph.node if n.op_type == "Cast"] + all_casts = [n for n in graph.node if n.op_type == "Cast" and not n.domain and not n.overload] tail_to_chain: dict[str, list] = {} for node in all_casts: @@ -882,15 +883,13 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx ): continue last_output = ch[-1].output[0] - last_to = get_node_attr_i(ch[-1], "to", 1) # ONNX Cast default: FLOAT - new_cast = helper.make_node( "Cast", inputs=[first_input], outputs=[last_output], - to=last_to, name=cache.unique_name(f"{ch[-1].name or 'cast'}_chain"), ) + new_cast.attribute.extend(ch[-1].attribute) idx = node_idx[id(ch[0])] replacements[idx] = (ch[:], [new_cast]) cnt += len(ch) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index ee15e261358..d6a936e6c42 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -42,6 +42,7 @@ import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_graph_cleanup as _mod from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache @@ -220,24 +221,33 @@ def test_canonicalize_3d_input_duplicate_unsqueezes_deduped(): ) -def test_canonicalize_4d_input_unchanged(): - """4-D graph input → step 1 is a noop.""" - relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") +@pytest.mark.parametrize("axes_input", [False, True]) +def test_canonicalize_4d_input_unchanged(axes_input): + """Keep 4D data inputs and optional Unsqueeze control inputs unchanged.""" + inputs = [_vi("X", TensorProto.FLOAT, [1, 2, 3, 4])] + output_shape = [1, 2, 3, 4] + node = helper.make_node("Relu", ["X"], ["Y"], name="relu") + feeds = {"X": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + if axes_input: + inputs.append(_vi("axes", TensorProto.INT64, [1])) + node = helper.make_node("Unsqueeze", ["X", "axes"], ["Y"], name="uq") + output_shape = [1, 1, 2, 3, 4] + feeds["axes"] = np.array([0], np.int64) graph = helper.make_graph( - [relu], + [node], "g", - [_vi("X", TensorProto.FLOAT, [1, 2, 3, 4])], - [_vi("Y", TensorProto.FLOAT, [1, 2, 3, 4])], + inputs, + [_vi("Y", TensorProto.FLOAT, output_shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) - orig_nodes = len(model.graph.node) - - tr = copy.deepcopy(model) - _prepare_model(tr) - tr = _canonicalize_inputs(tr) - - assert len(tr.graph.node) == orig_nodes, "4-D input: no nodes should be added" + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + before = model.SerializeToString() + tr = _canonicalize_inputs(copy.deepcopy(model)) + assert tr.SerializeToString() == before + onnx.checker.check_model(tr, full_check=True) + np.testing.assert_array_equal(_run_ort(tr, feeds)["Y"], expected["Y"]) # ══════════════════════════════════════════════════════════════════════════════ @@ -463,12 +473,18 @@ def test_intermediary_squeeze_replaced_with_reshape(): assert shape == [1, 2, 6, 8], f"expected [1,2,6,8], got {shape}" -def test_intermediary_unsqueeze_replaced_with_reshape(): - """Unsqueeze [2,4,6] → [2,4,1,6] (axis=2) is replaced by Reshape → [2,4,1,6]. - - pad([2,4,6]) = [1,2,4,6], pad([2,4,1,6]) = [2,4,1,6] — shapes differ → Reshape inserted. - """ - unsq_axes = numpy_helper.from_array(np.array([2], dtype=np.int64), name="unsq_ax") +@pytest.mark.parametrize( + ("axes", "output_shape"), + [ + ([2], [2, 4, 1, 6]), + ([-1, -3], [2, 4, 1, 6, 1]), + ([2, -1], [2, 4, 1, 6, 1]), + ([-5, -3], [1, 2, 1, 4, 6]), + ], +) +def test_intermediary_unsqueeze_replaced_with_reshape(axes, output_shape): + """Preserve multi-axis Unsqueeze layout when output shape metadata is absent.""" + unsq_axes = numpy_helper.from_array(np.array(axes, dtype=np.int64), name="unsq_ax") relu = helper.make_node("Relu", ["X"], ["X_relu"], name="relu") unsqueeze = helper.make_node("Unsqueeze", ["X_relu", "unsq_ax"], ["Uq_out"], name="uq") relu2 = helper.make_node("Relu", ["Uq_out"], ["Y"], name="relu2") @@ -476,14 +492,16 @@ def test_intermediary_unsqueeze_replaced_with_reshape(): [relu, unsqueeze, relu2], "g", [_vi("X", TensorProto.FLOAT, [2, 4, 6])], - [_vi("Y", TensorProto.FLOAT, [2, 4, 1, 6])], + [_vi("Y", TensorProto.FLOAT, output_shape)], initializer=[unsq_axes], + value_info=[_vi("X_relu", TensorProto.FLOAT, [2, 4, 6])], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + feeds = {"X": _RNG.standard_normal([2, 4, 6]).astype(np.float32)} + expected = _run_ort(model, feeds) tr = copy.deepcopy(model) - _prepare_model(tr) tr = _replace_sq_unsq(tr) assert not any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze must be removed" @@ -492,7 +510,9 @@ def test_intermediary_unsqueeze_replaced_with_reshape(): shape_init = next((i for i in tr.graph.initializer if i.name == "uq_4d_shape"), None) assert shape_init is not None, "Reshape shape initializer must exist" shape = numpy_helper.to_array(shape_init).tolist() - assert shape == [2, 4, 1, 6], f"expected [2,4,1,6], got {shape}" + assert shape == output_shape + onnx.checker.check_model(tr, full_check=True) + np.testing.assert_array_equal(_run_ort(tr, feeds)["Y"], expected["Y"]) def test_intermediary_squeeze_at_graph_output_replaced(): @@ -1302,6 +1322,88 @@ def test_cast_chain_three_nodes(): assert to_val == TensorProto.INT32 +@pytest.mark.parametrize( + ("opset", "target", "attributes", "values"), + [ + (19, TensorProto.FLOAT8E4M3FN, {"saturate": 0}, [-700, 700, 1.25, 1.75]), + ( + 24, + TensorProto.FLOAT8E8M0, + {"saturate": 0, "round_mode": "down"}, + [1.25, 1.75, 3.25, 0.6], + ), + ], +) +def test_cast_chain_preserves_final_attributes(opset, target, attributes, values): + """Retain float8 saturation and rounding semantics when removing a lossless cast.""" + first = helper.make_node("Cast", ["X"], ["wide"], to=TensorProto.DOUBLE) + last = helper.make_node("Cast", ["wide"], ["Y"], to=target, **attributes) + model = helper.make_model( + helper.make_graph( + [first, last], + "cast_attributes", + [_vi("X", TensorProto.FLOAT, [4])], + [_vi("Y", target, [4])], + value_info=[_vi("wide", TensorProto.DOUBLE, [4])], + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.array(values, np.float32)} + with np.errstate(over="ignore", invalid="ignore"): + expected = ReferenceEvaluator(model).run(None, feeds)[0] + transformed = _remove_cast_chains(copy.deepcopy(model)) + actual = ReferenceEvaluator(transformed).run(None, feeds)[0] + np.testing.assert_array_equal(actual.astype(np.float32), expected.astype(np.float32)) + assert len(transformed.graph.node) == 1 + assert list(transformed.graph.node[0].attribute) == list(last.attribute) + onnx.checker.check_model(transformed, full_check=True) + + +def test_custom_domain_cast_is_not_collapsed(): + """Treat a custom-domain Cast as opaque instead of assuming standard Cast semantics.""" + function = helper.make_function( + "test.modelopt", + "Cast", + ["X"], + ["Y"], + [ + helper.make_node("Cast", ["X"], ["wide"], to=TensorProto.DOUBLE), + helper.make_node("Neg", ["wide"], ["Y"]), + ], + [helper.make_opsetid("", _OPSET_VERSION)], + attributes=["to"], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Cast", ["X"], ["wide"], domain="test.modelopt", to=TensorProto.DOUBLE + ), + helper.make_node("Cast", ["wide"], ["Y"], to=TensorProto.FLOAT), + ], + "custom_cast", + [_vi("X", TensorProto.FLOAT, [4])], + [_vi("Y", TensorProto.FLOAT, [4])], + value_info=[_vi("wide", TensorProto.DOUBLE, [4])], + ), + opset_imports=[ + helper.make_opsetid("", _OPSET_VERSION), + helper.make_opsetid("test.modelopt", 1), + ], + functions=[function], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.arange(1, 5, dtype=np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds)[0] + before = model.SerializeToString() + transformed = _remove_cast_chains(copy.deepcopy(model)) + actual = ReferenceEvaluator(transformed).run(None, feeds)[0] + np.testing.assert_array_equal(actual, expected) + assert transformed.SerializeToString() == before + onnx.checker.check_model(transformed, full_check=True) + + def test_single_cast_unchanged(): """A lone Cast node is not modified.""" cast1 = helper.make_node("Cast", ["X"], ["Y"], name="cast1", to=TensorProto.INT32) From 7e917432262327f8e4b4f68965ec357fa150c326 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 12:25:01 +0000 Subject: [PATCH 4/9] Preserve subgraph captures during DLA cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 42 ++++++++- .../dla_transforms/test_dla_graph_cleanup.py | 93 +++++++++++++++++++ 2 files changed, 130 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index b4d028e41d2..c92cca1b953 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -30,6 +30,7 @@ from ...logging_config import logger from ..utils.graph_utils import ( GraphCache, + _captured_names, _node_inputs, add_unique_initializers, add_value_info, @@ -108,6 +109,19 @@ ) +def _subgraph_captures(graph: onnx.GraphProto) -> set[str]: + """Collect outer-scope values referenced inside this graph's nested subgraphs.""" + captured: set[str] = set() + for node in graph.node: + for attribute in node.attribute: + if attribute.type == onnx.AttributeProto.GRAPH: + captured.update(_captured_names(attribute.g)) + elif attribute.type == onnx.AttributeProto.GRAPHS: + for subgraph in attribute.graphs: + captured.update(_captured_names(subgraph)) + return captured + + def _rank_sensitive_consumers( graph: onnx.GraphProto, cache: GraphCache ) -> tuple[set[str], set[tuple[int, int]]]: @@ -145,6 +159,7 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode nodes_to_insert: list[tuple] = [] # (new_node, reference_node) nodes_to_remove: list = [] promotion_safety = None + captured: set[str] = set() cnt = 0 for graph_input in list(graph.input): @@ -162,6 +177,7 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode if promotion_safety is None: promotion_safety = _rank_sensitive_consumers(graph, cache) + captured = _subgraph_captures(graph) sensitive_values, sensitive_edges = promotion_safety # Read concrete input shape (may be None for symbolic dims) @@ -181,6 +197,8 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode ): continue if c_node.op_type == "Unsqueeze" and c_idx == 0: + if c_node.output[0] in captured: + continue old_shape = cache.get_shape(c_node.output[0]) if old_shape is None or ( len(old_shape) < 4 @@ -474,17 +492,18 @@ def _replace_intermediary_squeeze_unsqueeze( def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: """Remove noop Reshapes then collapse consecutive Reshape chains. - Pass 1 — Remove every noop Reshape (``in_shape == out_shape``) completely: + Pass 1 — Remove noop Reshapes (``in_shape == out_shape``) without captured outputs: consumers are wired directly to the Reshape input; at graph-output boundaries an ``Identity`` node is inserted to keep the output name alive. Pass 2 — Collapse consecutive Reshape chains (each member's output consumed - only by the next Reshape) into a single Reshape; chains that are noops as - a whole are removed the same way. + only by the next Reshape, without subgraph captures) into a single Reshape; + chains that are noops as a whole are removed the same way. """ graph = model.graph graph_output_names = cache.graph_output_names noop_count = 0 + captured = _subgraph_captures(graph) chain_count = 0 chain_nodes_count = 0 # ── Pass 1: Remove standalone noop Reshapes ─────────────────────────────── @@ -501,6 +520,8 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo # Noop reshape: rewire downstream node-consumers to bypass this node inp = node.input[0] out_tensor = node.output[0] + if out_tensor in captured: + continue for c_node, c_idx in cache.get_consumers(out_tensor): if c_node.input[c_idx] == out_tensor: c_node.input[c_idx] = inp @@ -543,7 +564,11 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo if inp_tensor and inp_tensor in tail_to_chain: # Extend existing chain only if this node is the sole consumer consumers = cache.get_consumers(inp_tensor) - if len(consumers) == 1 and inp_tensor not in cache.graph_output_names: + if ( + len(consumers) == 1 + and inp_tensor not in cache.graph_output_names + and inp_tensor not in captured + ): chain = tail_to_chain.pop(inp_tensor) chain.append(node) tail_to_chain[node.output[0]] = chain @@ -562,6 +587,8 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo inp = ch[0].input[0] out_tensor = ch[-1].output[0] if list(in_shape) == list(out_shape): + if out_tensor in captured: + continue # Entire chain is a noop for c_node, c_idx in cache.get_consumers(out_tensor): if c_node.input[c_idx] == out_tensor: @@ -838,13 +865,18 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx # Build chains: tail_output → [cast_node, ...] all_casts = [n for n in graph.node if n.op_type == "Cast" and not n.domain and not n.overload] + captured = _subgraph_captures(graph) tail_to_chain: dict[str, list] = {} for node in all_casts: inp_tensor = node.input[0] if node.input else None if inp_tensor and inp_tensor in tail_to_chain: consumers = cache.get_consumers(inp_tensor) - if len(consumers) == 1 and inp_tensor not in cache.graph_output_names: + if ( + len(consumers) == 1 + and inp_tensor not in cache.graph_output_names + and inp_tensor not in captured + ): chain = tail_to_chain.pop(inp_tensor) chain.append(node) tail_to_chain[node.output[0]] = chain diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index d6a936e6c42..dc1bd6c12c3 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -702,6 +702,99 @@ def test_single_reshape_unchanged(): assert len(tr.graph.node) == orig_count, "single Reshape must not be changed" +@pytest.mark.parametrize( + "kind", + ["reshape_noop", "reshape_chain", "reshape_noop_chain", "cast_chain", "unsqueeze", "control"], +) +def test_cleanup_preserves_subgraph_captures(kind): + """Keep chain values alive when If branches capture them, including direct controls.""" + dtype = TensorProto.BOOL if kind == "control" else TensorProto.FLOAT + input_shape = [] if kind == "control" else [6] + target_shape = input_shape + captured = "Y" + captured_type = dtype + initializers = [] + value_info = [] + if kind == "cast_chain": + nodes = [ + helper.make_node("Cast", ["X"], ["T"], to=TensorProto.DOUBLE), + helper.make_node("Cast", ["T"], ["Y"], to=TensorProto.FLOAT), + ] + captured = "T" + captured_type = TensorProto.DOUBLE + value_info.append(_vi("T", captured_type, input_shape)) + elif kind == "unsqueeze": + input_shape = [2, 3] + target_shape = [1, 1, 2, 3] + initializers.append(numpy_helper.from_array(np.array([0, 1], np.int64), name="axes")) + nodes = [helper.make_node("Unsqueeze", ["X", "axes"], ["Y"])] + else: + shapes = [[2, 3], [6]] if "chain" in kind else [input_shape] + if kind == "reshape_chain": + shapes[-1] = [1, 2, 3] + captured = "T" + nodes = [] + previous = "X" + for index, shape in enumerate(shapes): + name = f"shape_{index}" + initializers.append(numpy_helper.from_array(np.array(shape, np.int64), name=name)) + output = "Y" if index == len(shapes) - 1 else "T" + nodes.append(helper.make_node("Reshape", [previous, name], [output])) + value_info.append(_vi(output, dtype, shape)) + previous = output + target_shape = shapes[-1] + captured_shape = [2, 3] if captured == "T" and kind != "cast_chain" else target_shape + if kind in ("cast_chain", "unsqueeze"): + value_info.append(_vi("Y", dtype, target_shape)) + inputs = [_vi("X", dtype, input_shape)] + condition = "Y" if kind == "control" else "condition" + if kind != "control": + inputs.append(_vi(condition, TensorProto.BOOL, [])) + branch = helper.make_graph( + [helper.make_node("Identity", [captured], ["branch_out"])], + "capture", + [], + [_vi("branch_out", captured_type, captured_shape)], + ) + nodes.extend( + [ + helper.make_node("Identity", ["Y"], ["main"]), + helper.make_node( + "If", [condition], ["captured"], then_branch=branch, else_branch=branch + ), + ] + ) + model = helper.make_model( + helper.make_graph( + nodes, + "captured_chain", + inputs, + [_vi("main", dtype, target_shape), _vi("captured", captured_type, captured_shape)], + initializer=initializers, + value_info=value_info, + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + feeds = { + "X": np.array(True) + if kind == "control" + else np.arange(6, dtype=np.float32).reshape(input_shape) + } + if kind != "control": + feeds[condition] = np.array(True) + expected = _run_ort(model, feeds) + transform = _remove_cast_chains if kind == "cast_chain" else _remove_reshapes + if kind == "unsqueeze": + transform = _canonicalize_inputs + transformed = transform(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + actual = _run_ort(transformed, feeds) + for name in expected: + np.testing.assert_array_equal(actual[name], expected[name]) + + # ══════════════════════════════════════════════════════════════════════════════ # Step 4 — Fold constants # ══════════════════════════════════════════════════════════════════════════════ From dcd4d307b382a3ed96e717956d572c38e0ba3638 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 12:38:52 +0000 Subject: [PATCH 5/9] Preserve PRelu slope ranks in DLA cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 2 +- .../dla_transforms/test_dla_graph_cleanup.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index c92cca1b953..73d1d67d75d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -142,7 +142,7 @@ def _rank_sensitive_consumers( if not rank_agnostic: sensitive_values.update(_node_inputs(node)) for index, name in enumerate(node.input): - if name and (not rank_agnostic or (node.op_type == "Clip" and index > 0)): + if name and (not rank_agnostic or (node.op_type in ("Clip", "PRelu") and index > 0)): sensitive_values.add(name) sensitive_edges.add((id(node), index)) return sensitive_values, sensitive_edges diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index dc1bd6c12c3..cb19eebfc51 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -256,7 +256,9 @@ def test_canonicalize_4d_input_unchanged(axes_input): @pytest.mark.parametrize("stage", ["input", "squeeze", "unsqueeze"]) -@pytest.mark.parametrize("consumer", ["Transpose", "Softmax", "MatMul", "transpose_via_relu"]) +@pytest.mark.parametrize( + "consumer", ["Transpose", "Softmax", "MatMul", "transpose_via_relu", "prelu_slope"] +) def test_rank_promotion_preserves_rank_sensitive_consumers(stage, consumer): """Preserve rank-sensitive behavior, including consumers behind an elementwise node.""" input_shape = [1, 2, 3] if stage == "squeeze" else [2, 3] @@ -287,6 +289,15 @@ def test_rank_promotion_preserves_rank_sensitive_consumers(stage, consumer): elif consumer == "Softmax": nodes.append(helper.make_node(consumer, [source], ["Y"], axis=len(source_shape) - 2)) output_shape = source_shape + elif consumer == "prelu_slope": + initializers.append( + numpy_helper.from_array( + -np.arange(1, np.prod(source_shape) + 1, dtype=np.float32).reshape(source_shape), + "data", + ) + ) + nodes.append(helper.make_node("PRelu", ["data", source], ["Y"])) + output_shape = source_shape else: initializers.append( numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "W") From d4d5684fce0217963f1717759fec4d3701e34668 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 13:00:06 +0000 Subject: [PATCH 6/9] Honor model opsets in DLA cleanup axes builders Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 49 ++++---- .../dla_transforms/test_dla_graph_cleanup.py | 112 +++++++++++------- 2 files changed, 96 insertions(+), 65 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index 73d1d67d75d..363ad11f8b4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -109,6 +109,21 @@ ) +def _set_axes(model, cache: GraphCache, node: onnx.NodeProto, axes, preferred_name: str) -> None: + """Attach axes using the model's default-domain Squeeze/Unsqueeze schema.""" + opset = next((o.version for o in model.opset_import if o.domain == ""), 1) + if opset < 13: + node.attribute.append( + helper.make_attribute("axes", list(axes), attr_type=onnx.AttributeProto.INTS) + ) + else: + axes_name = cache.unique_name(preferred_name) + model.graph.initializer.append( + numpy_helper.from_array(np.asarray(axes, dtype=np.int64), name=axes_name) + ) + node.input.append(axes_name) + + def _subgraph_captures(graph: onnx.GraphProto) -> set[str]: """Collect outer-scope values referenced inside this graph's nested subgraphs.""" captured: set[str] = set() @@ -212,18 +227,14 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode if not consumers: continue - axes_name = cache.unique_name(f"{graph_input.name}_unsqueeze_axes") out_name = cache.unique_name(f"{graph_input.name}_4d") - add_unique_initializers( - graph, unique_init, [numpy_helper.from_array(axes_arr, name=axes_name)] - ) - unsq = helper.make_node( "Unsqueeze", - inputs=[graph_input.name, axes_name], + inputs=[graph_input.name], outputs=[out_name], name=cache.unique_name(f"{graph_input.name}_unsqueeze"), ) + _set_axes(model, cache, unsq, axes_arr, f"{graph_input.name}_unsqueeze_axes") add_value_info(graph, out_name, tt.elem_type, canonical_4d) # Find the first consumer node in graph order for insertion point @@ -771,22 +782,13 @@ def _replace_graph_output_reshape_with_squeeze( # Update the Reshape node to target the 4D shape node.input[1] = rs_shape_name node.output[0] = rs_4d_out - sq_axes_name = cache.unique_name(f"{pfx}_sq0_axes") - add_unique_initializers( - graph, - unique_init, - [ - numpy_helper.from_array( - np.array(squeeze_axes, dtype=np.int64), name=sq_axes_name - ) - ], - ) sq = helper.make_node( "Squeeze", - inputs=[rs_4d_out, sq_axes_name], + inputs=[rs_4d_out], outputs=[out_name], name=cache.unique_name(f"{pfx}_squeeze_leading"), ) + _set_axes(model, cache, sq, squeeze_axes, f"{pfx}_sq0_axes") go_batch[go_node_idx[id(node)]] = ([node], [node, sq]) cnt += 1 logger.debug( @@ -799,20 +801,17 @@ def _replace_graph_output_reshape_with_squeeze( continue squeeze_axes.sort() # ascending axis order required by ONNX Squeeze + if not squeeze_axes: + continue # Build Squeeze node - axes_name = cache.unique_name(f"{pfx}_squeeze_axes") - add_unique_initializers( - graph, - unique_init, - [numpy_helper.from_array(np.array(squeeze_axes, dtype=np.int64), name=axes_name)], - ) sq = helper.make_node( "Squeeze", - inputs=[in_name, axes_name], + inputs=[in_name], outputs=[out_name], name=cache.unique_name(f"{pfx}_squeeze"), ) + _set_axes(model, cache, sq, squeeze_axes, f"{pfx}_squeeze_axes") go_batch[go_node_idx[id(node)]] = ([node], [sq]) cnt += 1 logger.debug( @@ -849,7 +848,7 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: return model -# ── Step 7: Remove consecutive Cast chains ────────────────────────────────── +# ── Step 6: Remove consecutive Cast chains ────────────────────────────────── def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx.ModelProto: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index cb19eebfc51..b58f3a655d0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -13,30 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_graph_cleanup`. - """Structural and ORT parity tests for ``dla_graph_cleanup``. -The combined transform applies five steps: - 1. Canonicalize non-4D inputs (adds Unsqueeze, deduplicates existing ones) - 2. Replace intermediary Squeeze/Unsqueeze with 4-D Reshape - 3. Collapse Reshape chains (noop → remove; multi → single Reshape) - 4. Fold constants via onnx-graphsurgeon - 5. Remove unused initializers - -Steps 1 and 2 intentionally change tensor shapes, so ORT parity only works -for end-to-end models that happen to keep graph I/O shapes the same. -Tests for those steps verify structural changes (node types, counts). -Steps 3-5 preserve graph I/O shapes, so ORT parity is verified for them. +The combined transform applies seven stages in order: + 1. Replace intermediary Squeeze/Unsqueeze with 4-D Reshape. + 2. Canonicalize eligible non-4D inputs. + 3. Collapse Reshape chains. + 4. Adapt graph-output Reshapes with Squeeze. + 5. Fold constants via onnx-graphsurgeon. + 6. Collapse lossless Cast chains. + 7. Remove unused initializers. + +Numerical cases use ORT; float8 and custom-function cases use ONNX ReferenceEvaluator. """ from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx @@ -52,8 +45,6 @@ import onnxruntime as ort from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "graph_cleanup_test_models" _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -158,7 +149,8 @@ def test_canonicalize_3d_input_unsqueeze_added(): np.testing.assert_array_equal(actual["Z"], expected["Z"]) -def test_canonicalize_2d_input_unsqueeze_added(): +@pytest.mark.parametrize("opset", [12, 13, 17]) +def test_canonicalize_2d_input_unsqueeze_added(opset): """Promote a 2D input when broadcasting already fixes downstream rank at four.""" add = helper.make_node("Add", ["X", "W"], ["sum"], name="add") transpose = helper.make_node("Transpose", ["sum"], ["Y"], perm=[0, 1, 3, 2]) @@ -170,7 +162,14 @@ def test_canonicalize_2d_input_unsqueeze_added(): [numpy_helper.from_array(np.zeros([1, 1, 4, 5], np.float32), "W")], value_info=[_vi("sum", TensorProto.FLOAT, [1, 1, 4, 5])], ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("test.modelopt", 99), + helper.make_opsetid("ai.onnx", 17), + helper.make_opsetid("", opset), + ], + ) _clamp_ir(model) feeds = {"X": _RNG.standard_normal([4, 5]).astype(np.float32)} expected = _run_ort(model, feeds) @@ -181,8 +180,13 @@ def test_canonicalize_2d_input_unsqueeze_added(): unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] assert len(unsq_nodes) == 1 - axes_init = next(i for i in tr.graph.initializer if i.name.endswith("_unsqueeze_axes")) - axes = numpy_helper.to_array(axes_init).tolist() + if opset < 13: + axes = list(next(attr.ints for attr in unsq_nodes[0].attribute if attr.name == "axes")) + assert len(unsq_nodes[0].input) == 1 + assert not any(i.name.endswith("_unsqueeze_axes") for i in tr.graph.initializer) + else: + axes_init = next(i for i in tr.graph.initializer if i.name == unsq_nodes[0].input[1]) + axes = numpy_helper.to_array(axes_init).tolist() assert axes == [0, 1], f"axes should be [0,1] for 2-D input, got {axes}" add_node = next(node for node in tr.graph.node if node.op_type == "Add") assert add_node.input[0] == unsq_nodes[0].output[0] @@ -1098,18 +1102,26 @@ def test_interleaved_reshape_chains_collapsed(names): # ══════════════════════════════════════════════════════════════════════════════ -def test_graph_output_reshape_replaced_with_squeeze(): - """[1,3,4,1] → Reshape → [3,4] at graph output → replaced by Squeeze(axes=[0,3]); ORT parity.""" - sh_init = numpy_helper.from_array(np.array([3, 4], dtype=np.int64), name="rs_shape") +@pytest.mark.parametrize("opset", [12, 13, 17]) +@pytest.mark.parametrize("output_shape", [[3, 4], [1, 3, 4, 1]]) +def test_graph_output_reshape_replaced_with_squeeze(opset, output_shape): + """Use opset-compatible output Squeezes while retaining identity Reshapes.""" + sh_init = numpy_helper.from_array(np.array(output_shape, dtype=np.int64), name="rs_shape") reshape = helper.make_node("Reshape", ["X", "rs_shape"], ["Y"], name="rs") graph = helper.make_graph( [reshape], "g", [_vi("X", TensorProto.FLOAT, [1, 3, 4, 1])], - [_vi("Y", TensorProto.FLOAT, [3, 4])], + [_vi("Y", TensorProto.FLOAT, output_shape)], initializer=[sh_init], ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("test.modelopt", 99), + helper.make_opsetid("", opset), + ], + ) _clamp_ir(model) x = _RNG.standard_normal([1, 3, 4, 1]).astype(np.float32) @@ -1119,21 +1131,27 @@ def test_graph_output_reshape_replaced_with_squeeze(): tr = copy.deepcopy(model) _prepare_model(tr) tr = _reshape_to_squeeze(tr) + onnx.checker.check_model(tr, full_check=True) _prepare_model(tr) got = _run_ort(tr, {"X": x}) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + if output_shape == [1, 3, 4, 1]: + assert any(node.op_type == "Reshape" for node in tr.graph.node) + assert not any(node.op_type == "Squeeze" for node in tr.graph.node) + return assert not any(n.op_type == "Reshape" for n in tr.graph.node), "Reshape must be replaced" sq_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] assert len(sq_nodes) == 1, "One Squeeze must be inserted" - # Verify axes initializer contains [0, 3] - axes_init = next((i for i in tr.graph.initializer if "squeeze_axes" in i.name), None) - assert axes_init is not None, "Squeeze axes initializer must exist" - assert numpy_helper.to_array(axes_init).tolist() == [0, 3], ( - f"Expected axes [0,3], got {numpy_helper.to_array(axes_init).tolist()}" - ) - - np.testing.assert_array_equal(ref["Y"], got["Y"]) + if opset < 13: + axes = list(next(attr.ints for attr in sq_nodes[0].attribute if attr.name == "axes")) + assert len(sq_nodes[0].input) == 1 + assert not any("squeeze_axes" in i.name for i in tr.graph.initializer) + else: + axes_init = next(i for i in tr.graph.initializer if i.name == sq_nodes[0].input[1]) + axes = numpy_helper.to_array(axes_init).tolist() + assert axes == [0, 3] def test_graph_output_reshape_squeeze_axes_left_to_right(): @@ -1177,7 +1195,8 @@ def test_graph_output_reshape_squeeze_axes_left_to_right(): @pytest.mark.parametrize("output_shape", [[6, 4], [24]]) -def test_graph_output_reshape_not_squeezable_padded_to_4d(output_shape): +@pytest.mark.parametrize("opset", [12, 13, 17]) +def test_graph_output_reshape_not_squeezable_padded_to_4d(output_shape, opset): """[2,3,4] → Reshape → [6,4] at graph output: not a pure squeeze, output is 2D (<4D). Pad the Reshape target to rank four and squeeze every prepended unit dimension. @@ -1192,7 +1211,13 @@ def test_graph_output_reshape_not_squeezable_padded_to_4d(output_shape): [_vi("Y", TensorProto.FLOAT, output_shape)], initializer=[sh_init], ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("test.modelopt", 99), + helper.make_opsetid("", opset), + ], + ) _clamp_ir(model) x = _RNG.standard_normal([2, 3, 4]).astype(np.float32) @@ -1216,8 +1241,15 @@ def test_graph_output_reshape_not_squeezable_padded_to_4d(output_shape): numpy_helper.to_array(rs_shape_init).tolist() == [1] * (4 - len(output_shape)) + output_shape ) - axes = next(i for i in tr.graph.initializer if i.name == sq_nodes[0].input[1]) - assert numpy_helper.to_array(axes).tolist() == list(range(4 - len(output_shape))) + if opset < 13: + axes = list(next(attr.ints for attr in sq_nodes[0].attribute if attr.name == "axes")) + assert len(sq_nodes[0].input) == 1 + assert not any("sq0_axes" in i.name for i in tr.graph.initializer) + else: + axes_init = next(i for i in tr.graph.initializer if i.name == sq_nodes[0].input[1]) + axes = numpy_helper.to_array(axes_init).tolist() + assert axes == list(range(4 - len(output_shape))) + onnx.checker.check_model(tr, full_check=True) np.testing.assert_array_equal(ref["Y"], got["Y"]) From 49507cfa4ff3af11171c8476d68ba58a09d66bc5 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 13:19:43 +0000 Subject: [PATCH 7/9] Preserve ONNX domain semantics in DLA cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 17 ++-- .../dla_transforms/test_dla_graph_cleanup.py | 83 ++++++++++++++++++- 2 files changed, 93 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index 363ad11f8b4..46fe75d6315 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -111,7 +111,8 @@ def _set_axes(model, cache: GraphCache, node: onnx.NodeProto, axes, preferred_name: str) -> None: """Attach axes using the model's default-domain Squeeze/Unsqueeze schema.""" - opset = next((o.version for o in model.opset_import if o.domain == ""), 1) + versions = {o.domain: o.version for o in model.opset_import} + opset = versions.get("", versions.get("ai.onnx", 1)) if opset < 13: node.attribute.append( helper.make_attribute("axes", list(axes), attr_type=onnx.AttributeProto.INTS) @@ -148,7 +149,7 @@ def _rank_sensitive_consumers( sensitive_values: set[str] = set() sensitive_edges: set[tuple[int, int]] = set() for node in reversed(graph.node): - rank_agnostic = node.domain in ("", "ai.onnx") and node.op_type in _RANK_AGNOSTIC_OPS + rank_agnostic = not node.domain and not node.overload and node.op_type in _RANK_AGNOSTIC_OPS if rank_agnostic: rank_agnostic = not any( output in sensitive_values and len(cache.get_shape(output) or []) < 4 @@ -206,6 +207,8 @@ def _canonicalize_inputs(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mode consumers = [] for c_node, c_idx in consumer_map.get(graph_input.name, []): + if c_node.domain or c_node.overload: + continue if (c_node.op_type == "Squeeze" and "Gather" in c_node.name) or c_node.op_type in ( "GroupQueryAttention", "Gather", @@ -442,7 +445,7 @@ def _replace_intermediary_squeeze_unsqueeze( cnt_replace = 0 for node in list(graph.node): - if node.op_type not in ("Squeeze", "Unsqueeze"): + if node.domain or node.overload or node.op_type not in ("Squeeze", "Unsqueeze"): continue if not node.input or not node.output: continue @@ -522,6 +525,8 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo pass1_batch: dict[int, tuple[list, list]] = {} for node in cache.nodes_by_op("Reshape"): + if node.domain or node.overload: + continue in_shape = cache.get_shape(node.input[0]) out_shape = cache.get_shape(node.output[0]) if in_shape is None or out_shape is None: @@ -565,7 +570,9 @@ def _remove_reshape_chains(model: onnx.ModelProto, cache: GraphCache) -> onnx.Mo # members (e.g. r1→r2→r5→r3 where r5 is not a Reshape: r1,r2,r3 still # form a chain as long as each Reshape is the sole consumer of the previous # Reshape's output). - all_reshapes = [n for n in graph.node if n.op_type == "Reshape"] + all_reshapes = [ + n for n in graph.node if n.op_type == "Reshape" and not n.domain and not n.overload + ] # output_name → chain (list of Reshape nodes whose last output == output_name) tail_to_chain: dict[str, list] = {} @@ -711,7 +718,7 @@ def _replace_graph_output_reshape_with_squeeze( cnt = 0 for node in list(graph.node): - if node.op_type != "Reshape": + if node.domain or node.overload or node.op_type != "Reshape": continue if not node.output or node.output[0] not in graph_output_names: continue diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index b58f3a655d0..86834c10625 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -150,7 +150,8 @@ def test_canonicalize_3d_input_unsqueeze_added(): @pytest.mark.parametrize("opset", [12, 13, 17]) -def test_canonicalize_2d_input_unsqueeze_added(opset): +@pytest.mark.parametrize("alias_only", [False, True]) +def test_canonicalize_2d_input_unsqueeze_added(opset, alias_only): """Promote a 2D input when broadcasting already fixes downstream rank at four.""" add = helper.make_node("Add", ["X", "W"], ["sum"], name="add") transpose = helper.make_node("Transpose", ["sum"], ["Y"], perm=[0, 1, 3, 2]) @@ -170,6 +171,9 @@ def test_canonicalize_2d_input_unsqueeze_added(opset): helper.make_opsetid("", opset), ], ) + if alias_only: + model.ClearField("opset_import") + model.opset_import.append(helper.make_opsetid("ai.onnx", opset)) _clamp_ir(model) feeds = {"X": _RNG.standard_normal([4, 5]).astype(np.float32)} expected = _run_ort(model, feeds) @@ -190,7 +194,9 @@ def test_canonicalize_2d_input_unsqueeze_added(opset): assert axes == [0, 1], f"axes should be [0,1] for 2-D input, got {axes}" add_node = next(node for node in tr.graph.node if node.op_type == "Add") assert add_node.input[0] == unsq_nodes[0].output[0] - onnx.checker.check_model(tr, full_check=True) + # ONNX 1.21 requires an empty-domain import; ORT also accepts alias-only imports. + if not alias_only: + onnx.checker.check_model(tr, full_check=True) np.testing.assert_array_equal(_run_ort(tr, feeds)["Y"], expected["Y"]) @@ -1540,6 +1546,79 @@ def test_custom_domain_cast_is_not_collapsed(): onnx.checker.check_model(transformed, full_check=True) +@pytest.mark.parametrize( + ("stage", "op_type", "input_shape", "output_shape", "parameters"), + [ + ("input", "Unsqueeze", [2, 3], [1, 1, 2, 3], [0, 1]), + ("intermediary", "Unsqueeze", [2, 3], [1, 1, 2, 3], [0, 1]), + ("intermediary", "Squeeze", [1, 1, 2, 3], [2, 3], [0, 1]), + ("reshape", "Reshape", [2, 3], [2, 3], [2, 3]), + ("chain", "Reshape", [6], [2, 3], [2, 3]), + ("output", "Reshape", [1, 2, 3, 1], [2, 3], [2, 3]), + ], +) +def test_custom_domain_shape_operations_are_preserved( + stage, op_type, input_shape, output_shape, parameters +): + """Preserve local functions whose names match standard shape operations.""" + function = helper.make_function( + "test.modelopt", + op_type, + ["X", "parameters"], + ["Y"], + [ + helper.make_node(op_type, ["X", "parameters"], ["adapted"]), + helper.make_node("Neg", ["adapted"], ["Y"]), + ], + [helper.make_opsetid("", _OPSET_VERSION)], + ) + nodes, value_info = [], [] + source = "X" + if stage == "intermediary": + nodes.append(helper.make_node("Identity", ["X"], ["middle"])) + value_info.append(_vi("middle", TensorProto.FLOAT, input_shape)) + source = "middle" + output = "T" if stage == "chain" else "Y" + nodes.append( + helper.make_node(op_type, [source, "parameters"], [output], domain="test.modelopt") + ) + initializers = [numpy_helper.from_array(np.array(parameters, np.int64), name="parameters")] + if stage == "chain": + value_info.append(_vi("T", TensorProto.FLOAT, output_shape)) + output_shape = [1, 2, 3] + initializers.append(numpy_helper.from_array(np.array(output_shape, np.int64), name="final")) + nodes.append(helper.make_node("Reshape", ["T", "final"], ["Y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "custom_shape_operation", + [_vi("X", TensorProto.FLOAT, input_shape)], + [_vi("Y", TensorProto.FLOAT, output_shape)], + initializer=initializers, + value_info=value_info, + ), + opset_imports=[ + helper.make_opsetid("", _OPSET_VERSION), + helper.make_opsetid("test.modelopt", 1), + ], + functions=[function], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.arange(1, 7, dtype=np.float32).reshape(input_shape)} + expected = ReferenceEvaluator(model).run(None, feeds)[0] + transform = { + "input": _canonicalize_inputs, + "intermediary": _replace_sq_unsq, + "reshape": _remove_reshapes, + "chain": _remove_reshapes, + "output": _reshape_to_squeeze, + }[stage] + transformed = transform(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(transformed).run(None, feeds)[0], expected) + assert transformed.SerializeToString() == model.SerializeToString() + + def test_single_cast_unchanged(): """A lone Cast node is not modified.""" cast1 = helper.make_node("Cast", ["X"], ["Y"], name="cast1", to=TensorProto.INT32) From d85975bb05c42a4b0bb5d352b5d49a020e960da4 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 13:39:19 +0000 Subject: [PATCH 8/9] Preserve unresolved Squeeze axes in DLA cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 2 + .../dla_transforms/test_dla_graph_cleanup.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index 46fe75d6315..fe81fc218ca 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -352,6 +352,8 @@ def _compute_sq_unsq_out_shape( if node.op_type == "Squeeze": if axes is None: + if len(node.input) > 1 and node.input[1]: + return None return [d for d in in_shape if d != 1] out = list(in_shape) for ax in sorted([a if a >= 0 else len(out) + a for a in axes], reverse=True): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index 86834c10625..a5e3b2e096d 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -459,6 +459,43 @@ def test_squeeze_fallback_uses_selected_axes(opset, axes): ) +@pytest.mark.parametrize("axes_mode", ["dynamic", "omitted"]) +def test_squeeze_fallback_distinguishes_unknown_and_omitted_axes(axes_mode): + """Retain unresolved axes and preserve the remove-all fallback for omitted axes.""" + inputs = [_vi("X", TensorProto.FLOAT, [1, 2, 1, 3])] + node_inputs = ["X"] + if axes_mode == "dynamic": + inputs.append(_vi("axes", TensorProto.INT64, [1])) + node_inputs.append("axes") + squeeze = helper.make_node("Squeeze", node_inputs, ["middle"], name="sq") + output_shape = [None, None, None] if axes_mode == "dynamic" else [2, 3] + model = helper.make_model( + helper.make_graph( + [squeeze, helper.make_node("Relu", ["middle"], ["Y"])], + "axes_resolution", + inputs, + [_vi("Y", TensorProto.FLOAT, output_shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + transformed = _replace_sq_unsq(copy.deepcopy(model)) + if axes_mode == "dynamic": + assert transformed.SerializeToString() == model.SerializeToString() + else: + transformed = _mod.infer_shapes(transformed) + onnx.checker.check_model(transformed, full_check=True) + for axes in ([0], [2]) if axes_mode == "dynamic" else [None]: + feeds = {"X": np.arange(6, dtype=np.float32).reshape(1, 2, 1, 3)} + if axes is not None: + feeds["axes"] = np.array(axes, np.int64) + expected = _run_ort(model, feeds)["Y"] + if axes_mode != "dynamic": + expected = expected.reshape(1, 1, 2, 3) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) + + def test_intermediary_squeeze_replaced_with_reshape(): """Squeeze [2,1,6,8] → [2,6,8] (axis=1) is replaced by Reshape → [1,2,6,8]. From 4aa8e97472791d1ff735c70643823589b9dce579 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 14:18:47 +0000 Subject: [PATCH 9/9] Preserve cleanup outputs around 4D promotion Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 62 ++++++++++++++++++- .../dla_transforms/test_dla_graph_cleanup.py | 57 ++++++++++++++--- 2 files changed, 109 insertions(+), 10 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index fe81fc218ca..0ebe278d46c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -23,6 +23,8 @@ from __future__ import annotations +import copy + import numpy as np import onnx from onnx import helper, numpy_helper @@ -943,11 +945,65 @@ def _remove_consecutive_casts(model: onnx.ModelProto, cache: GraphCache) -> onnx return model +def _restore_promoted_graph_outputs(model, cache: GraphCache, original_outputs: dict) -> bool: + """Restore declared output ranks while retaining leading-unit promotion inside the graph.""" + graph = model.graph + plans = [] + for output in graph.output: + original = original_outputs[output.name] + current_type = output.type.tensor_type + original_type = original.type.tensor_type + if not current_type.HasField("shape") or not original_type.HasField("shape"): + continue + padding = len(current_type.shape.dim) - len(original_type.shape.dim) + if padding == 0: + continue + producer = cache.get_producer(output.name) + if ( + padding < 0 + or any(dim.dim_value != 1 for dim in current_type.shape.dim[:padding]) + or producer is None + ): + raise ValueError( + f"Cleanup cannot restore non-unit shape changes for output {output.name!r}" + ) + plans.append((output, original, producer, padding)) + if not plans: + return False + + positions = {id(node): index for index, node in enumerate(graph.node)} + insertions = {} + for output, original, producer, padding in plans: + promoted_name = cache.unique_name(f"{output.name}_4d_output") + promoted_info = copy.deepcopy(output) + promoted_info.name = promoted_name + graph.value_info.append(promoted_info) + for index, name in enumerate(producer.output): + if name == output.name: + producer.output[index] = promoted_name + for consumer, index in cache.get_consumers(output.name): + consumer.input[index] = promoted_name + squeeze = helper.make_node( + "Squeeze", + [promoted_name], + [output.name], + name=cache.unique_name(f"{output.name}_restore_rank"), + ) + _set_axes(model, cache, squeeze, range(padding), f"{output.name}_restore_axes") + output.CopyFrom(original) + insertions.setdefault(positions[id(producer)] + 1, []).append(squeeze) + for position in sorted(insertions, reverse=True): + for node in reversed(insertions[position]): + graph.node.insert(position, node) + return True + + # ── Master transform ────────────────────────────────────────────────────────── def _apply_graph_cleanup(model: onnx.ModelProto) -> onnx.ModelProto: - """Apply the seven DLA cleanup stages while retaining rank-sensitive paths.""" + """Apply seven DLA cleanup stages while retaining output signatures and rank-sensitive paths.""" + original_outputs = {output.name: copy.deepcopy(output) for output in model.graph.output} model = infer_shapes(model) cache = GraphCache(model.graph) @@ -959,6 +1015,8 @@ def _apply_graph_cleanup(model: onnx.ModelProto) -> onnx.ModelProto: model = _canonicalize_inputs(model, cache) model = infer_shapes(model) cache.rebuild() + if _restore_promoted_graph_outputs(model, cache, original_outputs): + cache.rebuild() # Step 3 — collapse Reshape chains model = _remove_reshape_chains(model, cache) @@ -976,5 +1034,7 @@ def _apply_graph_cleanup(model: onnx.ModelProto) -> onnx.ModelProto: # Step 7 — remove unused initializers model = _remove_unused_initializers(model) + for output in model.graph.output: + output.CopyFrom(original_outputs[output.name]) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py index a5e3b2e096d..9f7fbb17411 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_graph_cleanup.py @@ -115,19 +115,27 @@ def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: # ══════════════════════════════════════════════════════════════════════════════ -def test_canonicalize_3d_input_unsqueeze_added(): - """Promote the safe branch of a shared input while preserving a Transpose branch.""" +@pytest.mark.parametrize("opset", [12, 17]) +@pytest.mark.parametrize("input_shape", [[1], [6], [1, 6], [2, 3, 4]]) +def test_canonicalize_input_promotes_safe_branch_and_preserves_outputs(opset, input_shape): + """Keep 4D interior promotion while preserving both pointwise and Transpose outputs.""" relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") - transpose = helper.make_node("Transpose", ["X"], ["Z"], perm=[2, 1, 0]) + transpose = helper.make_node( + "Transpose", ["X"], ["Z"], perm=list(reversed(range(len(input_shape)))) + ) graph = helper.make_graph( [relu, transpose], "g", - [_vi("X", TensorProto.FLOAT, [2, 3, 4])], - [_vi("Y", TensorProto.FLOAT, [2, 3, 4]), _vi("Z", TensorProto.FLOAT, [4, 3, 2])], + [_vi("X", TensorProto.FLOAT, input_shape)], + [ + _vi("Y", TensorProto.FLOAT, input_shape), + _vi("Z", TensorProto.FLOAT, list(reversed(input_shape))), + ], ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) _clamp_ir(model) - feeds = {"X": _RNG.standard_normal([2, 3, 4]).astype(np.float32)} + onnx.checker.check_model(model, full_check=True) + feeds = {"X": _RNG.standard_normal(input_shape).astype(np.float32)} expected = _run_ort(model, feeds) tr = copy.deepcopy(model) @@ -135,20 +143,51 @@ def test_canonicalize_3d_input_unsqueeze_added(): onnx.checker.check_model(tr, full_check=True) unsq_nodes = [n for n in tr.graph.node if n.op_type == "Unsqueeze"] - assert len(unsq_nodes) == 1, "one Unsqueeze must be inserted for the 3-D input" + assert len(unsq_nodes) == 1 unsq = unsq_nodes[0] assert unsq.input[0] == "X", "Unsqueeze reads the graph input" # Relu must now read the Unsqueeze output, not X directly relu_node = next(n for n in tr.graph.node if n.op_type == "Relu") assert relu_node.input[0] == unsq.output[0], "Relu must be redirected to Unsqueeze output" + assert len(GraphCache(tr.graph).get_shape(relu_node.output[0])) == 4 transpose_node = next(n for n in tr.graph.node if n.op_type == "Transpose") assert transpose_node.input[0] == "X" actual = _run_ort(tr, feeds) - np.testing.assert_array_equal(actual["Y"], expected["Y"].reshape(1, 2, 3, 4)) + assert list(tr.graph.output) == list(model.graph.output) + np.testing.assert_array_equal(actual["Y"], expected["Y"]) np.testing.assert_array_equal(actual["Z"], expected["Z"]) +@pytest.mark.parametrize("opset", [12, 17]) +def test_cleanup_preserves_scalar_output_after_intermediary_promotion(opset): + """Restore a scalar result after an intermediary Squeeze is promoted to 4D.""" + squeeze = helper.make_node("Squeeze", ["X"], ["middle"], axes=[0, 1]) + initializers = [] + if opset >= 13: + squeeze.ClearField("attribute") + squeeze.input.append("axes") + initializers.append(numpy_helper.from_array(np.array([0, 1], np.int64), name="axes")) + model = helper.make_model( + helper.make_graph( + [squeeze, helper.make_node("Relu", ["middle"], ["Y"])], + "scalar_boundary", + [_vi("X", TensorProto.FLOAT, [1, 1])], + [_vi("Y", TensorProto.FLOAT, [])], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.array([[2]], np.float32)} + expected = _run_ort(model, feeds)["Y"] + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert list(transformed.graph.output) == list(model.graph.output) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) + + @pytest.mark.parametrize("opset", [12, 13, 17]) @pytest.mark.parametrize("alias_only", [False, True]) def test_canonicalize_2d_input_unsqueeze_added(opset, alias_only):