From f6bf9554e15334f8e3ade87f4040d27a232e37c6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:09 +0000 Subject: [PATCH 1/6] [DLA_transforms_7] [7/17] Logged casts, dequantization and reshape stages Signed-off-by: mgohil --- .../dla_transforms/dla_remove_deqlin.py | 214 +++++++ .../dla_transforms/dla_reshape_patterns.py | 264 +++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 15 + .../dla_transforms/test_dla_remove_deqlin.py | 545 ++++++++++++++++++ 4 files changed, 1038 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py new file mode 100644 index 00000000000..fef37a6675f --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py @@ -0,0 +1,214 @@ +# 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. + +"""Fold ``DequantizeLinear`` nodes whose data input is an initializer into a float32 initializer. + +Skip conditions — DQ node is left in place when its output feeds: + +* ``Conv`` weight (``input[1]``) with INT8 / UINT8 / INT4 / UINT4 zero-point. +* ``ConvTranspose`` weight with INT8 / UINT8 zero-point. +* ``MatMul`` / ``Gemm`` weight that will be converted to a ``ConvTranspose`` chain by + ``dla_matmul_to_transpose_conv_transpose`` (detected via ``check_to_apply_transpose``). + Pre-fusing the float32 weight would corrupt the per-channel quantization axis that the + conversion adjusts. +""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + check_to_apply_transpose, + get_node_attr_i, +) + +# ── Zero-point types that trigger the skip rule ────────────────────────────── +_CONV_SKIP_ZP = frozenset( + {TensorProto.INT8, TensorProto.UINT8, TensorProto.INT4, TensorProto.UINT4} +) +_CONVT_SKIP_ZP = frozenset({TensorProto.INT8, TensorProto.UINT8}) +_MATMUL_SKIP_ZP = _CONV_SKIP_ZP + + +# ── Helpers ─────────────────────────────────────────────────────────────────── + + +def _dequantize( + x_arr: np.ndarray, + scale_arr: np.ndarray, + zp_arr: np.ndarray, + axis: int, +) -> np.ndarray: + """Compute ``(x - zero_point) * scale`` with per-channel broadcasting. + + For scalar (rank-0) ``x`` the broadcast reshape is skipped. + """ + x = x_arr.astype(np.int32) + zp = zp_arr.astype(np.int32) + scale = scale_arr.astype(np.float32) + + if x.ndim > 0: + # Reshape scale / zp for per-channel broadcasting along ``axis``. + broadcast_shape = [1] * x.ndim + broadcast_shape[axis] = -1 + zp = zp.reshape(broadcast_shape) + scale = scale.reshape(broadcast_shape) + + return ((x - zp) * scale).astype(np.float32) + + +def _should_skip( + dq_output: str, + zp_dtype: int, + consumer_map: dict, + model: onnx.ModelProto, + shape_map: dict, +) -> bool: + """Return True if the DQ node must be kept (downstream op handles quantised weight).""" + consumers = consumer_map.get(dq_output, []) + if len(consumers) != 1: + return False # multiple consumers → always dequantize + + consumer_node, input_idx = consumers[0] + is_weight = input_idx == 1 + op = consumer_node.op_type + + if op == "Conv" and is_weight and zp_dtype in _CONV_SKIP_ZP: + return True + if op == "ConvTranspose" and is_weight and zp_dtype in _CONVT_SKIP_ZP: + return True + if op in ("MatMul", "Gemm") and is_weight and zp_dtype in _MATMUL_SKIP_ZP: + try: + return check_to_apply_transpose(consumer_node, model, shape_map) + except Exception: + return False + return False + + +# ── Main transform ───────────────────────────────────────────────────────────── + + +def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: + """Replace eligible DequantizeLinear nodes with pre-computed float32 initializers.""" + graph = model.graph + cache = GraphCache(graph) + + # Snapshot for check_to_apply_transpose (expects a plain dict) + shape_map_snapshot: dict = {} + for vi in list(graph.value_info) + list(graph.input) + list(graph.output): + s = cache.get_shape(vi.name) + if s is not None: + shape_map_snapshot[vi.name] = s + + consumer_map = cache.consumer_map + initializer_names = cache.init_names + unique_init: set[str] = set(cache.init_names) + nodes_to_remove: list = [] + cnt = 0 + + for node in cache.nodes_by_op("DequantizeLinear"): + if get_node_attr_i(node, "block_size", 0): + continue + if not node.input or node.input[0] not in initializer_names: + continue # data input is not a static initializer — skip + + dq_output = node.output[0] + + # Fetch scale and zero-point initializers (inputs 1 and 2) + scale_init = cache.get_init(node.input[1]) if len(node.input) > 1 else None + zp_init = cache.get_init(node.input[2]) if len(node.input) > 2 else None + + if scale_init is None: + logger.debug( + "DequantizeLinear %r: scale not found as initializer; skipping.", + node.name, + ) + continue + + if zp_init is None: + logger.debug( + "DequantizeLinear %r: zero_point not found as initializer; using 0 as zero_point.", + node.name, + ) + + zp_dtype: int = zp_init.data_type if zp_init is not None else scale_init.data_type + + if _should_skip(dq_output, zp_dtype, consumer_map, model, shape_map_snapshot): + logger.debug( + "DequantizeLinear %r: skip (downstream op keeps quantised weight).", node.name + ) + continue + + # ── Dequantize ──────────────────────────────────────────────────────── + x_init = cache.get_init(node.input[0]) + x_arr = numpy_helper.to_array(x_init) + scale_arr = numpy_helper.to_array(scale_init) + zp_arr = ( + numpy_helper.to_array(zp_init) + if zp_init is not None + else np.zeros(scale_arr.shape, dtype=np.int32) + ) + + axis = get_node_attr_i(node, "axis", 0) + + dequantized = _dequantize(x_arr, scale_arr, zp_arr, axis) + new_name = node.input[0] + "_dequantized" + previous = next((tensor for tensor in graph.initializer if tensor.name == new_name), None) + if previous is None: + new_name = cache.unique_name(new_name) + if previous is not None and ( + list(previous.dims) != list(dequantized.shape) + or previous.data_type != TensorProto.FLOAT + or numpy_helper.to_array(previous).tobytes() != dequantized.tobytes() + ): + new_name = cache.unique_name(new_name) + new_init = numpy_helper.from_array(dequantized, name=new_name) + + # ── Rewire all consumers ────────────────────────────────────────────── + for consumer_node, _ in consumer_map.get(dq_output, []): + for i, inp in enumerate(consumer_node.input): + if inp == dq_output: + consumer_node.input[i] = new_name + # no break — update every occurrence of dq_output in this node's inputs + + add_unique_initializers(graph, unique_init, [new_init]) + if dq_output in cache.graph_output_names: + position = next(index for index, current in enumerate(graph.node) if current is node) + graph.node.insert( + position, + helper.make_node( + "Identity", + [new_name], + [dq_output], + name=cache.unique_name(node.name + "_output"), + ), + ) + + nodes_to_remove.append(node) + cnt += 1 + logger.debug("DequantizeLinear %r → float32 initializer %r.", node.name, new_name) + + for node in nodes_to_remove: + graph.node.remove(node) + + logger.debug( + "remove_deqlin: folded %d DequantizeLinear node(s) into float32 initializers.", cnt + ) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py new file mode 100644 index 00000000000..21ddca1dcc4 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 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. + +"""The three separately ordered five-dimensional reshape pattern passes.""" + +import numpy as np +from onnx import TensorProto, helper + +from ..utils.graph_utils import GraphCache, _attribute, add_initializer +from .dla_5d_reshape_to_4d import _build_reduce_replacement_or_raise + +__all__ = [] + + +def _follower(cache, output, op_type): + consumers = cache.get_consumers(output) + if output in cache.graph_output_names or len(consumers) != 1: + return None + return next( + ( + node + for node, port in cache.get_consumers(output) + if port == 0 and node.op_type == op_type + ), + None, + ) + + +def _replace(graph, replacements): + nodes = list(graph.node) + del graph.node[:] + for node in nodes: + graph.node.extend(replacements.get(id(node), [node])) + + +def _apply_reshape_reduce_sum(model): + return _reshape_reduce_sum(model, quantized=False) + + +def _apply_reshape_qdq_reduce_sum(model): + return _reshape_reduce_sum(model, quantized=True) + + +def _reshape_reduce_sum(model, *, quantized): + graph = model.graph + cache = GraphCache(graph) + replacements = {} + count = 0 + for reshape in graph.node: + if reshape.op_type != "Reshape": + continue + shape = cache.get_init_array(reshape.input[1]) + if shape is None or shape.size != 5: + continue + middle = [] + output = reshape.output[0] + if quantized: + quantize = _follower(cache, output, "QuantizeLinear") + if quantize is None: + continue + dequantize = _follower(cache, quantize.output[0], "DequantizeLinear") + if dequantize is None: + continue + middle = [quantize, dequantize] + output = dequantize.output[0] + reduce = _follower(cache, output, "ReduceSum") + if reduce is None: + continue + # This reference pattern splits the final width of eight into two fours. + input_shape = cache.get_shape(reshape.input[0]) + axes_values = cache.get_init_array(reduce.input[1]) if len(reduce.input) > 1 else None + keepdims = _attribute(reduce, "keepdims") + fast_path = ( + input_shape is not None + and len(input_shape) == 4 + and input_shape[-1] == 8 + and shape.tolist() == [*input_shape[:-1], 2, 4] + and axes_values is not None + and axes_values.reshape(-1).tolist() == [-1] + and keepdims is not None + and keepdims.i == 0 + ) + if not fast_path: + while any( + name.startswith("u5d_rs_") and name.endswith(f"_{count}") + for name in cache.init_names + ): + count += 1 + removed, added = _build_reduce_replacement_or_raise( + model, reshape, middle, reduce, set(cache.init_names), count, cache + ) + replacements.update({id(node): [] for node in removed}) + replacements[id(reshape)] = added + count += 1 + continue + prefix = f"reshape_{'qdq_' if quantized else ''}reduce_{count}" + axes = add_initializer( + graph, + prefix + "_slice_axes", + np.asarray([3], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + nodes, outputs = [], [] + for branch in range(2): + start = add_initializer( + graph, + f"{prefix}_{branch}_starts", + np.asarray([branch * 4], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + end = add_initializer( + graph, + f"{prefix}_{branch}_ends", + np.asarray([(branch + 1) * 4], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + current = f"{prefix}_{branch}_slice" + nodes.append( + helper.make_node( + "Slice", [reshape.input[0], start, end, axes], [current], name=current + ) + ) + for index, original in enumerate(middle): + name = f"{prefix}_{branch}_qdq_{index}" + # The C++ reference uses the quantizer's scale and zero point for both. + nodes.append( + helper.make_node( + original.op_type, [current, *middle[0].input[1:3]], [name], name=name + ) + ) + current = name + output = f"{prefix}_{branch}_reduce" + nodes.append( + helper.make_node("ReduceSum", [current, *reduce.input[1:]], [output], name=output) + ) + outputs.append(output) + nodes.append( + helper.make_node( + "Concat", outputs, list(reduce.output), name=prefix + "_axis_updated", axis=3 + ) + ) + for original in [*middle, reduce]: + replacements[id(original)] = [] + replacements[id(reshape)] = nodes + count += 1 + _replace(graph, replacements) + return model + + +def _apply_reshape_transpose_reshape(model): + graph = model.graph + cache = GraphCache(graph) + replacements = {} + count = 0 + for first in graph.node: + if first.op_type != "Reshape" or "_transformed" in first.name: + continue + shape = cache.get_init_array(first.input[1]) + if shape is None or shape.size != 5: + continue + transpose = _follower(cache, first.output[0], "Transpose") + if transpose is None: + continue + perm = _attribute(transpose, "perm") + if perm is None: + continue + last = _follower(cache, transpose.output[0], "Reshape") + if last is None: + continue + target = cache.get_init_array(last.input[1]) + if target is None: + continue + input_shape = cache.get_shape(first.input[0]) + if not input_shape or input_shape[0] != int(shape[0]): + raise ValueError(f"{first.name!r}: leading input dimension must equal 5D slice count") + if shape[0] <= 0 or not perm.ints or perm.ints[0] != 0: + raise ValueError(f"{first.name!r}: unsupported leading dimension/permutation") + prefix = f"reshape_transpose_{count}" + axes = add_initializer( + graph, + prefix + "_slice_axes", + np.asarray([0], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + middle_shape = add_initializer( + graph, + prefix + "_4d_shape", + np.asarray(shape.tolist()[1:], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + final_shape = add_initializer( + graph, + prefix + "_final_shape", + np.asarray(target, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + nodes, outputs = [], [] + for branch in range(int(shape[0])): + start = add_initializer( + graph, + f"{prefix}_{branch}_start", + np.asarray([branch], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + end = add_initializer( + graph, + f"{prefix}_{branch}_end", + np.asarray([branch + 1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + sliced, reshaped, transposed = [ + f"{prefix}_{branch}_{s}" for s in ("slice", "reshape", "transpose") + ] + nodes.extend( + [ + helper.make_node( + "Slice", [first.input[0], start, end, axes], [sliced], name=sliced + ), + helper.make_node("Reshape", [sliced, middle_shape], [reshaped], name=reshaped), + helper.make_node( + "Transpose", + [reshaped], + [transposed], + name=transposed, + perm=[value - 1 for value in perm.ints[1:]], + ), + ] + ) + outputs.append(transposed) + concat = prefix + "_concat" + nodes.extend( + [ + helper.make_node( + "Concat", outputs, [concat], name=prefix + "_axis_updated", axis=0 + ), + helper.make_node( + "Reshape", [concat, final_shape], list(last.output), name=prefix + "_final" + ), + ] + ) + replacements.update({id(first): nodes, id(transpose): [], id(last): []}) + count += 1 + _replace(graph, replacements) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 80d645854c9..883996fa32c 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -32,8 +32,14 @@ from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32 from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers +from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes +from .dla_transforms.dla_reshape_patterns import ( + _apply_reshape_qdq_reduce_sum, + _apply_reshape_reduce_sum, + _apply_reshape_transpose_reshape, +) from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, infer_shapes, @@ -45,6 +51,10 @@ __all__ = [] +def _remove_deqlin_and_unused_initializers(model): + return _remove_unused_initializers(_apply_remove_deqlin(model)) + + def _apply_remove_qdq(model, **options): """Run QDQ removal while retaining the model result for ordered stage dispatch.""" _remove_qdq_from_model(model, **options) @@ -61,6 +71,11 @@ def _apply_constants_to_initializers(model, **options): ("RemoveQDQ", _apply_remove_qdq), ("ConstantsToInitializers", _apply_constants_to_initializers), ("Cast", _apply_cast_to_fp32), + ("RemoveDeqLin", _remove_deqlin_and_unused_initializers), + ("ShapeInference", infer_shapes), + ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), + ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), + ("ReshapeReduceSum", _apply_reshape_reduce_sum), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py new file mode 100644 index 00000000000..e4d5c5813bc --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py @@ -0,0 +1,545 @@ +# 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_remove_deqlin`. + +"""ORT CPU parity tests for ``dla_remove_deqlin``. + +Transform rule +-------------- +A ``DequantizeLinear`` node whose first input (data) is a graph initializer is +replaced by a pre-computed float32 initializer: + + ``DequantizeLinear(x_init, scale, zp)`` → ``x_init_dequantized (float32)`` + +The original quantised initializer and the DQ node are removed. + +Skip conditions (DQ node kept) +------------------------------- +* Conv / ConvTranspose weight input with a recognised quantised zero-point type. +* MatMul / Gemm weight that ``check_to_apply_transpose`` flags for conversion. +* DQ whose data input is NOT a static initializer (e.g. a graph input). + +Coverage +-------- +uint16 DQ on init → dequantized, ORT parity +int8 DQ on init feeding non-Conv → dequantized +DQ → Conv weight INT8 zp → kept (skipped) +DQ → Conv weight UINT8 zp → kept +DQ → ConvTranspose weight INT8 zp → kept +DQ → MatMul weight (will convert) → kept +DQ data is a graph input → never removed +DQ with two consumers → dequantized, both consumers rewired +No DQ nodes → graph unchanged, ORT parity +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms import dla_remove_deqlin as _mod + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "remove_deqlin_test_models" +_OPSET_VERSION = 21 +_MAX_IR_VERSION_FOR_ORT = 10 + + +_apply = _mod._apply_remove_deqlin + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model: onnx.ModelProto) -> None: + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + inferred = SymbolicShapeInference.infer_shapes(model) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + except Exception: + pass + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + import onnxruntime as ort # type: ignore[import-untyped] + + _clamp_ir(model) + sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +def _has_dq(model: onnx.ModelProto) -> bool: + return any(n.op_type == "DequantizeLinear" for n in model.graph.node) + + +def _init_names(model: onnx.ModelProto) -> set[str]: + return {i.name for i in model.graph.initializer} + + +_RNG = np.random.default_rng(2024) + + +# ── Helpers to build DQ sub-graphs ─────────────────────────────────────────── + + +def _make_dq_init( + name: str, + data: np.ndarray, + scale: np.ndarray, + zp: np.ndarray, +) -> tuple[onnx.TensorProto, onnx.TensorProto, onnx.TensorProto]: + """Return (data_init, scale_init, zp_init) ONNX initializers.""" + return ( + numpy_helper.from_array(data, name=f"{name}_x"), + numpy_helper.from_array(scale.astype(np.float32), name=f"{name}_scale"), + numpy_helper.from_array(zp, name=f"{name}_zp"), + ) + + +def _make_dq_node(name: str, output_name: str) -> onnx.NodeProto: + return helper.make_node( + "DequantizeLinear", + inputs=[f"{name}_x", f"{name}_scale", f"{name}_zp"], + outputs=[output_name], + name=f"dq_{name}", + ) + + +# ── 1. uint16 DQ on initializer → dequantized, ORT parity ──────────────────── + + +def test_uint16_dq_init_dequantized(): + """DQ(uint16 init) feeding Relu → folded into float32 init; ORT parity.""" + shape = [4, 8] + scale_val = np.float32(0.01) + zp_val = np.uint16(0) + x_data = _RNG.integers(0, 65535, shape, dtype=np.int64).astype(np.uint16) + + x_init, scale_init, zp_init = _make_dq_init("w", x_data, np.array(scale_val), np.array(zp_val)) + dq_node = _make_dq_node("w", "w_float") + relu = helper.make_node("Relu", ["X_add"], ["Y"], name="relu") + add = helper.make_node("Add", ["X", "w_float"], ["X_add"], name="add") + + graph = helper.make_graph( + [dq_node, add, relu], + "u16_dq", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[x_init, scale_init, zp_init], + value_info=[ + _vi("w_float", TensorProto.FLOAT, shape), + _vi("X_add", TensorProto.FLOAT, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not _has_dq(tr), "DQ node must be removed" + assert "w_x_dequantized" in _init_names(tr), "float32 initializer must be added" + # original quantised init is left for _remove_unused_initializers to clean up + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-4, atol=1e-4) + + +# ── 2. int8 DQ feeding non-Conv op → dequantized ───────────────────────────── + + +def test_int8_dq_non_conv_dequantized(): + """DQ(int8) feeding a Mul (not Conv) → folded into float32.""" + shape = [2, 4] + x_data = _RNG.integers(-127, 127, shape, dtype=np.int8) + scale_val = np.array(0.1, dtype=np.float32) + zp_val = np.array(0, dtype=np.int8) + + x_init, scale_init, zp_init = _make_dq_init("b", x_data, scale_val, zp_val) + dq_node = _make_dq_node("b", "b_float") + mul = helper.make_node("Mul", ["X", "b_float"], ["Y"], name="mul") + + graph = helper.make_graph( + [dq_node, mul], + "int8_dq_mul", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[x_init, scale_init, zp_init], + value_info=[_vi("b_float", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not _has_dq(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-4, atol=1e-4) + + +def test_dq_without_zero_point_dequantized(): + """DQ(data, scale) with omitted zero_point uses zero and folds into float32.""" + shape = [2, 4] + x_data = _RNG.integers(-127, 127, shape, dtype=np.int8) + scale_val = np.array(0.2, dtype=np.float32) + + x_init = numpy_helper.from_array(x_data, name="no_zp_x") + scale_init = numpy_helper.from_array(scale_val, name="no_zp_scale") + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["no_zp_x", "no_zp_scale"], + outputs=["no_zp_float"], + name="dq_no_zp", + ) + add = helper.make_node("Add", ["X", "no_zp_float"], ["Y"], name="add") + + graph = helper.make_graph( + [dq_node, add], + "dq_no_zero_point", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[x_init, scale_init], + value_info=[_vi("no_zp_float", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not _has_dq(tr) + assert "no_zp_x_dequantized" in _init_names(tr) + expected_init = next(i for i in tr.graph.initializer if i.name == "no_zp_x_dequantized") + np.testing.assert_allclose( + numpy_helper.to_array(expected_init), x_data.astype(np.float32) * scale_val + ) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-4, atol=1e-4) + + +# ── 3. DQ → Conv weight INT8 zp: SKIPPED ───────────────────────────────────── + + +def test_dq_conv_weight_int8_skipped(): + """DQ(int8 zp) feeding Conv input[1] must be kept.""" + in_c, out_c, k = 4, 8, 3 + x_data = _RNG.integers(-127, 127, [out_c, in_c, k, k], dtype=np.int8) + scale_val = np.ones(out_c, dtype=np.float32) * 0.1 + zp_val = np.zeros(out_c, dtype=np.int8) + + x_init, scale_init, zp_init = _make_dq_init("w", x_data, scale_val, zp_val) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["w_x", "w_scale", "w_zp"], + outputs=["w_float"], + name="dq_w", + axis=0, + ) + conv = helper.make_node("Conv", ["X", "w_float"], ["Y"], name="conv", pads=[1, 1, 1, 1]) + + graph = helper.make_graph( + [dq_node, conv], + "dq_conv_skip", + [_vi("X", TensorProto.FLOAT, [1, in_c, 8, 8])], + [_vi("Y", TensorProto.FLOAT, [1, out_c, 8, 8])], + initializer=[x_init, scale_init, zp_init], + value_info=[_vi("w_float", TensorProto.FLOAT, [out_c, in_c, k, k])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _has_dq(tr), "DQ node must be kept for Conv INT8 weight" + + +# ── 4. DQ → Conv weight UINT8 zp: SKIPPED ──────────────────────────────────── + + +def test_dq_conv_weight_uint8_skipped(): + """DQ(uint8 zp) feeding Conv input[1] must be kept.""" + in_c, out_c, k = 4, 8, 3 + x_data = _RNG.integers(0, 255, [out_c, in_c, k, k], dtype=np.uint8) + scale_val = np.ones(out_c, dtype=np.float32) * 0.1 + zp_val = np.full(out_c, 128, dtype=np.uint8) + + x_init, scale_init, zp_init = _make_dq_init("w", x_data, scale_val, zp_val) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["w_x", "w_scale", "w_zp"], + outputs=["w_float"], + name="dq_w", + axis=0, + ) + conv = helper.make_node("Conv", ["X", "w_float"], ["Y"], name="conv", pads=[1, 1, 1, 1]) + + graph = helper.make_graph( + [dq_node, conv], + "dq_conv_u8_skip", + [_vi("X", TensorProto.FLOAT, [1, in_c, 8, 8])], + [_vi("Y", TensorProto.FLOAT, [1, out_c, 8, 8])], + initializer=[x_init, scale_init, zp_init], + value_info=[_vi("w_float", TensorProto.FLOAT, [out_c, in_c, k, k])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _has_dq(tr), "DQ node must be kept for Conv UINT8 weight" + + +# ── 5. DQ → ConvTranspose weight INT8 zp: SKIPPED ──────────────────────────── + + +def test_dq_convtranspose_weight_int8_skipped(): + """DQ(int8 zp) feeding ConvTranspose input[1] must be kept.""" + in_c, out_c, k = 8, 4, 3 + x_data = _RNG.integers(-127, 127, [in_c, out_c, k, k], dtype=np.int8) + scale_val = np.ones(out_c, dtype=np.float32) * 0.05 + zp_val = np.zeros(out_c, dtype=np.int8) + + x_init, scale_init, zp_init = _make_dq_init("w", x_data, scale_val, zp_val) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["w_x", "w_scale", "w_zp"], + outputs=["w_float"], + name="dq_w", + axis=1, + ) + convt = helper.make_node( + "ConvTranspose", ["X", "w_float"], ["Y"], name="convt", pads=[1, 1, 1, 1] + ) + + graph = helper.make_graph( + [dq_node, convt], + "dq_convt_skip", + [_vi("X", TensorProto.FLOAT, [1, in_c, 6, 6])], + [_vi("Y", TensorProto.FLOAT, [1, out_c, 6, 6])], + initializer=[x_init, scale_init, zp_init], + value_info=[_vi("w_float", TensorProto.FLOAT, [in_c, out_c, k, k])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _has_dq(tr), "DQ node must be kept for ConvTranspose INT8 weight" + + +# ── 6. DQ → MatMul weight (will convert): SKIPPED ──────────────────────────── + + +def test_dq_matmul_weight_will_convert_skipped(): + """DQ(int8) feeding MatMul input[1] where check_to_apply_transpose=True → DQ kept. + + Weight shape [K, N] (2D) and activation [M, K] (2D) → check_to_apply_transpose + returns True, meaning dla_matmul_to_transpose_conv_transpose will handle this node. + """ + m, k, n = 4, 8, 16 + w_data = _RNG.integers(-127, 127, [k, n], dtype=np.int8) + scale_val = np.ones(n, dtype=np.float32) * 0.05 + zp_val = np.zeros(n, dtype=np.int8) + + x_init, scale_init, zp_init = _make_dq_init("w", w_data, scale_val, zp_val) + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["w_x", "w_scale", "w_zp"], + outputs=["w_float"], + name="dq_w", + axis=0, + ) + matmul = helper.make_node("MatMul", ["X", "w_float"], ["Y"], name="matmul") + + graph = helper.make_graph( + [dq_node, matmul], + "dq_matmul_skip", + [_vi("X", TensorProto.FLOAT, [m, k])], + [_vi("Y", TensorProto.FLOAT, [m, n])], + initializer=[x_init, scale_init, zp_init], + value_info=[_vi("w_float", TensorProto.FLOAT, [k, n])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _has_dq(tr), "DQ node must be kept for MatMul that will be converted to ConvTranspose" + + +# ── 7. DQ data is a graph input → never removed ─────────────────────────────── + + +def test_dq_dynamic_input_not_removed(): + """DQ whose input[0] is a graph input (not an initializer) must be left in place.""" + shape = [4, 8] + scale_init = numpy_helper.from_array(np.array(0.1, dtype=np.float32), name="scale") + zp_init = numpy_helper.from_array(np.array(0, dtype=np.uint16), name="zp") + + dq_node = helper.make_node( + "DequantizeLinear", + inputs=["X_quant", "scale", "zp"], + outputs=["Y"], + name="dq", + ) + graph = helper.make_graph( + [dq_node], + "dq_dynamic", + [_vi("X_quant", TensorProto.UINT16, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale_init, zp_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_node_count = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert len(tr.graph.node) == original_node_count, "Graph must be unchanged" + assert _has_dq(tr), "DQ node must remain when data is a graph input" + + +# ── 8. DQ with two consumers → both rewired ─────────────────────────────────── + + +def test_dq_multiple_consumers_rewired(): + """DQ folded and BOTH downstream consumers rewired to the new float init.""" + shape = [2, 4] + x_data = _RNG.integers(0, 255, shape, dtype=np.uint16) + scale_val = np.array(0.01, dtype=np.float32) + zp_val = np.array(0, dtype=np.uint16) + + x_init, scale_init, zp_init = _make_dq_init("c", x_data, scale_val, zp_val) + dq_node = _make_dq_node("c", "c_float") + add1 = helper.make_node("Add", ["X", "c_float"], ["add1_out"], name="add1") + add2 = helper.make_node("Add", ["add1_out", "c_float"], ["Y"], name="add2") + + graph = helper.make_graph( + [dq_node, add1, add2], + "dq_multi_consumer", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[x_init, scale_init, zp_init], + value_info=[ + _vi("c_float", TensorProto.FLOAT, shape), + _vi("add1_out", TensorProto.FLOAT, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert not _has_dq(tr), "DQ must be removed" + # Both add1 and add2 must reference the new float init + new_name = "c_x_dequantized" + for node in tr.graph.node: + if node.name in ("add1", "add2"): + assert new_name in node.input, f"{node.name} was not rewired to {new_name!r}" + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-4, atol=1e-4) + + +# ── 9. No DQ nodes: graph unchanged ────────────────────────────────────────── + + +def test_no_dq_nodes_noop(): + """Graph with no DQ nodes must be left completely unchanged.""" + shape = [2, 4] + w_init = numpy_helper.from_array(_RNG.standard_normal(shape).astype(np.float32), name="W") + add = helper.make_node("Add", ["X", "W"], ["Y"], name="add") + + graph = helper.make_graph( + [add], + "no_dq", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[w_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_count = len(model.graph.node) + + x = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert len(tr.graph.node) == original_count + np.testing.assert_array_equal(ref["Y"], got["Y"]) From f082a3ab1cb943ec3e850a40fcbf600583149438 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 19:11:16 +0000 Subject: [PATCH 2/6] Preserve DLA dequantization and reshape-stage contracts Signed-off-by: mgohil --- .../dla_transforms/dla_remove_deqlin.py | 113 ++++-- .../dla_transforms/dla_reshape_patterns.py | 349 ++++++++---------- .../onnx/graph_surgery/make_dla_compatible.py | 1 + .../dla_transforms/test_dla_remove_deqlin.py | 332 +++++++++++++++++ .../test_dla_reshape_patterns.py | 230 ++++++++++++ 5 files changed, 807 insertions(+), 218 deletions(-) create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py index fef37a6675f..9142e5a9846 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py @@ -60,18 +60,34 @@ def _dequantize( For scalar (rank-0) ``x`` the broadcast reshape is skipped. """ - x = x_arr.astype(np.int32) - zp = zp_arr.astype(np.int32) + int32_data = x_arr.dtype == np.int32 + zp = zp_arr.astype(np.int32 if int32_data else np.float32) scale = scale_arr.astype(np.float32) - - if x.ndim > 0: - # Reshape scale / zp for per-channel broadcasting along ``axis``. - broadcast_shape = [1] * x.ndim + if zp.size == 1: + zp = zp.reshape(()) + if scale.size == 1: + scale = scale.reshape(()) + if x_arr.ndim > 0 and (zp.size > 1 or scale.size > 1): + broadcast_shape = [1] * x_arr.ndim broadcast_shape[axis] = -1 zp = zp.reshape(broadcast_shape) scale = scale.reshape(broadcast_shape) - - return ((x - zp) * scale).astype(np.float32) + if int32_data: + # Keep established INT32 rounding without a whole-weight float64 temporary. + result = np.empty(x_arr.shape, dtype=np.float32) + with np.nditer( + [x_arr, zp, scale, result], + flags=["external_loop", "buffered", "zerosize_ok"], + op_flags=[["readonly"], ["readonly"], ["readonly"], ["writeonly"]], + buffersize=2**18, + ) as chunks: + for values, zeros, scales, output in chunks: + output[...] = (values - zeros) * scales + return result + result = x_arr.astype(np.float32) + np.subtract(result, zp, out=result) + np.multiply(result, scale, out=result) + return result def _should_skip( @@ -81,24 +97,21 @@ def _should_skip( model: onnx.ModelProto, shape_map: dict, ) -> bool: - """Return True if the DQ node must be kept (downstream op handles quantised weight).""" - consumers = consumer_map.get(dq_output, []) - if len(consumers) != 1: - return False # multiple consumers → always dequantize - - consumer_node, input_idx = consumers[0] - is_weight = input_idx == 1 - op = consumer_node.op_type - - if op == "Conv" and is_weight and zp_dtype in _CONV_SKIP_ZP: - return True - if op == "ConvTranspose" and is_weight and zp_dtype in _CONVT_SKIP_ZP: - return True - if op in ("MatMul", "Gemm") and is_weight and zp_dtype in _MATMUL_SKIP_ZP: - try: - return check_to_apply_transpose(consumer_node, model, shape_map) - except Exception: - return False + """Retain weights required or potentially required by any downstream consumer.""" + for consumer_node, input_idx in consumer_map.get(dq_output, []): + if input_idx != 1 or consumer_node.domain not in {"", "ai.onnx"} or consumer_node.overload: + continue + op = consumer_node.op_type + if op == "Conv" and zp_dtype in _CONV_SKIP_ZP: + return True + if op == "ConvTranspose" and zp_dtype in _CONVT_SKIP_ZP: + return True + if op in ("MatMul", "Gemm") and zp_dtype in _MATMUL_SKIP_ZP: + try: + if check_to_apply_transpose(consumer_node, model, shape_map): + return True + except (ValueError, RuntimeError): + return True return False @@ -109,6 +122,10 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: """Replace eligible DequantizeLinear nodes with pre-computed float32 initializers.""" graph = model.graph cache = GraphCache(graph) + local_functions = { + (function.domain, function.name, function.overload) for function in model.functions + } + opsets = {entry.domain: entry.version for entry in model.opset_import} # Snapshot for check_to_apply_transpose (expects a plain dict) shape_map_snapshot: dict = {} @@ -124,6 +141,14 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: cnt = 0 for node in cache.nodes_by_op("DequantizeLinear"): + if node.overload or (node.domain, node.op_type, node.overload) in local_functions: + continue + if node.domain not in {"", "ai.onnx"} and not ( + node.domain == "com.microsoft" and opsets.get(node.domain) == 1 + ): + continue + if any(name in cache.graph_input_names for name in node.input): + continue if get_node_attr_i(node, "block_size", 0): continue if not node.input or node.input[0] not in initializer_names: @@ -135,6 +160,19 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: scale_init = cache.get_init(node.input[1]) if len(node.input) > 1 else None zp_init = cache.get_init(node.input[2]) if len(node.input) > 2 else None + if len(node.input) > 2 and node.input[2] and zp_init is None: + continue + x_init = cache.get_init(node.input[0]) + if x_init is None: + continue + if ( + x_init.data_type + not in _CONV_SKIP_ZP | {TensorProto.INT16, TensorProto.UINT16, TensorProto.INT32} + or (scale_init is not None and scale_init.data_type != TensorProto.FLOAT) + or get_node_attr_i(node, "output_dtype", 0) not in {0, TensorProto.FLOAT} + ): + continue + if scale_init is None: logger.debug( "DequantizeLinear %r: scale not found as initializer; skipping.", @@ -148,7 +186,7 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: node.name, ) - zp_dtype: int = zp_init.data_type if zp_init is not None else scale_init.data_type + zp_dtype: int = zp_init.data_type if zp_init is not None else x_init.data_type if _should_skip(dq_output, zp_dtype, consumer_map, model, shape_map_snapshot): logger.debug( @@ -157,7 +195,6 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: continue # ── Dequantize ──────────────────────────────────────────────────────── - x_init = cache.get_init(node.input[0]) x_arr = numpy_helper.to_array(x_init) scale_arr = numpy_helper.to_array(scale_init) zp_arr = ( @@ -166,20 +203,27 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: else np.zeros(scale_arr.shape, dtype=np.int32) ) - axis = get_node_attr_i(node, "axis", 0) + axis = get_node_attr_i(node, "axis", 1) dequantized = _dequantize(x_arr, scale_arr, zp_arr, axis) + del x_arr, scale_arr, zp_arr new_name = node.input[0] + "_dequantized" - previous = next((tensor for tensor in graph.initializer if tensor.name == new_name), None) + previous = cache.get_init(new_name) if previous is None: new_name = cache.unique_name(new_name) if previous is not None and ( - list(previous.dims) != list(dequantized.shape) + new_name in cache.graph_input_names + or list(previous.dims) != list(dequantized.shape) or previous.data_type != TensorProto.FLOAT or numpy_helper.to_array(previous).tobytes() != dequantized.tobytes() ): new_name = cache.unique_name(new_name) - new_init = numpy_helper.from_array(dequantized, name=new_name) + new_init = ( + numpy_helper.from_array(dequantized, name=new_name) + if new_name not in unique_init + else None + ) + del dequantized # ── Rewire all consumers ────────────────────────────────────────────── for consumer_node, _ in consumer_map.get(dq_output, []): @@ -188,7 +232,10 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: consumer_node.input[i] = new_name # no break — update every occurrence of dq_output in this node's inputs - add_unique_initializers(graph, unique_init, [new_init]) + if new_init is not None: + add_unique_initializers(graph, unique_init, [new_init]) + cache.init_map[new_name] = graph.initializer[-1] + del new_init if dq_output in cache.graph_output_names: position = next(index for index, current in enumerate(graph.node) if current is node) graph.node.insert( diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py index 21ddca1dcc4..5aa562cf28f 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -18,247 +18,226 @@ import numpy as np from onnx import TensorProto, helper -from ..utils.graph_utils import GraphCache, _attribute, add_initializer -from .dla_5d_reshape_to_4d import _build_reduce_replacement_or_raise +from ..utils.graph_utils import GraphCache, _attribute, add_initializer, batch_replace_nodes +from .dla_5d_reshape_to_4d import ( + _build_reduce_replacement_or_raise, + _build_transpose_replacement_or_raise, + _duplicate_clip_qdq_chain, + _is_5d_reshape_candidate, + _known_non_scalar_control, +) +from .dla_reorder_nodes import priority_topological_order __all__ = [] def _follower(cache, output, op_type): + """Match an unshared standard-domain data follower without replacing public tensors.""" consumers = cache.get_consumers(output) if output in cache.graph_output_names or len(consumers) != 1: return None return next( ( node - for node, port in cache.get_consumers(output) - if port == 0 and node.op_type == op_type + for node, port in consumers + if port == 0 + and node.op_type == op_type + and node.domain in {"", "ai.onnx"} + and not node.overload ), None, ) def _replace(graph, replacements): - nodes = list(graph.node) - del graph.node[:] - for node in nodes: - graph.node.extend(replacements.get(id(node), [node])) + """Apply planned replacements without copying untouched payloads and order late controls.""" + removed_outputs = { + output for node in graph.node if id(node) in replacements for output in node.output + } + positions = {output: index for index, node in enumerate(graph.node) for output in node.output} + needs_reorder = any( + positions.get(name, -1) > index and name not in removed_outputs + for index, node in enumerate(graph.node) + for added in replacements.get(id(node), []) + for name in added.input + ) + batch_replace_nodes( + graph, + { + index: ([node], replacements[id(node)]) + for index, node in enumerate(graph.node) + if id(node) in replacements + }, + ) + if needs_reorder: + order, _, _ = priority_topological_order(graph) + nodes = list(graph.node) + ranks = {id(nodes[index]): rank for rank, index in enumerate(order)} + graph.node.sort(key=lambda node: ranks[id(node)]) def _apply_reshape_reduce_sum(model): + """Lower the separately ordered unquantized reshape/reduction stage.""" return _reshape_reduce_sum(model, quantized=False) def _apply_reshape_qdq_reduce_sum(model): + """Lower the separately ordered per-tensor QDQ reshape/reduction stage.""" return _reshape_reduce_sum(model, quantized=True) def _reshape_reduce_sum(model, *, quantized): + """Plan disjoint reductions while preserving controls and restoring initializers on refusal.""" graph = model.graph cache = GraphCache(graph) replacements = {} count = 0 - for reshape in graph.node: - if reshape.op_type != "Reshape": - continue - shape = cache.get_init_array(reshape.input[1]) - if shape is None or shape.size != 5: - continue - middle = [] - output = reshape.output[0] - if quantized: - quantize = _follower(cache, output, "QuantizeLinear") - if quantize is None: + initializer_count = len(graph.initializer) + try: + for reshape in graph.node: + if not _is_5d_reshape_candidate(cache, reshape): + continue + if id(reshape) in replacements: + raise ValueError(f"{reshape.name!r}: overlapping reshape patterns") + shape = cache.get_init_array(reshape.input[1]) + if shape is None or shape.size != 5: continue - dequantize = _follower(cache, quantize.output[0], "DequantizeLinear") - if dequantize is None: + middle = [] + output = reshape.output[0] + if quantized: + quantize = _follower(cache, output, "QuantizeLinear") + if quantize is None: + continue + dequantize = _follower(cache, quantize.output[0], "DequantizeLinear") + if dequantize is None: + continue + middle = [quantize, dequantize] + output = dequantize.output[0] + reduce = _follower(cache, output, "ReduceSum") + if reduce is None: continue - middle = [quantize, dequantize] - output = dequantize.output[0] - reduce = _follower(cache, output, "ReduceSum") - if reduce is None: - continue - # This reference pattern splits the final width of eight into two fours. - input_shape = cache.get_shape(reshape.input[0]) - axes_values = cache.get_init_array(reduce.input[1]) if len(reduce.input) > 1 else None - keepdims = _attribute(reduce, "keepdims") - fast_path = ( - input_shape is not None - and len(input_shape) == 4 - and input_shape[-1] == 8 - and shape.tolist() == [*input_shape[:-1], 2, 4] - and axes_values is not None - and axes_values.reshape(-1).tolist() == [-1] - and keepdims is not None - and keepdims.i == 0 - ) - if not fast_path: - while any( - name.startswith("u5d_rs_") and name.endswith(f"_{count}") - for name in cache.init_names + if any( + name in cache.graph_input_names for name in [reshape.input[1], *reduce.input[1:]] ): - count += 1 - removed, added = _build_reduce_replacement_or_raise( - model, reshape, middle, reduce, set(cache.init_names), count, cache - ) - replacements.update({id(node): [] for node in removed}) - replacements[id(reshape)] = added - count += 1 - continue - prefix = f"reshape_{'qdq_' if quantized else ''}reduce_{count}" - axes = add_initializer( - graph, - prefix + "_slice_axes", - np.asarray([3], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - nodes, outputs = [], [] - for branch in range(2): - start = add_initializer( - graph, - f"{prefix}_{branch}_starts", - np.asarray([branch * 4], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, + raise ValueError( + f"{reshape.name!r}: shape and axes must be fixed graph initializers" + ) + if any( + any(attr.name in {"axis", "block_size"} for attr in node.attribute) + or any(_known_non_scalar_control(cache, name) for name in node.input[1:]) + for node in middle + ): + raise ValueError(f"{reshape.name!r}: only per-tensor QDQ middleware is supported") + # This pattern splits the final width of eight into two fours. + input_shape = cache.get_shape(reshape.input[0]) + axes_values = cache.get_init_array(reduce.input[1]) if len(reduce.input) > 1 else None + keepdims = _attribute(reduce, "keepdims") + fast_path = ( + input_shape is not None + and len(input_shape) == 4 + and input_shape[-1] == 8 + and shape.tolist() == [*input_shape[:-1], 2, 4] + and axes_values is not None + and axes_values.reshape(-1).tolist() == [-1] + and keepdims is not None + and keepdims.i == 0 ) - end = add_initializer( + if not fast_path: + removed, added = _build_reduce_replacement_or_raise( + model, reshape, middle, reduce, set(cache.init_names), count, cache + ) + replacements.update({id(node): [] for node in removed}) + replacements[id(reshape)] = added + count += 1 + continue + prefix = f"reshape_{'qdq_' if quantized else ''}reduce_{count}" + axes = add_initializer( graph, - f"{prefix}_{branch}_ends", - np.asarray([(branch + 1) * 4], dtype=np.int64), + cache.unique_name(prefix + "_slice_axes"), + np.asarray([3], dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - current = f"{prefix}_{branch}_slice" - nodes.append( - helper.make_node( - "Slice", [reshape.input[0], start, end, axes], [current], name=current + nodes, outputs = [], [] + for branch in range(2): + start = add_initializer( + graph, + cache.unique_name(f"{prefix}_{branch}_starts"), + np.asarray([branch * 4], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, ) - ) - for index, original in enumerate(middle): - name = f"{prefix}_{branch}_qdq_{index}" - # The C++ reference uses the quantizer's scale and zero point for both. + end = add_initializer( + graph, + cache.unique_name(f"{prefix}_{branch}_ends"), + np.asarray([(branch + 1) * 4], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + current = cache.unique_name(f"{prefix}_{branch}_slice") nodes.append( helper.make_node( - original.op_type, [current, *middle[0].input[1:3]], [name], name=name + "Slice", [reshape.input[0], start, end, axes], [current], name=current ) ) - current = name - output = f"{prefix}_{branch}_reduce" + duplicated, current = _duplicate_clip_qdq_chain( + middle, current, f"{prefix}_{branch}", cache + ) + nodes.extend(duplicated) + output = cache.unique_name(f"{prefix}_{branch}_reduce") + nodes.append( + helper.make_node( + "ReduceSum", [current, *reduce.input[1:]], [output], name=output + ) + ) + outputs.append(output) nodes.append( - helper.make_node("ReduceSum", [current, *reduce.input[1:]], [output], name=output) - ) - outputs.append(output) - nodes.append( - helper.make_node( - "Concat", outputs, list(reduce.output), name=prefix + "_axis_updated", axis=3 + helper.make_node( + "Concat", + outputs, + list(reduce.output), + name=cache.unique_name(prefix + "_axis_updated"), + axis=3, + ) ) - ) - for original in [*middle, reduce]: - replacements[id(original)] = [] - replacements[id(reshape)] = nodes - count += 1 + for original in [*middle, reduce]: + replacements[id(original)] = [] + replacements[id(reshape)] = nodes + count += 1 + except Exception: + del graph.initializer[initializer_count:] + raise _replace(graph, replacements) return model def _apply_reshape_transpose_reshape(model): + """Reuse the validated transpose builder for the separately ordered direct pattern.""" graph = model.graph cache = GraphCache(graph) replacements = {} count = 0 - for first in graph.node: - if first.op_type != "Reshape" or "_transformed" in first.name: - continue - shape = cache.get_init_array(first.input[1]) - if shape is None or shape.size != 5: - continue - transpose = _follower(cache, first.output[0], "Transpose") - if transpose is None: - continue - perm = _attribute(transpose, "perm") - if perm is None: - continue - last = _follower(cache, transpose.output[0], "Reshape") - if last is None: - continue - target = cache.get_init_array(last.input[1]) - if target is None: - continue - input_shape = cache.get_shape(first.input[0]) - if not input_shape or input_shape[0] != int(shape[0]): - raise ValueError(f"{first.name!r}: leading input dimension must equal 5D slice count") - if shape[0] <= 0 or not perm.ints or perm.ints[0] != 0: - raise ValueError(f"{first.name!r}: unsupported leading dimension/permutation") - prefix = f"reshape_transpose_{count}" - axes = add_initializer( - graph, - prefix + "_slice_axes", - np.asarray([0], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - middle_shape = add_initializer( - graph, - prefix + "_4d_shape", - np.asarray(shape.tolist()[1:], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - final_shape = add_initializer( - graph, - prefix + "_final_shape", - np.asarray(target, dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - nodes, outputs = [], [] - for branch in range(int(shape[0])): - start = add_initializer( - graph, - f"{prefix}_{branch}_start", - np.asarray([branch], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - end = add_initializer( - graph, - f"{prefix}_{branch}_end", - np.asarray([branch + 1], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - sliced, reshaped, transposed = [ - f"{prefix}_{branch}_{s}" for s in ("slice", "reshape", "transpose") - ] - nodes.extend( - [ - helper.make_node( - "Slice", [first.input[0], start, end, axes], [sliced], name=sliced - ), - helper.make_node("Reshape", [sliced, middle_shape], [reshaped], name=reshaped), - helper.make_node( - "Transpose", - [reshaped], - [transposed], - name=transposed, - perm=[value - 1 for value in perm.ints[1:]], - ), - ] + initializer_count = len(graph.initializer) + try: + for first in graph.node: + if not _is_5d_reshape_candidate(cache, first) or "_transformed" in first.name: + continue + if id(first) in replacements: + raise ValueError(f"{first.name!r}: overlapping reshape patterns") + transpose = _follower(cache, first.output[0], "Transpose") + if transpose is None: + continue + last = _follower(cache, transpose.output[0], "Reshape") + if last is None: + continue + removed, added = _build_transpose_replacement_or_raise( + model, first, [], transpose, set(cache.init_names), count, cache ) - outputs.append(transposed) - concat = prefix + "_concat" - nodes.extend( - [ - helper.make_node( - "Concat", outputs, [concat], name=prefix + "_axis_updated", axis=0 - ), - helper.make_node( - "Reshape", [concat, final_shape], list(last.output), name=prefix + "_final" - ), - ] - ) - replacements.update({id(first): nodes, id(transpose): [], id(last): []}) - count += 1 + replacements.update({id(node): [] for node in removed}) + replacements[id(first)] = added + count += 1 + except Exception: + del graph.initializer[initializer_count:] + raise _replace(graph, replacements) return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 883996fa32c..100480015ed 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -52,6 +52,7 @@ def _remove_deqlin_and_unused_initializers(model): + """Fold eligible constant dequantizations and discard unused initializers.""" return _remove_unused_initializers(_apply_remove_deqlin(model)) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py index e4d5c5813bc..71ee9b91008 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py @@ -57,6 +57,7 @@ import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_remove_deqlin as _mod @@ -75,15 +76,19 @@ def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Create typed tensor metadata.""" return helper.make_tensor_value_info(name, elem_type, shape) 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 _prepare_model(model: onnx.ModelProto) -> None: + """Infer intermediate metadata when the optional runtime supports the graph.""" try: + # ONNX Runtime is optional for these execution tests. from onnxruntime.tools.symbolic_shape_infer import ( SymbolicShapeInference, # type: ignore[import-untyped] ) @@ -97,6 +102,8 @@ def _prepare_model(model: onnx.ModelProto) -> None: def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + """Execute the graph outputs with the CPU runtime.""" + # ONNX Runtime is optional for these execution tests. import onnxruntime as ort # type: ignore[import-untyped] _clamp_ir(model) @@ -105,10 +112,12 @@ def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: def _has_dq(model: onnx.ModelProto) -> bool: + """Check whether the graph retains a dequantization node.""" return any(n.op_type == "DequantizeLinear" for n in model.graph.node) def _init_names(model: onnx.ModelProto) -> set[str]: + """Collect graph initializer names.""" return {i.name for i in model.graph.initializer} @@ -133,6 +142,7 @@ def _make_dq_init( def _make_dq_node(name: str, output_name: str) -> onnx.NodeProto: + """Build a dequantization node using the named controls.""" return helper.make_node( "DequantizeLinear", inputs=[f"{name}_x", f"{name}_scale", f"{name}_zp"], @@ -543,3 +553,325 @@ def test_no_dq_nodes_noop(): assert len(tr.graph.node) == original_count np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +@pytest.mark.parametrize( + "case", + [ + "default-axis", + "scalar-rank-one", + "scalar-rank-zero", + "runtime-zero-point", + "overridable-data", + "overridable-scale", + "overridable-zero-point", + "fp16-scale", + ], +) +def test_constant_folding_preserves_runtime_and_type_contracts(case): + """Fold only supported fixed controls, retaining runtime overrides and output types.""" + data = np.asarray([[2, 4, 6], [8, 10, 12]], dtype=np.uint8) + if case == "scalar-rank-zero": + data = np.asarray(5, np.uint8) + if case == "scalar-rank-one": + data = data.reshape(-1) + scale = np.asarray( + [0.25, 0.5, 1.0] if case == "default-axis" else 0.5, + dtype=np.float16 if case == "fp16-scale" else np.float32, + ) + if case == "scalar-rank-zero": + scale = np.asarray([0.5], np.float32) + initializers = [ + numpy_helper.from_array(data, name="data"), + numpy_helper.from_array(scale, name="scale"), + ] + inputs, feeds = [], {} + node_inputs = ["data", "scale"] + if "zero-point" in case: + node_inputs.append("zero") + inputs.append(_vi("zero", TensorProto.UINT8, [])) + feeds["zero"] = np.asarray(2, np.uint8) + if case == "overridable-zero-point": + initializers.append(numpy_helper.from_array(np.asarray(0, np.uint8), name="zero")) + for name, value in [("data", data), ("scale", scale)]: + if case == f"overridable-{name}": + inputs.append( + _vi( + name, + TensorProto.UINT8 if name == "data" else TensorProto.FLOAT, + list(value.shape), + ) + ) + feeds[name] = value + np.asarray(2, dtype=value.dtype) + out_type = TensorProto.FLOAT16 if case == "fp16-scale" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node("DequantizeLinear", node_inputs, ["Y"], name="dq")], + case, + inputs, + [_vi("Y", out_type, list(data.shape))], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds)[0] + before = model.SerializeToString() + _apply(model) + onnx.checker.check_model(model, full_check=True) + actual = ReferenceEvaluator(model).run(None, feeds)[0] + np.testing.assert_array_equal(actual, expected) + assert actual.dtype == expected.dtype + if case in {"default-axis", "scalar-rank-one", "scalar-rank-zero"}: + assert not _has_dq(model) + else: + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shared", [False, True]) +@pytest.mark.parametrize("omit_zero", [False, True]) +def test_quantized_conv_weight_retained_for_all_consumers(shared, omit_zero): + """Retain a required quantized weight with shared uses or an omitted zero point.""" + initializers = [ + numpy_helper.from_array(np.ones([1, 1, 1, 1], dtype=np.uint8), name="weight"), + numpy_helper.from_array(np.asarray(0.25, np.float32), name="scale"), + ] + inputs = ["weight", "scale"] + if not omit_zero: + inputs.append("zero") + initializers.append(numpy_helper.from_array(np.asarray(0, np.uint8), name="zero")) + nodes = [ + helper.make_node("DequantizeLinear", inputs, ["dequantized"], name="dq"), + helper.make_node("Conv", ["X", "dequantized"], ["Y"], name="conv"), + ] + outputs = [_vi("Y", TensorProto.FLOAT, [1, 1, 1, 1])] + if shared: + nodes.append(helper.make_node("Identity", ["dequantized"], ["shared"], name="shared")) + outputs.append(_vi("shared", TensorProto.FLOAT, [1, 1, 1, 1])) + model = helper.make_model( + helper.make_graph( + nodes, + "quantized_weight", + [_vi("X", TensorProto.FLOAT, [1, 1, 1, 1])], + outputs, + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.asarray([[[[2.0]]]], np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + for actual, reference in zip(ReferenceEvaluator(model).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("overridable", [False, True]) +def test_reused_float_initializer_preserves_override_contract(overridable): + """Reuse equal fixed weights while avoiding initializer defaults callers may override.""" + data = np.asarray([[2, 4, 6]], np.uint8) + expected = data.astype(np.float32) * np.float32(0.5) + initializers = [ + numpy_helper.from_array(data, name="data"), + numpy_helper.from_array(np.float32(0.5), name="scale"), + numpy_helper.from_array(expected, name="data_dequantized"), + ] + inputs = [_vi("data_dequantized", TensorProto.FLOAT, [1, 3])] if overridable else [] + model = helper.make_model( + helper.make_graph( + [helper.make_node("DequantizeLinear", ["data", "scale"], ["Y"], name="dq")], + "reused", + inputs, + [_vi("Y", TensorProto.FLOAT, [1, 3])], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"data_dequantized": expected + 10} if overridable else {} + reference = ReferenceEvaluator(model).run(None, feeds) + _apply(model) + onnx.checker.check_model(model, full_check=True) + assert len(model.graph.initializer) == (4 if overridable else 3) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], reference[0]) + + +@pytest.mark.parametrize("domain", ["custom", "com.microsoft"]) +def test_custom_dequantization_function_remains_unchanged(domain): + """Preserve custom functions that share a built-in operation name.""" + function = helper.make_function( + domain, + "DequantizeLinear", + ["data", "scale"], + ["Y"], + [ + helper.make_node("Cast", ["data"], ["cast"], to=TensorProto.FLOAT), + helper.make_node("Neg", ["cast"], ["Y"]), + ], + opset_imports=[helper.make_opsetid("", 21)], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", ["data", "scale"], ["Y"], domain=domain, name="custom_dq" + ) + ], + "custom", + [], + [_vi("Y", TensorProto.FLOAT, [2])], + initializer=[ + numpy_helper.from_array(np.asarray([2, 4], np.uint8), name="data"), + numpy_helper.from_array(np.float32(0.5), name="scale"), + ], + ), + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid(domain, 1)], + functions=[function], + ) + onnx.checker.check_model(model, full_check=True) + reference = ReferenceEvaluator(model).run(None, {}) + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, {})[0], reference[0]) + + +@pytest.mark.parametrize("dtype", [np.int8, np.uint8, np.int16, np.uint16, np.int32]) +@pytest.mark.parametrize("axis", [1, -1]) +def test_supported_integer_axis_folding_is_exact(dtype, axis): + """Keep exact float32 results across supported integer widths and negative axes.""" + bounds = np.iinfo(np.int16 if dtype == np.int32 else dtype) + data = np.asarray([[bounds.min, 0, bounds.max], [1, bounds.max, bounds.min]], dtype=dtype) + scale = np.asarray([0.1234567, 0.25, 0.75], np.float32) + zero = ( + np.zeros(3, dtype=dtype) + if dtype == np.int32 + else np.asarray([bounds.max, 0, 1], dtype=dtype) + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", ["data", "scale", "zero"], ["Y"], axis=axis, name="dq" + ) + ], + "axis_dtype", + [], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + initializer=[ + numpy_helper.from_array(data, name="data"), + numpy_helper.from_array(scale, name="scale"), + numpy_helper.from_array(zero, name="zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + reference = ReferenceEvaluator(model).run(None, {}) + _apply(model) + onnx.checker.check_model(model, full_check=True) + assert not _has_dq(model) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, {})[0], reference[0]) + + +def test_supported_runtime_dequantization_schema_folds(): + """Execute the supported runtime contrib schema before and after constant folding.""" + data = np.asarray([[2, 4, 6], [8, 10, 12]], np.uint16) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", + ["data", "scale", "zero"], + ["Y"], + name="dq", + domain="com.microsoft", + ) + ], + "runtime_dq", + [], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + initializer=[ + numpy_helper.from_array(data, name="data"), + numpy_helper.from_array(np.float32(0.5), name="scale"), + numpy_helper.from_array(np.uint16(2), name="zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, {}) + _apply(model) + onnx.checker.check_model(model, full_check=True) + assert not _has_dq(model) + np.testing.assert_array_equal(_run_ort(model, {})["Y"], expected["Y"]) + + +@pytest.mark.parametrize("axis", [1, -1]) +def test_large_int32_folding_preserves_established_rounding(axis): + """Retain existing INT32 rounding while evaluating intermediate arithmetic in bounded chunks.""" + data = np.asarray( + [[16777217, -1356484993, 1129686157], [16777219, -18865023, 17258856]], np.int32 + ) + scale = np.asarray([0.1234567, 0.25, 0.0001], np.float32) + zero = np.zeros(3, np.int32) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", ["data", "scale", "zero"], ["Y"], name="dq", axis=axis + ) + ], + "large_int32", + [], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + initializer=[ + numpy_helper.from_array(data, name="data"), + numpy_helper.from_array(scale, name="scale"), + numpy_helper.from_array(zero, name="zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + reference = ReferenceEvaluator(model).run(None, {})[0] + expected = ((data - zero) * scale).astype(np.float32) + _apply(model) + onnx.checker.check_model(model, full_check=True) + actual = ReferenceEvaluator(model).run(None, {})[0] + assert not _has_dq(model) + np.testing.assert_array_equal(actual, expected) + np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=0) + + +def test_unknown_matmul_layout_keeps_quantized_weight(): + """Keep quantized weights until later inference can establish conversion eligibility.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "DequantizeLinear", ["weight", "scale", "zero"], ["dequantized"], name="dq" + ), + helper.make_node("MatMul", ["X", "dequantized"], ["Y"], name="matmul"), + ], + "unknown_layout", + [_vi("X", TensorProto.FLOAT, ["M", "K"])], + [_vi("Y", TensorProto.FLOAT, ["M", 4])], + initializer=[ + numpy_helper.from_array(np.arange(12, dtype=np.uint8).reshape(3, 4), name="weight"), + numpy_helper.from_array(np.float32(0.5), name="scale"), + numpy_helper.from_array(np.uint8(0), name="zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.ones([2, 3], np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected[0]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py new file mode 100644 index 00000000000..3838e8aa349 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -0,0 +1,230 @@ +# 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. + +"""Execute and validate the ordered DLA reshape pattern stages.""" + +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_reshape_patterns import ( + _apply_reshape_qdq_reduce_sum, + _apply_reshape_reduce_sum, + _apply_reshape_transpose_reshape, +) + +ort = pytest.importorskip("onnxruntime", reason="ORT required for numerical stage checks") + + +def _make_model( + pattern, + *, + changed_controls=False, + per_axis=False, + zeros=None, + collision=False, + late_control=False, +): + """Build valid transpose or fast/fallback reduction inputs with optional QDQ controls.""" + transpose = pattern == "transpose" + quantized = pattern == "qdq" + input_shape = [2, 3, 4, 5] if transpose else [2, 1, 3, 12 if pattern == "fallback" else 8] + shape5 = [2, 1, 3, 4, 5] if transpose else [2, 1, 3, 3 if pattern == "fallback" else 2, 4] + if zeros == "first": + shape5 = [2, 1, 0, 3, 5] + initializers = [numpy_helper.from_array(np.asarray(shape5, np.int64), name="shape5")] + nodes = [helper.make_node("Reshape", ["X", "shape5"], ["reshaped"], name="reshape")] + current = "reshaped" + if quantized: + scale = np.asarray([0.25] * 4 if per_axis else 0.25, np.float32) + zero = np.asarray([2] * 4 if per_axis else 2, np.uint8) + initializers.extend( + [ + numpy_helper.from_array(scale, name="q_scale"), + numpy_helper.from_array(zero, name="q_zero"), + ] + ) + if changed_controls: + initializers.extend( + [ + numpy_helper.from_array(scale * 2, name="dq_scale"), + numpy_helper.from_array(zero + 2, name="dq_zero"), + ] + ) + q_scale = "q_scale" + if late_control: + nodes.append( + helper.make_node("Identity", [q_scale], ["runtime_scale"], name="scale_producer") + ) + q_scale = "runtime_scale" + attrs = {"axis": 4} if per_axis else {} + nodes.extend( + [ + helper.make_node( + "QuantizeLinear", + [current, q_scale, "q_zero"], + ["quantized"], + name="quantize", + **attrs, + ), + helper.make_node( + "DequantizeLinear", + [ + "quantized", + "dq_scale" if changed_controls else q_scale, + "dq_zero" if changed_controls else "q_zero", + ], + ["dequantized"], + name="dequantize", + **attrs, + ), + ] + ) + current = "dequantized" + if transpose: + target = [2, -1, 0, 4] if zeros == "final" else [2, 3, 5, 4] + output_shape = [2, 15, 1, 4] if zeros == "final" else [2, 3, 5, 4] + initializers.append(numpy_helper.from_array(np.asarray(target, np.int64), name="target")) + nodes.extend( + [ + helper.make_node( + "Transpose", [current], ["transposed"], name="transpose", perm=[0, 2, 1, 4, 3] + ), + helper.make_node("Reshape", ["transposed", "target"], ["Y"], name="final"), + ] + ) + else: + output_shape = [2, 1, 3, shape5[-2]] + initializers.append(numpy_helper.from_array(np.asarray([-1], np.int64), name="axes")) + nodes.append( + helper.make_node("ReduceSum", [current, "axes"], ["Y"], name="reduce", keepdims=0) + ) + outputs = [helper.make_tensor_value_info("Y", TensorProto.FLOAT, output_shape)] + if collision: + name = "reshape_qdq_reduce_0_0_slice" + nodes.insert(0, helper.make_node("Identity", ["X"], [name], name="existing_output")) + outputs.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, input_shape)) + model = helper.make_model( + helper.make_graph( + nodes, + pattern, + [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], + outputs, + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + model.ir_version = 10 + onnx.checker.check_model(model, full_check=True) + feeds = {"X": (np.arange(np.prod(input_shape), dtype=np.float32).reshape(input_shape) - 17) / 8} + return model, feeds + + +def _execute(model, feeds, runtime): + """Execute all graph outputs with the selected independent engine.""" + if runtime == "reference": + return ReferenceEvaluator(model).run(None, feeds) + session = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return session.run(None, feeds) + + +def _transform(pattern): + """Select the registered stage for each synthetic pattern.""" + return { + "transpose": _apply_reshape_transpose_reshape, + "qdq": _apply_reshape_qdq_reduce_sum, + }.get(pattern, _apply_reshape_reduce_sum) + + +@pytest.mark.parametrize("pattern", ["plain", "qdq", "fallback", "transpose"]) +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +def test_registered_reshape_paths_execute(pattern, runtime): + """Run fast, QDQ, fallback and nontrivial transpose stages through real engines.""" + model, feeds = _make_model(pattern) + expected = _execute(model, feeds, runtime) + transformed = _transform(pattern)(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert not any( + node.op_type == "Reshape" and node.input[1] == "shape5" for node in transformed.graph.node + ) + if pattern == "qdq": + assert sum(node.op_type == "QuantizeLinear" for node in transformed.graph.node) == 2 + assert sum(node.op_type == "DequantizeLinear" for node in transformed.graph.node) == 2 + for actual, reference in zip(_execute(transformed, feeds, runtime), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +def test_fast_qdq_keeps_each_nodes_controls(runtime): + """Preserve different quantizer and dequantizer controls in every branch.""" + model, feeds = _make_model("qdq", changed_controls=True) + expected = _execute(model, feeds, runtime) + transformed = _apply_reshape_qdq_reduce_sum(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + for actual, reference in zip(_execute(transformed, feeds, runtime), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_fast_per_axis_qdq_refuses_before_mutation(): + """Refuse rank-sensitive per-axis middleware before adding branch initializers.""" + model, feeds = _make_model("qdq", per_axis=True) + ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="per-tensor"): + _apply_reshape_qdq_reduce_sum(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("zeros", ["first", "final"]) +def test_transpose_resolves_copied_dimensions(zeros): + """Resolve copied dimensions against their original input before changing rank.""" + model, feeds = _make_model("transpose", zeros=zeros) + expected = ReferenceEvaluator(model).run(None, feeds) + transformed = _apply_reshape_transpose_reshape(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + for actual, reference in zip(ReferenceEvaluator(transformed).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("option", ["collision", "late_control"]) +def test_qdq_rewrite_preserves_names_and_control_order(option): + """Retain unrelated public tensors and controls produced after the original reshape.""" + model, feeds = _make_model("qdq", **{option: True}) + expected = ReferenceEvaluator(model).run(None, feeds) + transformed = _apply_reshape_qdq_reduce_sum(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + for actual, reference in zip(ReferenceEvaluator(transformed).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_later_refusal_restores_earlier_planning_initializers(): + """Leave the whole model unchanged when an unsupported second pattern is found.""" + first, _ = _make_model("qdq") + second, _ = _make_model("qdq", per_axis=True) + first = onnx.compose.add_prefix(first, "first_") + second = onnx.compose.add_prefix(second, "second_") + first.graph.node.extend(second.graph.node) + first.graph.initializer.extend(second.graph.initializer) + first.graph.input.extend(second.graph.input) + first.graph.output.extend(second.graph.output) + onnx.checker.check_model(first, full_check=True) + before = first.SerializeToString() + with pytest.raises(ValueError, match="per-tensor"): + _apply_reshape_qdq_reduce_sum(first) + assert first.SerializeToString() == before From 8d2da63e6140896a4886c91a41e344f793c7bb8a Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 19:51:00 +0000 Subject: [PATCH 3/6] Preserve DLA control-flow captures and complete pattern matching Signed-off-by: mgohil --- .../dla_transforms/dla_remove_deqlin.py | 4 + .../dla_transforms/dla_reshape_patterns.py | 17 ++- .../dla_transforms/test_dla_remove_deqlin.py | 77 +++++++++++- .../test_dla_reshape_patterns.py | 117 ++++++++++++++++++ 4 files changed, 208 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py index 9142e5a9846..e0f37e92f81 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_deqlin.py @@ -38,6 +38,7 @@ check_to_apply_transpose, get_node_attr_i, ) +from .dla_graph_cleanup import _subgraph_captures # ── Zero-point types that trigger the skip rule ────────────────────────────── _CONV_SKIP_ZP = frozenset( @@ -122,6 +123,7 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: """Replace eligible DequantizeLinear nodes with pre-computed float32 initializers.""" graph = model.graph cache = GraphCache(graph) + captured = _subgraph_captures(graph) local_functions = { (function.domain, function.name, function.overload) for function in model.functions } @@ -155,6 +157,8 @@ def _apply_remove_deqlin(model: onnx.ModelProto) -> onnx.ModelProto: continue # data input is not a static initializer — skip dq_output = node.output[0] + if dq_output in captured: + continue # Fetch scale and zero-point initializers (inputs 1 and 2) scale_init = cache.get_init(node.input[1]) if len(node.input) > 1 else None diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py index 5aa562cf28f..23700566ac6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -26,6 +26,7 @@ _is_5d_reshape_candidate, _known_non_scalar_control, ) +from .dla_graph_cleanup import _subgraph_captures from .dla_reorder_nodes import priority_topological_order __all__ = [] @@ -51,9 +52,17 @@ def _follower(cache, output, op_type): def _replace(graph, replacements): """Apply planned replacements without copying untouched payloads and order late controls.""" + if not replacements: + return removed_outputs = { output for node in graph.node if id(node) in replacements for output in node.output } + added_outputs = { + output for nodes in replacements.values() for node in nodes for output in node.output + } + removed_captures = (removed_outputs - added_outputs) & _subgraph_captures(graph) + if removed_captures: + raise ValueError(f"Cannot rewrite captured tensors: {sorted(removed_captures)}") positions = {output: index for index, node in enumerate(graph.node) for output in node.output} needs_reorder = any( positions.get(name, -1) > index and name not in removed_outputs @@ -204,10 +213,10 @@ def _reshape_reduce_sum(model, *, quantized): replacements[id(original)] = [] replacements[id(reshape)] = nodes count += 1 + _replace(graph, replacements) except Exception: del graph.initializer[initializer_count:] raise - _replace(graph, replacements) return model @@ -222,22 +231,22 @@ def _apply_reshape_transpose_reshape(model): for first in graph.node: if not _is_5d_reshape_candidate(cache, first) or "_transformed" in first.name: continue - if id(first) in replacements: - raise ValueError(f"{first.name!r}: overlapping reshape patterns") transpose = _follower(cache, first.output[0], "Transpose") if transpose is None: continue last = _follower(cache, transpose.output[0], "Reshape") if last is None: continue + if id(first) in replacements: + raise ValueError(f"{first.name!r}: overlapping reshape patterns") removed, added = _build_transpose_replacement_or_raise( model, first, [], transpose, set(cache.init_names), count, cache ) replacements.update({id(node): [] for node in removed}) replacements[id(first)] = added count += 1 + _replace(graph, replacements) except Exception: del graph.initializer[initializer_count:] raise - _replace(graph, replacements) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py index 71ee9b91008..95ae4edf377 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_deqlin.py @@ -51,7 +51,6 @@ from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx @@ -63,8 +62,6 @@ pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "remove_deqlin_test_models" _OPSET_VERSION = 21 _MAX_IR_VERSION_FOR_ORT = 10 @@ -875,3 +872,77 @@ def test_unknown_matmul_layout_keeps_quantized_weight(): _apply(model) assert model.SerializeToString() == before np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("depth", [1, 2]) +@pytest.mark.parametrize("direct_consumer", [False, True]) +def test_dequantization_preserves_nested_captures(depth, direct_consumer): + """Preserve DQ outputs captured by nested control flow with or without direct uses.""" + shape = [2, 3] + branch = helper.make_graph( + [helper.make_node("Identity", ["dequantized"], ["branch_output"])], + "branch", + [], + [_vi("branch_output", TensorProto.FLOAT, shape)], + ) + if depth == 2: + branch = helper.make_graph( + [ + helper.make_node( + "If", + ["condition"], + ["nested"], + then_branch=branch, + else_branch=copy.deepcopy(branch), + ) + ], + "nested_branch", + [], + [_vi("nested", TensorProto.FLOAT, shape)], + ) + nodes = [ + helper.make_node("DequantizeLinear", ["data", "scale", "zero"], ["dequantized"], name="dq") + ] + outputs = [_vi("Y", TensorProto.FLOAT, shape)] + if direct_consumer: + nodes.append(helper.make_node("Identity", ["dequantized"], ["direct"], name="direct")) + outputs.append(_vi("direct", TensorProto.FLOAT, shape)) + nodes.append( + helper.make_node( + "If", + ["condition"], + ["Y"], + name="capture", + then_branch=branch, + else_branch=copy.deepcopy(branch), + ) + ) + model = helper.make_model( + helper.make_graph( + nodes, + "captured_dq", + [_vi("condition", TensorProto.BOOL, [])], + outputs, + initializer=[ + numpy_helper.from_array(np.arange(6, dtype=np.uint8).reshape(shape), name="data"), + numpy_helper.from_array(np.asarray(0.5, np.float32), name="scale"), + numpy_helper.from_array(np.asarray(0, np.uint8), name="zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + references = { + condition: ReferenceEvaluator(model).run(None, {"condition": np.asarray(condition)}) + for condition in [False, True] + } + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + onnx.checker.check_model(model, full_check=True) + for condition, expected in references.items(): + for actual, reference in zip( + ReferenceEvaluator(model).run(None, {"condition": np.asarray(condition)}), expected + ): + np.testing.assert_array_equal(actual, reference) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py index 3838e8aa349..62f67dfb4ef 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -228,3 +228,120 @@ def test_later_refusal_restores_earlier_planning_initializers(): with pytest.raises(ValueError, match="per-tensor"): _apply_reshape_qdq_reduce_sum(first) assert first.SerializeToString() == before + + +def _attach_if_capture(model, tensor, shape, dtype=TensorProto.FLOAT): + """Attach a checked control-flow output that captures the requested outer tensor.""" + branch = helper.make_graph( + [helper.make_node("Identity", [tensor], ["branch_output"])], + "branch", + [], + [helper.make_tensor_value_info("branch_output", dtype, shape)], + ) + model.graph.input.append(helper.make_tensor_value_info("condition", TensorProto.BOOL, [])) + model.graph.node.append( + helper.make_node( + "If", + ["condition"], + ["capture_output"], + name="capture", + then_branch=branch, + else_branch=copy.deepcopy(branch), + ) + ) + model.graph.output.append(helper.make_tensor_value_info("capture_output", dtype, shape)) + onnx.checker.check_model(model, full_check=True) + + +@pytest.mark.parametrize( + ("pattern", "tensor"), + [ + ("plain", "reshaped"), + ("fallback", "reshaped"), + ("qdq", "reshaped"), + ("qdq", "quantized"), + ("qdq", "dequantized"), + ("transpose", "reshaped"), + ("transpose", "transposed"), + ], +) +def test_reshape_refuses_captured_intermediates(pattern, tensor): + """Reject removal of captured intermediate names and restore all planned initializers.""" + model, feeds = _make_model(pattern) + shape = ( + [2, 1, 3, 4, 5] + if pattern == "transpose" + else [2, 1, 3, 3 if pattern == "fallback" else 2, 4] + ) + if tensor == "transposed": + shape = [2, 3, 1, 5, 4] + _attach_if_capture( + model, tensor, shape, TensorProto.UINT8 if tensor == "quantized" else TensorProto.FLOAT + ) + feeds["condition"] = np.asarray(True) + ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="captur"): + _transform(pattern)(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("pattern", ["plain", "fallback", "qdq", "transpose"]) +def test_reshape_preserves_captured_terminal_outputs(pattern): + """Allow a captured terminal output when the replacement recreates the same name.""" + model, feeds = _make_model(pattern) + shape = [dim.dim_value for dim in model.graph.output[0].type.tensor_type.shape.dim] + _attach_if_capture(model, "Y", shape) + expected = { + condition: ReferenceEvaluator(model).run( + None, {**feeds, "condition": np.asarray(condition)} + ) + for condition in [False, True] + } + transformed = _transform(pattern)(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + for condition, references in expected.items(): + for actual, reference in zip( + ReferenceEvaluator(transformed).run( + None, {**feeds, "condition": np.asarray(condition)} + ), + references, + ): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("chained", [False, True]) +def test_terminal_five_dimensional_target_is_not_an_overlap(chained): + """Distinguish an unmatched terminal 5D target from two genuinely overlapping triples.""" + model, feeds = _make_model("transpose") + target = next(tensor for tensor in model.graph.initializer if tensor.name == "target") + target.CopyFrom(numpy_helper.from_array(np.asarray([2, 3, 1, 5, 4], np.int64), name="target")) + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 1, 5, 4]) + ) + if chained: + model.graph.node.extend( + [ + helper.make_node("Transpose", ["Y"], ["again"], name="again", perm=[0, 2, 1, 4, 3]), + helper.make_node("Reshape", ["again", "last_target"], ["Z"], name="last"), + ] + ) + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([2, 3, 5, 4], np.int64), name="last_target") + ) + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Z", TensorProto.FLOAT, [2, 3, 5, 4]) + ) + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds) + before = model.SerializeToString() + if chained: + with pytest.raises(ValueError, match="overlap"): + _apply_reshape_transpose_reshape(model) + assert model.SerializeToString() == before + else: + transformed = _apply_reshape_transpose_reshape(model) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal( + ReferenceEvaluator(transformed).run(None, feeds)[0], expected[0] + ) From 0c37970915ee0af45ec0a3cb841daeafe106e1e2 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:52:46 +0000 Subject: [PATCH 4/6] Declare computed scalar QDQ control ranks in the stage fixture Signed-off-by: mgohil --- .../dla_transforms/test_dla_reshape_patterns.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py index 62f67dfb4ef..5fb2de37233 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -127,6 +127,11 @@ def _make_model( [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)], outputs, initializer=initializers, + value_info=[ + helper.make_tensor_value_info("runtime_scale", TensorProto.FLOAT, list(scale.shape)) + ] + if quantized and late_control + else [], ), opset_imports=[helper.make_opsetid("", 21)], ) From 8325db82c36bb28f6decbe7e1ebd08b051923b53 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 14:25:41 +0000 Subject: [PATCH 5/6] Retain generated transpose shapes for ordered reductions Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 20 +++++++- .../dla_transforms/dla_reshape_patterns.py | 12 ++++- .../test_dla_reshape_patterns.py | 49 +++++++++++++++++++ 3 files changed, 79 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 572fa6f19b9..635654ee10a 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 @@ -250,7 +250,15 @@ 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 + model, + reshape_node, + mid_a, + tr, + unique_initializer: set, + counter: int, + cache: GraphCache, + *, + inferred_value_info: list | None = None, ) -> tuple[list, list]: """Build the slice-based transpose decomposition from reshape targets and permutation. @@ -437,6 +445,16 @@ def _build_transpose_replacement_or_raise( ) ) to_remove = [reshape_node, *mid_a, tr, *mid_b, r2] + if inferred_value_info is not None: + concat_shape = [int(reshape_4d_vals[index]) for index in new_transpose_perm] + concat_shape[concat_axis] *= num_slices + inferred_value_info.append( + helper.make_tensor_value_info( + concat_out, + cache.get_dtype(output_name) or cache.get_dtype(input_tensor) or TensorProto.FLOAT, + concat_shape, + ) + ) return to_remove, nodes_to_add diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py index 23700566ac6..60a260b1cf4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -225,6 +225,7 @@ def _apply_reshape_transpose_reshape(model): graph = model.graph cache = GraphCache(graph) replacements = {} + inferred_value_info = [] count = 0 initializer_count = len(graph.initializer) try: @@ -240,12 +241,21 @@ def _apply_reshape_transpose_reshape(model): if id(first) in replacements: raise ValueError(f"{first.name!r}: overlapping reshape patterns") removed, added = _build_transpose_replacement_or_raise( - model, first, [], transpose, set(cache.init_names), count, cache + model, + first, + [], + transpose, + set(cache.init_names), + count, + cache, + inferred_value_info=inferred_value_info, ) replacements.update({id(node): [] for node in removed}) replacements[id(first)] = added count += 1 _replace(graph, replacements) + # Ordered reduction stages read new Concat shapes before the next global inference. + graph.value_info.extend(inferred_value_info) except Exception: del graph.initializer[initializer_count:] raise diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py index 5fb2de37233..a7fb53af5bc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -28,6 +28,7 @@ _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, ) +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible ort = pytest.importorskip("onnxruntime", reason="ORT required for numerical stage checks") @@ -175,6 +176,54 @@ def test_registered_reshape_paths_execute(pattern, runtime): np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("entrypoint", ["stages", "pipeline"]) +def test_transpose_then_reduction_recovers_generated_concat_shape(runtime, quantized, entrypoint): + """Execute reduction of a new 5D transpose result without a full inference pass between stages.""" + model, feeds = _make_model("transpose") + target = next(tensor for tensor in model.graph.initializer if tensor.name == "target") + target.CopyFrom(numpy_helper.from_array(np.array([2, 3, 1, 5, 4], np.int64), "target")) + model.graph.node[-1].output[0] = "wide" + current = "wide" + if quantized: + model.graph.initializer.extend( + [ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.uint8(128), "zero"), + ] + ) + model.graph.node.extend( + [ + helper.make_node("QuantizeLinear", [current, "scale", "zero"], ["quantized"]), + helper.make_node( + "DequantizeLinear", ["quantized", "scale", "zero"], ["dequantized"] + ), + ] + ) + current = "dequantized" + model.graph.initializer.append(numpy_helper.from_array(np.array([4], np.int64), "reduce_axes")) + model.graph.node.append( + helper.make_node("ReduceSum", [current, "reduce_axes"], ["Y"], keepdims=0) + ) + model.graph.output[0].CopyFrom( + helper.make_tensor_value_info("Y", TensorProto.FLOAT, [2, 3, 1, 5]) + ) + model = onnx.shape_inference.infer_shapes(model) + onnx.checker.check_model(model, full_check=True) + expected = _execute(model, feeds, runtime) + if entrypoint == "pipeline": + transformed = _transform_make_dla_compatible( + copy.deepcopy(model), verbose=False, qdq_quantized_dtypes=[] + ) + else: + transformed = _apply_reshape_transpose_reshape(copy.deepcopy(model)) + transformed = _apply_reshape_qdq_reduce_sum(transformed) + transformed = _apply_reshape_reduce_sum(transformed) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(_execute(transformed, feeds, runtime)[0], expected[0]) + + @pytest.mark.parametrize("runtime", ["reference", "ort"]) def test_fast_qdq_keeps_each_nodes_controls(runtime): """Preserve different quantizer and dequantizer controls in every branch.""" From 20d3ac1dc09a85ca7898f684061d7032c5e52d01 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 15:01:10 +0000 Subject: [PATCH 6/6] Retain generated branch layouts for later rank stages Signed-off-by: mgohil --- .../dla_transforms/dla_5d_reshape_to_4d.py | 12 ++++++++++++ 1 file changed, 12 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 635654ee10a..b718c71d5ca 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 @@ -425,6 +425,18 @@ def _build_transpose_replacement_or_raise( ) nodes_to_add.extend(br_b) slice_outputs.append(t_b) + if inferred_value_info is not None: + assert tr_in_shape is not None + sliced_shape = list(tr_in_shape) + sliced_shape[slice_axis] = 1 + transposed_shape = [int(reshape_4d_vals[index]) for index in new_transpose_perm] + dtype = cache.get_dtype(input_tensor) or TensorProto.FLOAT + for name, shape in ( + (slice_out, sliced_shape), + (reshape_4d_out, reshape_4d_vals.tolist()), + (transpose_branch_out, transposed_shape), + ): + inferred_value_info.append(helper.make_tensor_value_info(name, dtype, shape)) concat_out = cache.unique_name(f"u5d_tr_concat_{counter}") nodes_to_add.append(