diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_constants_to_initializers.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_constants_to_initializers.py new file mode 100644 index 00000000000..7e525fca4fa --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_constants_to_initializers.py @@ -0,0 +1,115 @@ +# 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 ``Constant`` ops into graph ``initializer`` entries.""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import add_unique_initializers + + +def _tensor_from_constant_node(node: onnx.NodeProto) -> onnx.TensorProto: + """Build a dense ``TensorProto`` initializer from a ``Constant`` node's attributes.""" + if not node.output: + raise ValueError(f"Constant node {node.name!r} has no output") + out_name = node.output[0] + + for attr in node.attribute: + if attr.name == "value": + arr = numpy_helper.to_array(attr.t) + return numpy_helper.from_array(arr, name=out_name) + if attr.name == "value_float": + arr = np.asarray(attr.f, dtype=np.float32).reshape(()) + return numpy_helper.from_array(arr, name=out_name) + if attr.name == "value_floats": + arr = np.asarray(list(attr.floats), dtype=np.float32) + return numpy_helper.from_array(arr, name=out_name) + if attr.name == "value_int": + arr = np.asarray(np.int64(attr.i), dtype=np.int64).reshape(()) + return numpy_helper.from_array(arr, name=out_name) + if attr.name == "value_ints": + arr = np.asarray(list(attr.ints), dtype=np.int64) + return numpy_helper.from_array(arr, name=out_name) + if attr.name == "sparse_value": + sparse = attr.sparse_tensor + values = numpy_helper.to_array(sparse.values) + indices = numpy_helper.to_array(sparse.indices) + arr = np.zeros(sparse.dims, dtype=values.dtype) + if indices.ndim == 1: + arr.reshape(-1)[indices] = values + else: + arr[tuple(indices.T)] = values + return numpy_helper.from_array(arr, name=out_name) + + attr_names = [a.name for a in node.attribute] + raise ValueError( + f"Constant node {node.name!r} has no supported value attribute " + f"(expected value, value_float(s), value_int(s), or sparse_value); got {attr_names}" + ) + + +def _add_legacy_initializer_inputs(model: onnx.ModelProto) -> None: + """Declare initializer defaults as graph inputs for ONNX IR versions through 3.""" + if model.ir_version > 3: + return + graph = model.graph + input_names = {value.name for value in graph.input} + for init in graph.initializer: + if init.name not in input_names: + graph.input.append( + helper.make_tensor_value_info(init.name, init.data_type, list(init.dims)) + ) + input_names.add(init.name) + + +def _constants_to_initializers_from_model( + model: onnx.ModelProto, + *, + verbose: bool, +) -> int: + """Fold supported Constant nodes in place and return their conversion count.""" + graph = model.graph + unique: set[str] = {i.name for i in graph.initializer} + nodes_to_remove: list[onnx.NodeProto] = [] + + for node in list(graph.node): + if node.op_type != "Constant" or node.domain not in {"", "ai.onnx"} or node.overload: + continue + if not node.output: + continue + try: + init = _tensor_from_constant_node(node) + except ValueError as exc: + if verbose: + logger.warning("%s", exc) + continue + before = len(unique) + add_unique_initializers(graph, unique, [init]) + if len(unique) > before: + nodes_to_remove.append(node) + + for node in nodes_to_remove: + graph.node.remove(node) + if nodes_to_remove: + _add_legacy_initializer_inputs(model) + + if verbose: + logger.debug("Converted %d Constant nodes to initializers", len(nodes_to_remove)) + return len(nodes_to_remove) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_handle_qdq.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_handle_qdq.py new file mode 100644 index 00000000000..82de3409b23 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_handle_qdq.py @@ -0,0 +1,156 @@ +# 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. + +"""Wrap every QuantizeLinear / DequantizeLinear node with Unsqueeze → op → Squeeze. + +Both Q and DQ are handled independently regardless of whether their data input +is a static initializer or a dynamic tensor. Scale and zero-point inputs are +left untouched. The ``axis`` attribute is shifted by ``4 - orig_rank`` to +keep per-channel quantization referencing the correct dimension. +Blocked quantization is left unchanged because its scales track the data rank. + +Pattern applied to each node:: + + data [*orig_shape] + → Unsqueeze → data_4d [1,…,1, *orig_shape] + → Q / DQ → out_4d [1,…,1, *orig_shape] + → Squeeze → out [*orig_shape] +""" + +from __future__ import annotations + +import onnx +from onnx import helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_value_info, + axes_for_rank, + get_node_attr_i, + pad4d, + set_node_attr_i, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _set_axes +from .dla_remove_qdq import _qdq_output_dtype + +_axes_for_rank = axes_for_rank + + +def _apply_handle_qdq(model: onnx.ModelProto) -> onnx.ModelProto: + """Wrap eligible Q / DQ nodes with Unsqueeze → node → Squeeze.""" + graph = model.graph + cache = GraphCache(graph) + versions = {opset.domain: opset.version for opset in model.opset_import} + opset = versions.get("", versions.get("ai.onnx", 1)) + + replacements: dict[int, list[onnx.NodeProto]] = {} + cnt = [0] + + for node_pos, node in enumerate(list(graph.node)): + if node.domain not in {"", "ai.onnx", "com.microsoft"} or node.overload: + continue + if node.op_type not in ("QuantizeLinear", "DequantizeLinear"): + continue + if not node.input or not node.output: + continue + if get_node_attr_i(node, "block_size", 0) > 0: + continue + + data_name = node.input[0] + orig_shape: list | None = cache.get_shape(data_name) + if orig_shape is None: + info = cache.vi_map.get(data_name) + if info is not None and info.type.tensor_type.HasField("shape"): + orig_shape = [ + dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None + for dim in info.type.tensor_type.shape.dim + ] + + if orig_shape is None: + logger.debug("%s %r: data input shape unknown; skipping.", node.op_type, node.name) + continue + if len(orig_shape) not in (1, 2, 3): + continue # already 4-D, nothing to do + + orig_rank = len(orig_shape) + delta = 4 - orig_rank + axes = _axes_for_rank(orig_rank) + pfx = node.name or f"qdq_{cnt[0]}" + c = cnt[0] + cnt[0] += 1 + + # ── Unsqueeze data input to 4D ──────────────────────────────────────── + in4d = cache.unique_name(f"{pfx}_in4d_{c}") + unsq = helper.make_node( + "Unsqueeze", + inputs=[data_name], + outputs=[in4d], + name=cache.unique_name(f"{pfx}_unsq_{c}"), + ) + _set_axes(model, cache, unsq, axes, f"{pfx}_unsq_axes_{c}") + in_dtype = cache.get_dtype(data_name) or onnx.TensorProto.FLOAT + add_value_info(graph, in4d, in_dtype, pad4d(orig_shape)) + node.input[0] = in4d + + # ── Update axis attribute (ONNX Q/DQ default: 1) ───────────────────── + has_axis = any(a.name == "axis" for a in node.attribute) + if (opset >= 13 or node.domain == "com.microsoft") and ( + has_axis or (node.domain != "com.microsoft" and len(orig_shape) > 1) + ): + orig_axis = get_node_attr_i(node, "axis", 1) + if orig_axis >= 0: + set_node_attr_i(node, "axis", orig_axis + delta) + + # ── Squeeze output back to original rank ────────────────────────────── + orig_out = node.output[0] + out4d = cache.unique_name(f"{pfx}_out4d_{c}") + out_dtype = cache.get_dtype(orig_out) or _qdq_output_dtype(node, cache) + if out_dtype: + add_value_info(graph, out4d, out_dtype, pad4d(orig_shape)) + node.output[0] = out4d + + c2 = cnt[0] + cnt[0] += 1 + sq = helper.make_node( + "Squeeze", + inputs=[out4d], + outputs=[orig_out], + name=cache.unique_name(f"{pfx}_sq_{c2}"), + ) + _set_axes(model, cache, sq, axes, f"{pfx}_sq_axes_{c2}") + + replacements[node_pos] = [unsq, node, sq] + logger.debug( + "%s %r: wrapped Unsqueeze→op→Squeeze (rank %d→4, axis delta=%d).", + node.op_type, + node.name, + orig_rank, + delta, + ) + + if replacements: + nodes = [ + replacement + for index, node in enumerate(graph.node) + for replacement in replacements.get(index, [node]) + ] + del graph.node[:] + graph.node.extend(nodes) + _add_legacy_initializer_inputs(model) + + logger.debug("handle_qdq: wrapped %d Q/DQ node(s) with Unsqueeze/Squeeze.", len(replacements)) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_qdq.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_qdq.py new file mode 100644 index 00000000000..5ef7a1240c9 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_qdq.py @@ -0,0 +1,313 @@ +# 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. + +"""Remove paired QuantizeLinear / DequantizeLinear nodes for selected quantized dtypes. + +Pairs are removed when the quantizer's storage element type matches one of the +user-configurable storage types (default: UINT16 and INT16). See :mod:`utils.graph_utils` for +name/NumPy mappings. +Scope-captured pairs are preserved. Input clipping requires matching data/output +types and scalar static bounds; unsupported clipping is rejected before mutation. +""" + +from __future__ import annotations + +from collections import defaultdict + +import numpy as np +import onnx +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + calculate_clip_range, + get_node_attr_i, + insert_nodes_at_position, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures + + +def _qdq_output_dtype(node: onnx.NodeProto, cache: GraphCache) -> int | None: + """Resolve Q/DQ output types from explicit attributes and typed control inputs.""" + declared = get_node_attr_i(node, "output_dtype", 0) + if declared: + return declared + index = 2 if node.op_type == "QuantizeLinear" else 1 + if len(node.input) > index and node.input[index]: + return cache.get_dtype(node.input[index]) + return onnx.TensorProto.UINT8 if node.op_type == "QuantizeLinear" else None + + +def _remove_qdq_from_model( + model: onnx.ModelProto, + *, + keep_clip_after_inputs: bool, + max_chained_qdq_pairs: int, + quant_types_to_remove: frozenset[int], +) -> int: + """Mutate ``model`` in place. Returns number of QDQ pairs removed.""" + q_output_to_q_node_map: dict = {} + dq_input_to_dq_node_map: dict = defaultdict(list) + graph = model.graph + cache = GraphCache(graph) + versions = {opset.domain: opset.version for opset in model.opset_import} + opset = versions.get("", versions.get("ai.onnx", 1)) + unique_initializer: set[str] = set() + node_datatype_map: dict[str, int] = {} + captured = _subgraph_captures(graph) + + for node in graph.node: + if node.domain not in {"", "ai.onnx", "com.microsoft"} or node.overload: + continue + if node.op_type == "QuantizeLinear": + if node.input[0] in captured or node.output[0] in captured: + continue + if not node.output[0] or len(node.output) > 1: + continue + q_output_to_q_node_map[node.output[0]] = node + zp_dtype = _qdq_output_dtype(node, cache) + if zp_dtype is not None: + node_datatype_map[node.output[0]] = zp_dtype + in_dtype = cache.get_dtype(node.input[0]) + if in_dtype is not None: + node_datatype_map[node.input[0]] = in_dtype + elif node.op_type == "DequantizeLinear": + if node.output[0] in captured: + continue + if not node.input[0]: + continue + dq_input_to_dq_node_map[node.input[0]].append(node) + out_dtype = cache.get_dtype(node.output[0]) or _qdq_output_dtype(node, cache) + if out_dtype is not None: + node_datatype_map[node.output[0]] = out_dtype + + qdq_node_pair_output_to_input_map: dict[str, str] = {} + cnt = 0 + input_names = {inp.name for inp in model.graph.input} + aliases = {} + input_clips = {} + clip_types = ( + { + kind + for constraint in onnx.defs.get_schema("Clip", opset).type_constraints + for kind in constraint.allowed_type_strs + } + if keep_clip_after_inputs + else set() + ) + for q_output, q_node in q_output_to_q_node_map.items(): + if node_datatype_map.get(q_output) not in quant_types_to_remove: + continue + for dq_node in dq_input_to_dq_node_map.get(q_output, []): + in_dtype = node_datatype_map.get(q_node.input[0]) + out_dtype = node_datatype_map.get(dq_node.output[0]) + if in_dtype is None or out_dtype is None: + raise ValueError( + f"Cannot remove QDQ pair ({q_node.name!r}, {dq_node.name!r}): " + "dtype is unknown. Ensure shape inference has been run." + ) + if keep_clip_after_inputs and q_node.input[0] in input_names: + if ( + in_dtype != out_dtype + or f"tensor({onnx.TensorProto.DataType.Name(out_dtype).lower()})" + not in clip_types + ): + raise NotImplementedError( + "Input clipping requires matching types supported by the model's Clip schema" + ) + zero_name = dq_node.input[2] if len(dq_node.input) > 2 else "" + if zero_name and cache.get_init(zero_name) is None: + raise NotImplementedError( + "Fold zero-point constants before preserving input clips" + ) + if ( + not zero_name + and cache.get_dtype(dq_node.input[0]) is None + and node_datatype_map[q_output] != onnx.TensorProto.UINT8 + ): + raise ValueError( + "Infer quantized tensor types before preserving input clips without a zero point" + ) + clip_min, clip_max = calculate_clip_range(dq_node, model) + if clip_min.size != 1 or clip_max.size != 1: + raise NotImplementedError( + "Per-axis input clipping requires an additional transformation" + ) + dtype = onnx.helper.tensor_dtype_to_np_dtype(out_dtype) + input_clips[dq_node.output[0]] = tuple( + np.asarray(bound, dtype=dtype).reshape(()) for bound in (clip_min, clip_max) + ) + continue + if in_dtype == out_dtype: + aliases[dq_node.output[0]] = q_node.input[0] + for output in aliases: + current, visited = output, set() + while current in aliases: + if current in visited or len(visited) > max_chained_qdq_pairs: + raise ValueError( + f"More than {max_chained_qdq_pairs} chained QDQ pairs at tensor {output!r}" + ) + visited.add(current) + current = aliases[current] + + for q_output, q_node in q_output_to_q_node_map.items(): + if q_output not in dq_input_to_dq_node_map: + continue + q_dtype = node_datatype_map.get(q_output) + if q_dtype is None or q_dtype not in quant_types_to_remove: + continue + + for dq_node in dq_input_to_dq_node_map[q_output]: + if keep_clip_after_inputs and q_node.input[0] in input_names: + clip_min, clip_max = input_clips[dq_node.output[0]] + clip_inputs = [q_node.input[0]] + clip_options = {} + if opset < 11: + clip_options = {"min": float(clip_min), "max": float(clip_max)} + else: + for suffix, bound in (("min", clip_min), ("max", clip_max)): + name = cache.unique_name(f"{dq_node.name}_clip_{suffix}") + add_unique_initializers( + graph, unique_initializer, [numpy_helper.from_array(bound, name)] + ) + clip_inputs.append(name) + clip_node = helper.make_node( + "Clip", + inputs=clip_inputs, + outputs=[dq_node.output[0]], + name=cache.unique_name(f"{dq_node.name}_clip"), + **clip_options, + ) + insert_nodes_at_position(graph, [clip_node], q_node) + else: + # The preparation pass validated these dtypes before any mutations. + in_dt = node_datatype_map.get(q_node.input[0]) + out_dt = node_datatype_map.get(dq_node.output[0]) + if in_dt != out_dt: + cast_node = helper.make_node( + "Cast", + inputs=[q_node.input[0]], + outputs=[dq_node.output[0]], + name=f"{dq_node.name}_cast", + to=out_dt, + ) + insert_nodes_at_position(graph, [cast_node], q_node) + else: + qdq_node_pair_output_to_input_map[dq_node.output[0]] = q_node.input[0] + + graph.node.remove(dq_node) + cnt += 1 + + dq_ids = {id(node) for node in dq_input_to_dq_node_map[q_output]} + if q_output not in cache.graph_output_names and all( + id(node) in dq_ids for node, _ in cache.get_consumers(q_output) + ): + graph.node.remove(q_node) + + for node in graph.node: + for i, input_name in enumerate(node.input): + if input_name not in qdq_node_pair_output_to_input_map: + continue + qdq_node_pair_output = input_name + num_connected = 0 + while True: + qdq_node_pair_input = qdq_node_pair_output_to_input_map[qdq_node_pair_output] + if qdq_node_pair_input not in qdq_node_pair_output_to_input_map: + node.input[i] = qdq_node_pair_input + break + qdq_node_pair_output = qdq_node_pair_input + num_connected += 1 + if num_connected > max_chained_qdq_pairs: + raise ValueError( + f"More than {max_chained_qdq_pairs} chained QDQ pairs at tensor " + f"{input_name!r}; raise max_chained_qdq_pairs if intentional." + ) + + consumer_node_map: dict[str, list] = defaultdict(list) + producer_node_map: dict[str, onnx.NodeProto] = {} + for node in graph.node: + for input_ in node.input: + consumer_node_map[input_].append(node) + for output in node.output: + producer_node_map[output] = node + + renamed_outputs: dict[str, str] = {} + graph_output_names = {output.name for output in graph.output} + for graph_output in graph.output: + if graph_output.name not in qdq_node_pair_output_to_input_map: + continue + qdq_node_pair_output = graph_output.name + num_connected = 0 + while True: + qdq_node_pair_input = qdq_node_pair_output_to_input_map[qdq_node_pair_output] + if qdq_node_pair_input not in qdq_node_pair_output_to_input_map: + node = producer_node_map.get(qdq_node_pair_input) + if qdq_node_pair_input in graph_output_names: + graph.node.append( + helper.make_node( + "Identity", + [qdq_node_pair_input], + [graph_output.name], + name=f"identity_output_{graph_output.name}", + ) + ) + break + if qdq_node_pair_input in renamed_outputs: + graph.node.append( + helper.make_node( + "Identity", + [renamed_outputs[qdq_node_pair_input]], + [graph_output.name], + name=f"identity_fanout_{graph_output.name}", + ) + ) + break + if node is None: + # Terminal tensor is a graph input — insert Identity to bridge it. + identity = helper.make_node( + "Identity", + inputs=[qdq_node_pair_input], + outputs=[graph_output.name], + name=f"identity_bypass_{graph_output.name}", + ) + graph.node.append(identity) + else: + for i, output_name in enumerate(node.output): + if output_name == qdq_node_pair_input: + node.output[i] = graph_output.name + renamed_outputs[qdq_node_pair_input] = graph_output.name + break + consumers = consumer_node_map.get(qdq_node_pair_input, []) + for consumer in consumers: + for j, input_name in enumerate(consumer.input): + if input_name == qdq_node_pair_input: + consumer.input[j] = graph_output.name + break + break + qdq_node_pair_output = qdq_node_pair_input + num_connected += 1 + if num_connected > max_chained_qdq_pairs: + raise ValueError( + f"More than {max_chained_qdq_pairs} chained QDQ pairs for graph output " + f"{graph_output.name!r}." + ) + + if cnt: + _add_legacy_initializer_inputs(model) + logger.debug("Removed %d QuantizeLinear and DequantizeLinear pairs", cnt) + return cnt diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 8431e2040e6..80d645854c9 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -30,7 +30,9 @@ from ..logging_config import logger 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_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, @@ -41,7 +43,23 @@ if TYPE_CHECKING: from collections.abc import Iterable __all__ = [] + + +def _apply_remove_qdq(model, **options): + """Run QDQ removal while retaining the model result for ordered stage dispatch.""" + _remove_qdq_from_model(model, **options) + return model + + +def _apply_constants_to_initializers(model, **options): + """Fold constants while retaining the model result for ordered stage dispatch.""" + _constants_to_initializers_from_model(model, **options) + return model + + DLA_TRANSFORMS = ( + ("RemoveQDQ", _apply_remove_qdq), + ("ConstantsToInitializers", _apply_constants_to_initializers), ("Cast", _apply_cast_to_fp32), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py new file mode 100644 index 00000000000..60650dc6a9b --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_constants_to_initializers.py @@ -0,0 +1,445 @@ +# 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_constants_to_initializers`. + +"""ORT CPU parity tests for ``constants_to_initializers`` / ``_constants_to_initializers_from_model``. + +Transform rule +-------------- +Every ``Constant`` node is removed and replaced by a graph initializer carrying +the same tensor data. Downstream consumers are unaffected because they +reference the tensor by name, which is preserved. + +Supported Constant payload attributes +-------------------------------------- +``value`` — dense TensorProto (any dtype / any shape) +``value_float`` — scalar float32 +``value_floats`` — 1-D float32 list +``value_int`` — scalar int64 +``value_ints`` — 1-D int64 list + +Coverage +-------- +value (float32 tensor) → node removed, initializer added, ORT parity +value_float (scalar) → node removed, initializer added, ORT parity +value_floats (1-D list) → node removed, initializer added, ORT parity +value_int (scalar) → node removed, initializer added +value_ints (1-D list) → node removed, initializer added +no Constant nodes → graph unchanged, ORT output unchanged +multiple Constants → all converted in one pass, ORT parity +""" + +from __future__ import annotations + +import copy + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms import dla_constants_to_initializers as _mod + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._constants_to_initializers_from_model + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Build a typed tensor descriptor for a small test graph.""" + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + """Keep the test model compatible with the supported ORT versions.""" + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model: onnx.ModelProto) -> None: + """Refresh intermediate metadata and surface invalid graph or inference errors.""" + inferred = onnx.shape_inference.infer_shapes(model, strict_mode=True) + onnx.checker.check_model(inferred, full_check=True) + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + """Execute the test model on the CPU and return named outputs.""" + _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 _initializer_names(model: onnx.ModelProto) -> set[str]: + """Collect initializer names to check that folded values remain addressable.""" + return {i.name for i in model.graph.initializer} + + +def _constant_node_names(model: onnx.ModelProto) -> set[str]: + """Collect Constant output names that have not been folded.""" + return {n.output[0] for n in model.graph.node if n.op_type == "Constant"} + + +_RNG = np.random.default_rng(2024) + + +# ── 1. value attribute (dense float32 tensor) ──────────────────────────────── + + +@pytest.mark.parametrize(("ir_version", "opset"), [(3, 7), (4, 9), (9, _OPSET_VERSION)]) +def test_constant_value_tensor_converted(ir_version, opset): + """Constant with 'value' (TensorProto) replaced by initializer; ORT output unchanged.""" + shape = [2, 4] + const_data = _RNG.standard_normal(shape).astype(np.float32) + const_init = numpy_helper.from_array(const_data, name="const_w") + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["const_w"], + name="c_w", + value=const_init, + ) + # Add const_w to X element-wise + add_node = helper.make_node("Add", ["X", "const_w"], ["Y"], name="add") + + graph = helper.make_graph( + [const_node, add_node], + "const_value", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + value_info=[_vi("const_w", TensorProto.FLOAT, shape)], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", opset)], ir_version=ir_version + ) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + onnx.checker.check_model(tr, full_check=True) + assert [value.name for value in tr.graph.input] == ( + ["X", "const_w"] if ir_version <= 3 else ["X"] + ) + got = _run_ort(tr, {"X": x}) + + assert cnt == 1 + assert "const_w" in _initializer_names(tr) + assert "const_w" not in _constant_node_names(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) + + +# ── 2. value_float attribute (scalar float32) ──────────────────────────────── + + +@pytest.mark.parametrize("layout", ["linear", "coordinates", "empty"]) +def test_sparse_constant_is_densified_without_changing_values(layout): + """Sparse Constant layouts become dense initializers with identical tensor values.""" + shape = [2, 3] + values = np.array([] if layout == "empty" else [2.0, -4.0], dtype=np.float32) + indices = np.array( + [[0, 1], [1, 2]] if layout == "coordinates" else ([] if layout == "empty" else [1, 5]), + dtype=np.int64, + ) + sparse = helper.make_sparse_tensor( + numpy_helper.from_array(values, "K"), numpy_helper.from_array(indices, "indices"), shape + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Constant", [], ["K"], sparse_value=sparse), + helper.make_node("Add", ["X", "K"], ["Y"]), + ], + "sparse_constant", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + expected = np.zeros(shape, dtype=np.float32) + if layout != "empty": + expected.reshape(-1)[[1, 5]] = values + assert _apply(model, verbose=False) == 1 + assert not _constant_node_names(model) + initializer = next(value for value in model.graph.initializer if value.name == "K") + np.testing.assert_array_equal(numpy_helper.to_array(initializer), expected) + onnx.checker.check_model(model, full_check=True) + x = np.arange(6, dtype=np.float32).reshape(shape) + np.testing.assert_array_equal(_run_ort(model, {"X": x})["Y"], x + expected) + + +def test_constant_value_float_converted(): + """Constant with 'value_float' attribute converted and ORT output matches.""" + shape = [3, 4] + scalar = 2.5 + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["scalar_c"], + name="c_f", + value_float=scalar, + ) + mul_node = helper.make_node("Mul", ["X", "scalar_c"], ["Y"], name="mul") + + graph = helper.make_graph( + [const_node, mul_node], + "const_float", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + value_info=[_vi("scalar_c", TensorProto.FLOAT, [])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 1 + assert "scalar_c" in _initializer_names(tr) + assert "scalar_c" not in _constant_node_names(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) + + +# ── 3. value_floats attribute (1-D float32 list) ────────────────────────────── + + +def test_constant_value_floats_converted(): + """Constant with 'value_floats' (list) converted; ORT output matches.""" + floats = [1.0, 2.0, 3.0, 4.0] + shape_x = [2, 4] + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["bias"], + name="c_fs", + value_floats=floats, + ) + add_node = helper.make_node("Add", ["X", "bias"], ["Y"], name="add") + + graph = helper.make_graph( + [const_node, add_node], + "const_floats", + [_vi("X", TensorProto.FLOAT, shape_x)], + [_vi("Y", TensorProto.FLOAT, shape_x)], + value_info=[_vi("bias", TensorProto.FLOAT, [len(floats)])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape_x).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 1 + assert "bias" in _initializer_names(tr) + assert "bias" not in _constant_node_names(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) + + +# ── 4. value_int attribute (scalar int64) ──────────────────────────────────── + + +def test_constant_value_int_converted(): + """Constant with 'value_int' converted to initializer (int64 scalar). + + Uses Gather with a scalar index, which is the natural consumer of a + scalar int64 constant (ReduceSum opset-13 requires a 1-D axes tensor). + """ + shape = [4, 8] + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["idx"], + name="c_i", + value_int=1, + ) + # Gather selects row 1 from X → output shape [8] + gather_node = helper.make_node("Gather", ["X", "idx"], ["Y"], name="gather", axis=0) + + graph = helper.make_graph( + [const_node, gather_node], + "const_int", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, [shape[1]])], + value_info=[_vi("idx", TensorProto.INT64, [])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 1 + assert "idx" in _initializer_names(tr) + assert "idx" not in _constant_node_names(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) + + +# ── 5. value_ints attribute (1-D int64 list) ────────────────────────────────── + + +def test_constant_value_ints_converted(): + """Constant with 'value_ints' (shape tensor) converted; Reshape ORT parity.""" + shape_in = [2, 4, 8] + target_shape = [2, 32] + const_node = helper.make_node( + "Constant", + inputs=[], + outputs=["new_shape"], + name="c_is", + value_ints=target_shape, + ) + reshape_node = helper.make_node("Reshape", ["X", "new_shape"], ["Y"], name="reshape") + + graph = helper.make_graph( + [const_node, reshape_node], + "const_ints", + [_vi("X", TensorProto.FLOAT, shape_in)], + [_vi("Y", TensorProto.FLOAT, target_shape)], + value_info=[_vi("new_shape", TensorProto.INT64, [len(target_shape)])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape_in).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 1 + assert "new_shape" in _initializer_names(tr) + assert "new_shape" not in _constant_node_names(tr) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) + + +# ── 6. No Constant nodes: graph unchanged ──────────────────────────────────── + + +def test_no_constants_noop(): + """Model with no Constant nodes: graph and ORT output unchanged.""" + shape = [2, 4] + w_init = numpy_helper.from_array(_RNG.standard_normal(shape).astype(np.float32), name="W") + add_node = helper.make_node("Add", ["X", "W"], ["Y"], name="add") + + graph = helper.make_graph( + [add_node], + "no_const", + [_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_node_count = len(model.graph.node) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 0 + assert len(tr.graph.node) == original_node_count + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 7. Multiple Constants in one graph ─────────────────────────────────────── + + +def test_multiple_constants_all_converted(): + """All Constant nodes in the graph are converted in a single pass; ORT parity.""" + shape = [3, 4] + w_data = _RNG.standard_normal(shape).astype(np.float32) + b_data = _RNG.standard_normal([shape[1]]).astype(np.float32) + + w_const = helper.make_node( + "Constant", + inputs=[], + outputs=["W"], + name="c_w", + value=numpy_helper.from_array(w_data, name="W"), + ) + b_const = helper.make_node( + "Constant", + inputs=[], + outputs=["B"], + name="c_b", + value=numpy_helper.from_array(b_data, name="B"), + ) + matmul = helper.make_node("MatMul", ["X", "W"], ["mm_out"], name="mm") + add = helper.make_node("Add", ["mm_out", "B"], ["Y"], name="add") + + graph = helper.make_graph( + [w_const, b_const, matmul, add], + "multi_const", + [_vi("X", TensorProto.FLOAT, [shape[0], shape[0]])], + [_vi("Y", TensorProto.FLOAT, [shape[0], shape[1]])], + value_info=[ + _vi("W", TensorProto.FLOAT, shape), + _vi("B", TensorProto.FLOAT, [shape[1]]), + _vi("mm_out", TensorProto.FLOAT, [shape[0], shape[1]]), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal([shape[0], shape[0]]).astype(np.float32) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + cnt = _apply(tr, verbose=False) + got = _run_ort(tr, {"X": x}) + + assert cnt == 2 + assert "W" in _initializer_names(tr) + assert "B" in _initializer_names(tr) + assert not any(n.op_type == "Constant" for n in tr.graph.node) + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-5) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py new file mode 100644 index 00000000000..7359554d03b --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_handle_qdq.py @@ -0,0 +1,471 @@ +# 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_handle_qdq`. + +"""ORT CPU parity tests for ``dla_handle_qdq``. + +Transform rule +-------------- +Every ``QuantizeLinear`` and ``DequantizeLinear`` node whose data input is +non-4D gets wrapped independently: + + data [*orig_shape] + → Unsqueeze → data_4d [1,…,1, *orig_shape] + → Q / DQ → out_4d [1,…,1, *orig_shape] + → Squeeze → out [*orig_shape] + +Scale / zero-point inputs are untouched. The ``axis`` attribute is shifted +by ``4 - orig_rank``. Nodes whose data input is already 4-D are skipped. + +Coverage +-------- +QuantizeLinear 3-D input (dynamic) → wrapped, axis adjusted, ORT parity +DequantizeLinear 3-D input (static init) → wrapped, axis adjusted, ORT parity +DequantizeLinear 2-D input (dynamic) → wrapped, ORT parity +Q already 4-D → skipped (no Unsqueeze/Squeeze inserted) +""" + +from __future__ import annotations + +import copy + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator + +from modelopt.onnx.graph_surgery.dla_transforms import dla_handle_qdq as _mod + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_OPSET_VERSION = 21 # INT4/UINT4 require ≥ opset 21 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_handle_qdq + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Build a typed tensor descriptor for a small test graph.""" + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + """Keep the test model compatible with the supported ORT versions.""" + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model: onnx.ModelProto) -> None: + """Refresh intermediate metadata and surface invalid graph or inference errors.""" + inferred = onnx.shape_inference.infer_shapes(model, strict_mode=True) + onnx.checker.check_model(inferred, full_check=True) + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + """Execute the test model on the CPU and return named outputs.""" + _clamp_ir(model) + options = ort.SessionOptions() + # Preserve the wrappers under test; ORT's QDQ propagation can introduce invalid signed Q nodes. + options.add_session_config_entry( + "optimization.disable_specified_optimizers", "QDQPropagationTransformer" + ) + sess = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +def _assert_wrapped(tr, op_type: str, orig_name: str) -> None: + """Check Unsqueeze→op→Squeeze chain is present and op has the right name.""" + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + f"Unsqueeze must be inserted before {op_type}" + ) + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + f"Squeeze must be inserted after {op_type}" + ) + assert any(n.op_type == op_type for n in tr.graph.node), f"{op_type} node must still be present" + + +_RNG = np.random.default_rng(2024) + + +@pytest.mark.parametrize("op_type", ["QuantizeLinear", "DequantizeLinear"]) +def test_runtime_per_tensor_qdq_keeps_axis_absent(op_type): + """Runtime per-tensor QDQ retains scalar controls and its absent axis when wrapped.""" + quantize = op_type == "QuantizeLinear" + shape = [2, 3] + input_dtype = TensorProto.FLOAT if quantize else TensorProto.UINT8 + output_dtype = TensorProto.UINT8 if quantize else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node(op_type, ["X", "scale", "zero"], ["Y"], domain="com.microsoft")], + "runtime_per_tensor", + [_vi("X", input_dtype, shape)], + [_vi("Y", output_dtype, shape)], + initializer=[ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 10), helper.make_opsetid("com.microsoft", 1)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + data = np.arange(6, dtype=np.float32 if quantize else np.uint8).reshape(shape) + expected = _run_ort(model, {"X": data})["Y"] + _apply(model) + onnx.checker.check_model(model, full_check=True) + transformed = next(node for node in model.graph.node if node.op_type == op_type) + assert not any(attribute.name == "axis" for attribute in transformed.attribute) + _assert_wrapped(model, op_type, transformed.name) + np.testing.assert_array_equal(_run_ort(model, {"X": data})["Y"], expected) + + +# ── 1. QuantizeLinear — 3-D dynamic input, per-channel axis ────────────────── + + +@pytest.mark.parametrize( + ("opset", "domain", "ir_version"), [(21, "", 9), (10, "com.microsoft", 9), (13, "", 3)] +) +@pytest.mark.parametrize("shape", [[2, 4, 8], ["batch", 4, None]]) +def test_quantize_linear_3d_dynamic_wrapped(opset, domain, ir_version, shape): + """Q [2,4,8] axis=1 → Unsqueeze→Q(axis=2)→Squeeze; output shape and values match base.""" + scale = numpy_helper.from_array(np.full((4,), 0.1, dtype=np.float32), name="scale") + zp = numpy_helper.from_array(np.zeros((4,), dtype=np.int8), name="zp") + + q = helper.make_node( + "QuantizeLinear", ["X", "scale", "zp"], ["Y"], name="q", axis=1, domain=domain + ) + graph = helper.make_graph( + [q], + "g", + [ + _vi("X", TensorProto.FLOAT, shape), + *( + [_vi("scale", TensorProto.FLOAT, [4]), _vi("zp", TensorProto.INT8, [4])] + if ir_version <= 3 + else [] + ), + ], + [_vi("Y", TensorProto.INT8, shape)], + initializer=[scale, zp], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", opset), + *([helper.make_opsetid(domain, 1)] if domain else []), + ], + ir_version=ir_version, + ) + _clamp_ir(model) + + x = _RNG.standard_normal([2, 4, 8]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + q_node = next(n for n in tr.graph.node if n.op_type == "QuantizeLinear") + for name in [q_node.input[0], q_node.output[0]]: + info = next(value for value in tr.graph.value_info if value.name == name) + actual = [ + d.dim_value if d.HasField("dim_value") else d.dim_param or None + for d in info.type.tensor_type.shape.dim + ] + assert actual == [1, *shape] + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + _assert_wrapped(tr, "QuantizeLinear", "q") + + # axis must have been shifted from 1 to 2 (4-3=1 delta) + q_node = next(n for n in tr.graph.node if n.op_type == "QuantizeLinear") + axis_attr = next((a for a in q_node.attribute if a.name == "axis"), None) + assert axis_attr is not None and axis_attr.i == 2, ( + f"axis must be shifted to 2, got {axis_attr.i if axis_attr else 'None'}" + ) + + # scale and zp must be untouched (still 1-D) + s_init = next(i for i in tr.graph.initializer if i.name == "scale") + z_init = next(i for i in tr.graph.initializer if i.name == "zp") + assert list(s_init.dims) == [4], f"scale must stay 1-D [4], got {list(s_init.dims)}" + assert list(z_init.dims) == [4], f"zp must stay 1-D [4], got {list(z_init.dims)}" + + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 2. DequantizeLinear — 3-D static initializer data input, per-channel ───── + + +@pytest.mark.parametrize(("opset", "domain"), [(21, ""), (10, "com.microsoft")]) +def test_dequantize_linear_3d_static_init_wrapped(opset, domain): + """DQ with static int8 initializer [2,4,8] axis=1 → wrapped; float output matches base.""" + shape = [2, 4, 8] + data = numpy_helper.from_array((_RNG.integers(-5, 6, shape)).astype(np.int8), name="data") + scale = numpy_helper.from_array(np.full((4,), 0.1, dtype=np.float32), name="dq_scale") + zp = numpy_helper.from_array(np.zeros((4,), dtype=np.int8), name="dq_zp") + + dq = helper.make_node( + "DequantizeLinear", ["data", "dq_scale", "dq_zp"], ["Y"], name="dq", axis=1, domain=domain + ) + graph = helper.make_graph( + [dq], + "g", + [], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[data, scale, zp], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", opset), + *([helper.make_opsetid(domain, 1)] if domain else []), + ], + ) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {}) + + _assert_wrapped(tr, "DequantizeLinear", "dq") + + dq_node = next(n for n in tr.graph.node if n.op_type == "DequantizeLinear") + axis_attr = next((a for a in dq_node.attribute if a.name == "axis"), None) + assert axis_attr is not None and axis_attr.i == 2, ( + f"axis must be shifted to 2, got {axis_attr.i if axis_attr else 'None'}" + ) + + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-6) + + +# ── 3. DequantizeLinear — 2-D dynamic input ─────────────────────────────────── + + +def test_dequantize_linear_2d_dynamic_wrapped(): + """DQ [4,8] dynamic input (per-tensor, no axis) → Unsqueeze→DQ→Squeeze; ORT parity.""" + shape = [4, 8] + scale = numpy_helper.from_array(np.array(0.1, dtype=np.float32), name="s2") + zp = numpy_helper.from_array(np.array(0, dtype=np.int8), name="z2") + + dq = helper.make_node("DequantizeLinear", ["X", "s2", "z2"], ["Y"], name="dq2") + graph = helper.make_graph( + [dq], + "g", + [_vi("X", TensorProto.INT8, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.integers(-5, 6, shape).astype(np.int8) + _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_wrapped(tr, "DequantizeLinear", "dq2") + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-6) + + +# ── 4. QuantizeLinear already 4-D → skipped ────────────────────────────────── + + +def test_quantize_linear_4d_skipped(): + """Q with 4-D input is already 4-D → no Unsqueeze/Squeeze inserted, graph unchanged.""" + shape = [1, 2, 4, 8] + scale = numpy_helper.from_array(np.full((2,), 0.1, dtype=np.float32), name="s4d") + zp = numpy_helper.from_array(np.zeros((2,), dtype=np.int8), name="z4d") + + q = helper.make_node("QuantizeLinear", ["X", "s4d", "z4d"], ["Y"], name="q4d", axis=1) + graph = helper.make_graph( + [q], + "g", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.INT8, shape)], + initializer=[scale, zp], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_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) == orig_count, "4-D input: no nodes should be added" + assert not any(n.op_type == "Unsqueeze" for n in tr.graph.node) + assert not any(n.op_type == "Squeeze" for n in tr.graph.node) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +@pytest.mark.parametrize( + ("opset", "collision"), + [ + (10, None), + (12, None), + (21, None), + (21, "node0_unsq_axes_0"), + (21, "node0_in4d_0"), + (21, "node0_out4d_0"), + (21, "node0_unsq_0"), + ], +) +def test_mixed_qdq_nodes_use_their_snapshot_positions(opset, collision): + """Mixed-rank Q/DQ nodes retain values when only selected positions are wrapped.""" + specs = [ + ("QuantizeLinear", [3]), + ("DequantizeLinear", [1, 1, 2, 2]), + ("QuantizeLinear", [2, 2, 3]), + ("DequantizeLinear", [2, 3]), + ] + nodes, inputs, outputs, feeds = [], [], [], {} + for index, (op_type, shape) in enumerate(specs): + input_name, output_name = f"X{index}", f"Y{index}" + quantize = op_type == "QuantizeLinear" + input_dtype = TensorProto.FLOAT if quantize else TensorProto.UINT8 + output_dtype = TensorProto.UINT8 if quantize else TensorProto.FLOAT + nodes.append( + helper.make_node( + op_type, [input_name, "scale", "zero"], [output_name], name=f"node{index}" + ) + ) + inputs.append(_vi(input_name, input_dtype, shape)) + outputs.append(_vi(output_name, output_dtype, shape)) + feeds[input_name] = np.arange( + np.prod(shape), dtype=np.float32 if quantize else np.uint8 + ).reshape(shape) + initializers = [ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + ] + if collision == "node0_unsq_axes_0": + initializers.append(numpy_helper.from_array(np.array([0], dtype=np.int64), collision)) + elif collision == "node0_unsq_0": + nodes.append(helper.make_node("Identity", ["X1"], ["protected"], name=collision)) + outputs.append(_vi("protected", TensorProto.UINT8, [1, 1, 2, 2])) + elif collision: + inputs.append(_vi(collision, TensorProto.FLOAT, [1])) + feeds[collision] = np.array([1], dtype=np.float32) + model = helper.make_model( + helper.make_graph( + nodes, + "mixed_qdq", + inputs, + outputs, + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", opset)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + unchanged = next(node for node in transformed.graph.node if node.name == "node1") + assert list(unchanged.input) == ["X1", "scale", "zero"] + for name, output in _run_ort(transformed, feeds).items(): + np.testing.assert_array_equal(output, expected[name]) + + +@pytest.mark.parametrize("op_type", ["QuantizeLinear", "DequantizeLinear"]) +def test_untyped_qdq_output_keeps_its_inferred_dtype(op_type): + """An internal output without value_info remains correctly typed after wrapping.""" + quantize = op_type == "QuantizeLinear" + input_dtype = TensorProto.FLOAT if quantize else TensorProto.INT8 + output_dtype = TensorProto.INT8 if quantize else TensorProto.FLOAT + shape = [2, 3] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node(op_type, ["X", "scale", "zero"], ["hidden"], name="qdq"), + helper.make_node("Identity", ["hidden"], ["Y"]), + ], + "untyped_output", + [_vi("X", input_dtype, shape)], + [_vi("Y", output_dtype, shape)], + initializer=[ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.int8(0), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + x = np.arange(6, dtype=np.float32 if quantize else np.int8).reshape(shape) + expected = _run_ort(model, {"X": x}) + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(_run_ort(transformed, {"X": x})["Y"], expected["Y"]) + + +@pytest.mark.parametrize("op_type", ["QuantizeLinear", "DequantizeLinear"]) +def test_blocked_qdq_is_preserved_until_scale_rank_is_supported(op_type): + """Blocked quantization keeps data and scale ranks consistent by remaining unchanged.""" + quantize = op_type == "QuantizeLinear" + input_dtype = TensorProto.FLOAT if quantize else TensorProto.INT8 + output_dtype = TensorProto.INT8 if quantize else TensorProto.FLOAT + shape = [2, 4] + model = helper.make_model( + helper.make_graph( + [helper.make_node(op_type, ["X", "scale", "zero"], ["Y"], axis=1, block_size=2)], + "blocked", + [_vi("X", input_dtype, shape)], + [_vi("Y", output_dtype, shape)], + initializer=[ + numpy_helper.from_array(np.full((2, 2), 0.25, dtype=np.float32), "scale"), + numpy_helper.from_array(np.zeros((2, 2), dtype=np.int8), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + x = np.arange(8, dtype=np.float32 if quantize else np.int8).reshape(shape) + expected = ReferenceEvaluator(model).run(None, {"X": x})[0] + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, {"X": x})[0], expected) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_preprocessing_domains.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_preprocessing_domains.py new file mode 100644 index 00000000000..fbd9d545fe2 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_preprocessing_domains.py @@ -0,0 +1,149 @@ +# 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. + +"""DLA preprocessing must preserve unrelated custom operator functions.""" + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_constants_to_initializers import ( + _constants_to_initializers_from_model, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_handle_qdq import _apply_handle_qdq +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_qdq import _remove_qdq_from_model + +ort = pytest.importorskip( + "onnxruntime", reason="onnxruntime required for custom-function execution" +) + + +@pytest.mark.parametrize("stage", ["constants", "handle_q", "handle_dq", "remove_qdq"]) +def test_preprocessing_keeps_custom_functions_with_standard_names(stage): + """Custom functions retain their semantics even when they reuse recognized operator names.""" + functions, initializers, value_info = [], [], [] + shape = [2, 3] + inputs = [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)] + if stage == "constants": + body = helper.make_node("Constant", [], ["value"]) + attribute = helper.make_attribute_ref("value", onnx.AttributeProto.TENSOR) + attribute.ref_attr_name = "value" + body.attribute.append(attribute) + functions.append( + helper.make_function( + "local", + "Constant", + [], + ["Y"], + [body, helper.make_node("Neg", ["value"], ["Y"])], + [helper.make_opsetid("", 21)], + attributes=["value"], + ) + ) + nodes = [ + helper.make_node( + "Constant", + [], + ["Y"], + domain="local", + value=numpy_helper.from_array(np.ones(shape, dtype=np.float32), "seed"), + ) + ] + inputs, feeds = [], {} + else: + initializers = [ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.uint16(0), "zero"), + ] + feeds = {"X": np.arange(6, dtype=np.float32).reshape(shape)} + if stage == "remove_qdq": + functions = [ + helper.make_function( + "local", + "QuantizeLinear", + ["X", "S", "Z"], + ["Y"], + [ + helper.make_node( + "Constant", [], ["one"], value=numpy_helper.from_array(np.float32(1)) + ), + helper.make_node("Add", ["X", "one"], ["Y"]), + ], + [helper.make_opsetid("", 21)], + ), + helper.make_function( + "local", + "DequantizeLinear", + ["X", "S", "Z"], + ["Y"], + [helper.make_node("Identity", ["X"], ["Y"])], + [helper.make_opsetid("", 21)], + ), + ] + nodes = [ + helper.make_node("QuantizeLinear", ["X", "scale", "zero"], ["q"], domain="local"), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["Y"], domain="local"), + ] + value_info = [helper.make_tensor_value_info("q", TensorProto.FLOAT, shape)] + else: + op_type = "QuantizeLinear" if stage == "handle_q" else "DequantizeLinear" + functions = [ + helper.make_function( + "local", + op_type, + ["X", "S", "Z"], + ["Y"], + [helper.make_node("Transpose", ["X"], ["Y"], perm=[1, 0])], + [helper.make_opsetid("", 21)], + ) + ] + nodes = [helper.make_node(op_type, ["X", "scale", "zero"], ["Y"], domain="local")] + shape = [3, 2] + model = helper.make_model( + helper.make_graph( + nodes, + "custom", + inputs, + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, shape)], + initializer=initializers, + value_info=value_info, + ), + functions=functions, + opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("local", 1)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + expected = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds)[0] + before = model.SerializeToString() + if stage == "constants": + _constants_to_initializers_from_model(model, verbose=False) + elif stage == "remove_qdq": + _remove_qdq_from_model( + model, + keep_clip_after_inputs=False, + max_chained_qdq_pairs=5, + quant_types_to_remove=frozenset([TensorProto.UINT16]), + ) + else: + _apply_handle_qdq(model) + onnx.checker.check_model(model, full_check=True) + actual = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds)[0] + np.testing.assert_array_equal(actual, expected) + assert model.SerializeToString() == before diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py new file mode 100644 index 00000000000..860638cee53 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_qdq.py @@ -0,0 +1,671 @@ +# 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. + +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_remove_qdq`. + +"""ORT CPU parity tests for ``remove_qdq`` / ``_remove_qdq_from_model``. + +Transform rule +-------------- +A contiguous QuantizeLinear → DequantizeLinear pair is removed when the +quantizer zero-point element type belongs to ``qdq_quantized_dtypes`` +(default: UINT16 and INT16). + +After removal the DQ output is rewired directly to the Q input; downstream +nodes continue to see a float32 tensor with no quantisation error. + +ORT note: ORT does not reliably execute UINT16 / INT16 QDQ ops across all +opset / runtime versions, so the base model is *not* executed with ORT. +Instead we build numpy references that model the expected post-transform +output (the float computation without any quantisation round-trip) and +compare against ORT on the transformed (float-only) graph. + +Coverage +-------- +single uint16 Q→DQ → pair removed, Y = relu(X) +single uint8 Q→DQ → NOT removed (not in default types) +custom dtype uint8 remove → pair removed when uint8 specified +chained uint16+uint16 → both pairs removed (consecutive Q→DQ→Q→DQ) +chained uint16+uint8 → only uint16 pair removed +graph-output Q→DQ → producer rewired to graph output name +no QDQ nodes → graph unchanged +int16 Q→DQ → removed by default +""" + +from __future__ import annotations + +import copy +import importlib + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_qdq import _remove_qdq_from_model + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_OPSET_VERSION = 21 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_onnx_dtypes_mod = importlib.import_module("modelopt.onnx.graph_surgery.utils.graph_utils") +_DEFAULT_QUANT_TYPES = _onnx_dtypes_mod.DEFAULT_QDQ_REMOVE_QUANT_TYPES +_parse_dtype_list = _onnx_dtypes_mod.parse_qdq_quantized_dtype_list + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Build a typed tensor descriptor for a small test graph.""" + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + """Keep the test model compatible with the supported ORT versions.""" + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model: onnx.ModelProto) -> onnx.ModelProto: + """Refresh intermediate metadata and surface invalid graph or inference errors.""" + inferred = onnx.shape_inference.infer_shapes(model, strict_mode=True) + onnx.checker.check_model(inferred, full_check=True) + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + return model + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + """Execute the test model on the CPU and return named outputs.""" + _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 _apply(model, quant_types=None, **kw): + """Run _remove_qdq_from_model with sensible defaults.""" + qt = _DEFAULT_QUANT_TYPES if quant_types is None else _parse_dtype_list(quant_types) + return _remove_qdq_from_model( + model, + keep_clip_after_inputs=kw.get("keep_clip_after_inputs", False), + max_chained_qdq_pairs=kw.get("max_chained_qdq_pairs", 5), + quant_types_to_remove=qt, + ) + + +def _make_qdq_init(name_prefix, scale_val, zp_val, zp_np_dtype): + """Build a scale and zero point for the selected quantized storage type.""" + scale = numpy_helper.from_array(np.float32(scale_val), name=f"{name_prefix}_scale") + zp = numpy_helper.from_array(zp_np_dtype(zp_val), name=f"{name_prefix}_zp") + return scale, zp + + +def _make_qdq_nodes(name_prefix, input_name, output_name, onnx_zp_type): + """Return (Q node, DQ node, intermediate quantized tensor name).""" + q_out = f"{name_prefix}_q_out" + q = helper.make_node( + "QuantizeLinear", + inputs=[input_name, f"{name_prefix}_scale", f"{name_prefix}_zp"], + outputs=[q_out], + name=f"{name_prefix}_q", + ) + dq = helper.make_node( + "DequantizeLinear", + inputs=[q_out, f"{name_prefix}_scale", f"{name_prefix}_zp"], + outputs=[output_name], + name=f"{name_prefix}_dq", + ) + return q, dq, q_out + + +_RNG = np.random.default_rng(2024) + + +# ── 1. Single uint16 Q→DQ: pair removed, output is relu(X) ─────────────────── + + +def test_remove_single_uint16_qdq(): + """uint16 Q→DQ sandwiching Relu: pair removed; ORT output matches relu(X).""" + shape = [2, 4, 8] + scale, zp = _make_qdq_init("u16", 0.01, 0, np.uint16) + relu_in = helper.make_node("Relu", ["X"], ["relu_out"], name="relu_pre") + q, dq, q_out = _make_qdq_nodes("u16", "relu_out", "Y", TensorProto.UINT16) + + graph = helper.make_graph( + [relu_in, q, dq], + "single_u16", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + value_info=[ + _vi("relu_out", TensorProto.FLOAT, shape), + _vi(q_out, TensorProto.UINT16, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = np.maximum(x, 0.0) # After removing QDQ: Y = relu(X) + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 1, f"Expected 1 pair removed, got {pairs_removed}" + assert got["Y"].shape == tuple(shape) + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + assert not any(n.op_type in ("QuantizeLinear", "DequantizeLinear") for n in tr.graph.node) + + +# ── 2. Single uint8 Q→DQ: NOT removed (not in default types) ───────────────── + + +def test_remove_uint8_not_removed_by_default(): + """uint8 Q→DQ must NOT be removed with default quant_types (UINT16/INT16 only).""" + shape = [3, 8] + scale, zp = _make_qdq_init("u8", 0.1, 128, np.uint8) + q, dq, q_out = _make_qdq_nodes("u8", "X", "Y", TensorProto.UINT8) + + graph = helper.make_graph( + [q, dq], + "u8_graph", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + value_info=[ + _vi(q_out, TensorProto.UINT8, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_count = len(model.graph.node) + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + + assert pairs_removed == 0, f"Expected 0 pairs removed, got {pairs_removed}" + assert len(tr.graph.node) == original_count + assert any(n.op_type == "QuantizeLinear" for n in tr.graph.node) + assert any(n.op_type == "DequantizeLinear" for n in tr.graph.node) + + +# ── 3. Custom dtype: uint8 explicitly added → removed ──────────────────────── + + +def test_remove_uint8_with_custom_dtype(): + """uint8 Q→DQ removed when uint8 explicitly in qdq_quantized_dtypes.""" + shape = [3, 8] + scale, zp = _make_qdq_init("u8c", 0.1, 128, np.uint8) + q, dq, q_out = _make_qdq_nodes("u8c", "X", "Y", TensorProto.UINT8) + + graph = helper.make_graph( + [q, dq], + "u8_custom", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + value_info=[ + _vi(q_out, TensorProto.UINT8, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = x # After removing QDQ: Y = X + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr, quant_types=["uint8"]) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 1 + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + assert not any(n.op_type in ("QuantizeLinear", "DequantizeLinear") for n in tr.graph.node) + + +# ── 4. Chained consecutive uint16 Q→DQ→Q→DQ: both pairs removed ───────────── + + +def test_remove_chained_uint16_qdq(): + """Two consecutive uint16 Q→DQ pairs in sequence: both removed, Y = X.""" + shape = [4, 8] + scale1, zp1 = _make_qdq_init("ch1", 0.01, 0, np.uint16) + scale2, zp2 = _make_qdq_init("ch2", 0.02, 0, np.uint16) + + # X → Q1 → DQ1("mid") → Q2 → DQ2 → Y + q1, dq1, q1_out = _make_qdq_nodes("ch1", "X", "mid", TensorProto.UINT16) + q2, dq2, q2_out = _make_qdq_nodes("ch2", "mid", "Y", TensorProto.UINT16) + + graph = helper.make_graph( + [q1, dq1, q2, dq2], + "chained_u16", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale1, zp1, scale2, zp2], + value_info=[ + _vi(q1_out, TensorProto.UINT16, shape), + _vi("mid", TensorProto.FLOAT, shape), + _vi(q2_out, TensorProto.UINT16, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = x # After removing both QDQ: Y = X + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 2, f"Expected 2 pairs removed, got {pairs_removed}" + assert got["Y"].shape == tuple(shape) + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + assert not any(n.op_type in ("QuantizeLinear", "DequantizeLinear") for n in tr.graph.node) + + +# ── 5. Chained uint16+uint8: only uint16 pair removed ──────────────────────── + + +def test_remove_chained_mixed_dtypes(): + """Chain uint16 Q→DQ → Relu → uint8 Q→DQ: only uint16 pair removed.""" + shape = [4, 8] + scale16, zp16 = _make_qdq_init("m16", 0.01, 0, np.uint16) + scale8, zp8 = _make_qdq_init("m8", 0.1, 128, np.uint8) + + # X → Q16 → DQ16("after_16") → Relu("relu_out") → Q8 → DQ8 → Y + q16, dq16, q16_out = _make_qdq_nodes("m16", "X", "after_16", TensorProto.UINT16) + relu = helper.make_node("Relu", ["after_16"], ["relu_out"], name="relu") + q8, dq8, q8_out = _make_qdq_nodes("m8", "relu_out", "Y", TensorProto.UINT8) + + graph = helper.make_graph( + [q16, dq16, relu, q8, dq8], + "mixed_chain", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale16, zp16, scale8, zp8], + value_info=[ + _vi(q16_out, TensorProto.UINT16, shape), + _vi("after_16", TensorProto.FLOAT, shape), + _vi("relu_out", TensorProto.FLOAT, shape), + _vi(q8_out, TensorProto.UINT8, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + + assert pairs_removed == 1, f"Expected 1 pair removed, got {pairs_removed}" + + # uint16 nodes gone, uint8 nodes still present + q_nodes = [n for n in tr.graph.node if n.op_type == "QuantizeLinear"] + dq_nodes = [n for n in tr.graph.node if n.op_type == "DequantizeLinear"] + assert len(q_nodes) == 1 and q_nodes[0].name == "m8_q", "Only uint8 Q node should remain" + assert len(dq_nodes) == 1 and dq_nodes[0].name == "m8_dq", "Only uint8 DQ node should remain" + + # Relu must now consume X directly (uint16 bypassed) + relu_node = next(n for n in tr.graph.node if n.op_type == "Relu") + assert relu_node.input[0] == "X", f"Relu should consume X, got {relu_node.input[0]!r}" + + +# ── 6. Q→DQ directly to graph output ───────────────────────────────────────── + + +def test_remove_qdq_to_graph_output(): + """Q→DQ whose DQ output is a graph output: producer output rewired to graph output name.""" + shape = [2, 4] + scale, zp = _make_qdq_init("go", 0.01, 0, np.uint16) + + relu = helper.make_node("Relu", ["X"], ["relu_out"], name="relu") + q, dq, q_out = _make_qdq_nodes("go", "relu_out", "Y", TensorProto.UINT16) + + graph = helper.make_graph( + [relu, q, dq], + "graph_out", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + value_info=[ + _vi("relu_out", TensorProto.FLOAT, shape), + _vi(q_out, TensorProto.UINT16, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = np.maximum(x, 0.0) # After removing QDQ: Y = relu(X) + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 1 + assert got["Y"].shape == tuple(shape) + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + assert not any(n.op_type in ("QuantizeLinear", "DequantizeLinear") for n in tr.graph.node) + + +# ── 7. No QDQ nodes: graph unchanged ───────────────────────────────────────── + + +def test_remove_qdq_noop(): + """Model without any QDQ nodes: graph and ORT output unchanged.""" + shape = [2, 4, 8] + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "no_qdq", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + 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) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 0 + assert len(tr.graph.node) == original_count + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 8. int16 Q→DQ: removed by default ──────────────────────────────────────── + + +def test_remove_int16_qdq(): + """int16 Q→DQ is in the default remove types and must be removed.""" + shape = [2, 8] + scale, zp = _make_qdq_init("i16", 0.02, 0, np.int16) + q, dq, q_out = _make_qdq_nodes("i16", "X", "Y", TensorProto.INT16) + + graph = helper.make_graph( + [q, dq], + "i16_graph", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale, zp], + value_info=[ + _vi(q_out, TensorProto.INT16, shape), + ], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = x # After removing QDQ: Y = X + + tr = copy.deepcopy(model) + _prepare_model(tr) + pairs_removed = _apply(tr) + got = _run_ort(tr, {"X": x}) + + assert pairs_removed == 1 + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + assert not any(n.op_type in ("QuantizeLinear", "DequantizeLinear") for n in tr.graph.node) + + +@pytest.mark.parametrize("extra_outputs", [[], ["a"], ["a", "b"], ["activation", "a", "b"]]) +def test_remove_qdq_fanout_preserves_every_graph_output(extra_outputs): + """Preserve all output names when dequantizers share one quantizer.""" + shape = [1, 1, 1, 3] + nodes = [ + helper.make_node("Relu", ["X"], ["activation"]), + helper.make_node("QuantizeLinear", ["activation", "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["a"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["b"]), + helper.make_node("Add", ["a", "b"], ["Y"]), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "fanout", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi(name, TensorProto.FLOAT, shape) for name in [*extra_outputs, "Y"]], + initializer=[ + numpy_helper.from_array(np.float32(1), "scale"), + numpy_helper.from_array(np.uint16(0), "zero"), + ], + value_info=[_vi("activation", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + _clamp_ir(model) + assert _apply(model) == 2 + onnx.checker.check_model(model) + data = np.array([-1, 2, 3], dtype=np.float32).reshape(shape) + results = _run_ort(model, {"X": data}) + for name in extra_outputs: + np.testing.assert_array_equal(results[name], np.maximum(data, 0)) + np.testing.assert_array_equal(results["Y"], np.maximum(data, 0) * 2) + + +@pytest.mark.parametrize("empty_zero_point", [False, True]) +def test_remove_default_uint8_qdq_without_explicit_zero_point(empty_zero_point): + """Recognize the default UINT8 zero point when removing a selected QDQ pair.""" + shape = [1, 1, 1, 3] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "QuantizeLinear", ["X", "scale"] + ([""] if empty_zero_point else []), ["q"] + ), + helper.make_node( + "DequantizeLinear", ["q", "scale"] + ([""] if empty_zero_point else []), ["Y"] + ), + ], + "default_zero", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[numpy_helper.from_array(np.float32(1), "scale")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + assert _apply(model, quant_types=[TensorProto.UINT8]) == 1 + data = np.arange(3, dtype=np.float32).reshape(shape) + np.testing.assert_array_equal(_run_ort(model, {"X": data})["Y"], data) + + +@pytest.mark.parametrize("storage_dtype", [TensorProto.UINT16, TensorProto.INT16]) +@pytest.mark.parametrize("empty_zero_point", [False, True]) +def test_output_dtype_selects_quantizers_without_a_zero_point(storage_dtype, empty_zero_point): + """Explicit quantized output types control selection when the zero-point slot is absent.""" + extra = [""] if empty_zero_point else [] + shape = [2, 3] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "QuantizeLinear", ["X", "scale", *extra], ["q"], output_dtype=storage_dtype + ), + helper.make_node("DequantizeLinear", ["q", "scale", *extra], ["Y"]), + ], + "explicit_storage", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[numpy_helper.from_array(np.float32(0.25), "scale")], + value_info=[_vi("q", storage_dtype, shape)], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + excluded = copy.deepcopy(model) + before = excluded.SerializeToString() + assert _apply(excluded, quant_types=["uint8"]) == 0 + assert excluded.SerializeToString() == before + assert _apply(model) == 1 + onnx.checker.check_model(model, full_check=True) + x = np.arange(6, dtype=np.float32).reshape(shape) + np.testing.assert_array_equal(_run_ort(model, {"X": x})["Y"], x) + + +@pytest.mark.parametrize("captured", ["activation", "q", "D"]) +def test_remove_qdq_keeps_values_captured_by_control_flow(captured): + """QDQ removal leaves lexically captured tensors defined with their original types.""" + shape = [3] + dtype = TensorProto.UINT8 if captured == "q" else TensorProto.FLOAT + branch = helper.make_graph( + [helper.make_node("Identity", [captured], ["branch_out"])], + "branch", + [], + [_vi("branch_out", dtype, shape)], + ) + public_output = "Y" if captured == "D" else "D" + bridge = [helper.make_node("Identity", ["D"], ["Y"])] if captured == "D" else [] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Relu", ["X"], ["activation"]), + helper.make_node("QuantizeLinear", ["activation", "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["D"]), + *bridge, + helper.make_node( + "If", ["condition"], ["Z"], then_branch=branch, else_branch=branch + ), + ], + "captured", + [_vi("X", TensorProto.FLOAT, shape), _vi("condition", TensorProto.BOOL, [])], + [_vi(public_output, TensorProto.FLOAT, shape), _vi("Z", dtype, shape)], + initializer=[ + numpy_helper.from_array(np.float32(1), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + ], + value_info=[ + _vi("activation", TensorProto.FLOAT, shape), + _vi("q", TensorProto.UINT8, shape), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.array([-1, 1, 2], dtype=np.float32), "condition": np.array(True)} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + count = _apply(model, quant_types=["uint8"]) + onnx.checker.check_model(model, full_check=True) + assert count == 0 and model.SerializeToString() == before + for name, actual in _run_ort(model, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) + + +@pytest.mark.parametrize( + ("opset", "dtype", "ir_version"), + [ + (10, TensorProto.FLOAT, 9), + (21, TensorProto.FLOAT16, 9), + (21, TensorProto.FLOAT, 9), + (21, TensorProto.FLOAT, 3), + ], +) +def test_input_clip_respects_opset_and_output_dtype(opset, dtype, ir_version): + """Input clipping uses a valid schema and returns values in the declared output dtype.""" + np_dtype = np.float16 if dtype == TensorProto.FLOAT16 else np.float32 + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("QuantizeLinear", ["X", "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["Y"]), + ], + "input_clip", + [ + _vi("X", dtype, [3]), + *( + [_vi("scale", dtype, []), _vi("zero", TensorProto.UINT8, [])] + if ir_version <= 3 + else [] + ), + ], + [_vi("Y", dtype, [3])], + initializer=[ + numpy_helper.from_array(np_dtype(0.25), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", opset)], + ir_version=ir_version, + ) + onnx.checker.check_model(model, full_check=True) + assert _apply(model, quant_types=["uint8"], keep_clip_after_inputs=True) == 1 + onnx.checker.check_model(model, full_check=True) + x = np.array([-10, 1, 100], dtype=np_dtype) + np.testing.assert_array_equal(_run_ort(model, {"X": x})["Y"], np.clip(x, 0, 63.75)) + + +@pytest.mark.parametrize("unsupported", ["per_axis", "dtype_change", "untyped_storage"]) +def test_input_clip_refusal_preserves_all_pairs(unsupported): + """An unsupported later pair leaves earlier valid pairs and caller state unchanged.""" + shape = [2, 3] + output_dtype = TensorProto.FLOAT16 if unsupported == "dtype_change" else TensorProto.FLOAT + per_axis = unsupported == "per_axis" + missing_zero = unsupported == "untyped_storage" + scale = np.full(3, 0.25, dtype=np.float32) if per_axis else np.float32(0.25) + zero = np.zeros(3, dtype=np.uint8) if per_axis else np.uint8(0) + second_inputs = ["W", "second_scale"] + ([] if missing_zero else ["second_zero"]) + options = {"axis": 1} if per_axis else {} + q_options = {"output_dtype": TensorProto.UINT16} if missing_zero else options + dq_scale = "half_scale" if unsupported == "dtype_change" else "second_scale" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("QuantizeLinear", ["X", "scale", "zero"], ["first_q"]), + helper.make_node("DequantizeLinear", ["first_q", "scale", "zero"], ["Y"]), + helper.make_node("QuantizeLinear", second_inputs, ["second_q"], **q_options), + helper.make_node( + "DequantizeLinear", ["second_q", dq_scale, *second_inputs[2:]], ["Z"], **options + ), + ], + "clip_refusal", + [_vi(name, TensorProto.FLOAT, shape) for name in ("X", "W")], + [_vi("Y", TensorProto.FLOAT, shape), _vi("Z", output_dtype, shape)], + initializer=[ + numpy_helper.from_array(np.float32(0.25), "scale"), + numpy_helper.from_array(np.uint8(0), "zero"), + numpy_helper.from_array(scale, "second_scale"), + *([] if missing_zero else [numpy_helper.from_array(zero, "second_zero")]), + *( + [numpy_helper.from_array(np.float16(0.25), "half_scale")] + if unsupported == "dtype_change" + else [] + ), + ], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises((NotImplementedError, ValueError)): + _apply(model, quant_types=["uint8", "uint16"], keep_clip_after_inputs=True) + assert model.SerializeToString() == before