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..0f5427eae83 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_5d_reshape_to_4d.py @@ -0,0 +1,734 @@ +# 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``. 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, +) +from .dla_graph_cleanup import _subgraph_captures +from .dla_reorder_nodes import priority_topological_order + +_REDUCE_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) + + +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 + 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 None + or not info.type.tensor_type.HasField("shape") + or bool(info.type.tensor_type.shape.dim) + ) + + +def _parse_middleware_strict( + 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: + 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}" + ) + 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.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": + 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") + q_consumers = cache.get_consumers(n.output[0]) + dq_nodes = [ + node + for node, port in q_consumers + if port == 0 and node.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 dq.domain not in {"", "ai.onnx"} or dq.overload: + raise ValueError( + f"{path_context}: custom DequantizeLinear middleware is unsupported" + ) + if ( + 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:]} + ) + ): + 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: + """Identify Reshape nodes with a constant five-dimensional target.""" + 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: + return False + shape_vals = numpy_helper.to_array(ini) + return len(shape_vals) == 5 + + +def _classify_5d_reshape_follower(model, reshape_node, cache: GraphCache) -> tuple: + r"""Return (``\"transpose\"``, mid_a, transpose_node) or (``\"reduce\"``, middle, reduce_node).""" + ctx = f"5D Reshape {reshape_node.name!r}" + r_out = reshape_node.output[0] + stop = frozenset({"Transpose", *_REDUCE_OPS}) + 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: + 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, 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 + out = x_in + new_nodes: list = [] + i = 0 + m = 0 + while m < len(middleware): + n = middleware[m] + if n.op_type == "Clip": + 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=cache.unique_name(f"{n.name}_u5d_{suffix}_{m}"), + **{a.name: helper.get_attribute_value(a) for a in n.attribute}, + ) + ) + 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 = 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=cache.unique_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=cache.unique_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 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( + 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: + """Concatenate branches along the preserved leading axis.""" + 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]: + """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 + duplicated per branch. + """ + 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": + 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) + tr_out = tr.output[0] + mid_b, r2 = _parse_middleware_strict(tr_out, frozenset({"Reshape"}), cache, path_context=ctx) + ini2 = cache.get_init(r2.input[1]) + 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]) + 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) + 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 + 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: + 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=cache.unique_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=cache.unique_name(f"u5d_tr_{i}_starts_{counter}"), + data_type=TensorProto.INT64, + dims=[1], + vals=[i], + ) + ) + slice_ends.append( + helper.make_tensor( + name=cache.unique_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=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=cache.unique_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 = cache.unique_name(f"u5d_tr_slice_{i}_{counter}") + nodes_to_add.append( + helper.make_node( + "Slice", + [ + input_tensor, + slice_starts[i].name, + slice_ends[i].name, + slice_axes.name, + ], + [slice_out], + name=slice_out, + ) + ) + reshape_4d_out = cache.unique_name(f"u5d_tr_r4o_{i}_{counter}") + nodes_to_add.append( + helper.make_node( + "Reshape", + [slice_out, reshape_4d_shape.name], + [reshape_4d_out], + 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", cache) + nodes_to_add.extend(br_a) + 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=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", cache + ) + nodes_to_add.extend(br_b) + slice_outputs.append(t_b) + + 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=cache.unique_name(f"u5d_tr_concatn_{counter}"), + ) + ) + nodes_to_add.append( + helper.make_node( + "Reshape", + [concat_out, reshape_final_shape.name], + [output_name], + name=cache.unique_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. + + Supports 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 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: + 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]) + 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()}" + 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 = 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 + ) + # 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 = 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=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 = 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] + ) + 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 = 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] + ) + 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 = 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=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}", cache) + nodes_to_add.extend(br_nodes) + + 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=cache.unique_name(f"u5d_rs_{op_type}_{i}_{counter}"), + **reduce_branch_kw, + ) + ) + reduce_outs.append(rs_out) + + if original_keepdims: + 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] + ) + add_unique_initializers(graph, unique_initializer, [unsq_ini]) + nodes_to_add.append( + helper.make_node( + "Concat", + reduce_outs, + [concat_out], + axis=split_ax, + name=cache.unique_name(f"u5d_rs_concat_{counter}"), + ) + ) + nodes_to_add.append( + helper.make_node( + "Unsqueeze", + [concat_out, usq_name], + [output_name], + name=cache.unique_name(f"u5d_rs_unsqueeze_{counter}"), + ) + ) + else: + nodes_to_add.append( + helper.make_node( + "Concat", + reduce_outs, + [output_name], + axis=split_ax, + name=cache.unique_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, cache) + 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): + """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.graph_utils.infer_shapes` + first so intermediate tensor shapes are available where needed. + + 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. + """ + graph = model.graph + 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], + ) + batch: dict[int, tuple[list, list]] = {} + consumed: set[int] = set() + needs_reorder = False + 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) + captured = _subgraph_captures(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): + 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 + ) + 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 + 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: + 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) + 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)", + 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..866c8c3a4f7 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_5d_reshape_to_4d.py @@ -0,0 +1,844 @@ +# 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 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 + +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. +_REDUCE_OPSET = 18 +_GENERIC_OPSET = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +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]: + """Execute all graph outputs with the CPU runtime.""" + 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 + + +@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 = [ + _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): + """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 {} + 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): + """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 = _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) + 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)}" + ) + 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 = _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 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))) + + +@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 + + +@pytest.mark.parametrize("branch", ["transpose", "reduce"]) +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +@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( + "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).reshape(control_shape)), + _init("zero", np.asarray(7, np.uint8).reshape(control_shape)), + ] + 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} + 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) == ( + 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)] + ) + 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.""" + 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"] + ) + 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]) + def test_late_dynamic_controls_preserve_values(self, path, middleware, skip_inference): + """Keep runtime controls ordered and refuse QDQ controls whose rank is unproven.""" + 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) + 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): + 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.""" + 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)), + ] + ) + if transpose: + model.graph.value_info.extend( + [ + helper.make_tensor_value_info(scale_name, TensorProto.FLOAT, []), + helper.make_tensor_value_info(zero_name, TensorProto.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.""" + 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 + + @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 + 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))) + 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): + """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"): + _apply_5d_reshape_to_4d(model) + assert model.SerializeToString() == before + + +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"] + + 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)