From 317c87831e811a7c4c55957e8991b0704772d881 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:07 +0000 Subject: [PATCH 01/15] [DLA_transforms_6] [6/17] Standalone five-dimensional reshape conversion Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 634 ++++++++++++++++++ .../test_dla_5d_reshape_to_4d.py | 347 ++++++++++ 2 files changed, 981 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py new file mode 100644 index 00000000000..7569b6675f0 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -0,0 +1,634 @@ +# 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. + +"""5D ``Reshape`` → 4D subgraph rewrites for DLA-style graphs. + +:func:`_apply_5d_reshape_to_4d` runs ORT symbolic shape inference first +(unless ``skip_symbolic_shape_inference``), then classifies every 5D ``Reshape`` follower chain and applies +all subgraph replacements in discovery order. + +Reduce followers use merge-(n-1,n) → ``Slice`` x N → optional ``Clip`` / Q→DQ **per branch** → +per-branch reduce (axes ``[n-1]``, ``keepdims=1``) → ``Concat``, matching +``scratch_space/verify_reduce_5d_vs_4d_decompose.py``. Reduce axis ``n`` must be ``1..4`` (initializer); +axis ``0`` is rejected. ``keepdims=1`` on the original op adds ``Unsqueeze`` at axis ``n``. + +Unsupported followers raise ``ValueError`` with ``op_type`` and node ``name``. +""" + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + batch_replace_nodes, + infer_shapes, + resolve_reshape_shape, +) + +_REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) + + +def _consumers_first_input(graph, tensor_name: str) -> list: + return [n for n in graph.node if n.input and len(n.input) > 0 and n.input[0] == tensor_name] + + +def _parse_middleware_strict( + graph, start_out: str, stop_op_types: frozenset, *, path_context: str +) -> tuple: + """Walk ``Clip`` / Q→DQ middleware then a terminal op; raise ``ValueError`` on dead ends or bad ops.""" + middle: list = [] + cur = start_out + while True: + cands = _consumers_first_input(graph, cur) + uses = sum(name == cur for node in graph.node for name in node.input) + if uses != 1 or any(value.name == cur for value in graph.output): + raise ValueError( + f"{path_context}: cannot replace shared or graph-output tensor {cur!r}" + ) + if not cands: + msg = ( + f"{path_context}: no consumer for tensor {cur!r}; " + f"expected middleware (Clip / Q→DQ) or terminal in {sorted(stop_op_types)}" + ) + raise ValueError(msg) + n = cands[0] + if n.op_type in stop_op_types: + return middle, n + if n.op_type == "Clip": + middle.append(n) + cur = n.output[0] + continue + if n.op_type == "QuantizeLinear": + if any(a.name in {"axis", "block_size"} for a in n.attribute): + raise ValueError(f"{path_context}: only per-tensor QDQ middleware is supported") + dq_nodes = [ + x + for x in _consumers_first_input(graph, n.output[0]) + if x.op_type == "DequantizeLinear" + ] + if not dq_nodes: + msg = ( + f"{path_context}: QuantizeLinear node {n.name!r} (op_type={n.op_type!r}) " + f"has no DequantizeLinear consumer" + ) + raise ValueError(msg) + dq = dq_nodes[0] + if ( + sum(name == n.output[0] for node in graph.node for name in node.input) != 1 + or any(value.name == n.output[0] for value in graph.output) + or any(a.name in {"axis", "block_size"} for a in dq.attribute) + ): + raise ValueError(f"{path_context}: QDQ must be unshared and per-tensor") + middle.extend([n, dq]) + cur = dq.output[0] + continue + msg = ( + f"{path_context}: unsupported follower op_type={n.op_type!r} name={n.name!r}; " + f"allowed middleware: Clip, QuantizeLinear→DequantizeLinear; " + f"terminals: {', '.join(sorted(stop_op_types))}" + ) + raise ValueError(msg) + + +def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: + if node.op_type != "Reshape": + return False + ini = cache.get_init(node.input[1]) + if ini is None: + return False + shape_vals = numpy_helper.to_array(ini) + return len(shape_vals) == 5 + + +def _classify_5d_reshape_follower(model, reshape_node) -> tuple: + r"""Return (``\"transpose\"``, mid_a, transpose_node) or (``\"reduce\"``, middle, reduce_node).""" + graph = model.graph + ctx = f"5D Reshape {reshape_node.name!r}" + r_out = reshape_node.output[0] + stop = frozenset({"Transpose", *_REDUCE_OPS}) + middle, term = _parse_middleware_strict(graph, r_out, stop, path_context=ctx) + if term.op_type == "Transpose": + return ("transpose", middle, term) + elif term.op_type in _REDUCE_OPS: + return ("reduce", middle, term) + else: + msg = f"unsupported terminal op {term.op_type!r} in 5D Reshape {reshape_node.name!r}" + raise ValueError(msg) + + +def _duplicate_clip_qdq_chain(middleware: list, x_in: str, suffix: str) -> tuple[list, str]: + """Rebuild Clip / Q→DQ sequence on ``x_in`` with fresh tensor names. Returns ``(new_nodes, output_name)``.""" + if not middleware: + return [], x_in + out = x_in + new_nodes: list = [] + i = 0 + m = 0 + while m < len(middleware): + n = middleware[m] + if n.op_type == "Clip": + new_out = f"_u5d_{suffix}_c{m}_{i}" + inputs = [out, *list(n.input[1:])] + new_nodes.append( + helper.make_node("Clip", inputs, [new_out], name=f"{n.name}_u5d_{suffix}_{m}") + ) + out = new_out + m += 1 + i += 1 + elif n.op_type == "QuantizeLinear": + if m + 1 >= len(middleware) or middleware[m + 1].op_type != "DequantizeLinear": + msg = "QuantizeLinear without following DequantizeLinear in middleware chain" + raise ValueError(msg) + q, dq = middleware[m], middleware[m + 1] + q_out = f"_u5d_{suffix}_q{m}_{i}" + dq_out = f"_u5d_{suffix}_dq{m}_{i}" + new_nodes.append( + helper.make_node( + "QuantizeLinear", + [out, *q.input[1:]], + [q_out], + name=f"{q.name}_u5d_{suffix}_{m}", + **{a.name: helper.get_attribute_value(a) for a in q.attribute}, + ) + ) + new_nodes.append( + helper.make_node( + "DequantizeLinear", + [q_out, *dq.input[1:]], + [dq_out], + name=f"{dq.name}_u5d_{suffix}_{m}", + **{a.name: helper.get_attribute_value(a) for a in dq.attribute}, + ) + ) + out = dq_out + m += 2 + i += 1 + else: + msg = f"unsupported middleware op {n.op_type!r} in unified 5D reshape path" + raise ValueError(msg) + return new_nodes, out + + +def _infer_transpose_path_slice_axis( + cache: GraphCache, input_tensor: str, num_slices: int, *, ctx: str +) -> int: + """Require the axis-zero layout used by the reference decomposition.""" + in_shape = cache.get_shape(input_tensor) + if in_shape is None or not in_shape or in_shape[0] != num_slices: + raise ValueError( + f"{ctx}: leading input dimension must equal slice count {num_slices}, got {in_shape}" + ) + return 0 + + +def _transpose_path_concat_axis_4d(slice_axis: int, transpose_perm: list[int]) -> int: + """Concat axis for 4D branch outputs (scratch: axis 0 when perm[0]==0 and input is sliced on axis 0).""" + if slice_axis != 0: + msg = ( + "5D reshape→transpose path: non-zero Slice axis is not supported " + f"(slice_axis={slice_axis}); only leading batch-dim slicing is supported" + ) + raise ValueError(msg) + if int(transpose_perm[0]) != 0: + raise ValueError("internal: concat axis derivation requires transpose perm[0] == 0") + return 0 + + +def _build_transpose_replacement_or_raise( + model, reshape_node, mid_a, tr, unique_initializer: set, counter: int, cache: GraphCache +) -> tuple[list, list]: + """Slice-based decomposition aligned with scratch_space ``transform_reshape_transpose_reshape``. + + Shapes and permutation come from graph initializers: slice count ``shape1[0]``, 4D reshape ``shape1[1:]``, + final reshape ``shape2`` (any rank), 4D transpose perm ``perm[1:]-1``. Optional ``Clip`` / Q→DQ middleware is + duplicated per branch. + """ + graph = model.graph + ctx = f"5D Reshape {reshape_node.name!r} → Transpose {tr.name!r}" + transpose_perm = None + for attr in tr.attribute: + if attr.name == "perm": + transpose_perm = [int(x) for x in attr.ints] + break + if transpose_perm is None: + msg = f"{ctx}: missing `perm` attribute on Transpose" + raise ValueError(msg) + if transpose_perm[0] != 0: + msg = ( + f"{ctx}: Transpose perm[0] must be 0, got perm={transpose_perm} " + f"(op_type={tr.op_type!r}, name={tr.name!r})" + ) + raise ValueError(msg) + if "_transformed" in tr.name: + msg = f"{ctx}: Transpose node name contains '_transformed' ({tr.name!r})" + raise ValueError(msg) + + tr_out = tr.output[0] + mid_b, r2 = _parse_middleware_strict(graph, tr_out, frozenset({"Reshape"}), path_context=ctx) + if "_transformed" in r2.name: + msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" + raise ValueError(msg) + ini2 = cache.get_init(r2.input[1]) + assert ini2 is not None, f"initializer for node {r2.name} is None" + ini1 = cache.get_init(reshape_node.input[1]) + assert ini1 is not None, f"initializer for node {reshape_node.name} is None" + + shape1_raw = np.asarray(numpy_helper.to_array(ini1), dtype=np.int64).reshape(-1) + shape2 = np.asarray(numpy_helper.to_array(ini2), dtype=np.int64).reshape(-1) + if shape1_raw.size != 5: + msg = ( + f"{ctx}: first Reshape shape initializer must have length 5, got {shape1_raw.tolist()}" + ) + raise ValueError(msg) + # Resolve -1/0 in the shape using the actual input tensor shape + input_tensor_tr = reshape_node.input[0] + tr_in_shape = cache.get_shape(input_tensor_tr) + if tr_in_shape is not None and (-1 in shape1_raw or 0 in shape1_raw): + shape1 = np.array(resolve_reshape_shape(shape1_raw, list(tr_in_shape)), dtype=np.int64) + else: + shape1 = shape1_raw + num_slices = int(shape1[0]) + reshape_4d_vals = shape1[1:] + if reshape_4d_vals.size != 4: + msg = f"{ctx}: expected four tail dims after leading slice count, got {reshape_4d_vals.tolist()}" + raise ValueError(msg) + + new_transpose_perm = [transpose_perm[i + 1] - 1 for i in range(4)] + + input_tensor = reshape_node.input[0] + output_name = r2.output[0] + slice_axis = _infer_transpose_path_slice_axis(cache, input_tensor, num_slices, ctx=ctx) + concat_axis = _transpose_path_concat_axis_4d(slice_axis, transpose_perm) + + slice_axes = helper.make_tensor( + name=f"u5d_tr_slice_axes_{counter}", + data_type=TensorProto.INT64, + dims=[1], + vals=[slice_axis], + ) + slice_starts = [] + slice_ends = [] + for i in range(num_slices): + slice_starts.append( + helper.make_tensor( + name=f"u5d_tr_{i}_starts_{counter}", data_type=TensorProto.INT64, dims=[1], vals=[i] + ) + ) + slice_ends.append( + helper.make_tensor( + name=f"u5d_tr_{i}_ends_{counter}", + data_type=TensorProto.INT64, + dims=[1], + vals=[i + 1], + ) + ) + add_unique_initializers(graph, unique_initializer, [slice_axes, *slice_starts, *slice_ends]) + + reshape_4d_shape = helper.make_tensor( + name=f"u5d_tr_r4_{counter}", + data_type=TensorProto.INT64, + dims=[4], + vals=reshape_4d_vals.tolist(), + ) + final_rank = int(shape2.size) + reshape_final_shape = helper.make_tensor( + name=f"u5d_tr_r1_{counter}", + data_type=TensorProto.INT64, + dims=[final_rank], + vals=shape2.tolist(), + ) + add_unique_initializers(graph, unique_initializer, [reshape_4d_shape, reshape_final_shape]) + + nodes_to_add: list = [] + slice_outputs: list[str] = [] + for i in range(num_slices): + slice_out = f"u5d_tr_slice_{i}_{counter}" + nodes_to_add.append( + helper.make_node( + "Slice", + [ + input_tensor, + f"u5d_tr_{i}_starts_{counter}", + f"u5d_tr_{i}_ends_{counter}", + f"u5d_tr_slice_axes_{counter}", + ], + [slice_out], + name=f"u5d_tr_slice_{i}_{counter}", + ) + ) + reshape_4d_out = f"u5d_tr_r4o_{i}_{counter}" + nodes_to_add.append( + helper.make_node( + "Reshape", + [slice_out, f"u5d_tr_r4_{counter}"], + [reshape_4d_out], + name=f"u5d_tr_r4n_{i}_{counter}", + ) + ) + br_a, t_a = _duplicate_clip_qdq_chain(mid_a, reshape_4d_out, f"tr{counter}_{i}a") + nodes_to_add.extend(br_a) + transpose_branch_out = f"u5d_tr_t_{i}_{counter}" + nodes_to_add.append( + helper.make_node( + "Transpose", + [t_a], + [transpose_branch_out], + perm=new_transpose_perm, + name=f"u5d_tr_tn_{i}_{counter}", + ) + ) + br_b, t_b = _duplicate_clip_qdq_chain(mid_b, transpose_branch_out, f"tr{counter}_{i}b") + nodes_to_add.extend(br_b) + slice_outputs.append(t_b) + + concat_out = f"u5d_tr_concat_{counter}" + nodes_to_add.append( + helper.make_node( + "Concat", + slice_outputs, + [concat_out], + axis=concat_axis, + name=f"u5d_tr_concatn_{counter}", + ) + ) + nodes_to_add.append( + helper.make_node( + "Reshape", + [concat_out, f"u5d_tr_r1_{counter}"], + [output_name], + name=f"u5d_tr_final_{counter}", + ) + ) + to_remove = [reshape_node, *mid_a, tr, *mid_b, r2] + return to_remove, nodes_to_add + + +def _build_reduce_replacement_or_raise( + model, reshape_node, middle, red, unique_initializer: set, counter: int, cache: GraphCache +) -> tuple[list, list]: + """Replace 5D Reshape → … → Reduce with merge-(n-1,n) → Slice x N → optional Clip/Q→DQ x N → Reduce x N → Concat. + + ``Clip`` / Q→DQ middleware is duplicated **per branch** (after each Slice, before each Reduce) + so that quantisation is applied to each chunk individually, preserving the original semantics. + + Matches ``scratch_space/verify_reduce_5d_vs_4d_decompose.py`` for reduce axis ``n`` in ``{1,2,3,4}``. + Axis ``0`` is not supported (no observed models); raises ``ValueError``. + """ + graph = model.graph + ctx = f"5D Reshape {reshape_node.name!r} → {red.op_type} {red.name!r}" + if len(red.input) < 2: + msg = f"{ctx}: missing axes input (op_type={red.op_type!r})" + raise ValueError(msg) + + axes_name = red.input[1] + ini_axes = cache.get_init(axes_name) + if ini_axes is None: + msg = f"{ctx}: reduce axes {axes_name!r} must be a graph initializer for decomposition" + raise ValueError(msg) + axes_arr = np.asarray(numpy_helper.to_array(ini_axes), dtype=np.int64).reshape(-1) + if axes_arr.size != 1: + msg = f"{ctx}: expected a single reduce axis, got {axes_arr.tolist()}" + raise ValueError(msg) + n = int(axes_arr[0]) + if n < 0: + n += 5 + if n <= 0 or n >= 5: + msg = ( + f"{ctx}: reduce axis must be 1..4 on the 5D Reshape target (axis 0 not supported), " + f"got raw initializer {numpy_helper.to_array(ini_axes).tolist()}" + ) + raise ValueError(msg) + + ini1 = cache.get_init(reshape_node.input[1]) + assert ini1 is not None, f"initializer for node {reshape_node.name} is None" + shape1_raw = np.asarray(numpy_helper.to_array(ini1), dtype=np.int64).reshape(-1) + if shape1_raw.size != 5: + msg = f"{ctx}: first Reshape shape must have length 5, got {shape1_raw.tolist()}" + raise ValueError(msg) + + input_tensor = reshape_node.input[0] + reshape_in_shape = cache.get_shape(input_tensor) + if reshape_in_shape is None: + msg = f"{ctx}: shape of Reshape input {input_tensor!r} is unknown" + raise ValueError(msg) + + # Resolve -1/0 in the shape using the actual input tensor shape + if -1 in shape1_raw or 0 in shape1_raw: + shape1 = np.array(resolve_reshape_shape(shape1_raw, list(reshape_in_shape)), dtype=np.int64) + else: + shape1 = shape1_raw + + r = [int(x) for x in shape1.tolist()] + split_ax = n - 1 + num_splits = r[n - 1] + chunk = r[n] + merged = [*r[:split_ax], r[n - 1] * r[n], *r[n + 1 :]] + if num_splits <= 0 or chunk <= 0: + msg = f"{ctx}: invalid 5D shape {r}" + raise ValueError(msg) + + if 0 not in reshape_in_shape and -1 not in r: + prod_in = int(np.prod(reshape_in_shape)) + prod_r = int(np.prod(r)) + if prod_in != prod_r: + msg = ( + f"{ctx}: Reshape input element count {prod_in} != product of 5D shape {prod_r} " + f"(input shape {list(reshape_in_shape)}, 5D {r})" + ) + raise ValueError(msg) + + output_name = red.output[0] + op_type = red.op_type + reduce_kw: dict = {} + for a in red.attribute: + if a.name == "keepdims": + reduce_kw["keepdims"] = a.i + elif a.name == "noop_with_empty_axes": + reduce_kw["noop_with_empty_axes"] = a.i + + original_keepdims = bool(reduce_kw.get("keepdims", 1)) # ONNX default: 1 + reduce_branch_kw = dict(reduce_kw) + reduce_branch_kw["keepdims"] = 1 + + msh_name = f"u5d_rs_msh_{counter}" + ax4_name = f"u5d_rs_ax4_{counter}" + merged_shape_ini = helper.make_tensor( + name=msh_name, data_type=TensorProto.INT64, dims=[4], vals=merged + ) + # Use negative axis so the reduce op is robust to leading-dim padding. + # For a 4-D tensor, axis k is equivalent to axis k-4 (e.g. 3 → -1). + axes_4d_ini = helper.make_tensor( + name=ax4_name, data_type=TensorProto.INT64, dims=[1], vals=[split_ax - 4] + ) + add_unique_initializers( + graph, + unique_initializer, + [merged_shape_ini, axes_4d_ini], + ) + + merged_out = f"u5d_rs_m4_{counter}" + nodes_to_add: list = [] + nodes_to_add.append( + helper.make_node( + "Reshape", + [input_tensor, msh_name], + [merged_out], + name=f"u5d_rs_reshape4_{counter}", + ) + ) + + # Replace Split with individual Slice nodes — one per chunk along split_ax. + # Middleware (Clip / Q→DQ) is applied per-branch after each Slice, before the Reduce. + sl_axes_name = f"u5d_rs_slax_{counter}" + sl_axes_ini = helper.make_tensor( + name=sl_axes_name, data_type=TensorProto.INT64, dims=[1], vals=[split_ax] + ) + add_unique_initializers(graph, unique_initializer, [sl_axes_ini]) + + reduce_outs: list[str] = [] + for i in range(num_splits): + start = i * chunk + end = start + chunk + sl_starts_name = f"u5d_rs_slst_{i}_{counter}" + sl_ends_name = f"u5d_rs_slen_{i}_{counter}" + sl_starts_ini = helper.make_tensor( + name=sl_starts_name, data_type=TensorProto.INT64, dims=[1], vals=[start] + ) + sl_ends_ini = helper.make_tensor( + name=sl_ends_name, data_type=TensorProto.INT64, dims=[1], vals=[end] + ) + add_unique_initializers(graph, unique_initializer, [sl_starts_ini, sl_ends_ini]) + sl_out = f"u5d_rs_sp_{i}_{counter}" + nodes_to_add.append( + helper.make_node( + "Slice", + [merged_out, sl_starts_name, sl_ends_name, sl_axes_name], + [sl_out], + name=f"u5d_rs_slice_{i}_{counter}", + ) + ) + # Apply per-branch middleware (Clip / Q→DQ) between Slice and Reduce + br_nodes, reduce_in = _duplicate_clip_qdq_chain(middle, sl_out, f"rs{counter}_{i}") + nodes_to_add.extend(br_nodes) + + rs_out = f"u5d_rs_o_{i}_{counter}" + nodes_to_add.append( + helper.make_node( + op_type, + [reduce_in, ax4_name], + [rs_out], + name=f"u5d_rs_{op_type}_{i}_{counter}", + **reduce_branch_kw, + ) + ) + reduce_outs.append(rs_out) + + if original_keepdims: + concat_out = f"u5d_rs_cat_{counter}" + usq_name = f"u5d_rs_usq_{counter}" + unsq_ini = helper.make_tensor( + name=usq_name, data_type=TensorProto.INT64, dims=[1], vals=[n] + ) + add_unique_initializers(graph, unique_initializer, [unsq_ini]) + nodes_to_add.append( + helper.make_node( + "Concat", + reduce_outs, + [concat_out], + axis=split_ax, + name=f"u5d_rs_concat_{counter}", + ) + ) + nodes_to_add.append( + helper.make_node( + "Unsqueeze", + [concat_out, usq_name], + [output_name], + name=f"u5d_rs_unsqueeze_{counter}", + ) + ) + else: + nodes_to_add.append( + helper.make_node( + "Concat", + reduce_outs, + [output_name], + axis=split_ax, + name=f"u5d_rs_concat_{counter}", + ) + ) + + to_remove = [reshape_node, *middle, red] + return to_remove, nodes_to_add + + +def _dispatch_5d_reshape_candidate( + model, reshape_node, unique_initializer: set, counter: int, cache: GraphCache +) -> tuple[list, list]: + """Classify a single 5D Reshape and return ``(nodes_to_remove, nodes_to_add)``.""" + branch, a, b = _classify_5d_reshape_follower(model, reshape_node) + if branch == "transpose": + return _build_transpose_replacement_or_raise( + model, reshape_node, a, b, unique_initializer, counter, cache + ) + return _build_reduce_replacement_or_raise( + model, reshape_node, a, b, unique_initializer, counter, cache + ) + + +def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = False): + """Parse the graph without mutating it, then replace every matching 5D ``Reshape`` subgraph. + + Unless ``skip_symbolic_shape_inference`` is true, runs + :func:`~modelopt.onnx.graph_surgery.utils.symbolic_shape_inference.run_ort_symbolic_shape_inference` + first so intermediate tensor shapes are available where needed. + + Phase 1 walks a snapshot of ``graph.node``, dispatches each 5D ``Reshape``, and records + ``(nodes_to_remove, nodes_to_add)``. Phase 2 runs :func:`replace_nodes_with_topological_order` + for each plan in that same order. + + Classification: after optional ``Clip`` / Q→DQ middleware, the first terminal must be ``Transpose`` + (reshape→transpose→reshape pipeline) or a reduce op in :data:`_REDUCE_OPS`. Any other op raises + ``ValueError`` with ``op_type`` and node ``name``. Pattern validation errors include the same. + """ + if not skip_symbolic_shape_inference: + model = infer_shapes(model) + graph = model.graph + cache = GraphCache(graph) + node_idx = {id(n): i for i, n in enumerate(graph.node)} + unique_initializer: set = set() + counter = 0 + batch: dict[int, tuple[list, list]] = {} + + for node in list(graph.node): + if not _is_5d_reshape_candidate(cache, node): + continue + nodes_to_remove, nodes_to_add = _dispatch_5d_reshape_candidate( + model, node, unique_initializer, counter, cache + ) + idx = node_idx[id(nodes_to_remove[0])] if nodes_to_remove else len(graph.node) + batch[idx] = (nodes_to_remove, nodes_to_add) + counter += 1 + + batch_replace_nodes(graph, batch) + + logger.debug( + "5D reshape to 4D: classified and rewrote %d 5D Reshape subgraph(s) (transpose or reduce path)", + len(batch), + ) + + return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py new file mode 100644 index 00000000000..528e0205ac8 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -0,0 +1,347 @@ +# 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_5d_reshape_to_4d`. + +"""Tests for 5D Reshape to 4D graph surgery (reduce and transpose paths).""" + +from __future__ import annotations + +import copy + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_5d_reshape_to_4d import _apply_5d_reshape_to_4d +from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes + +pytest.importorskip("onnxruntime", reason="onnxruntime required for numerical parity checks") + +# ReduceMin with axes-as-input requires opset 18+. +# Transpose / NoOp models use opset 13 (axes as attribute) — IR version clamped separately. +_REDUCE_OPSET = 18 +_GENERIC_OPSET = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + return numpy_helper.from_array(arr, name=name) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + import onnxruntime as ort + + m = copy.deepcopy(model) + _clamp_ir(m) + sess = ort.InferenceSession( + m.SerializeToString(), + providers=["CPUExecutionProvider"], + ) + out = {} + for o in m.graph.output: + out[o.name] = sess.run([o.name], feeds)[0] + return out + + +def _make_reduce_min_axis4_model() -> onnx.ModelProto: + """X [6,4,4,16] -> Reshape(2,4,12,4,4) -> ReduceMin(axis=4) -> [2,4,12,4].""" + inits = [ + _init("shape5", np.array([2, 4, 12, 4, 4], dtype=np.int64)), + _init("axes_reduce", np.array([4], dtype=np.int64)), + ] + nodes = [ + helper.make_node("Reshape", ["X", "shape5"], ["T"], name="reshape_5d"), + helper.make_node( + "ReduceMin", + ["T", "axes_reduce"], + ["Y"], + name="reduce_min", + keepdims=0, + ), + ] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [6, 4, 4, 16])] + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 4, 12, 4])] + graph = helper.make_graph(nodes, "reduce5d", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _REDUCE_OPSET)]) + onnx.checker.check_model(model) + return model + + +def _make_reduce_with_clip_model() -> onnx.ModelProto: + """Reshape 5D -> Clip -> ReduceMin(axis=4).""" + inits = [ + _init("shape5", np.array([2, 4, 12, 4, 4], dtype=np.int64)), + _init("axes_reduce", np.array([4], dtype=np.int64)), + _init("clip_min", np.array(-1e6, dtype=np.float32).reshape(1)), + _init("clip_max", np.array(1e6, dtype=np.float32).reshape(1)), + ] + nodes = [ + helper.make_node("Reshape", ["X", "shape5"], ["T"], name="reshape_5d"), + helper.make_node("Clip", ["T", "clip_min", "clip_max"], ["Tc"], name="clip_mid"), + helper.make_node( + "ReduceMin", + ["Tc", "axes_reduce"], + ["Y"], + name="reduce_min", + keepdims=0, + ), + ] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [6, 4, 4, 16])] + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 4, 12, 4])] + graph = helper.make_graph(nodes, "reduce_clip", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _REDUCE_OPSET)]) + onnx.checker.check_model(model) + return model + + +def _make_transpose_reshape_model() -> onnx.ModelProto: + """X [6,4,4,16] -> Reshape(2,4,12,4,4) -> Transpose(perm[0]==0) -> Reshape(2,12,16,4).""" + inits = [ + _init("shape5", np.array([2, 4, 12, 4, 4], dtype=np.int64)), + _init("shape_out", np.array([2, 12, 16, 4], dtype=np.int64)), + ] + nodes = [ + helper.make_node("Reshape", ["X", "shape5"], ["T5"], name="reshape_5d"), + helper.make_node( + "Transpose", + ["T5"], + ["Tv"], + name="transpose_mid", + perm=[0, 2, 1, 4, 3], + ), + helper.make_node("Reshape", ["Tv", "shape_out"], ["Y"], name="reshape_out"), + ] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [6, 4, 4, 16])] + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 12, 16, 4])] + graph = helper.make_graph(nodes, "tr5d", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _GENERIC_OPSET)]) + onnx.checker.check_model(model) + return model + + +def _make_reduce_axis0_model() -> onnx.ModelProto: + """Invalid for this pass: ReduceMin axis 0 on 5D.""" + inits = [ + _init("shape5", np.array([2, 4, 12, 4, 4], dtype=np.int64)), + _init("axes_reduce", np.array([0], dtype=np.int64)), + ] + nodes = [ + helper.make_node("Reshape", ["X", "shape5"], ["T"], name="reshape_5d"), + helper.make_node( + "ReduceMin", + ["T", "axes_reduce"], + ["Y"], + name="reduce_min", + keepdims=0, + ), + ] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [6, 4, 4, 16])] + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [4, 12, 4, 4])] + graph = helper.make_graph(nodes, "bad_axis", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _REDUCE_OPSET)]) + onnx.checker.check_model(model) + return model + + +def _make_only_4d_reshape_model() -> onnx.ModelProto: + """Reshape to 4D target — not a 5D candidate.""" + inits = [_init("shape4", np.array([24, 4, 16], dtype=np.int64))] + nodes = [helper.make_node("Reshape", ["X", "shape4"], ["Y"], name="reshape_4d")] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, [6, 4, 4, 16])] + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [24, 4, 16])] + graph = helper.make_graph(nodes, "r4", inputs, outputs, initializer=inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _GENERIC_OPSET)]) + onnx.checker.check_model(model) + return model + + +class TestDla5dReshapeTo4dReduce: + @pytest.mark.parametrize("unsupported", ["axis", "shared"]) + def test_unsafe_qdq_middleware_refuses_before_mutation(self, unsupported): + model = _make_reduce_min_axis4_model() + reshape, reduce = model.graph.node + attrs = {"axis": 4} if unsupported == "axis" else {} + quantize = helper.make_node( + "QuantizeLinear", [reshape.output[0], "scale", "zero"], ["q"], **attrs + ) + dequantize = helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["dq"]) + reduce.input[0] = "dq" + nodes = [reshape, quantize, dequantize, reduce] + if unsupported == "shared": + nodes.append(helper.make_node("Identity", ["q"], ["other"])) + del model.graph.node[:] + model.graph.node.extend(nodes) + model.graph.initializer.extend( + [ + numpy_helper.from_array(np.float32(0.1), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + ] + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"per-tensor|unshared"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + + def test_reduce_min_axis4_ort_matches_reference(self): + orig = _make_reduce_min_axis4_model() + x = np.random.default_rng(0).standard_normal((6, 4, 4, 16)).astype(np.float32) + y_ref = _run_ort(orig, {"X": x})["Y"] + + rewritten = copy.deepcopy(orig) + rewritten = infer_shapes(rewritten) + rewritten = _apply_5d_reshape_to_4d(rewritten) + y_new = _run_ort(rewritten, {"X": x})["Y"] + + assert y_ref.shape == y_new.shape == (2, 4, 12, 4) + np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) + + # Verify Split is NOT present — replaced by Slice nodes + assert not any(n.op_type == "Split" for n in rewritten.graph.node), ( + "Split must be replaced by Slice nodes" + ) + slice_nodes = [n for n in rewritten.graph.node if n.op_type == "Slice"] + assert len(slice_nodes) == 4, ( + f"Expected 4 Slice nodes (one per chunk), got {len(slice_nodes)}" + ) + + def test_reduce_with_clip_ort_matches_reference(self): + orig = _make_reduce_with_clip_model() + x = np.random.default_rng(1).standard_normal((6, 4, 4, 16)).astype(np.float32) + y_ref = _run_ort(orig, {"X": x})["Y"] + + rewritten = copy.deepcopy(orig) + rewritten = infer_shapes(rewritten) + rewritten = _apply_5d_reshape_to_4d(rewritten) + y_new = _run_ort(rewritten, {"X": x})["Y"] + + assert y_ref.shape == y_new.shape + np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) + + assert not any(n.op_type == "Split" for n in rewritten.graph.node), ( + "Split must be replaced by Slice nodes" + ) + + def test_clip_duplicated_per_branch(self): + """Clip is applied once per Slice branch (before each Reduce), not on the merged tensor.""" + orig = _make_reduce_with_clip_model() + x = np.random.default_rng(3).standard_normal((6, 4, 4, 16)).astype(np.float32) + y_ref = _run_ort(orig, {"X": x})["Y"] + + rewritten = copy.deepcopy(orig) + rewritten = infer_shapes(rewritten) + rewritten = _apply_5d_reshape_to_4d(rewritten) + y_new = _run_ort(rewritten, {"X": x})["Y"] + + assert y_ref.shape == y_new.shape + np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) + # num_splits for the test model: shape [2,4,12,4,4], axis=4 → split_ax=3, num_splits=4 + clip_nodes = [n for n in rewritten.graph.node if n.op_type == "Clip"] + num_splits = 4 + assert len(clip_nodes) == num_splits, ( + f"Clip must appear once per branch ({num_splits}), got {len(clip_nodes)}" + ) + + def test_reduce_nodes_use_negative_axis(self): + """Newly created ReduceMin nodes must use a negative axis value (e.g. -1 not 3). + + For the test model: shape [2,4,12,4,4], reduce axis=4 → split_ax=3 → negative axis=-1. + """ + orig = _make_reduce_min_axis4_model() + rewritten = copy.deepcopy(orig) + rewritten = infer_shapes(rewritten) + rewritten = _apply_5d_reshape_to_4d(rewritten) + + reduce_nodes = [n for n in rewritten.graph.node if n.op_type == "ReduceMin"] + assert len(reduce_nodes) > 0, "ReduceMin nodes must be present after rewrite" + + for rn in reduce_nodes: + # axis is stored as input[1] (opset 18 style) — find its initializer value + axes_name = rn.input[1] if len(rn.input) > 1 else None + assert axes_name, f"ReduceMin {rn.name!r} must have axes as input[1]" + axes_init = next((i for i in rewritten.graph.initializer if i.name == axes_name), None) + assert axes_init is not None, f"axes initializer {axes_name!r} not found" + axis_val = int(numpy_helper.to_array(axes_init).flat[0]) + assert axis_val < 0, f"ReduceMin {rn.name!r}: expected negative axis, got {axis_val}" + + def test_reduce_axis_zero_raises(self): + orig = _make_reduce_axis0_model() + with pytest.raises(ValueError, match=r"reduce axis must be 1\.\.4"): + _apply_5d_reshape_to_4d(infer_shapes(copy.deepcopy(orig))) + + +class TestDla5dReshapeTo4dTranspose: + def test_transpose_pipeline_rewritten(self): + """The supported leading-slice layout preserves values and output shape.""" + orig = _make_transpose_reshape_model() + dims = orig.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value = 2 + dims[1].dim_value = 12 + values = np.random.default_rng(5).standard_normal((2, 12, 4, 16)).astype(np.float32) + expected = _run_ort(orig, {"X": values})["Y"] + rewritten = _apply_5d_reshape_to_4d(infer_shapes(copy.deepcopy(orig))) + actual = _run_ort(rewritten, {"X": values})["Y"] + np.testing.assert_array_equal(actual, expected) + + # 5D Reshape must be removed + assert not any(n.name == "reshape_5d" for n in rewritten.graph.node), ( + "5D Reshape must be removed by the transform" + ) + # Slice nodes inserted by decomposition + assert any(n.op_type == "Slice" for n in rewritten.graph.node), ( + "Slice nodes must be inserted" + ) + # Transpose preserved (one per slice) + assert any(n.op_type == "Transpose" for n in rewritten.graph.node), ( + "Transpose nodes must remain" + ) + # Concat stitches results + assert any(n.op_type == "Concat" for n in rewritten.graph.node), ( + "Concat node must be present" + ) + + def test_incompatible_leading_slice_layout_refuses_before_mutation(self): + model = infer_shapes(_make_transpose_reshape_model()) + before = model.SerializeToString() + with pytest.raises(ValueError, match="leading input dimension"): + _apply_5d_reshape_to_4d(model) + assert model.SerializeToString() == before + + +class TestDla5dReshapeTo4dNoOp: + def test_no_5d_reshape_leaves_graph_unchanged(self): + orig = _make_only_4d_reshape_model() + x = np.random.default_rng(4).standard_normal((6, 4, 4, 16)).astype(np.float32) + y_ref = _run_ort(orig, {"X": x})["Y"] + + before = [n.SerializeToString() for n in orig.graph.node] + rewritten = copy.deepcopy(orig) + rewritten = infer_shapes(rewritten) + rewritten = _apply_5d_reshape_to_4d(rewritten) + after = [n.SerializeToString() for n in rewritten.graph.node] + + y_new = _run_ort(rewritten, {"X": x})["Y"] + + assert before == after + assert y_ref.shape == y_new.shape + np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) From 5cc4b429f286d5dbf9cc490c422f227f3a17241f Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 16:12:14 +0000 Subject: [PATCH 02/15] Fix DLA 5D reshape planning and output shapes Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 71 +++++++--- .../test_dla_5d_reshape_to_4d.py | 133 +++++++++++------- 2 files changed, 133 insertions(+), 71 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 7569b6675f0..3e979ebf0fc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -43,6 +43,7 @@ def _consumers_first_input(graph, tensor_name: str) -> list: + """Find consumers using the tensor as their data input.""" return [n for n in graph.node if n.input and len(n.input) > 0 and n.input[0] == tensor_name] @@ -105,6 +106,7 @@ def _parse_middleware_strict( def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: + """Identify Reshape nodes with a constant five-dimensional target.""" if node.op_type != "Reshape": return False ini = cache.get_init(node.input[1]) @@ -243,7 +245,8 @@ def _build_transpose_replacement_or_raise( msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" raise ValueError(msg) ini2 = cache.get_init(r2.input[1]) - assert ini2 is not None, f"initializer for node {r2.name} is None" + if ini2 is None: + raise ValueError(f"{ctx}: final Reshape shape {r2.input[1]!r} must be a graph initializer") ini1 = cache.get_init(reshape_node.input[1]) assert ini1 is not None, f"initializer for node {reshape_node.name} is None" @@ -261,6 +264,14 @@ def _build_transpose_replacement_or_raise( shape1 = np.array(resolve_reshape_shape(shape1_raw, list(tr_in_shape)), dtype=np.int64) else: shape1 = shape1_raw + if 0 in shape2: + if any(a.name == "allowzero" and a.i for a in r2.attribute): + raise ValueError(f"{ctx}: literal zero dimensions are not supported") + # Copied dimensions refer to the original 5D Transpose, before Concat merges axes. + shape2 = np.array( + [shape1[transpose_perm[i]] if dim == 0 else dim for i, dim in enumerate(shape2)], + dtype=np.int64, + ) num_slices = int(shape1[0]) reshape_4d_vals = shape1[1:] if reshape_4d_vals.size != 4: @@ -591,38 +602,56 @@ def _dispatch_5d_reshape_candidate( def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = False): - """Parse the graph without mutating it, then replace every matching 5D ``Reshape`` subgraph. + """Replace matching 5D ``Reshape`` subgraphs, restoring the model if planning fails. Unless ``skip_symbolic_shape_inference`` is true, runs :func:`~modelopt.onnx.graph_surgery.utils.symbolic_shape_inference.run_ort_symbolic_shape_inference` first so intermediate tensor shapes are available where needed. - Phase 1 walks a snapshot of ``graph.node``, dispatches each 5D ``Reshape``, and records - ``(nodes_to_remove, nodes_to_add)``. Phase 2 runs :func:`replace_nodes_with_topological_order` - for each plan in that same order. + Plan disjoint replacements before applying them with :func:`batch_replace_nodes`. Classification: after optional ``Clip`` / Q→DQ middleware, the first terminal must be ``Transpose`` (reshape→transpose→reshape pipeline) or a reduce op in :data:`_REDUCE_OPS`. Any other op raises ``ValueError`` with ``op_type`` and node ``name``. Pattern validation errors include the same. """ - if not skip_symbolic_shape_inference: - model = infer_shapes(model) graph = model.graph - cache = GraphCache(graph) - node_idx = {id(n): i for i, n in enumerate(graph.node)} - unique_initializer: set = set() - counter = 0 - batch: dict[int, tuple[list, list]] = {} - - for node in list(graph.node): - if not _is_5d_reshape_candidate(cache, node): - continue - nodes_to_remove, nodes_to_add = _dispatch_5d_reshape_candidate( - model, node, unique_initializer, counter, cache + initializer_count = len(graph.initializer) + metadata = None + if not skip_symbolic_shape_inference: + metadata = ( + list(graph.value_info), + [output.type.SerializeToString() for output in graph.output], ) - idx = node_idx[id(nodes_to_remove[0])] if nodes_to_remove else len(graph.node) - batch[idx] = (nodes_to_remove, nodes_to_add) - counter += 1 + batch: dict[int, tuple[list, list]] = {} + consumed: set[int] = set() + try: + if metadata is not None: + infer_shapes(model) + for output, output_type in zip(graph.output, metadata[1]): + if output.type.HasField("tensor_type") and not output.type.tensor_type.HasField( + "shape" + ): + output.type.ParseFromString(output_type) + cache = GraphCache(graph) + unique_initializer: set = set() + for idx, node in enumerate(graph.node): + if not _is_5d_reshape_candidate(cache, node): + continue + if id(node) in consumed: + raise ValueError(f"5D Reshape {node.name!r} overlaps another 5D Reshape rewrite") + nodes_to_remove, nodes_to_add = _dispatch_5d_reshape_candidate( + model, node, unique_initializer, len(batch), cache + ) + consumed.update(id(removed) for removed in nodes_to_remove) + batch[idx] = (nodes_to_remove, nodes_to_add) + except Exception: + del graph.initializer[initializer_count:] + if metadata is not None: + del graph.value_info[:] + graph.value_info.extend(metadata[0]) + for output, output_type in zip(graph.output, metadata[1]): + output.type.ParseFromString(output_type) + raise batch_replace_nodes(graph, batch) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 528e0205ac8..ec36a836ba7 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -32,7 +32,7 @@ from modelopt.onnx.graph_surgery.dla_transforms.dla_5d_reshape_to_4d import _apply_5d_reshape_to_4d from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes -pytest.importorskip("onnxruntime", reason="onnxruntime required for numerical parity checks") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for numerical parity checks") # ReduceMin with axes-as-input requires opset 18+. # Transpose / NoOp models use opset 13 (axes as attribute) — IR version clamped separately. @@ -42,16 +42,17 @@ def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + """Create a named initializer from an array.""" return numpy_helper.from_array(arr, name=name) def _clamp_ir(model: onnx.ModelProto) -> None: + """Use an IR version supported by the test runtime.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _run_ort(model: onnx.ModelProto, feeds: dict[str, np.ndarray]) -> dict[str, np.ndarray]: - import onnxruntime as ort - + """Execute all graph outputs with the CPU runtime.""" m = copy.deepcopy(model) _clamp_ir(m) sess = ort.InferenceSession( @@ -179,6 +180,7 @@ def _make_only_4d_reshape_model() -> onnx.ModelProto: class TestDla5dReshapeTo4dReduce: @pytest.mark.parametrize("unsupported", ["axis", "shared"]) def test_unsafe_qdq_middleware_refuses_before_mutation(self, unsupported): + """Refuse middleware whose branching or axes cannot be preserved.""" model = _make_reduce_min_axis4_model() reshape, reduce = model.graph.node attrs = {"axis": 4} if unsupported == "axis" else {} @@ -204,13 +206,14 @@ def test_unsafe_qdq_middleware_refuses_before_mutation(self, unsupported): assert model.SerializeToString() == before def test_reduce_min_axis4_ort_matches_reference(self): + """Preserve reduction values while generating slices and negative axes.""" orig = _make_reduce_min_axis4_model() x = np.random.default_rng(0).standard_normal((6, 4, 4, 16)).astype(np.float32) y_ref = _run_ort(orig, {"X": x})["Y"] rewritten = copy.deepcopy(orig) - rewritten = infer_shapes(rewritten) rewritten = _apply_5d_reshape_to_4d(rewritten) + onnx.checker.check_model(rewritten) y_new = _run_ort(rewritten, {"X": x})["Y"] assert y_ref.shape == y_new.shape == (2, 4, 12, 4) @@ -224,73 +227,101 @@ def test_reduce_min_axis4_ort_matches_reference(self): assert len(slice_nodes) == 4, ( f"Expected 4 Slice nodes (one per chunk), got {len(slice_nodes)}" ) + initializers = {init.name: init for init in rewritten.graph.initializer} + for node in rewritten.graph.node: + if node.op_type == "ReduceMin": + assert numpy_helper.to_array(initializers[node.input[1]]).tolist() == [-1] def test_reduce_with_clip_ort_matches_reference(self): + """Preserve clipped reduction values with one Clip per branch.""" orig = _make_reduce_with_clip_model() x = np.random.default_rng(1).standard_normal((6, 4, 4, 16)).astype(np.float32) y_ref = _run_ort(orig, {"X": x})["Y"] rewritten = copy.deepcopy(orig) - rewritten = infer_shapes(rewritten) rewritten = _apply_5d_reshape_to_4d(rewritten) + onnx.checker.check_model(rewritten) y_new = _run_ort(rewritten, {"X": x})["Y"] assert y_ref.shape == y_new.shape np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) - assert not any(n.op_type == "Split" for n in rewritten.graph.node), ( - "Split must be replaced by Slice nodes" - ) - - def test_clip_duplicated_per_branch(self): - """Clip is applied once per Slice branch (before each Reduce), not on the merged tensor.""" - orig = _make_reduce_with_clip_model() - x = np.random.default_rng(3).standard_normal((6, 4, 4, 16)).astype(np.float32) - y_ref = _run_ort(orig, {"X": x})["Y"] - - rewritten = copy.deepcopy(orig) - rewritten = infer_shapes(rewritten) - rewritten = _apply_5d_reshape_to_4d(rewritten) - y_new = _run_ort(rewritten, {"X": x})["Y"] - - assert y_ref.shape == y_new.shape - np.testing.assert_allclose(y_ref, y_new, rtol=0, atol=0) - # num_splits for the test model: shape [2,4,12,4,4], axis=4 → split_ax=3, num_splits=4 - clip_nodes = [n for n in rewritten.graph.node if n.op_type == "Clip"] - num_splits = 4 - assert len(clip_nodes) == num_splits, ( - f"Clip must appear once per branch ({num_splits}), got {len(clip_nodes)}" - ) - - def test_reduce_nodes_use_negative_axis(self): - """Newly created ReduceMin nodes must use a negative axis value (e.g. -1 not 3). - - For the test model: shape [2,4,12,4,4], reduce axis=4 → split_ax=3 → negative axis=-1. - """ - orig = _make_reduce_min_axis4_model() - rewritten = copy.deepcopy(orig) - rewritten = infer_shapes(rewritten) - rewritten = _apply_5d_reshape_to_4d(rewritten) - - reduce_nodes = [n for n in rewritten.graph.node if n.op_type == "ReduceMin"] - assert len(reduce_nodes) > 0, "ReduceMin nodes must be present after rewrite" - - for rn in reduce_nodes: - # axis is stored as input[1] (opset 18 style) — find its initializer value - axes_name = rn.input[1] if len(rn.input) > 1 else None - assert axes_name, f"ReduceMin {rn.name!r} must have axes as input[1]" - axes_init = next((i for i in rewritten.graph.initializer if i.name == axes_name), None) - assert axes_init is not None, f"axes initializer {axes_name!r} not found" - axis_val = int(numpy_helper.to_array(axes_init).flat[0]) - assert axis_val < 0, f"ReduceMin {rn.name!r}: expected negative axis, got {axis_val}" + assert sum(node.op_type == "Clip" for node in rewritten.graph.node) == 4 def test_reduce_axis_zero_raises(self): + """Refuse the unsupported batch-axis reduction.""" orig = _make_reduce_axis0_model() with pytest.raises(ValueError, match=r"reduce axis must be 1\.\.4"): _apply_5d_reshape_to_4d(infer_shapes(copy.deepcopy(orig))) class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize("shape", [[0, -1], [0, 0, 4, -1], [0, 0, 0, -1]]) + def test_final_reshape_zero_dimensions_preserve_values(self, shape): + """Resolve copied dimensions against the original five-dimensional input.""" + model = _make_transpose_reshape_model() + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + model.graph.initializer[1].CopyFrom(_init("shape_out", np.array(shape, np.int64))) + resolved = [2, 12, 4, 4, 4] + target = [resolved[i] if dim == 0 else dim for i, dim in enumerate(shape)] + target = np.empty((2, 12, 4, 16)).reshape(target).shape + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Y", TensorProto.FLOAT, target) + ) + values = np.random.default_rng(7).standard_normal((2, 12, 4, 16)).astype(np.float32) + expected = _run_ort(model, {"X": values})["Y"] + rewritten = _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + onnx.checker.check_model(rewritten) + np.testing.assert_array_equal(_run_ort(rewritten, {"X": values})["Y"], expected) + + @pytest.mark.parametrize("unsupported", ["dynamic", "overlap", "late_failure"]) + @pytest.mark.parametrize("skip_inference", [True, False]) + def test_unsupported_patterns_leave_model_unchanged(self, unsupported, skip_inference): + """Refusal preserves nodes, initializers and inferred metadata, even on retry.""" + model = _make_transpose_reshape_model() + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + if unsupported == "dynamic": + del model.graph.initializer[1] + model.graph.input.append( + helper.make_tensor_value_info("shape_out", TensorProto.INT64, [4]) + ) + elif unsupported == "overlap": + model.graph.initializer[1].CopyFrom( + _init("shape_out", np.array([2, 12, 4, 4, 4], np.int64)) + ) + model.graph.initializer.append(_init("axes", np.array([4], np.int64))) + model.graph.node[-1].output[0] = "second_5d" + model.graph.node.append( + helper.make_node("ReduceMin", ["second_5d", "axes"], ["Y"], keepdims=0) + ) + model.opset_import[0].version = _REDUCE_OPSET + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 12, 4, 4]) + ) + else: + model.graph.initializer.append( + _init("late_shape", np.array([2, 4, 12, 4, 4], np.int64)) + ) + model.graph.node.extend( + [ + helper.make_node( + "Reshape", ["X", "late_shape"], ["late_5d"], name="late_reshape" + ), + helper.make_node("Identity", ["late_5d"], ["late_output"], name="unsupported"), + ] + ) + model.graph.output.append( + helper.make_tensor_value_info("late_output", TensorProto.FLOAT, [2, 4, 12, 4, 4]) + ) + onnx.checker.check_model(model) + before = model.SerializeToString() + for _ in range(2): + with pytest.raises(ValueError): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + assert model.SerializeToString() == before + def test_transpose_pipeline_rewritten(self): """The supported leading-slice layout preserves values and output shape.""" orig = _make_transpose_reshape_model() @@ -321,6 +352,7 @@ def test_transpose_pipeline_rewritten(self): ) def test_incompatible_leading_slice_layout_refuses_before_mutation(self): + """Refuse a leading dimension that does not match the slice count.""" model = infer_shapes(_make_transpose_reshape_model()) before = model.SerializeToString() with pytest.raises(ValueError, match="leading input dimension"): @@ -330,6 +362,7 @@ def test_incompatible_leading_slice_layout_refuses_before_mutation(self): class TestDla5dReshapeTo4dNoOp: def test_no_5d_reshape_leaves_graph_unchanged(self): + """Leave unrelated reshape nodes and their values unchanged.""" orig = _make_only_4d_reshape_model() x = np.random.default_rng(4).standard_normal((6, 4, 4, 16)).astype(np.float32) y_ref = _run_ort(orig, {"X": x})["Y"] From 8b2895b60f8133b68914af7271864acddf375227 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 16:35:52 +0000 Subject: [PATCH 03/15] Reserve generated names in DLA 5D reshape rewrites Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 99 ++++++++++--------- .../test_dla_5d_reshape_to_4d.py | 78 +++++++++++++++ 2 files changed, 132 insertions(+), 45 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 3e979ebf0fc..1336ef58a8f 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -132,7 +132,9 @@ def _classify_5d_reshape_follower(model, reshape_node) -> tuple: raise ValueError(msg) -def _duplicate_clip_qdq_chain(middleware: list, x_in: str, suffix: str) -> tuple[list, str]: +def _duplicate_clip_qdq_chain( + middleware: list, x_in: str, suffix: str, cache: GraphCache +) -> tuple[list, str]: """Rebuild Clip / Q→DQ sequence on ``x_in`` with fresh tensor names. Returns ``(new_nodes, output_name)``.""" if not middleware: return [], x_in @@ -143,10 +145,12 @@ def _duplicate_clip_qdq_chain(middleware: list, x_in: str, suffix: str) -> tuple while m < len(middleware): n = middleware[m] if n.op_type == "Clip": - new_out = f"_u5d_{suffix}_c{m}_{i}" + new_out = cache.unique_name(f"_u5d_{suffix}_c{m}_{i}") inputs = [out, *list(n.input[1:])] new_nodes.append( - helper.make_node("Clip", inputs, [new_out], name=f"{n.name}_u5d_{suffix}_{m}") + helper.make_node( + "Clip", inputs, [new_out], name=cache.unique_name(f"{n.name}_u5d_{suffix}_{m}") + ) ) out = new_out m += 1 @@ -156,14 +160,14 @@ def _duplicate_clip_qdq_chain(middleware: list, x_in: str, suffix: str) -> tuple msg = "QuantizeLinear without following DequantizeLinear in middleware chain" raise ValueError(msg) q, dq = middleware[m], middleware[m + 1] - q_out = f"_u5d_{suffix}_q{m}_{i}" - dq_out = f"_u5d_{suffix}_dq{m}_{i}" + q_out = cache.unique_name(f"_u5d_{suffix}_q{m}_{i}") + dq_out = cache.unique_name(f"_u5d_{suffix}_dq{m}_{i}") new_nodes.append( helper.make_node( "QuantizeLinear", [out, *q.input[1:]], [q_out], - name=f"{q.name}_u5d_{suffix}_{m}", + name=cache.unique_name(f"{q.name}_u5d_{suffix}_{m}"), **{a.name: helper.get_attribute_value(a) for a in q.attribute}, ) ) @@ -172,7 +176,7 @@ def _duplicate_clip_qdq_chain(middleware: list, x_in: str, suffix: str) -> tuple "DequantizeLinear", [q_out, *dq.input[1:]], [dq_out], - name=f"{dq.name}_u5d_{suffix}_{m}", + name=cache.unique_name(f"{dq.name}_u5d_{suffix}_{m}"), **{a.name: helper.get_attribute_value(a) for a in dq.attribute}, ) ) @@ -286,7 +290,7 @@ def _build_transpose_replacement_or_raise( concat_axis = _transpose_path_concat_axis_4d(slice_axis, transpose_perm) slice_axes = helper.make_tensor( - name=f"u5d_tr_slice_axes_{counter}", + name=cache.unique_name(f"u5d_tr_slice_axes_{counter}"), data_type=TensorProto.INT64, dims=[1], vals=[slice_axis], @@ -296,12 +300,15 @@ def _build_transpose_replacement_or_raise( for i in range(num_slices): slice_starts.append( helper.make_tensor( - name=f"u5d_tr_{i}_starts_{counter}", data_type=TensorProto.INT64, dims=[1], vals=[i] + name=cache.unique_name(f"u5d_tr_{i}_starts_{counter}"), + data_type=TensorProto.INT64, + dims=[1], + vals=[i], ) ) slice_ends.append( helper.make_tensor( - name=f"u5d_tr_{i}_ends_{counter}", + name=cache.unique_name(f"u5d_tr_{i}_ends_{counter}"), data_type=TensorProto.INT64, dims=[1], vals=[i + 1], @@ -310,14 +317,14 @@ def _build_transpose_replacement_or_raise( add_unique_initializers(graph, unique_initializer, [slice_axes, *slice_starts, *slice_ends]) reshape_4d_shape = helper.make_tensor( - name=f"u5d_tr_r4_{counter}", + name=cache.unique_name(f"u5d_tr_r4_{counter}"), data_type=TensorProto.INT64, dims=[4], vals=reshape_4d_vals.tolist(), ) final_rank = int(shape2.size) reshape_final_shape = helper.make_tensor( - name=f"u5d_tr_r1_{counter}", + name=cache.unique_name(f"u5d_tr_r1_{counter}"), data_type=TensorProto.INT64, dims=[final_rank], vals=shape2.tolist(), @@ -327,61 +334,63 @@ def _build_transpose_replacement_or_raise( nodes_to_add: list = [] slice_outputs: list[str] = [] for i in range(num_slices): - slice_out = f"u5d_tr_slice_{i}_{counter}" + slice_out = cache.unique_name(f"u5d_tr_slice_{i}_{counter}") nodes_to_add.append( helper.make_node( "Slice", [ input_tensor, - f"u5d_tr_{i}_starts_{counter}", - f"u5d_tr_{i}_ends_{counter}", - f"u5d_tr_slice_axes_{counter}", + slice_starts[i].name, + slice_ends[i].name, + slice_axes.name, ], [slice_out], - name=f"u5d_tr_slice_{i}_{counter}", + name=slice_out, ) ) - reshape_4d_out = f"u5d_tr_r4o_{i}_{counter}" + reshape_4d_out = cache.unique_name(f"u5d_tr_r4o_{i}_{counter}") nodes_to_add.append( helper.make_node( "Reshape", - [slice_out, f"u5d_tr_r4_{counter}"], + [slice_out, reshape_4d_shape.name], [reshape_4d_out], - name=f"u5d_tr_r4n_{i}_{counter}", + name=cache.unique_name(f"u5d_tr_r4n_{i}_{counter}"), ) ) - br_a, t_a = _duplicate_clip_qdq_chain(mid_a, reshape_4d_out, f"tr{counter}_{i}a") + br_a, t_a = _duplicate_clip_qdq_chain(mid_a, reshape_4d_out, f"tr{counter}_{i}a", cache) nodes_to_add.extend(br_a) - transpose_branch_out = f"u5d_tr_t_{i}_{counter}" + transpose_branch_out = cache.unique_name(f"u5d_tr_t_{i}_{counter}") nodes_to_add.append( helper.make_node( "Transpose", [t_a], [transpose_branch_out], perm=new_transpose_perm, - name=f"u5d_tr_tn_{i}_{counter}", + name=cache.unique_name(f"u5d_tr_tn_{i}_{counter}"), ) ) - br_b, t_b = _duplicate_clip_qdq_chain(mid_b, transpose_branch_out, f"tr{counter}_{i}b") + br_b, t_b = _duplicate_clip_qdq_chain( + mid_b, transpose_branch_out, f"tr{counter}_{i}b", cache + ) nodes_to_add.extend(br_b) slice_outputs.append(t_b) - concat_out = f"u5d_tr_concat_{counter}" + concat_out = cache.unique_name(f"u5d_tr_concat_{counter}") nodes_to_add.append( helper.make_node( "Concat", slice_outputs, [concat_out], axis=concat_axis, - name=f"u5d_tr_concatn_{counter}", + name=cache.unique_name(f"u5d_tr_concatn_{counter}"), ) ) nodes_to_add.append( helper.make_node( "Reshape", - [concat_out, f"u5d_tr_r1_{counter}"], + [concat_out, reshape_final_shape.name], [output_name], - name=f"u5d_tr_final_{counter}", + name=cache.unique_name(f"u5d_tr_final_{counter}"), ) ) to_remove = [reshape_node, *mid_a, tr, *mid_b, r2] @@ -475,8 +484,8 @@ def _build_reduce_replacement_or_raise( reduce_branch_kw = dict(reduce_kw) reduce_branch_kw["keepdims"] = 1 - msh_name = f"u5d_rs_msh_{counter}" - ax4_name = f"u5d_rs_ax4_{counter}" + msh_name = cache.unique_name(f"u5d_rs_msh_{counter}") + ax4_name = cache.unique_name(f"u5d_rs_ax4_{counter}") merged_shape_ini = helper.make_tensor( name=msh_name, data_type=TensorProto.INT64, dims=[4], vals=merged ) @@ -491,20 +500,20 @@ def _build_reduce_replacement_or_raise( [merged_shape_ini, axes_4d_ini], ) - merged_out = f"u5d_rs_m4_{counter}" + merged_out = cache.unique_name(f"u5d_rs_m4_{counter}") nodes_to_add: list = [] nodes_to_add.append( helper.make_node( "Reshape", [input_tensor, msh_name], [merged_out], - name=f"u5d_rs_reshape4_{counter}", + name=cache.unique_name(f"u5d_rs_reshape4_{counter}"), ) ) # Replace Split with individual Slice nodes — one per chunk along split_ax. # Middleware (Clip / Q→DQ) is applied per-branch after each Slice, before the Reduce. - sl_axes_name = f"u5d_rs_slax_{counter}" + sl_axes_name = cache.unique_name(f"u5d_rs_slax_{counter}") sl_axes_ini = helper.make_tensor( name=sl_axes_name, data_type=TensorProto.INT64, dims=[1], vals=[split_ax] ) @@ -514,8 +523,8 @@ def _build_reduce_replacement_or_raise( for i in range(num_splits): start = i * chunk end = start + chunk - sl_starts_name = f"u5d_rs_slst_{i}_{counter}" - sl_ends_name = f"u5d_rs_slen_{i}_{counter}" + sl_starts_name = cache.unique_name(f"u5d_rs_slst_{i}_{counter}") + sl_ends_name = cache.unique_name(f"u5d_rs_slen_{i}_{counter}") sl_starts_ini = helper.make_tensor( name=sl_starts_name, data_type=TensorProto.INT64, dims=[1], vals=[start] ) @@ -523,34 +532,34 @@ def _build_reduce_replacement_or_raise( name=sl_ends_name, data_type=TensorProto.INT64, dims=[1], vals=[end] ) add_unique_initializers(graph, unique_initializer, [sl_starts_ini, sl_ends_ini]) - sl_out = f"u5d_rs_sp_{i}_{counter}" + sl_out = cache.unique_name(f"u5d_rs_sp_{i}_{counter}") nodes_to_add.append( helper.make_node( "Slice", [merged_out, sl_starts_name, sl_ends_name, sl_axes_name], [sl_out], - name=f"u5d_rs_slice_{i}_{counter}", + name=cache.unique_name(f"u5d_rs_slice_{i}_{counter}"), ) ) # Apply per-branch middleware (Clip / Q→DQ) between Slice and Reduce - br_nodes, reduce_in = _duplicate_clip_qdq_chain(middle, sl_out, f"rs{counter}_{i}") + br_nodes, reduce_in = _duplicate_clip_qdq_chain(middle, sl_out, f"rs{counter}_{i}", cache) nodes_to_add.extend(br_nodes) - rs_out = f"u5d_rs_o_{i}_{counter}" + rs_out = cache.unique_name(f"u5d_rs_o_{i}_{counter}") nodes_to_add.append( helper.make_node( op_type, [reduce_in, ax4_name], [rs_out], - name=f"u5d_rs_{op_type}_{i}_{counter}", + name=cache.unique_name(f"u5d_rs_{op_type}_{i}_{counter}"), **reduce_branch_kw, ) ) reduce_outs.append(rs_out) if original_keepdims: - concat_out = f"u5d_rs_cat_{counter}" - usq_name = f"u5d_rs_usq_{counter}" + concat_out = cache.unique_name(f"u5d_rs_cat_{counter}") + usq_name = cache.unique_name(f"u5d_rs_usq_{counter}") unsq_ini = helper.make_tensor( name=usq_name, data_type=TensorProto.INT64, dims=[1], vals=[n] ) @@ -561,7 +570,7 @@ def _build_reduce_replacement_or_raise( reduce_outs, [concat_out], axis=split_ax, - name=f"u5d_rs_concat_{counter}", + name=cache.unique_name(f"u5d_rs_concat_{counter}"), ) ) nodes_to_add.append( @@ -569,7 +578,7 @@ def _build_reduce_replacement_or_raise( "Unsqueeze", [concat_out, usq_name], [output_name], - name=f"u5d_rs_unsqueeze_{counter}", + name=cache.unique_name(f"u5d_rs_unsqueeze_{counter}"), ) ) else: @@ -579,7 +588,7 @@ def _build_reduce_replacement_or_raise( reduce_outs, [output_name], axis=split_ax, - name=f"u5d_rs_concat_{counter}", + name=cache.unique_name(f"u5d_rs_concat_{counter}"), ) ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index ec36a836ba7..5d1f867ca53 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -28,6 +28,7 @@ import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms.dla_5d_reshape_to_4d import _apply_5d_reshape_to_4d from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes @@ -256,6 +257,83 @@ def test_reduce_axis_zero_raises(self): class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize("path", ["transpose", "reduce"]) + def test_generated_names_preserve_existing_controls_and_branches(self, path): + """Reserve generated names without retargeting existing data or QDQ controls.""" + transpose = path == "transpose" + model = _make_transpose_reshape_model() if transpose else _make_reduce_with_clip_model() + model.opset_import[0].version = 21 + prefix = "tr" if transpose else "rs" + if transpose: + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + clip = helper.make_node( + "Clip", ["T5", "clip_min", "clip_max"], ["clipped"], name="clip" + ) + model.graph.node.insert(1, clip) + model.graph.initializer.extend( + [ + _init("clip_min", np.array(-1, np.float32)), + _init("clip_max", np.array(1, np.float32)), + ] + ) + clip_output = model.graph.node[1].output[0] + model.graph.node[2].input[0] = "dequantized" + scale_name = "u5d_tr_r1_0" if transpose else "u5d_rs_slax_0" + suffix = "tr0_0a" if transpose else "rs0_0" + zero_name = f"_u5d_{suffix}_q1_1" + model.graph.node.insert( + 2, + helper.make_node( + "QuantizeLinear", [clip_output, scale_name, zero_name], ["quantized"], name="q" + ), + ) + model.graph.node.insert( + 3, + helper.make_node( + "DequantizeLinear", ["quantized", scale_name, zero_name], ["dequantized"], name="dq" + ), + ) + model.graph.initializer.extend( + [ + _init(scale_name, np.array(0.25, np.float32)), + _init(zero_name, np.array(0, np.uint8)), + ] + ) + renames = { + "X": "u5d_tr_r4_0" if transpose else "u5d_rs_msh_0", + "shape5": "u5d_tr_slice_axes_0" if transpose else "u5d_rs_ax4_0", + "clip_min": f"_u5d_{suffix}_c0_0", + } + for value in [*model.graph.input, *model.graph.initializer]: + value.name = renames.get(value.name, value.name) + for node in model.graph.node: + for index, name in enumerate(node.input): + node.input[index] = renames.get(name, name) + branch = "u5d_tr_r4o_0_0" if transpose else "u5d_rs_sp_0_0" + model.graph.node.append( + helper.make_node("Identity", [renames["X"]], [branch], name=f"u5d_{prefix}_slice_0_0") + ) + branch_info = copy.deepcopy(model.graph.input[0]) + branch_info.name = branch + model.graph.output.append(branch_info) + onnx.checker.check_model(model, full_check=True) + controls = {value.name: value.SerializeToString() for value in model.graph.initializer} + shape = [dim.dim_value for dim in model.graph.input[0].type.tensor_type.shape.dim] + values = np.random.default_rng(9).standard_normal(shape).astype(np.float32) + feeds = {renames["X"]: values} + expected = ReferenceEvaluator(model).run(None, feeds) + rewritten = _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + onnx.checker.check_model(rewritten, full_check=True) + assert len({init.name for init in rewritten.graph.initializer}) == len( + rewritten.graph.initializer + ) + for init in rewritten.graph.initializer: + if init.name in controls: + assert init.SerializeToString() == controls[init.name] + for actual, reference in zip(ReferenceEvaluator(rewritten).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + @pytest.mark.parametrize("shape", [[0, -1], [0, 0, 4, -1], [0, 0, 0, -1]]) def test_final_reshape_zero_dimensions_preserve_values(self, shape): """Resolve copied dimensions against the original five-dimensional input.""" From 5ac4316cc6890eec6a15bae37ec5c8177de95b23 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 16:39:42 +0000 Subject: [PATCH 04/15] Preserve Clip attributes in DLA reshape branches Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 6 +++++- .../dla_transforms/test_dla_5d_reshape_to_4d.py | 10 +++++++++- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 1336ef58a8f..f20f6b54977 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -149,7 +149,11 @@ def _duplicate_clip_qdq_chain( inputs = [out, *list(n.input[1:])] new_nodes.append( helper.make_node( - "Clip", inputs, [new_out], name=cache.unique_name(f"{n.name}_u5d_{suffix}_{m}") + "Clip", + inputs, + [new_out], + name=cache.unique_name(f"{n.name}_u5d_{suffix}_{m}"), + **{a.name: helper.get_attribute_value(a) for a in n.attribute}, ) ) out = new_out diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 5d1f867ca53..98ed02f164e 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -400,12 +400,20 @@ def test_unsupported_patterns_leave_model_unchanged(self, unsupported, skip_infe _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) assert model.SerializeToString() == before - def test_transpose_pipeline_rewritten(self): + @pytest.mark.parametrize("legacy_clip", [False, True]) + def test_transpose_pipeline_rewritten(self, legacy_clip): """The supported leading-slice layout preserves values and output shape.""" orig = _make_transpose_reshape_model() dims = orig.graph.input[0].type.tensor_type.shape.dim dims[0].dim_value = 2 dims[1].dim_value = 12 + if legacy_clip: + orig.opset_import[0].version = 10 + orig.graph.node[1].input[0] = "clipped" + orig.graph.node.insert( + 1, helper.make_node("Clip", ["T5"], ["clipped"], name="clip", min=-0.5, max=0.5) + ) + onnx.checker.check_model(orig, full_check=True) values = np.random.default_rng(5).standard_normal((2, 12, 4, 16)).astype(np.float32) expected = _run_ort(orig, {"X": values})["Y"] rewritten = _apply_5d_reshape_to_4d(infer_shapes(copy.deepcopy(orig))) From a86d0e7d89302e6e96af2a0de9a71ee94fd2fcb6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 17:00:16 +0000 Subject: [PATCH 05/15] Order DLA reshape branches after dynamic controls Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 19 ++++++++ .../test_dla_5d_reshape_to_4d.py | 47 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index f20f6b54977..f4fdc11af23 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -38,6 +38,7 @@ infer_shapes, resolve_reshape_shape, ) +from .dla_reorder_nodes import priority_topological_order _REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) @@ -637,6 +638,7 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals ) batch: dict[int, tuple[list, list]] = {} consumed: set[int] = set() + needs_reorder = False try: if metadata is not None: infer_shapes(model) @@ -646,6 +648,12 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals ): output.type.ParseFromString(output_type) cache = GraphCache(graph) + producer_positions = { + output: index + for index, node in enumerate(graph.node) + for output in node.output + if output + } unique_initializer: set = set() for idx, node in enumerate(graph.node): if not _is_5d_reshape_candidate(cache, node): @@ -655,6 +663,11 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals nodes_to_remove, nodes_to_add = _dispatch_5d_reshape_candidate( model, node, unique_initializer, len(batch), cache ) + needs_reorder |= any( + producer_positions.get(name, -1) > idx + for removed in nodes_to_remove + for name in removed.input[1:] + ) consumed.update(id(removed) for removed in nodes_to_remove) batch[idx] = (nodes_to_remove, nodes_to_add) except Exception: @@ -667,6 +680,12 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals raise batch_replace_nodes(graph, batch) + if needs_reorder: + order, _, _ = priority_topological_order(graph) + nodes = list(graph.node) + positions = {id(nodes[index]): rank for rank, index in enumerate(order)} + # Sorting the container in place avoids copying Constant tensor payloads. + graph.node.sort(key=lambda node: positions[id(node)]) logger.debug( "5D reshape to 4D: classified and rewrote %d 5D Reshape subgraph(s) (transpose or reduce path)", diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 98ed02f164e..ab0a3448962 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -257,6 +257,53 @@ def test_reduce_axis_zero_raises(self): class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize("path", ["transpose", "reduce"]) + @pytest.mark.parametrize("middleware", ["Clip", "QDQ"]) + @pytest.mark.parametrize("skip_inference", [False, True]) + def test_late_dynamic_controls_preserve_values(self, path, middleware, skip_inference): + """Keep runtime controls before their duplicated consumers after rewriting.""" + model = ( + _make_transpose_reshape_model() + if path == "transpose" + else _make_reduce_min_axis4_model() + ) + model.opset_import[0].version = 21 + if path == "transpose": + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + nodes = list(model.graph.node) + data = nodes[0].output[0] + control = helper.make_node("Identity", ["bound"], ["control"], name="late_control") + if middleware == "Clip": + middle = [helper.make_node("Clip", [data, "control"], ["middle"], name="clip")] + bound = np.float32(-0.5) + else: + middle = [ + helper.make_node( + "QuantizeLinear", [data, "control", "zero"], ["q"], name="quantize" + ), + helper.make_node( + "DequantizeLinear", ["q", "control", "zero"], ["middle"], name="dequantize" + ), + ] + model.graph.initializer.append(_init("zero", np.array(0, np.uint8))) + bound = np.float32(0.25) + nodes[1].input[0] = "middle" + del model.graph.node[:] + model.graph.node.extend([nodes[0], control, *middle, *nodes[1:]]) + model.graph.input.append(helper.make_tensor_value_info("bound", TensorProto.FLOAT, [])) + onnx.checker.check_model(model, full_check=True) + shape = [dim.dim_value for dim in model.graph.input[0].type.tensor_type.shape.dim] + feeds = { + "X": np.random.default_rng(11).standard_normal(shape).astype(np.float32), + "bound": bound, + } + expected = ReferenceEvaluator(model).run(None, feeds) + rewritten = _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + onnx.checker.check_model(rewritten, full_check=True) + for actual, reference in zip(ReferenceEvaluator(rewritten).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + @pytest.mark.parametrize("path", ["transpose", "reduce"]) def test_generated_names_preserve_existing_controls_and_branches(self, path): """Reserve generated names without retargeting existing data or QDQ controls.""" From 0736748aa87c4aae3dc3b5b208a6149ef351726d Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 17:29:17 +0000 Subject: [PATCH 06/15] Validate implicit per-axis controls in DLA reshape middleware Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 33 ++++++-- .../test_dla_5d_reshape_to_4d.py | 76 ++++++++++++++++++- 2 files changed, 102 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index f4fdc11af23..737b8494cd4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -48,8 +48,26 @@ def _consumers_first_input(graph, tensor_name: str) -> list: return [n for n in graph.node if n.input and len(n.input) > 0 and n.input[0] == tensor_name] +def _known_non_scalar_control(cache: GraphCache, name: str) -> bool: + """Identify controls whose known shape is not a fixed single element.""" + initializer = cache.get_init(name) + if initializer is not None and any(dim != 1 for dim in initializer.dims): + return True + if initializer is not None and name not in cache.graph_input_names: + return False + info = cache.vi_map.get(name) + return ( + info is not None + and info.type.tensor_type.HasField("shape") + and any( + not dim.HasField("dim_value") or dim.dim_value != 1 + for dim in info.type.tensor_type.shape.dim + ) + ) + + def _parse_middleware_strict( - graph, start_out: str, stop_op_types: frozenset, *, path_context: str + graph, start_out: str, stop_op_types: frozenset, cache: GraphCache, *, path_context: str ) -> tuple: """Walk ``Clip`` / Q→DQ middleware then a terminal op; raise ``ValueError`` on dead ends or bad ops.""" middle: list = [] @@ -93,6 +111,9 @@ def _parse_middleware_strict( sum(name == n.output[0] for node in graph.node for name in node.input) != 1 or any(value.name == n.output[0] for value in graph.output) or any(a.name in {"axis", "block_size"} for a in dq.attribute) + or any( + _known_non_scalar_control(cache, name) for name in {*n.input[1:], *dq.input[1:]} + ) ): raise ValueError(f"{path_context}: QDQ must be unshared and per-tensor") middle.extend([n, dq]) @@ -117,13 +138,13 @@ def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: return len(shape_vals) == 5 -def _classify_5d_reshape_follower(model, reshape_node) -> tuple: +def _classify_5d_reshape_follower(model, reshape_node, cache: GraphCache) -> tuple: r"""Return (``\"transpose\"``, mid_a, transpose_node) or (``\"reduce\"``, middle, reduce_node).""" graph = model.graph ctx = f"5D Reshape {reshape_node.name!r}" r_out = reshape_node.output[0] stop = frozenset({"Transpose", *_REDUCE_OPS}) - middle, term = _parse_middleware_strict(graph, r_out, stop, path_context=ctx) + middle, term = _parse_middleware_strict(graph, r_out, stop, cache, path_context=ctx) if term.op_type == "Transpose": return ("transpose", middle, term) elif term.op_type in _REDUCE_OPS: @@ -249,7 +270,9 @@ def _build_transpose_replacement_or_raise( raise ValueError(msg) tr_out = tr.output[0] - mid_b, r2 = _parse_middleware_strict(graph, tr_out, frozenset({"Reshape"}), path_context=ctx) + mid_b, r2 = _parse_middleware_strict( + graph, tr_out, frozenset({"Reshape"}), cache, path_context=ctx + ) if "_transformed" in r2.name: msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" raise ValueError(msg) @@ -605,7 +628,7 @@ def _dispatch_5d_reshape_candidate( model, reshape_node, unique_initializer: set, counter: int, cache: GraphCache ) -> tuple[list, list]: """Classify a single 5D Reshape and return ``(nodes_to_remove, nodes_to_add)``.""" - branch, a, b = _classify_5d_reshape_follower(model, reshape_node) + branch, a, b = _classify_5d_reshape_follower(model, reshape_node, cache) if branch == "transpose": return _build_transpose_replacement_or_raise( model, reshape_node, a, b, unique_initializer, counter, cache diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index ab0a3448962..1768394bd0a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -257,6 +257,71 @@ def test_reduce_axis_zero_raises(self): class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize( + "controls", ["initializer", "input", "symbolic", "overridable", "dequantizer", "zero_only"] + ) + def test_implicit_per_axis_qdq_refuses_before_mutation(self, controls): + """Refuse per-axis controls even when both QDQ nodes omit their axis.""" + model = _make_transpose_reshape_model() + model.opset_import[0].version = 21 + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + nodes = list(model.graph.node) + nodes[1].input[0] = "dequantized" + scale = np.array([0.125, 0.25, 0.5, 1], np.float32) + zero = np.zeros(4, np.uint8) + feeds = {"X": np.random.default_rng(13).standard_normal((2, 12, 4, 16)).astype(np.float32)} + leading = [] + if controls in {"input", "symbolic", "overridable"}: + shape = [4] if controls == "input" else ["channels"] + model.graph.input.extend( + [ + helper.make_tensor_value_info("scale", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("zero", TensorProto.UINT8, shape), + ] + ) + feeds.update(scale=scale, zero=zero) + if controls == "overridable": + model.graph.initializer.extend( + [ + _init("scale", np.array([0.25], np.float32)), + _init("zero", np.array([0], np.uint8)), + ] + ) + else: + model.graph.initializer.append(_init("zero", zero)) + if controls == "zero_only": + model.graph.input.append( + helper.make_tensor_value_info("scale_input", TensorProto.FLOAT, [4]) + ) + leading.append( + helper.make_node("Identity", ["scale_input"], ["scale"], name="scale_producer") + ) + feeds["scale_input"] = scale + else: + model.graph.initializer.append(_init("scale", scale)) + q_controls = ["scale", "zero"] + if controls == "dequantizer": + q_controls = ["scalar_scale", "scalar_zero"] + model.graph.initializer.extend( + [ + _init("scalar_scale", np.array(0.25, np.float32)), + _init("scalar_zero", np.array(0, np.uint8)), + ] + ) + quantize = helper.make_node("QuantizeLinear", ["T5", *q_controls], ["q"], name="quantize") + dequantize = helper.make_node( + "DequantizeLinear", ["q", "scale", "zero"], ["dequantized"], name="dequantize" + ) + del model.graph.node[:] + model.graph.node.extend([*leading, nodes[0], quantize, dequantize, *nodes[1:]]) + onnx.checker.check_model(model, full_check=True) + ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="per-tensor"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + @pytest.mark.parametrize("path", ["transpose", "reduce"]) @pytest.mark.parametrize("middleware", ["Clip", "QDQ"]) @pytest.mark.parametrize("skip_inference", [False, True]) @@ -343,10 +408,17 @@ def test_generated_names_preserve_existing_controls_and_branches(self, path): ) model.graph.initializer.extend( [ - _init(scale_name, np.array(0.25, np.float32)), - _init(zero_name, np.array(0, np.uint8)), + _init(scale_name, np.array([0.25] if transpose else 0.25, np.float32)), + _init(zero_name, np.array([0] if transpose else 0, np.uint8)), ] ) + if transpose: + model.graph.value_info.extend( + [ + helper.make_tensor_value_info(scale_name, TensorProto.FLOAT, ["scale_length"]), + helper.make_tensor_value_info(zero_name, TensorProto.UINT8, ["zero_length"]), + ] + ) renames = { "X": "u5d_tr_r4_0" if transpose else "u5d_rs_msh_0", "shape5": "u5d_tr_slice_axes_0" if transpose else "u5d_rs_ax4_0", From d62ab34f0ac612888dafe3be3c0a6c2ab5fb49a0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 17:42:29 +0000 Subject: [PATCH 07/15] Respect ONNX domains in DLA reshape pattern matching Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 11 +++- .../test_dla_5d_reshape_to_4d.py | 59 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 737b8494cd4..7931f728dcc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -86,6 +86,11 @@ def _parse_middleware_strict( ) raise ValueError(msg) n = cands[0] + if n.domain not in {"", "ai.onnx"} or n.overload: + raise ValueError( + f"{path_context}: unsupported follower {n.op_type!r} in domain {n.domain!r} " + f"with overload {n.overload!r}" + ) if n.op_type in stop_op_types: return middle, n if n.op_type == "Clip": @@ -107,6 +112,10 @@ def _parse_middleware_strict( ) raise ValueError(msg) dq = dq_nodes[0] + if dq.domain not in {"", "ai.onnx"} or dq.overload: + raise ValueError( + f"{path_context}: custom DequantizeLinear middleware is unsupported" + ) if ( sum(name == n.output[0] for node in graph.node for name in node.input) != 1 or any(value.name == n.output[0] for value in graph.output) @@ -129,7 +138,7 @@ def _parse_middleware_strict( def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: """Identify Reshape nodes with a constant five-dimensional target.""" - if node.op_type != "Reshape": + if node.op_type != "Reshape" or node.domain not in {"", "ai.onnx"} or node.overload: return False ini = cache.get_init(node.input[1]) if ini is None: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 1768394bd0a..fe89fb476a8 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -257,6 +257,65 @@ def test_reduce_axis_zero_raises(self): class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize("custom_op", ["Reshape", "Clip", "DequantizeLinear"]) + def test_custom_functions_are_not_lowered_as_builtin_ops(self, custom_op): + """Preserve custom candidates and refuse custom followers before mutation.""" + model = _make_transpose_reshape_model() + model.opset_import[0].version = 21 + model.opset_import.append(helper.make_opsetid("custom", 1)) + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + if custom_op == "Reshape": + model.graph.node[0].domain = "custom" + inputs = ["data", "shape"] + body = [ + helper.make_node("Reshape", inputs, ["reshaped"]), + helper.make_node("Neg", ["reshaped"], ["result"]), + ] + else: + model.graph.node[1].input[0] = "middle" + inputs = ["data", "scale", "zero"] + if custom_op == "Clip": + nodes = [ + helper.make_node("Clip", ["T5", "scale", "zero"], ["middle"], domain="custom") + ] + body = [helper.make_node("Identity", ["data"], ["result"])] + zero = np.array(0.5, np.float32) + else: + nodes = [ + helper.make_node("QuantizeLinear", ["T5", "scale", "zero"], ["q"]), + helper.make_node( + "DequantizeLinear", ["q", "scale", "zero"], ["middle"], domain="custom" + ), + ] + body = [helper.make_node("Cast", ["data"], ["result"], to=TensorProto.FLOAT)] + zero = np.array(0, np.uint8) + for index, node in enumerate(nodes, start=1): + model.graph.node.insert(index, node) + model.graph.initializer.extend( + [ + _init("scale", np.array(0.25, np.float32)), + _init("zero", zero), + ] + ) + model.functions.append( + helper.make_function( + "custom", custom_op, inputs, ["result"], body, [helper.make_opsetid("", 21)] + ) + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.random.default_rng(17).standard_normal((2, 12, 4, 16)).astype(np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + if custom_op == "Reshape": + rewritten = _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + for actual, reference in zip(ReferenceEvaluator(rewritten).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + else: + with pytest.raises(ValueError): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + @pytest.mark.parametrize( "controls", ["initializer", "input", "symbolic", "overridable", "dequantizer", "zero_only"] ) From dd36377bda3052239f9a417a912a3846f91f81e3 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 18:04:09 +0000 Subject: [PATCH 08/15] Guard DLA transpose rewriting against legacy ONNX schemas Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 12 ++++++- .../test_dla_5d_reshape_to_4d.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 7931f728dcc..bd904a3c89b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -138,7 +138,12 @@ def _parse_middleware_strict( def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: """Identify Reshape nodes with a constant five-dimensional target.""" - if node.op_type != "Reshape" or node.domain not in {"", "ai.onnx"} or node.overload: + if ( + node.op_type != "Reshape" + or node.domain not in {"", "ai.onnx"} + or node.overload + or len(node.input) < 2 + ): return False ini = cache.get_init(node.input[1]) if ini is None: @@ -260,6 +265,11 @@ def _build_transpose_replacement_or_raise( """ graph = model.graph ctx = f"5D Reshape {reshape_node.name!r} → Transpose {tr.name!r}" + opset = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 0 + ) + if opset < 10: + raise ValueError(f"{ctx}: input-based Slice requires ONNX opset 10 or newer, got {opset}") transpose_perm = None for attr in tr.attribute: if attr.name == "perm": diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index fe89fb476a8..2569afc9199 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -257,6 +257,38 @@ def test_reduce_axis_zero_raises(self): class TestDla5dReshapeTo4dTranspose: + @pytest.mark.parametrize( + ("opset", "skip_inference"), [(5, True), (7, False), (9, False), (9, True)] + ) + def test_unsupported_slice_schema_refuses_before_mutation(self, opset, skip_inference): + """Refuse versions without input-based Slice and restore inferred metadata.""" + model = _make_transpose_reshape_model() + model.opset_import[0].version = opset + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="opset"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + assert model.SerializeToString() == before + + def test_attribute_reshape_is_not_an_input_shape_candidate(self): + """Leave valid legacy attribute-form reshapes intact instead of indexing input 1.""" + model = _make_transpose_reshape_model() + model.opset_import[0].version = 1 + targets = { + init.name: numpy_helper.to_array(init).tolist() for init in model.graph.initializer + } + for node in model.graph.node: + if node.op_type == "Reshape": + node.attribute.append(helper.make_attribute("shape", targets[node.input[1]])) + del node.input[1:] + del model.graph.initializer[:] + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + @pytest.mark.parametrize("custom_op", ["Reshape", "Clip", "DequantizeLinear"]) def test_custom_functions_are_not_lowered_as_builtin_ops(self, custom_op): """Preserve custom candidates and refuse custom followers before mutation.""" From 40ac525ba914d7152250af36876cb9f12be5dfde Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 18:16:42 +0000 Subject: [PATCH 09/15] Reject runtime-overridable DLA reshape dimensions and axes Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 18 ++++++--- .../test_dla_5d_reshape_to_4d.py | 37 +++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index bd904a3c89b..2b712e3e2e3 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -296,10 +296,13 @@ def _build_transpose_replacement_or_raise( msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" raise ValueError(msg) ini2 = cache.get_init(r2.input[1]) - if ini2 is None: - raise ValueError(f"{ctx}: final Reshape shape {r2.input[1]!r} must be a graph initializer") + if ini2 is None or r2.input[1] in cache.graph_input_names: + raise ValueError( + f"{ctx}: final Reshape shape {r2.input[1]!r} must be a fixed graph initializer" + ) ini1 = cache.get_init(reshape_node.input[1]) - assert ini1 is not None, f"initializer for node {reshape_node.name} is None" + if ini1 is None or reshape_node.input[1] in cache.graph_input_names: + raise ValueError(f"{ctx}: Reshape shape must be a fixed graph initializer") shape1_raw = np.asarray(numpy_helper.to_array(ini1), dtype=np.int64).reshape(-1) shape2 = np.asarray(numpy_helper.to_array(ini2), dtype=np.int64).reshape(-1) @@ -463,8 +466,10 @@ def _build_reduce_replacement_or_raise( axes_name = red.input[1] ini_axes = cache.get_init(axes_name) - if ini_axes is None: - msg = f"{ctx}: reduce axes {axes_name!r} must be a graph initializer for decomposition" + if ini_axes is None or axes_name in cache.graph_input_names: + msg = ( + f"{ctx}: reduce axes {axes_name!r} must be a fixed graph initializer for decomposition" + ) raise ValueError(msg) axes_arr = np.asarray(numpy_helper.to_array(ini_axes), dtype=np.int64).reshape(-1) if axes_arr.size != 1: @@ -481,7 +486,8 @@ def _build_reduce_replacement_or_raise( raise ValueError(msg) ini1 = cache.get_init(reshape_node.input[1]) - assert ini1 is not None, f"initializer for node {reshape_node.name} is None" + if ini1 is None or reshape_node.input[1] in cache.graph_input_names: + raise ValueError(f"{ctx}: Reshape shape must be a fixed graph initializer") shape1_raw = np.asarray(numpy_helper.to_array(ini1), dtype=np.int64).reshape(-1) if shape1_raw.size != 5: msg = f"{ctx}: first Reshape shape must have length 5, got {shape1_raw.tolist()}" diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 2569afc9199..8791f46ea19 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -256,6 +256,43 @@ def test_reduce_axis_zero_raises(self): _apply_5d_reshape_to_4d(infer_shapes(copy.deepcopy(orig))) +@pytest.mark.parametrize("control", ["shape5", "shape_out", "axes_reduce"]) +@pytest.mark.parametrize("skip_inference", [False, True]) +def test_overridable_structure_refuses_before_mutation(control, skip_inference): + """Preserve runtime shape and axis overrides rather than baking in initializer defaults.""" + model = ( + _make_reduce_min_axis4_model() + if control == "axes_reduce" + else _make_transpose_reshape_model() + ) + if control != "axes_reduce": + dims = model.graph.input[0].type.tensor_type.shape.dim + dims[0].dim_value, dims[1].dim_value = 2, 12 + initializer = next(tensor for tensor in model.graph.initializer if tensor.name == control) + model.graph.input.append( + helper.make_tensor_value_info(control, TensorProto.INT64, list(initializer.dims)) + ) + overrides = {"shape5": [2, 12, 4, 4, 4], "shape_out": [2, 16, 12, 4], "axes_reduce": [3]} + if control == "shape_out": + for index, dim in enumerate(model.graph.output[0].type.tensor_type.shape.dim): + dim.ClearField("dim_value") + dim.dim_param = f"output_{index}" + onnx.checker.check_model(model, full_check=True) + data = np.arange(1536, dtype=np.float32).reshape( + [dim.dim_value for dim in model.graph.input[0].type.tensor_type.shape.dim] + ) + reference = ReferenceEvaluator(model) + expected = reference.run( + None, {"X": data, control: np.asarray(overrides[control], dtype=np.int64)} + )[0] + default = reference.run(None, {"X": data})[0] + assert not np.array_equal(default, expected) + before = model.SerializeToString() + with pytest.raises(ValueError, match="initializer"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + assert model.SerializeToString() == before + + class TestDla5dReshapeTo4dTranspose: @pytest.mark.parametrize( ("opset", "skip_inference"), [(5, True), (7, False), (9, False), (9, True)] From fadfa4ac771a42471ec42d3868745fd894f76278 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 18:22:27 +0000 Subject: [PATCH 10/15] Document DLA decomposition and verify singleton QDQ semantics Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 13 ++-- .../test_dla_5d_reshape_to_4d.py | 67 +++++++++++++++++++ 2 files changed, 73 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 2b712e3e2e3..abe4097560d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -20,8 +20,7 @@ all subgraph replacements in discovery order. Reduce followers use merge-(n-1,n) → ``Slice`` x N → optional ``Clip`` / Q→DQ **per branch** → -per-branch reduce (axes ``[n-1]``, ``keepdims=1``) → ``Concat``, matching -``scratch_space/verify_reduce_5d_vs_4d_decompose.py``. Reduce axis ``n`` must be ``1..4`` (initializer); +per-branch reduce (axes ``[n-1]``, ``keepdims=1``) → ``Concat``. Reduce axis ``n`` must be ``1..4`` (initializer); axis ``0`` is rejected. ``keepdims=1`` on the original op adds ``Unsqueeze`` at axis ``n``. Unsupported followers raise ``ValueError`` with ``op_type`` and node ``name``. @@ -232,7 +231,7 @@ def _duplicate_clip_qdq_chain( def _infer_transpose_path_slice_axis( cache: GraphCache, input_tensor: str, num_slices: int, *, ctx: str ) -> int: - """Require the axis-zero layout used by the reference decomposition.""" + """Require the leading input dimension to match the number of slices.""" in_shape = cache.get_shape(input_tensor) if in_shape is None or not in_shape or in_shape[0] != num_slices: raise ValueError( @@ -242,7 +241,7 @@ def _infer_transpose_path_slice_axis( def _transpose_path_concat_axis_4d(slice_axis: int, transpose_perm: list[int]) -> int: - """Concat axis for 4D branch outputs (scratch: axis 0 when perm[0]==0 and input is sliced on axis 0).""" + """Concatenate branches along the preserved leading axis.""" if slice_axis != 0: msg = ( "5D reshape→transpose path: non-zero Slice axis is not supported " @@ -257,7 +256,7 @@ def _transpose_path_concat_axis_4d(slice_axis: int, transpose_perm: list[int]) - def _build_transpose_replacement_or_raise( model, reshape_node, mid_a, tr, unique_initializer: set, counter: int, cache: GraphCache ) -> tuple[list, list]: - """Slice-based decomposition aligned with scratch_space ``transform_reshape_transpose_reshape``. + """Build the slice-based transpose decomposition from reshape targets and permutation. Shapes and permutation come from graph initializers: slice count ``shape1[0]``, 4D reshape ``shape1[1:]``, final reshape ``shape2`` (any rank), 4D transpose perm ``perm[1:]-1``. Optional ``Clip`` / Q→DQ middleware is @@ -455,7 +454,7 @@ def _build_reduce_replacement_or_raise( ``Clip`` / Q→DQ middleware is duplicated **per branch** (after each Slice, before each Reduce) so that quantisation is applied to each chunk individually, preserving the original semantics. - Matches ``scratch_space/verify_reduce_5d_vs_4d_decompose.py`` for reduce axis ``n`` in ``{1,2,3,4}``. + Supports reduce axis ``n`` in ``{1,2,3,4}``. Axis ``0`` is not supported (no observed models); raises ``ValueError``. """ graph = model.graph @@ -667,7 +666,7 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals """Replace matching 5D ``Reshape`` subgraphs, restoring the model if planning fails. Unless ``skip_symbolic_shape_inference`` is true, runs - :func:`~modelopt.onnx.graph_surgery.utils.symbolic_shape_inference.run_ort_symbolic_shape_inference` + :func:`~modelopt.onnx.graph_surgery.utils.graph_utils.infer_shapes` first so intermediate tensor shapes are available where needed. Plan disjoint replacements before applying them with :func:`batch_replace_nodes`. diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 8791f46ea19..8e4e1edfec4 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -293,6 +293,73 @@ def test_overridable_structure_refuses_before_mutation(control, skip_inference): assert model.SerializeToString() == before +@pytest.mark.parametrize("branch", ["transpose", "reduce"]) +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +def test_single_element_qdq_controls_survive_axis_length_change(branch, runtime): + """Execute single-element QDQ controls when lowering changes the default-axis length.""" + nodes = [ + helper.make_node("Reshape", ["X", "shape5"], ["reshaped"], name="reshape5"), + helper.make_node( + "QuantizeLinear", ["reshaped", "scale", "zero"], ["quantized"], name="quantize" + ), + helper.make_node( + "DequantizeLinear", ["quantized", "scale", "zero"], ["dequantized"], name="dequantize" + ), + ] + initializers = [ + _init("shape5", np.asarray([2, 1, 3, 4, 5], np.int64)), + _init("scale", np.asarray([0.25], np.float32)), + _init("zero", np.asarray([7], np.uint8)), + ] + if branch == "transpose": + nodes.extend( + [ + helper.make_node( + "Transpose", + ["dequantized"], + ["transposed"], + name="transpose", + perm=[0, 1, 3, 4, 2], + ), + helper.make_node("Reshape", ["transposed", "shape_out"], ["Y"], name="reshape_out"), + ] + ) + initializers.append(_init("shape_out", np.asarray([2, 1, 12, 5], np.int64))) + output_shape = [2, 1, 12, 5] + else: + nodes.append( + helper.make_node("ReduceMin", ["dequantized", "axes"], ["Y"], name="reduce", keepdims=0) + ) + initializers.append(_init("axes", np.asarray([2], np.int64))) + output_shape = [2, 1, 4, 5] + model = helper.make_model( + helper.make_graph( + nodes, + "single_element", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [2, 1, 3, 20])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 19)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.arange(120, dtype=np.float32).reshape(2, 1, 3, 20) / 16 - 2} + transformed = _apply_5d_reshape_to_4d(copy.deepcopy(model), skip_symbolic_shape_inference=True) + onnx.checker.check_model(transformed, full_check=True) + assert sum(node.op_type == "Slice" for node in transformed.graph.node) == ( + 2 if branch == "transpose" else 1 + ) + if runtime == "ort": + np.testing.assert_array_equal( + _run_ort(transformed, feeds)["Y"], _run_ort(model, feeds)["Y"] + ) + else: + np.testing.assert_array_equal( + ReferenceEvaluator(transformed).run(None, feeds)[0], + ReferenceEvaluator(model).run(None, feeds)[0], + ) + + class TestDla5dReshapeTo4dTranspose: @pytest.mark.parametrize( ("opset", "skip_inference"), [(5, True), (7, False), (9, False), (9, True)] From dea65abdfa22a4d1c7e986cd138dd059e1619597 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 18:35:57 +0000 Subject: [PATCH 11/15] Reuse DLA graph consumer indexes while planning reshape rewrites Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 31 +++++++------------ 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index abe4097560d..9fa4fb90a6d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -42,11 +42,6 @@ _REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) -def _consumers_first_input(graph, tensor_name: str) -> list: - """Find consumers using the tensor as their data input.""" - return [n for n in graph.node if n.input and len(n.input) > 0 and n.input[0] == tensor_name] - - def _known_non_scalar_control(cache: GraphCache, name: str) -> bool: """Identify controls whose known shape is not a fixed single element.""" initializer = cache.get_init(name) @@ -66,15 +61,15 @@ def _known_non_scalar_control(cache: GraphCache, name: str) -> bool: def _parse_middleware_strict( - graph, start_out: str, stop_op_types: frozenset, cache: GraphCache, *, path_context: str + start_out: str, stop_op_types: frozenset, cache: GraphCache, *, path_context: str ) -> tuple: """Walk ``Clip`` / Q→DQ middleware then a terminal op; raise ``ValueError`` on dead ends or bad ops.""" middle: list = [] cur = start_out while True: - cands = _consumers_first_input(graph, cur) - uses = sum(name == cur for node in graph.node for name in node.input) - if uses != 1 or any(value.name == cur for value in graph.output): + consumers = cache.get_consumers(cur) + cands = [node for node, port in consumers if port == 0] + if len(consumers) != 1 or cur in cache.graph_output_names: raise ValueError( f"{path_context}: cannot replace shared or graph-output tensor {cur!r}" ) @@ -99,10 +94,11 @@ def _parse_middleware_strict( if n.op_type == "QuantizeLinear": if any(a.name in {"axis", "block_size"} for a in n.attribute): raise ValueError(f"{path_context}: only per-tensor QDQ middleware is supported") + q_consumers = cache.get_consumers(n.output[0]) dq_nodes = [ - x - for x in _consumers_first_input(graph, n.output[0]) - if x.op_type == "DequantizeLinear" + node + for node, port in q_consumers + if port == 0 and node.op_type == "DequantizeLinear" ] if not dq_nodes: msg = ( @@ -116,8 +112,8 @@ def _parse_middleware_strict( f"{path_context}: custom DequantizeLinear middleware is unsupported" ) if ( - sum(name == n.output[0] for node in graph.node for name in node.input) != 1 - or any(value.name == n.output[0] for value in graph.output) + len(q_consumers) != 1 + or n.output[0] in cache.graph_output_names or any(a.name in {"axis", "block_size"} for a in dq.attribute) or any( _known_non_scalar_control(cache, name) for name in {*n.input[1:], *dq.input[1:]} @@ -153,11 +149,10 @@ def _is_5d_reshape_candidate(cache: GraphCache, node) -> bool: def _classify_5d_reshape_follower(model, reshape_node, cache: GraphCache) -> tuple: r"""Return (``\"transpose\"``, mid_a, transpose_node) or (``\"reduce\"``, middle, reduce_node).""" - graph = model.graph ctx = f"5D Reshape {reshape_node.name!r}" r_out = reshape_node.output[0] stop = frozenset({"Transpose", *_REDUCE_OPS}) - middle, term = _parse_middleware_strict(graph, r_out, stop, cache, path_context=ctx) + middle, term = _parse_middleware_strict(r_out, stop, cache, path_context=ctx) if term.op_type == "Transpose": return ("transpose", middle, term) elif term.op_type in _REDUCE_OPS: @@ -288,9 +283,7 @@ def _build_transpose_replacement_or_raise( raise ValueError(msg) tr_out = tr.output[0] - mid_b, r2 = _parse_middleware_strict( - graph, tr_out, frozenset({"Reshape"}), cache, path_context=ctx - ) + mid_b, r2 = _parse_middleware_strict(tr_out, frozenset({"Reshape"}), cache, path_context=ctx) if "_transformed" in r2.name: msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" raise ValueError(msg) From 534fbce12425e8f23891cce9d7bd499f42632617 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:41:47 +0000 Subject: [PATCH 12/15] Require scalar QDQ controls for five-dimensional lowering Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 27 ++++++++-------- .../test_dla_5d_reshape_to_4d.py | 31 +++++++++++++------ 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 9fa4fb90a6d..f0e599ad331 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -42,21 +42,21 @@ _REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) -def _known_non_scalar_control(cache: GraphCache, name: str) -> bool: - """Identify controls whose known shape is not a fixed single element.""" - initializer = cache.get_init(name) - if initializer is not None and any(dim != 1 for dim in initializer.dims): - return True - if initializer is not None and name not in cache.graph_input_names: +def _unproven_scalar_qdq_control(cache: GraphCache, name: str) -> bool: + """Refuse QDQ controls unless their rank is known to be zero.""" + if not name: return False + initializer = cache.get_init(name) + if initializer is not None: + if initializer.dims: + return True + if name not in cache.graph_input_names: + return False info = cache.vi_map.get(name) return ( - info is not None - and info.type.tensor_type.HasField("shape") - and any( - not dim.HasField("dim_value") or dim.dim_value != 1 - for dim in info.type.tensor_type.shape.dim - ) + info is None + or not info.type.tensor_type.HasField("shape") + or bool(info.type.tensor_type.shape.dim) ) @@ -116,7 +116,8 @@ def _parse_middleware_strict( or n.output[0] in cache.graph_output_names or any(a.name in {"axis", "block_size"} for a in dq.attribute) or any( - _known_non_scalar_control(cache, name) for name in {*n.input[1:], *dq.input[1:]} + _unproven_scalar_qdq_control(cache, name) + for name in {*n.input[1:], *dq.input[1:]} ) ): raise ValueError(f"{path_context}: QDQ must be unshared and per-tensor") diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 8e4e1edfec4..2d30f9391e2 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -295,8 +295,9 @@ def test_overridable_structure_refuses_before_mutation(control, skip_inference): @pytest.mark.parametrize("branch", ["transpose", "reduce"]) @pytest.mark.parametrize("runtime", ["reference", "ort"]) -def test_single_element_qdq_controls_survive_axis_length_change(branch, runtime): - """Execute single-element QDQ controls when lowering changes the default-axis length.""" +@pytest.mark.parametrize("control_shape", [[], [1]]) +def test_qdq_control_rank_when_lowering_changes_default_axis_length(branch, runtime, control_shape): + """Lower scalar QDQ controls and refuse per-axis singleton controls without mutation.""" nodes = [ helper.make_node("Reshape", ["X", "shape5"], ["reshaped"], name="reshape5"), helper.make_node( @@ -308,8 +309,8 @@ def test_single_element_qdq_controls_survive_axis_length_change(branch, runtime) ] initializers = [ _init("shape5", np.asarray([2, 1, 3, 4, 5], np.int64)), - _init("scale", np.asarray([0.25], np.float32)), - _init("zero", np.asarray([7], np.uint8)), + _init("scale", np.asarray(0.25, np.float32).reshape(control_shape)), + _init("zero", np.asarray(7, np.uint8).reshape(control_shape)), ] if branch == "transpose": nodes.extend( @@ -344,6 +345,12 @@ def test_single_element_qdq_controls_survive_axis_length_change(branch, runtime) ) onnx.checker.check_model(model, full_check=True) feeds = {"X": np.arange(120, dtype=np.float32).reshape(2, 1, 3, 20) / 16 - 2} + if control_shape: + before = model.SerializeToString() + with pytest.raises(ValueError, match="per-tensor"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + return transformed = _apply_5d_reshape_to_4d(copy.deepcopy(model), skip_symbolic_shape_inference=True) onnx.checker.check_model(transformed, full_check=True) assert sum(node.op_type == "Slice" for node in transformed.graph.node) == ( @@ -521,7 +528,7 @@ def test_implicit_per_axis_qdq_refuses_before_mutation(self, controls): @pytest.mark.parametrize("middleware", ["Clip", "QDQ"]) @pytest.mark.parametrize("skip_inference", [False, True]) def test_late_dynamic_controls_preserve_values(self, path, middleware, skip_inference): - """Keep runtime controls before their duplicated consumers after rewriting.""" + """Keep runtime controls ordered and refuse QDQ controls whose rank is unproven.""" model = ( _make_transpose_reshape_model() if path == "transpose" @@ -559,6 +566,12 @@ def test_late_dynamic_controls_preserve_values(self, path, middleware, skip_infe "bound": bound, } expected = ReferenceEvaluator(model).run(None, feeds) + if middleware == "QDQ" and skip_inference: + before = model.SerializeToString() + with pytest.raises(ValueError, match="per-tensor"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=True) + assert model.SerializeToString() == before + return rewritten = _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) onnx.checker.check_model(rewritten, full_check=True) for actual, reference in zip(ReferenceEvaluator(rewritten).run(None, feeds), expected): @@ -603,15 +616,15 @@ def test_generated_names_preserve_existing_controls_and_branches(self, path): ) model.graph.initializer.extend( [ - _init(scale_name, np.array([0.25] if transpose else 0.25, np.float32)), - _init(zero_name, np.array([0] if transpose else 0, np.uint8)), + _init(scale_name, np.array(0.25, np.float32)), + _init(zero_name, np.array(0, np.uint8)), ] ) if transpose: model.graph.value_info.extend( [ - helper.make_tensor_value_info(scale_name, TensorProto.FLOAT, ["scale_length"]), - helper.make_tensor_value_info(zero_name, TensorProto.UINT8, ["zero_length"]), + helper.make_tensor_value_info(scale_name, TensorProto.FLOAT, []), + helper.make_tensor_value_info(zero_name, TensorProto.UINT8, []), ] ) renames = { From 29aec9ccaec9d0a1565887a4fbcad1d88844e77d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:44:34 +0000 Subject: [PATCH 13/15] Retain the shared QDQ control helper import Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index f0e599ad331..8f47681a217 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -42,7 +42,7 @@ _REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) -def _unproven_scalar_qdq_control(cache: GraphCache, name: str) -> bool: +def _known_non_scalar_control(cache: GraphCache, name: str) -> bool: """Refuse QDQ controls unless their rank is known to be zero.""" if not name: return False @@ -116,8 +116,7 @@ def _parse_middleware_strict( or n.output[0] in cache.graph_output_names or any(a.name in {"axis", "block_size"} for a in dq.attribute) or any( - _unproven_scalar_qdq_control(cache, name) - for name in {*n.input[1:], *dq.input[1:]} + _known_non_scalar_control(cache, name) for name in {*n.input[1:], *dq.input[1:]} ) ): raise ValueError(f"{path_context}: QDQ must be unshared and per-tensor") From 03acd4968ff492ce319164fb500b23b13f2d732d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 06:30:33 +0000 Subject: [PATCH 14/15] Preserve nested captures during five-dimensional lowering Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 6 +++ .../test_dla_5d_reshape_to_4d.py | 44 +++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 8f47681a217..572fa6f19b9 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -37,6 +37,7 @@ infer_shapes, resolve_reshape_shape, ) +from .dla_graph_cleanup import _subgraph_captures from .dla_reorder_nodes import priority_topological_order _REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) @@ -688,6 +689,7 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals ): output.type.ParseFromString(output_type) cache = GraphCache(graph) + captured = _subgraph_captures(graph) producer_positions = { output: index for index, node in enumerate(graph.node) @@ -703,6 +705,10 @@ def _apply_5d_reshape_to_4d(model, *, skip_symbolic_shape_inference: bool = Fals nodes_to_remove, nodes_to_add = _dispatch_5d_reshape_candidate( model, node, unique_initializer, len(batch), cache ) + removed_outputs = {name for removed in nodes_to_remove for name in removed.output} + recreated_outputs = {name for added in nodes_to_add for name in added.output} + if lost := (removed_outputs - recreated_outputs) & captured: + raise ValueError(f"5D Reshape {node.name!r} cannot replace captured tensors {lost}") needs_reorder |= any( producer_positions.get(name, -1) > idx for removed in nodes_to_remove diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 2d30f9391e2..29ea92bdf6f 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -142,6 +142,50 @@ def _make_transpose_reshape_model() -> onnx.ModelProto: return model +@pytest.mark.parametrize("captured", ["T", "clip", "quantized", "dequantized", "Y"]) +@pytest.mark.parametrize("skip_inference", [False, True]) +def test_nested_capture_retains_native_producer(captured, skip_inference): + """Refuse removed middleware captures while allowing the recreated terminal output.""" + model = _make_reduce_min_axis4_model() + model.graph.initializer.extend( + [_init("scale", np.array(0.25, np.float32)), _init("zero", np.array(128, np.uint8))] + ) + model.graph.node[1].input[0] = "dequantized" + model.graph.node.insert(1, helper.make_node("Clip", ["T"], ["clip"])) + model.graph.node.insert( + 2, helper.make_node("QuantizeLinear", ["clip", "scale", "zero"], ["quantized"]) + ) + model.graph.node.insert( + 3, helper.make_node("DequantizeLinear", ["quantized", "scale", "zero"], ["dequantized"]) + ) + dtype = TensorProto.UINT8 if captured == "quantized" else TensorProto.FLOAT + shape = [2, 4, 12, 4] if captured == "Y" else [2, 4, 12, 4, 4] + branch = helper.make_graph( + [helper.make_node("Identity", [captured], ["view"])], + "capturing_branch", + [], + [helper.make_tensor_value_info("view", dtype, shape)], + ) + model.graph.initializer.append(_init("condition", np.array(True))) + model.graph.node.append( + helper.make_node("If", ["condition"], ["captured"], then_branch=branch, else_branch=branch) + ) + model.graph.output.append(helper.make_tensor_value_info("captured", dtype, shape)) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.arange(1536, dtype=np.float32).reshape(6, 4, 4, 16) / 128} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + if captured != "Y": + with pytest.raises(ValueError, match="captur"): + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + assert model.SerializeToString() == before + else: + _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) + onnx.checker.check_model(model, full_check=True) + for name, actual in _run_ort(model, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) + + def _make_reduce_axis0_model() -> onnx.ModelProto: """Invalid for this pass: ReduceMin axis 0 on 5D.""" inits = [ From 64a66d4af00583e853e2eaa652f2a658fcc2fea9 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 10:44:52 +0000 Subject: [PATCH 15/15] Accept valid reshape patterns regardless of node-name substrings Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 7 ------- .../dla_transforms/test_dla_5d_reshape_to_4d.py | 11 +++++++++-- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py index 572fa6f19b9..0f5427eae83 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -279,15 +279,8 @@ def _build_transpose_replacement_or_raise( f"(op_type={tr.op_type!r}, name={tr.name!r})" ) raise ValueError(msg) - if "_transformed" in tr.name: - msg = f"{ctx}: Transpose node name contains '_transformed' ({tr.name!r})" - raise ValueError(msg) - tr_out = tr.output[0] mid_b, r2 = _parse_middleware_strict(tr_out, frozenset({"Reshape"}), cache, path_context=ctx) - if "_transformed" in r2.name: - msg = f"{ctx}: final Reshape name contains '_transformed' ({r2.name!r})" - raise ValueError(msg) ini2 = cache.get_init(r2.input[1]) if ini2 is None or r2.input[1] in cache.graph_input_names: raise ValueError( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py index 29ea92bdf6f..866c8c3a4f7 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -771,10 +771,17 @@ def test_unsupported_patterns_leave_model_unchanged(self, unsupported, skip_infe _apply_5d_reshape_to_4d(model, skip_symbolic_shape_inference=skip_inference) assert model.SerializeToString() == before - @pytest.mark.parametrize("legacy_clip", [False, True]) - def test_transpose_pipeline_rewritten(self, legacy_clip): + @pytest.mark.parametrize( + ("legacy_clip", "marked_node"), + [(False, None), (True, None), (False, "transpose_mid"), (False, "reshape_out")], + ) + def test_transpose_pipeline_rewritten(self, legacy_clip, marked_node): """The supported leading-slice layout preserves values and output shape.""" orig = _make_transpose_reshape_model() + if marked_node is not None: + next( + node for node in orig.graph.node if node.name == marked_node + ).name += "_transformed" dims = orig.graph.input[0].type.tensor_type.shape.dim dims[0].dim_value = 2 dims[1].dim_value = 12