From 8c14beba222d1a0e69c42714f8d844c918d00e66 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:22 +0000 Subject: [PATCH 01/15] [DLA_transforms_16] [16/17] DLA public integration and reshape cleanup Signed-off-by: mgohil --- CHANGELOG.rst | 2 + modelopt/onnx/graph_surgery/__init__.py | 10 + modelopt/onnx/graph_surgery/__main__.py | 61 +++ .../dla_transforms/dla_remove_reshapes.py | 91 ++++ .../onnx/graph_surgery/make_dla_compatible.py | 48 +- .../test_dla_logged_pipeline.py | 130 ++++++ .../test_dla_make_compatible.py | 412 ++++++++++++++++++ .../dla_transforms/test_dla_rewrite_safety.py | 391 +++++++++++++++++ 8 files changed, 1144 insertions(+), 1 deletion(-) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index dcf30338bad..a41790cc9c3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,6 +34,8 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. +- Add the DLA graph-surgery pipeline with explicitly ordered transforms and a + ``--remove-all-qdq`` option. **Backward Breaking Changes** diff --git a/modelopt/onnx/graph_surgery/__init__.py b/modelopt/onnx/graph_surgery/__init__.py index 1b3e751445a..e0aec2af15e 100644 --- a/modelopt/onnx/graph_surgery/__init__.py +++ b/modelopt/onnx/graph_surgery/__init__.py @@ -23,6 +23,7 @@ - Converting model precision (e.g., FP16 to BF16) - Transposing DequantizeLinear weights for column-major storage optimization - Graph cleanup and optimization +- Making models compatible with the NVIDIA DLA accelerator Example usage: >>> from modelopt.onnx.graph_surgery import ( @@ -30,6 +31,7 @@ ... convert_fp16_to_bf16, ... transpose_dequantize_linear_weights, ... add_cross_kv_to_encoder, + ... make_dla_compatible, ... ) >>> # Replace attention with GQA for LLMs (FP16 model) >>> replace_attention_with_gqa( @@ -62,11 +64,17 @@ ... model_path="model_quantized.onnx", ... output_path="model_quantized_transposed.onnx", ... ) + >>> # Apply the ordered DLA compatibility pipeline + >>> make_dla_compatible( + ... model_path="model.onnx", + ... output_path="model_dla.onnx", + ... ) """ from .dq_transpose import transpose_dequantize_linear_weights from .encoder_cross_kv import add_cross_kv_to_encoder from .gqa_replacement import replace_attention_with_gqa +from .make_dla_compatible import dla_make_dla_compatible as make_dla_compatible from .utils.dtype_conversion import convert_fp16_to_bf16 _SURGERY_REGISTRY = { @@ -74,6 +82,7 @@ "add-cross-kv": add_cross_kv_to_encoder, "convert-bf16": convert_fp16_to_bf16, "transpose-dq": transpose_dequantize_linear_weights, + "make-dla-compatible": make_dla_compatible, } @@ -132,6 +141,7 @@ def run_graph_surgery( "add_cross_kv_to_encoder", "convert_fp16_to_bf16", "get_available_surgeries", + "make_dla_compatible", "replace_attention_with_gqa", "run_graph_surgery", "transpose_dequantize_linear_weights", diff --git a/modelopt/onnx/graph_surgery/__main__.py b/modelopt/onnx/graph_surgery/__main__.py index 98586e142d8..50730b38a0c 100644 --- a/modelopt/onnx/graph_surgery/__main__.py +++ b/modelopt/onnx/graph_surgery/__main__.py @@ -50,6 +50,12 @@ --input model_quantized.onnx \ --output model_quantized_transposed.onnx +Make a model compatible with the NVIDIA DLA accelerator:: + + python -m modelopt.onnx.graph_surgery make-dla-compatible \ + --input model.onnx \ + --output model_dla.onnx + Analyze attention pattern:: python -m modelopt.onnx.graph_surgery analyze \ @@ -84,6 +90,9 @@ def main(): Transpose DequantizeLinear weights: python -m modelopt.onnx.graph_surgery transpose-dq -i model_quantized.onnx -o model_transposed.onnx + Make a model DLA-compatible: + python -m modelopt.onnx.graph_surgery make-dla-compatible -i model.onnx -o model_dla.onnx + Analyze attention pattern: python -m modelopt.onnx.graph_surgery analyze -i model.onnx --layer 0 """, @@ -235,6 +244,44 @@ def main(): "-q", "--quiet", action="store_true", help="Suppress progress messages" ) + # Make DLA compatible subcommand + dla_parser = subparsers.add_parser( + "make-dla-compatible", + help="Make a model compatible with the NVIDIA DLA accelerator", + description="Apply the ordered DLA graph transforms and validate the resulting ONNX model.", + ) + dla_parser.add_argument("-i", "--input", required=True, help="Input ONNX model path") + dla_parser.add_argument("-o", "--output", required=True, help="Output ONNX model path") + dla_parser.add_argument( + "--no-external-data", + action="store_true", + help="Don't save weights as external data", + ) + dla_parser.add_argument( + "--external-data-name", + type=str, + default=None, + help="Name for external data file", + ) + qdq_group = dla_parser.add_mutually_exclusive_group() + qdq_group.add_argument( + "--qdq-quantized-dtypes", + nargs="+", + default=None, + metavar="DTYPE", + help=( + "Zero-point storage dtypes whose Q/DQ pairs are stripped in the early " + "QDQ-removal step (e.g. uint16 int16). Accepts TensorProto names or ints. " + "Default: UINT16 and INT16." + ), + ) + dla_parser.add_argument("-q", "--quiet", action="store_true", help="Suppress progress messages") + qdq_group.add_argument( + "--remove-all-qdq", + action="store_true", + help="Remove QDQ pairs for every dtype, matching the C++ remove-all mode", + ) + # Analyze subcommand analyze_parser = subparsers.add_parser( "analyze", @@ -305,6 +352,20 @@ def main(): verbose=not args.quiet, ) + elif args.command == "make-dla-compatible": + from . import run_graph_surgery + + run_graph_surgery( + "make-dla-compatible", + model_path=args.input, + output_path=args.output, + use_external_data=not args.no_external_data, + external_data_name=args.external_data_name, + verbose=not args.quiet, + qdq_quantized_dtypes=args.qdq_quantized_dtypes, + remove_all_qdq=args.remove_all_qdq, + ) + elif args.command == "analyze": from .gqa_replacement import analyze_attention_pattern diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py new file mode 100644 index 00000000000..7f7f28be6c1 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py @@ -0,0 +1,91 @@ +# 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. + +"""Collapse reshape chains at the final, inferred stage of the DLA pipeline.""" + +import numpy as np +from onnx import TensorProto, helper + +from ..utils.graph_utils import GraphCache, add_initializer + +__all__ = [] + + +def _apply_remove_reshapes(model): + graph = model.graph + cache = GraphCache(graph) + reshapes = cache.nodes_by_op("Reshape") + chains = [] + chain = [] + for node in reshapes: + if chain: + consumers = cache.get_consumers(chain[-1].output[0]) + if ( + chain[-1].output[0] in cache.graph_output_names + or len(consumers) != 1 + or consumers[0][0] is not node + ): + chains.append(chain) + chain = [] + chain.append(node) + if chain: + chains.append(chain) + replacements = {} + for chain in chains: + first, last = chain[0], chain[-1] + input_shape = cache.get_shape(first.input[0]) + output_shape = cache.get_shape(last.output[0]) + if ( + not input_shape + or not output_shape + or any(d <= 0 for d in [*input_shape, *output_shape]) + ): + continue + short, long = sorted((input_shape, output_shape), key=len) + removable = long[-len(short) :] == short + if removable and last.output[0] not in cache.graph_output_names: + for consumer, port in cache.get_consumers(last.output[0]): + consumer.input[port] = first.input[0] + elif input_shape == output_shape: + replacements[id(last)] = [ + helper.make_node( + "Identity", [first.input[0]], list(last.output), name=last.name + "_identity" + ) + ] + elif len(chain) > 1: + target = add_initializer( + graph, + cache.unique_name(last.name + "_reshape_shape"), + np.asarray(output_shape, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + replacements[id(last)] = [ + helper.make_node( + "Reshape", + [first.input[0], target], + list(last.output), + name=last.name + "_reshape", + ) + ] + else: + continue + for node in chain: + replacements.setdefault(id(node), []) + nodes = list(graph.node) + del graph.node[:] + for node in nodes: + graph.node.extend(replacements.get(id(node), [node])) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 17d0abba611..94403192dec 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -69,6 +69,7 @@ ) 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_remove_reshapes import _apply_remove_reshapes from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes from .dla_transforms.dla_reshape_patterns import ( _apply_reshape_qdq_reduce_sum, @@ -83,11 +84,13 @@ GraphCache, infer_shapes, parse_qdq_quantized_dtype_list, + save_onnx_model, ) if TYPE_CHECKING: from collections.abc import Iterable -__all__ = [] + +__all__ = ["DLA_TRANSFORMS", "dla_make_dla_compatible"] def _remove_deqlin_and_unused_initializers(model): @@ -148,6 +151,7 @@ def _apply_constants_to_initializers(model, **options): ("Greater", _apply_greater), ("ShapeInference", infer_shapes), ("TopK", partial(_apply_topk, rank_converted=True)), + ("RemoveReshapes", _apply_remove_reshapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) @@ -309,3 +313,45 @@ def _transform_make_dla_compatible( if verbose: logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started) return model + + +def dla_make_dla_compatible( + model_path: str, + output_path: str, + *, + use_external_data: bool = True, + external_data_name: str | None = None, + verbose: bool = True, + qdq_quantized_dtypes: Iterable[int | str] | None = None, + remove_all_qdq: bool = False, +) -> onnx.ModelProto: + """Load an ONNX model, apply the ordered DLA transforms, and save the result. + + Args: + model_path: Input file; external weights are loaded from its directory. + output_path: Destination ONNX file. + use_external_data: Store large tensor payloads in a sidecar file. + external_data_name: Sidecar filename; defaults to output filename plus "_data". + verbose: Log each transform stage. + qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed. + Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names. + remove_all_qdq: Remove QDQ pairs for every dtype, matching the C++ flag. + Cannot be combined with qdq_quantized_dtypes. + + Returns: + The transformed model. Invalid models and failed transformations raise. + """ + model = _transform_make_dla_compatible( + onnx.load(model_path, load_external_data=True), + verbose=verbose, + qdq_quantized_dtypes=qdq_quantized_dtypes, + remove_all_qdq=remove_all_qdq, + ) + save_onnx_model( + model, + output_path, + use_external_data=use_external_data, + external_data_name=external_data_name, + verbose=verbose, + ) + return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py new file mode 100644 index 00000000000..9df6d1cd813 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py @@ -0,0 +1,130 @@ +# 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. + +"""Protect logged pass order and numerical behavior at promoted boundaries.""" + +import copy + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.make_dla_compatible import ( + DLA_TRANSFORMS, + _transform_make_dla_compatible, +) + + +def test_logged_pass_order(): + # Recorded by the packaged C++ transform.log files, including both inference stages. + expected = [ + "RemoveQDQ", + "ConstantsToInitializers", + "Cast", + "RemoveDeqLin", + "ShapeInference", + "ReshapeTransposeReshape", + "ReshapeQDQReduceSum", + "ReshapeReduceSum", + "Non4dSliceAxis", + "Gather", + "DecomposeLstm", + "MatMulToTransposeConvTranspose", + "Non4dModelInputs", + "Non4dModelOutputs", + "Non4DReshape", + "ArgMax", + "StandaloneReduceOperations", + "RemoveIntermediarySqueezeAndUnsqueeze", + "GroupQueryAttention", + "GatherElements", + "Non4dInitializers", + "Non4DExpand", + "Non4DTile", + "Flatten", + "Non4DTranspose", + "Non4DLpNorm", + "FixInstanceNormChannelMismatch", + "UpdateAxis", + "Where", + "Not", + "Unsqueeze", + "Greater", + "ShapeInference", + "TopK", + "RemoveReshapes", + "RemoveUnusedInitializers", + ] + assert [name for name, _ in DLA_TRANSFORMS] == expected + + +def _run(model, feeds): + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + +@pytest.mark.parametrize("operation", ["ArgMax", "ReduceSum", "Concat", "Flatten"]) +def test_boundary_axis_conversion_preserves_values_and_shapes(operation): + input_shape = [2, 3, 4] + feeds = {"x": np.arange(24, dtype=np.float32).reshape(input_shape)} + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)] + initializers = [] + output_type = TensorProto.FLOAT + if operation == "ArgMax": + node = helper.make_node(operation, ["x"], ["y"], axis=2, keepdims=0) + output_shape, output_type = [2, 3], TensorProto.INT64 + expected = feeds["x"].argmax(axis=2) + elif operation == "ReduceSum": + node = helper.make_node(operation, ["x", "axes"], ["y"], keepdims=0) + initializers.append(numpy_helper.from_array(np.array([2], dtype=np.int64), "axes")) + output_shape = [2, 3] + expected = feeds["x"].sum(axis=2) + elif operation == "Concat": + inputs.append(helper.make_tensor_value_info("other", TensorProto.FLOAT, input_shape)) + feeds["other"] = feeds["x"] * 2 + node = helper.make_node(operation, ["x", "other"], ["y"], axis=2) + output_shape = [2, 3, 8] + expected = np.concatenate([feeds["x"], feeds["other"]], axis=2) + else: + node = helper.make_node(operation, ["x"], ["y"], axis=1) + output_shape = [2, 12] + expected = feeds["x"].reshape(output_shape) + model = helper.make_model( + helper.make_graph( + [node], + "boundaries", + inputs, + [helper.make_tensor_value_info("y", output_type, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + actual = _run(transformed, feeds) + assert actual[0].shape == tuple(output_shape) + np.testing.assert_array_equal(actual[0], expected) + + +def test_remove_all_qdq_rejects_conflicting_dtype_filter(): + with pytest.raises(ValueError, match="not both"): + _transform_make_dla_compatible( + None, remove_all_qdq=True, qdq_quantized_dtypes=[TensorProto.INT8] + ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py new file mode 100644 index 00000000000..62fbaac2aaa --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py @@ -0,0 +1,412 @@ +# 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. + +"""Focused assertions for the public DLA pipeline using small synthetic inputs.""" + +from __future__ import annotations + +import copy +import importlib + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required") +_dla_pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible + +_OPSET = 17 +_MAX_IR = 9 + + +def _vi(name, dtype, shape): + return helper.make_tensor_value_info(name, dtype, shape) + + +def _clamp_ir(model): + model.ir_version = min(model.ir_version, _MAX_IR) + + +def _run_pipeline(model): + """Deep-copy and run the DLA pipeline.""" + _clamp_ir(model) + return _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + + +def test_pipeline_reorders_nodes_after_final_unsqueeze(monkeypatch): + """The final reorder runs after the post-cleanup Unsqueeze transform.""" + + def fake_unsqueeze(model): + graph = model.graph + del graph.node[:] + graph.node.extend( + [ + helper.make_node( + "Constant", + [], + ["bias"], + name="const_after_unsqueeze", + value=numpy_helper.from_array(np.array([1.0], dtype=np.float32)), + ), + helper.make_node("Relu", ["X"], ["relu"], name="relu_after_unsqueeze"), + helper.make_node("Add", ["relu", "bias"], ["Y"], name="add_after_unsqueeze"), + ] + ) + return model + + model = helper.make_model( + helper.make_graph( + [helper.make_node("Identity", ["X"], ["Y"], name="identity")], + "g", + [_vi("X", TensorProto.FLOAT, [1])], + [_vi("Y", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + monkeypatch.setattr( + _dla_pipeline, + "DLA_TRANSFORMS", + tuple( + (name, fake_unsqueeze if name == "Unsqueeze" else transform) + for (name, transform) in _dla_pipeline.DLA_TRANSFORMS + ), + ) + result = _run_pipeline(model) + assert [node.name for node in result.graph.node] == [ + "relu_after_unsqueeze", + "const_after_unsqueeze", + "add_after_unsqueeze", + ] + + +def test_gather_with_cast_indices_promoted_to_4d(): + """Fold constant dequantization and preserve the dynamically selected row.""" + (seq_len, embed) = (5, 3) + x_data = np.random.default_rng(42).integers(-128, 127, (seq_len, embed), dtype=np.int8) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("DequantizeLinear", ["xq", "xs", "xz"], ["xf"], name="dq"), + helper.make_node("Cast", ["prev"], ["pc"], to=TensorProto.INT64, name="cast"), + helper.make_node("Gather", ["xf", "pc"], ["Y"], name="gather"), + ], + "g", + [_vi("prev", TensorProto.FLOAT, [1, 1])], + [_vi("Y", TensorProto.FLOAT, [1, 1, embed])], + initializer=[ + numpy_helper.from_array(x_data, "xq"), + numpy_helper.from_array(np.float32(0.003), "xs"), + numpy_helper.from_array(np.int8(0), "xz"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + gather = next(n for n in candidate.graph.node if n.op_type == "Gather") + assert next(a.i for a in gather.attribute if a.name == "axis") == 2 + assert all(n.op_type != "DequantizeLinear" for n in candidate.graph.node) + cast = next(n for n in candidate.graph.node if n.op_type == "Cast") + assert next(a.i for a in cast.attribute if a.name == "to") == TensorProto.INT32 + actual = ReferenceEvaluator(candidate).run(None, {"prev": np.array([[2]], dtype=np.float32)})[0] + np.testing.assert_allclose(actual, x_data[[[2]]].astype(np.float32) * np.float32(0.003)) + + +def test_matmul_add_softmax_2d(): + """MatMul [1,64]×[64,32] + Add bias + Softmax converted to Conv path.""" + w = np.random.default_rng(2).standard_normal((4, 3)).astype(np.float32) + b = np.random.default_rng(3).standard_normal((3,)).astype(np.float32) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("MatMul", ["X", "W"], ["mm"], name="matmul"), + helper.make_node("Add", ["mm", "B"], ["a"], name="add"), + helper.make_node("Softmax", ["a"], ["Y"], name="sm", axis=-1), + ], + "g", + [_vi("X", TensorProto.FLOAT, [1, 4])], + [_vi("Y", TensorProto.FLOAT, [1, 3])], + initializer=[numpy_helper.from_array(w, "W"), numpy_helper.from_array(b, "B")], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + conv = next(n for n in candidate.graph.node if n.op_type == "Conv") + weights = next(t for t in candidate.graph.initializer if t.name == conv.input[1]) + np.testing.assert_array_equal(numpy_helper.to_array(weights), w.T.reshape(3, 4, 1, 1)) + x = np.arange(4, dtype=np.float32).reshape(1, 4) + logits = x @ w + b + expected = np.exp(logits - logits.max(axis=-1, keepdims=True)) + expected /= expected.sum(axis=-1, keepdims=True) + actual = ReferenceEvaluator(candidate).run(None, {"X": x})[0] + np.testing.assert_allclose(actual, expected, rtol=1e-06, atol=1e-07) + + +def test_qdq_conv4d_relu_passthrough(): + """DQ(INT8→FP32) + Conv [1,3,8,8] + Relu — already 4D, should be mostly unchanged.""" + w = np.random.default_rng(4).integers(-128, 127, (16, 3, 3, 3), dtype=np.int8) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("DequantizeLinear", ["Wq", "Ws", "Wz"], ["Wf"], name="dq"), + helper.make_node( + "Conv", ["X", "Wf"], ["co"], name="conv", kernel_shape=[3, 3], pads=[1, 1, 1, 1] + ), + helper.make_node("Relu", ["co"], ["Y"], name="relu"), + ], + "g", + [_vi("X", TensorProto.FLOAT, [1, 3, 8, 8])], + [_vi("Y", TensorProto.FLOAT, [1, 16, 8, 8])], + initializer=[ + numpy_helper.from_array(w, "Wq"), + numpy_helper.from_array(np.float32(0.01), "Ws"), + numpy_helper.from_array(np.int8(0), "Wz"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + assert [n.op_type for n in candidate.graph.node] == ["DequantizeLinear", "Conv", "Relu"] + conv = next(n for n in candidate.graph.node if n.op_type == "Conv") + assert next(list(a.ints) for a in conv.attribute if a.name == "pads") == [1, 1, 1, 1] + assert candidate.graph.output[0].name == "Y" + assert [d.dim_value for d in candidate.graph.output[0].type.tensor_type.shape.dim] == [ + 1, + 16, + 8, + 8, + ] + + +def test_encoder_block_2d(): + """LayerNorm(2D) → MatMul → Add → Softmax → MatMul → Add (simplified attention/FFN).""" + h = 4 + rng = np.random.default_rng(10) + w1 = rng.standard_normal((h, h)).astype(np.float32) + w2 = rng.standard_normal((h, h)).astype(np.float32) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "LayerNormalization", + ["X", "ls", "lb"], + ["ln"], + name="ln", + axis=-1, + epsilon=1e-05, + ), + helper.make_node("MatMul", ["ln", "W1"], ["m1"], name="mm1"), + helper.make_node("Add", ["m1", "B1"], ["a1"], name="add1"), + helper.make_node("Softmax", ["a1"], ["sm"], name="sm", axis=-1), + helper.make_node("MatMul", ["sm", "W2"], ["m2"], name="mm2"), + helper.make_node("Add", ["m2", "B2"], ["Y"], name="add2"), + ], + "g", + [_vi("X", TensorProto.FLOAT, [1, h])], + [_vi("Y", TensorProto.FLOAT, [1, h])], + initializer=[ + numpy_helper.from_array(np.ones(h, dtype=np.float32), "ls"), + numpy_helper.from_array(np.zeros(h, dtype=np.float32), "lb"), + numpy_helper.from_array(w1, "W1"), + numpy_helper.from_array(np.zeros(h, dtype=np.float32), "B1"), + numpy_helper.from_array(w2, "W2"), + numpy_helper.from_array(np.zeros(h, dtype=np.float32), "B2"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + assert sum(n.op_type == "Conv" for n in candidate.graph.node) == 2 + assert all(n.op_type != "MatMul" for n in candidate.graph.node) + normalization = next(n for n in candidate.graph.node if n.op_type == "LayerNormalization") + assert next(a.i for a in normalization.attribute if a.name == "axis") == -1 + assert candidate.graph.output[0].name == "Y" + assert [d.dim_value for d in candidate.graph.output[0].type.tensor_type.shape.dim] == [1, h] + + +def test_gather_greater_where_chain(): + """Gather(2D data, dynamic idx) → Greater(threshold) → Where(mask, gathered, zeros). + + Exercises Gather 4D promotion, Greater int→float cast, Where→Mul rewrite. + Expected: Unsqueeze + Cast + Gather + Reshape + Greater + Cast + Mul + Squeeze. + """ + data = np.random.default_rng(30).standard_normal((10, 8)).astype(np.float32) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("Gather", ["D", "idx"], ["go"], name="gather", axis=0), + helper.make_node("Greater", ["go", "T"], ["mask"], name="gt"), + helper.make_node("Where", ["mask", "go", "Z"], ["Y"], name="where"), + ], + "g", + [_vi("idx", TensorProto.INT32, [1])], + [_vi("Y", TensorProto.FLOAT, [1, 8])], + initializer=[ + numpy_helper.from_array(data, "D"), + numpy_helper.from_array(np.float32(0.0), "T"), + numpy_helper.from_array(np.zeros((1, 8), dtype=np.float32), "Z"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + feeds = {"idx": np.array([1], dtype=np.int32)} + transformed_values = ReferenceEvaluator(candidate).run(None, feeds) + np.testing.assert_array_equal(transformed_values[0], np.maximum(data[[1]], np.float32(0))) + + +def test_split_dual_branch_add(): + """Split [1,4,8] axis=1 into [1,2,8]+[1,2,8] → Relu + Sigmoid → Add.""" + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("Split", ["X", "sp"], ["a", "b"], name="split", axis=1), + helper.make_node("Relu", ["a"], ["ra"], name="relu"), + helper.make_node("Sigmoid", ["b"], ["sb"], name="sig"), + helper.make_node("Add", ["ra", "sb"], ["Y"], name="add"), + ], + "g", + [_vi("X", TensorProto.FLOAT, [1, 4, 8])], + [_vi("Y", TensorProto.FLOAT, [1, 2, 8])], + initializer=[numpy_helper.from_array(np.array([2, 2], dtype=np.int64), "sp")], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + split = next(n for n in candidate.graph.node if n.op_type == "Split") + assert next(a.i for a in split.attribute if a.name == "axis") == 2 + assert len(split.output) == 2 + actual = ReferenceEvaluator(candidate).run(None, {"X": np.zeros((1, 4, 8), dtype=np.float32)})[ + 0 + ] + np.testing.assert_array_equal(actual, np.full((1, 2, 8), np.float32(0.5))) + + +def test_matmul_dynamic_then_5d_reshape_qdq_reduce(): + """MatMul(both graph inputs) → 5D Reshape → Q→DQ(int8) → ReduceSum → Q→DQ → Reshape.""" + scale = np.float32(0.05) + zp = np.int8(0) + inp = helper.make_model( + helper.make_graph( + [ + helper.make_node("MatMul", ["A", "B"], ["mm"], name="matmul"), + helper.make_node("Reshape", ["mm", "s5"], ["x5"], name="rs5d"), + helper.make_node("QuantizeLinear", ["x5", "qs", "qz"], ["xq"], name="q1"), + helper.make_node("DequantizeLinear", ["xq", "ds", "dz"], ["xdq"], name="dq1"), + helper.make_node("ReduceSum", ["xdq", "ax"], ["red"], name="rsum", keepdims=0), + helper.make_node("QuantizeLinear", ["red", "qs2", "qz2"], ["rq"], name="q2"), + helper.make_node("DequantizeLinear", ["rq", "ds2", "dz2"], ["rdq"], name="dq2"), + helper.make_node("Reshape", ["rdq", "s4"], ["Y"], name="rs4d"), + ], + "g", + [_vi("A", TensorProto.FLOAT, [1, 4, 8]), _vi("B", TensorProto.FLOAT, [1, 8, 6])], + [_vi("Y", TensorProto.FLOAT, [1, 4, 3, 1])], + initializer=[ + numpy_helper.from_array(np.array([1, 4, 2, 3, 1], dtype=np.int64), "s5"), + numpy_helper.from_array(np.array(scale), "qs"), + numpy_helper.from_array(np.array(zp), "qz"), + numpy_helper.from_array(np.array(scale), "ds"), + numpy_helper.from_array(np.array(zp), "dz"), + numpy_helper.from_array(np.array([2], dtype=np.int64), "ax"), + numpy_helper.from_array(np.array(scale), "qs2"), + numpy_helper.from_array(np.array(zp), "qz2"), + numpy_helper.from_array(np.array(scale), "ds2"), + numpy_helper.from_array(np.array(zp), "dz2"), + numpy_helper.from_array(np.array([1, 4, 3, 1], dtype=np.int64), "s4"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + candidate = _run_pipeline(inp) + assert sum(n.op_type == "Slice" for n in candidate.graph.node) == 4 + reductions = [n for n in candidate.graph.node if n.op_type == "ReduceSum"] + assert len(reductions) == 4 + assert all(next(a.i for a in n.attribute if a.name == "keepdims") == 1 for n in reductions) + feeds = {"A": np.full((1, 4, 8), np.float32(0.1)), "B": np.full((1, 8, 6), np.float32(0.2))} + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.inter_op_num_threads = 1 + actual = ort.InferenceSession( + candidate.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds)[0] + np.testing.assert_allclose(actual, np.full((1, 4, 3, 1), np.float32(0.3)), atol=1e-07) + + +def _count_op(graph, op_type): + return sum(1 for n in graph.node if n.op_type == op_type) + + +def _int8_qdq_relu_model(): + """X(4D) → QuantizeLinear(int8 zp) → DequantizeLinear(int8 zp) → Relu → Y.""" + return helper.make_model( + helper.make_graph( + [ + helper.make_node("QuantizeLinear", ["X", "s", "zp"], ["xq"], name="q"), + helper.make_node("DequantizeLinear", ["xq", "s", "zp"], ["xdq"], name="dq"), + helper.make_node("Relu", ["xdq"], ["Y"], name="relu"), + ], + "g", + [_vi("X", TensorProto.FLOAT, [1, 3, 8, 8])], + [_vi("Y", TensorProto.FLOAT, [1, 3, 8, 8])], + initializer=[ + numpy_helper.from_array(np.float32(0.02), "s"), + numpy_helper.from_array(np.int8(0), "zp"), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ) + + +def test_qdq_dtypes_default_keeps_int8_pair(): + """With the default dtypes (UINT16/INT16) an INT8 Q/DQ pair is left untouched.""" + model = _int8_qdq_relu_model() + _clamp_ir(model) + out = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert _count_op(out.graph, "QuantizeLinear") >= 1 + assert _count_op(out.graph, "DequantizeLinear") >= 1 + + +def test_qdq_dtypes_custom_removes_int8_pair(): + """Passing ``qdq_quantized_dtypes=['int8']`` strips the INT8 Q/DQ pair.""" + model = _int8_qdq_relu_model() + _clamp_ir(model) + out = _transform_make_dla_compatible( + copy.deepcopy(model), verbose=False, qdq_quantized_dtypes=["int8"] + ) + assert _count_op(out.graph, "QuantizeLinear") == 0 + assert _count_op(out.graph, "DequantizeLinear") == 0 + + +def test_qdq_dtypes_threaded_through_run_graph_surgery(tmp_path): + """The kwarg is forwarded by the dispatcher to the make-dla-compatible surgery.""" + import onnx + + from modelopt.onnx.graph_surgery import run_graph_surgery + + model = _int8_qdq_relu_model() + _clamp_ir(model) + in_path = tmp_path / "in.onnx" + out_path = tmp_path / "out.onnx" + onnx.save(model, str(in_path)) + result = run_graph_surgery( + "make-dla-compatible", + model_path=str(in_path), + output_path=str(out_path), + verbose=False, + qdq_quantized_dtypes=["int8"], + ) + assert _count_op(result.graph, "QuantizeLinear") == 0 + assert _count_op(result.graph, "DequantizeLinear") == 0 diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py new file mode 100644 index 00000000000..57bdfeb837e --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -0,0 +1,391 @@ +# 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. + +"""Model-level safety contracts spanning the ordered DLA graph rewrites.""" + +import copy + +import numpy as np +import onnx +import pytest +from onnx import TensorProto as T +from onnx import helper as h +from onnx import numpy_helper as nh +from onnx.reference import ReferenceEvaluator + +from modelopt.onnx import graph_surgery +from modelopt.onnx.graph_surgery.dla_transforms.dla_convert_ops_to_4d import ( + _apply_convert_ops_to_4d, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_graph_cleanup import ( + _remove_consecutive_casts, + _remove_unused_initializers, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_greater import _apply_greater +from modelopt.onnx.graph_surgery.dla_transforms.dla_indexing import _apply_gather +from modelopt.onnx.graph_surgery.dla_transforms.dla_matmul_to_transpose_conv_transpose import ( + _apply_matmul_to_transpose_conv_transpose, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_not import _apply_not +from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import ( + _apply_argmax, + _apply_non4d_lpnorm, + _apply_update_axis, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_deqlin import _apply_remove_deqlin +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_qdq import _remove_qdq_from_model +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_reshapes import _apply_remove_reshapes +from modelopt.onnx.graph_surgery.dla_transforms.dla_topk import _apply_topk +from modelopt.onnx.graph_surgery.dla_transforms.dla_where import _apply_where +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache + + +def _vi(name, dtype=T.FLOAT, shape=(2, 3)): + return h.make_tensor_value_info(name, dtype, shape) + + +def _model(nodes, inputs, outputs, initializers=(), value_info=()): + return h.make_model( + h.make_graph(nodes, "safety", inputs, outputs, initializers, value_info=value_info), + opset_imports=[h.make_opsetid("", 21)], + ) + + +def _same_values(model, transform, feeds): + model = onnx.shape_inference.infer_shapes(model) + expected = ReferenceEvaluator(model).run(None, feeds) + candidate = transform(copy.deepcopy(model)) + onnx.checker.check_model(candidate) + actual = ReferenceEvaluator(candidate).run(None, feeds) + for left, right in zip(expected, actual, strict=True): + np.testing.assert_allclose(right, left, rtol=1e-6, atol=1e-7) + return candidate + + +@pytest.mark.parametrize("static", [False, True]) +def test_shared_where_mask_keeps_both_polarities(static): + mask = np.array([[True, False, True]]) + initializers = [nh.from_array(np.float32(0), "zero")] + inputs = [_vi("x")] + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + if static: + initializers.append(nh.from_array(mask, "mask")) + else: + inputs.append(_vi("mask", T.BOOL, [1, 3])) + feeds["mask"] = mask + model = _model( + [ + h.make_node("Where", ["mask", "zero", "x"], ["a"]), + h.make_node("Where", ["mask", "x", "zero"], ["b"]), + ], + inputs, + [_vi("a"), _vi("b")], + initializers, + ) + _same_values(model, _apply_where, feeds) + + +def test_shared_integer_comparisons_have_one_cast_producer(): + model = _model( + [h.make_node("Greater", ["x", "x"], ["a"]), h.make_node("Greater", ["x", "z"], ["b"])], + [_vi("x", T.INT32)], + [_vi("a", T.BOOL), _vi("b", T.BOOL)], + [nh.from_array(np.int32(2), "z")], + ) + _same_values(model, _apply_greater, {"x": np.arange(6, dtype=np.int32).reshape(2, 3)}) + + +def test_quantized_graph_output_and_non_dq_consumer_survive(): + model = _model( + [ + h.make_node("QuantizeLinear", ["x", "s", "z"], ["q"]), + h.make_node("DequantizeLinear", ["q", "s", "z"], ["y"]), + h.make_node("Identity", ["q"], ["other"]), + ], + [_vi("x")], + [_vi("q", T.INT16), _vi("y"), _vi("other", T.INT16)], + [nh.from_array(np.float32(1), "s"), nh.from_array(np.int16(0), "z")], + ) + model = onnx.shape_inference.infer_shapes(model) + _remove_qdq_from_model( + model, + keep_clip_after_inputs=False, + max_chained_qdq_pairs=5, + quant_types_to_remove=frozenset({T.INT16}), + ) + onnx.checker.check_model(model) + assert any(node.op_type == "QuantizeLinear" for node in model.graph.node) + + +def test_shared_quantized_weight_with_different_scales_keeps_outputs(): + model = _model( + [ + h.make_node("DequantizeLinear", ["w", "s1", "z"], ["a"]), + h.make_node("DequantizeLinear", ["w", "s2", "z"], ["b"]), + ], + [], + [_vi("a", shape=[2]), _vi("b", shape=[2])], + [ + nh.from_array(np.array([2, 4], dtype=np.uint8), "w"), + nh.from_array(np.float32(0.5), "s1"), + nh.from_array(np.float32(0.25), "s2"), + nh.from_array(np.uint8(0), "z"), + ], + ) + _same_values(model, _apply_remove_deqlin, {}) + + +def test_blocked_dequantization_is_left_intact(): + model = _model( + [h.make_node("DequantizeLinear", ["w", "s", "z"], ["y"], axis=1, block_size=2)], + [], + [_vi("y", shape=[2, 4])], + [ + nh.from_array(np.arange(8, dtype=np.uint8).reshape(2, 4), "w"), + nh.from_array(np.ones((2, 2), dtype=np.float32), "s"), + nh.from_array(np.zeros((2, 2), dtype=np.uint8), "z"), + ], + ) + before = model.SerializeToString() + _apply_remove_deqlin(model) + assert model.SerializeToString() == before + + +def test_lossy_cast_roundtrip_keeps_its_rounding(): + model = _model( + [ + h.make_node("Cast", ["x"], ["half"], to=T.FLOAT16), + h.make_node("Cast", ["half"], ["y"], to=T.FLOAT), + ], + [_vi("x")], + [_vi("y")], + ) + _same_values( + model, + lambda m: _remove_consecutive_casts(m, GraphCache(m.graph)), + {"x": np.full((2, 3), 1.0003, dtype=np.float32)}, + ) + + +def test_symbolic_reshape_is_preserved(): + model = _model( + [h.make_node("Reshape", ["x", "shape"], ["y"])], + [_vi("x", shape=["N", 6])], + [_vi("y", shape=["N", 2, 3])], + [nh.from_array(np.array([0, 2, 3], dtype=np.int64), "shape")], + ) + _same_values( + model, _apply_remove_reshapes, {"x": np.arange(12, dtype=np.float32).reshape(2, 6)} + ) + + +def test_reshape_chain_keeps_exposed_intermediate(): + model = _model( + [h.make_node("Reshape", ["x", "s1"], ["a"]), h.make_node("Reshape", ["a", "s2"], ["b"])], + [_vi("x")], + [_vi("a", shape=[3, 2]), _vi("b", shape=[6])], + [ + nh.from_array(np.array([3, 2], dtype=np.int64), "s1"), + nh.from_array(np.array([6], dtype=np.int64), "s2"), + ], + ) + _same_values(model, _apply_remove_reshapes, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}) + + +@pytest.mark.parametrize("axis", [None, -2, 0, 1]) +@pytest.mark.parametrize("operation", ["ArgMax", "Softmax", "LpNormalization", "Gather"]) +def test_rank_promotion_preserves_axis_semantics(operation, axis): + attrs = {} if axis is None else {"axis": axis} + inputs = ["x"] + initializers = [] + dtype = T.FLOAT + shape = [2, 3] + if operation == "ArgMax": + dtype = T.INT64 + selected = 0 if axis is None else axis % 2 + shape = [1, 3] if selected == 0 else [2, 1] + elif operation == "Gather": + inputs.append("indices") + initializers.append(nh.from_array(np.array([1], dtype=np.int64), "indices")) + shape = [1, 3] if axis is None or axis % 2 == 0 else [2, 1] + model = _model( + [h.make_node(operation, inputs, ["y"], **attrs)], + [_vi("x")], + [_vi("y", dtype, shape)], + initializers, + ) + + def promote(candidate): + if operation == "Gather": + candidate = _apply_gather(candidate) + candidate = _apply_non4d_model_inputs(candidate) + candidate = _apply_non4d_model_outputs(candidate) + if operation == "ArgMax": + candidate = _apply_argmax(candidate) + elif operation == "LpNormalization": + candidate = _apply_non4d_lpnorm(candidate) + elif operation == "Softmax": + candidate = _apply_update_axis(candidate) + return candidate + + _same_values( + model, + promote, + {"x": np.array([[1, 5, 3], [4, 2, 6]], dtype=np.float32)}, + ) + + +def test_not_preserves_nonbinary_and_shared_boolean_uses(): + model = _model( + [ + h.make_node("Cast", ["x"], ["mask"], to=T.BOOL), + h.make_node("Not", ["mask"], ["y"]), + h.make_node("Identity", ["mask"], ["other"]), + ], + [_vi("x")], + [_vi("y", T.BOOL), _vi("other", T.BOOL)], + ) + _same_values(model, _apply_not, {"x": np.array([[0, -3, 0.5], [1, -1, 0]], dtype=np.float32)}) + + +@pytest.mark.parametrize("change", ["alpha", "bias"]) +def test_unsupported_gemm_refuses_before_mutation(change): + attrs = {"alpha": 2.0} if change == "alpha" else {} + bias = np.ones((2, 3) if change == "bias" else (3,), dtype=np.float32) + model = _model( + [h.make_node("Gemm", ["x", "w", "b"], ["y"], **attrs)], + [_vi("x")], + [_vi("y")], + [nh.from_array(np.eye(3, dtype=np.float32), "w"), nh.from_array(bias, "b")], + ) + before = model.SerializeToString() + with pytest.raises(ValueError): + _apply_matmul_to_transpose_conv_transpose(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("deferred", [False, True]) +def test_broadcast_matmul_restores_output_rank(deferred): + model = _model( + [h.make_node("MatMul", ["x", "w"], ["y"])], + [_vi("x")], + [_vi("y", shape=[1, 2, 3])], + [nh.from_array(np.eye(3, dtype=np.float32)[None], "w")], + ) + _same_values( + model, + (lambda model: _transform_make_dla_compatible(model, verbose=False)) + if deferred + else _apply_matmul_to_transpose_conv_transpose, + {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}, + ) + + +@pytest.mark.parametrize("axes", [None, [-1]]) +def test_legacy_reduce_restores_scalar_or_reduced_shape(axes): + inputs = ["x"] + initializers = [] + if axes is not None: + inputs.append("axes") + initializers.append(nh.from_array(np.array(axes, dtype=np.int64), "axes")) + model = _model( + [h.make_node("ReduceSum", inputs, ["y"], keepdims=0)], + [_vi("x")], + [_vi("y", shape=[] if axes is None else [2])], + initializers, + ) + _same_values( + model, _apply_convert_ops_to_4d, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + ) + + +def test_legacy_rank_dependent_op_is_left_intact(): + model = _model( + [h.make_node("Einsum", ["x"], ["y"], equation="ij->ji")], + [_vi("x")], + [_vi("y", shape=[3, 2])], + ) + _same_values( + model, _apply_convert_ops_to_4d, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + ) + + +def test_dispatcher_keeps_existing_return_and_keyword_contract(monkeypatch): + calls = [] + + def original(**kwargs): + calls.append(kwargs) + return {"converted": 3} + + monkeypatch.setitem(graph_surgery._SURGERY_REGISTRY, "convert-bf16", original) + result = graph_surgery.run_graph_surgery( + "convert-bf16", "input.onnx", "output.onnx", external_data=False, verbose=False + ) + assert result == {"converted": 3} + assert calls == [ + { + "model_path": "input.onnx", + "output_path": "output.onnx", + "external_data": False, + "verbose": False, + } + ] + + +def test_control_flow_capture_keeps_initializer(): + branch = h.make_graph( + [h.make_node("Identity", ["w"], ["v"])], "branch", [], [_vi("v", shape=[2])] + ) + model = _model( + [h.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch)], + [_vi("condition", T.BOOL, [])], + [_vi("y", shape=[2])], + [nh.from_array(np.array([1, 2], dtype=np.float32), "w")], + ) + _same_values(model, _remove_unused_initializers, {"condition": np.array(True)}) + + +def test_rank_converted_topk_chain_preserves_gathered_values(): + model = _model( + [ + h.make_node("TopK", ["scores", "k"], ["values", "indices"], axis=3), + h.make_node("Cast", ["indices"], ["cast1"], to=T.INT64), + h.make_node("Reshape", ["cast1", "shape"], ["reshaped"]), + h.make_node("Tile", ["reshaped", "repeats"], ["tiled"]), + h.make_node("Cast", ["tiled"], ["cast2"], to=T.INT32), + h.make_node("GatherElements", ["data", "cast2"], ["y"], axis=2), + ], + [_vi("scores", shape=[1, 1, 1, 4]), _vi("data", shape=[1, 1, 4, 3])], + [_vi("y", shape=[1, 1, 2, 3])], + [ + nh.from_array(np.array([2], dtype=np.int64), "k"), + nh.from_array(np.array([1, 1, 2, 1], dtype=np.int64), "shape"), + nh.from_array(np.array([1, 1, 1, 3], dtype=np.int64), "repeats"), + ], + ) + candidate = _same_values( + model, + lambda m: _apply_topk(m, rank_converted=True), + { + "scores": np.array([[[[1, 4, 2, 3]]]], dtype=np.float32), + "data": np.arange(12, dtype=np.float32).reshape(1, 1, 4, 3), + }, + ) + assert not any(node.op_type in {"Tile", "GatherElements"} for node in candidate.graph.node) From 8c8305fe1ed91a4f6a0d3b92119df21bafd4466a Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:27:18 +0000 Subject: [PATCH 02/15] Declare deferred output ranks in public rewrite fixtures Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_rewrite_safety.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py index 57bdfeb837e..a6d924b6f4c 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -233,10 +233,11 @@ def test_rank_promotion_preserves_axis_semantics(operation, axis): ) def promote(candidate): + """Declare deferred boundary ranks before applying the operation-specific stage.""" if operation == "Gather": candidate = _apply_gather(candidate) candidate = _apply_non4d_model_inputs(candidate) - candidate = _apply_non4d_model_outputs(candidate) + candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) if operation == "ArgMax": candidate = _apply_argmax(candidate) elif operation == "LpNormalization": From 82d9fb0e5ede3d63b19491cd4d24ee614e30030a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 02:20:54 +0000 Subject: [PATCH 03/15] Retain native rank views and runtime controls in final reshape cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_remove_reshapes.py | 41 ++-- .../test_dla_remove_reshapes.py | 203 ++++++++++++++++++ 2 files changed, 232 insertions(+), 12 deletions(-) create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py index 7f7f28be6c1..45be9754cd4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py @@ -18,22 +18,26 @@ import numpy as np from onnx import TensorProto, helper -from ..utils.graph_utils import GraphCache, add_initializer +from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes +from .dla_graph_cleanup import _subgraph_captures +from .dla_indexing import _operator_nodes __all__ = [] def _apply_remove_reshapes(model): + """Collapse fixed reshape chains while retaining rank-sensitive consumers and runtime controls.""" graph = model.graph cache = GraphCache(graph) - reshapes = cache.nodes_by_op("Reshape") + reshapes = _operator_nodes(model, {"Reshape"}) + captured = _subgraph_captures(graph) chains = [] chain = [] for node in reshapes: if chain: consumers = cache.get_consumers(chain[-1].output[0]) if ( - chain[-1].output[0] in cache.graph_output_names + chain[-1].output[0] in cache.graph_output_names | captured or len(consumers) != 1 or consumers[0][0] is not node ): @@ -45,6 +49,11 @@ def _apply_remove_reshapes(model): replacements = {} for chain in chains: first, last = chain[0], chain[-1] + if any( + node.input[1] not in cache.init_names or node.input[1] in cache.graph_input_names + for node in chain + ): + continue input_shape = cache.get_shape(first.input[0]) output_shape = cache.get_shape(last.output[0]) if ( @@ -53,15 +62,18 @@ def _apply_remove_reshapes(model): or any(d <= 0 for d in [*input_shape, *output_shape]) ): continue - short, long = sorted((input_shape, output_shape), key=len) - removable = long[-len(short) :] == short - if removable and last.output[0] not in cache.graph_output_names: + removable = input_shape == output_shape + if removable and last.output[0] not in cache.graph_output_names | captured: for consumer, port in cache.get_consumers(last.output[0]): consumer.input[port] = first.input[0] elif input_shape == output_shape: replacements[id(last)] = [ helper.make_node( - "Identity", [first.input[0]], list(last.output), name=last.name + "_identity" + "Identity", + [first.input[0]], + list(last.output), + name=cache.unique_name((last.name or last.output[0]) + "_identity"), + domain=last.domain, ) ] elif len(chain) > 1: @@ -77,15 +89,20 @@ def _apply_remove_reshapes(model): "Reshape", [first.input[0], target], list(last.output), - name=last.name + "_reshape", + name=cache.unique_name((last.name or last.output[0]) + "_reshape"), + domain=last.domain, ) ] else: continue for node in chain: replacements.setdefault(id(node), []) - nodes = list(graph.node) - del graph.node[:] - for node in nodes: - graph.node.extend(replacements.get(id(node), [node])) + batch_replace_nodes( + graph, + { + index: ([node], replacements[id(node)]) + for index, node in enumerate(graph.node) + if id(node) in replacements + }, + ) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py new file mode 100644 index 00000000000..3e2889a976d --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -0,0 +1,203 @@ +# 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. + +"""Native rank and payload contracts for the final reshape cleanup.""" + +import copy + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_reshapes import _apply_remove_reshapes + + +def _execute(model, feeds): + """Check and execute the actual CPU graph with native public outputs.""" + onnx.checker.check_model(model, full_check=True) + return ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]).run( + None, feeds + ) + + +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize("cast", [False, True]) +def test_native_gather_index_views_are_not_bypassed(legacy, cast): + """Keep the native vector view required by the Gather index port.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["source", "shape"], ["indices"]), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=2), + ], + "index_cleanup", + [ + helper.make_tensor_value_info("source", TensorProto.INT64, [1, 1, 1, 3]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 5, 4]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 3, 4])], + [numpy_helper.from_array(np.asarray([-1], np.int64), "shape")], + value_info=[helper.make_tensor_value_info("indices", TensorProto.INT64, [3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + if cast: + model.graph.node.insert( + 1, helper.make_node("Cast", ["indices"], ["cast_indices"], to=TensorProto.INT64) + ) + model.graph.node[2].input[1] = "cast_indices" + model.ir_version = 3 if legacy else 9 + if legacy: + model.graph.input.append(helper.make_tensor_value_info("shape", TensorProto.INT64, [1])) + feeds = { + "source": np.asarray([0, 2, 4], np.int64).reshape(1, 1, 1, 3), + "x": np.arange(20, dtype=np.float32).reshape(1, 1, 5, 4), + } + expected = _execute(model, feeds)[0] + converted = _apply_remove_reshapes(copy.deepcopy(model)) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +def test_reshape_cleanup_retains_untouched_constant_protobuf_payloads(): + """Leave an unrelated payload object in place while removing an identity Reshape.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("Reshape", ["x", "shape"], ["y"]), + ], + "payload_cleanup", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + [numpy_helper.from_array(np.asarray([1, 1, 2, 3], np.int64), "shape")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds) + constant = model.graph.node[0] + converted = _apply_remove_reshapes(model) + assert converted.graph.node[0] is constant + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_unnamed_public_identity_reshapes_allocate_distinct_node_names(): + """Preserve both public values when multiple unnamed shapes become Identity nodes.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "first_shape"], ["first"]), + helper.make_node("Reshape", ["z", "second_shape"], ["second"]), + ], + "public_names", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 1, 3, 2]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, [1, 1, 3, 2]), + ], + [ + numpy_helper.from_array(np.asarray([1, 1, 2, 3], np.int64), "first_shape"), + numpy_helper.from_array(np.asarray([1, 1, 3, 2], np.int64), "second_shape"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3), + "z": np.arange(6, dtype=np.float32).reshape(1, 1, 3, 2), + } + expected = _execute(model, feeds) + converted = _apply_remove_reshapes(copy.deepcopy(model)) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_nested_capture_prevents_reshape_chain_collapse(): + """Keep the intermediate native layout referenced inside a nested branch.""" + branch = helper.make_graph( + [helper.make_node("Identity", ["middle"], ["out"])], + "branch", + [], + [helper.make_tensor_value_info("out", TensorProto.FLOAT, [2, 3])], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "short_shape"], ["middle"]), + helper.make_node("Reshape", ["middle", "long_shape"], ["y"]), + helper.make_node( + "If", ["condition"], ["view"], then_branch=branch, else_branch=branch + ), + ], + "captured_shape", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("view", TensorProto.FLOAT, [2, 3]), + ], + [ + numpy_helper.from_array(np.asarray([2, 3], np.int64), "short_shape"), + numpy_helper.from_array(np.asarray([1, 1, 2, 3], np.int64), "long_shape"), + numpy_helper.from_array(np.asarray(True), "condition"), + ], + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds) + converted = _apply_remove_reshapes(copy.deepcopy(model)) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_padding_view_before_public_unary_output_keeps_its_rank(): + """Retain a real rank change even when its next consumer is rank-agnostic.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "shape"], ["padded"]), + helper.make_node("Relu", ["padded"], ["y"]), + ], + "padding_view", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 1, 3])], + [numpy_helper.from_array(np.asarray([1, 1, 1, 3], np.int64), "shape")], + value_info=[helper.make_tensor_value_info("padded", TensorProto.FLOAT, [1, 1, 1, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.asarray([-1, 0, 1], np.float32)} + expected = _execute(model, feeds)[0] + converted = _apply_remove_reshapes(copy.deepcopy(model)) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) From 4142a880cdf1125708522f07b46913cb878f0e45 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 05:02:34 +0000 Subject: [PATCH 04/15] Document DLA finite-input and integration test contracts Signed-off-by: mgohil --- .../onnx/graph_surgery/make_dla_compatible.py | 3 +++ .../test_dla_logged_pipeline.py | 4 ++++ .../test_dla_make_compatible.py | 4 ++++ .../dla_transforms/test_dla_rewrite_safety.py | 21 +++++++++++++++++++ 4 files changed, 32 insertions(+) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 94403192dec..0507c7259eb 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -327,6 +327,9 @@ def dla_make_dla_compatible( ) -> onnx.ModelProto: """Load an ONNX model, apply the ordered DLA transforms, and save the result. + Arithmetic lowering of zero-branch ``Where`` assumes finite FP32 branch tensors. + NaN and infinity values at those nodes are outside the supported input contract. + Args: model_path: Input file; external weights are loaded from its directory. output_path: Destination ONNX file. diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py index 9df6d1cd813..86f02724be9 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_logged_pipeline.py @@ -30,6 +30,7 @@ def test_logged_pass_order(): # Recorded by the packaged C++ transform.log files, including both inference stages. + """Keep every stage in the established DLA conversion order.""" expected = [ "RemoveQDQ", "ConstantsToInitializers", @@ -72,6 +73,7 @@ def test_logged_pass_order(): def _run(model, feeds): + """Execute a checked model with the CPU runtime.""" options = ort.SessionOptions() options.intra_op_num_threads = 1 options.inter_op_num_threads = 1 @@ -82,6 +84,7 @@ def _run(model, feeds): @pytest.mark.parametrize("operation", ["ArgMax", "ReduceSum", "Concat", "Flatten"]) def test_boundary_axis_conversion_preserves_values_and_shapes(operation): + """Preserve native axis semantics through boundary promotion.""" input_shape = [2, 3, 4] feeds = {"x": np.arange(24, dtype=np.float32).reshape(input_shape)} inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)] @@ -124,6 +127,7 @@ def test_boundary_axis_conversion_preserves_values_and_shapes(operation): def test_remove_all_qdq_rejects_conflicting_dtype_filter(): + """Refuse conflicting QDQ options before changing the graph.""" with pytest.raises(ValueError, match="not both"): _transform_make_dla_compatible( None, remove_all_qdq=True, qdq_quantized_dtypes=[TensorProto.INT8] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py index 62fbaac2aaa..2d8dabe2c82 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py @@ -34,10 +34,12 @@ def _vi(name, dtype, shape): + """Declare a typed tensor for a checked graph.""" return helper.make_tensor_value_info(name, dtype, shape) def _clamp_ir(model): + """Use the model IR supported by the CPU runtime.""" model.ir_version = min(model.ir_version, _MAX_IR) @@ -51,6 +53,7 @@ def test_pipeline_reorders_nodes_after_final_unsqueeze(monkeypatch): """The final reorder runs after the post-cleanup Unsqueeze transform.""" def fake_unsqueeze(model): + """Inject a deliberately unordered fixture for final reorder coverage.""" graph = model.graph del graph.node[:] graph.node.extend( @@ -346,6 +349,7 @@ def test_matmul_dynamic_then_5d_reshape_qdq_reduce(): def _count_op(graph, op_type): + """Count nodes matching an operator type.""" return sum(1 for n in graph.node if n.op_type == op_type) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py index a6d924b6f4c..499c933bf44 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -58,10 +58,12 @@ def _vi(name, dtype=T.FLOAT, shape=(2, 3)): + """Declare a typed tensor for a checked graph.""" return h.make_tensor_value_info(name, dtype, shape) def _model(nodes, inputs, outputs, initializers=(), value_info=()): + """Build a typed model for a focused safety contract.""" return h.make_model( h.make_graph(nodes, "safety", inputs, outputs, initializers, value_info=value_info), opset_imports=[h.make_opsetid("", 21)], @@ -69,6 +71,7 @@ def _model(nodes, inputs, outputs, initializers=(), value_info=()): def _same_values(model, transform, feeds): + """Check output values before and after a graph rewrite.""" model = onnx.shape_inference.infer_shapes(model) expected = ReferenceEvaluator(model).run(None, feeds) candidate = transform(copy.deepcopy(model)) @@ -81,6 +84,7 @@ def _same_values(model, transform, feeds): @pytest.mark.parametrize("static", [False, True]) def test_shared_where_mask_keeps_both_polarities(static): + """Preserve both selections when Where nodes share a condition mask.""" mask = np.array([[True, False, True]]) initializers = [nh.from_array(np.float32(0), "zero")] inputs = [_vi("x")] @@ -103,6 +107,7 @@ def test_shared_where_mask_keeps_both_polarities(static): def test_shared_integer_comparisons_have_one_cast_producer(): + """Reuse one private conversion across shared integer comparisons.""" model = _model( [h.make_node("Greater", ["x", "x"], ["a"]), h.make_node("Greater", ["x", "z"], ["b"])], [_vi("x", T.INT32)], @@ -113,6 +118,7 @@ def test_shared_integer_comparisons_have_one_cast_producer(): def test_quantized_graph_output_and_non_dq_consumer_survive(): + """Keep quantized owners needed by public outputs and other consumers.""" model = _model( [ h.make_node("QuantizeLinear", ["x", "s", "z"], ["q"]), @@ -135,6 +141,7 @@ def test_quantized_graph_output_and_non_dq_consumer_survive(): def test_shared_quantized_weight_with_different_scales_keeps_outputs(): + """Respect separate dequantization parameters on a shared weight.""" model = _model( [ h.make_node("DequantizeLinear", ["w", "s1", "z"], ["a"]), @@ -153,6 +160,7 @@ def test_shared_quantized_weight_with_different_scales_keeps_outputs(): def test_blocked_dequantization_is_left_intact(): + """Leave unsupported blocked dequantization layouts unchanged.""" model = _model( [h.make_node("DequantizeLinear", ["w", "s", "z"], ["y"], axis=1, block_size=2)], [], @@ -169,6 +177,7 @@ def test_blocked_dequantization_is_left_intact(): def test_lossy_cast_roundtrip_keeps_its_rounding(): + """Retain an intermediate cast whose rounding changes values.""" model = _model( [ h.make_node("Cast", ["x"], ["half"], to=T.FLOAT16), @@ -185,6 +194,7 @@ def test_lossy_cast_roundtrip_keeps_its_rounding(): def test_symbolic_reshape_is_preserved(): + """Keep reshape targets whose native dimensions are unresolved.""" model = _model( [h.make_node("Reshape", ["x", "shape"], ["y"])], [_vi("x", shape=["N", 6])], @@ -197,6 +207,7 @@ def test_symbolic_reshape_is_preserved(): def test_reshape_chain_keeps_exposed_intermediate(): + """Do not remove reshape results exposed outside a private chain.""" model = _model( [h.make_node("Reshape", ["x", "s1"], ["a"]), h.make_node("Reshape", ["a", "s2"], ["b"])], [_vi("x")], @@ -212,6 +223,7 @@ def test_reshape_chain_keeps_exposed_intermediate(): @pytest.mark.parametrize("axis", [None, -2, 0, 1]) @pytest.mark.parametrize("operation", ["ArgMax", "Softmax", "LpNormalization", "Gather"]) def test_rank_promotion_preserves_axis_semantics(operation, axis): + """Preserve native axes when inserting four-dimensional views.""" attrs = {} if axis is None else {"axis": axis} inputs = ["x"] initializers = [] @@ -254,6 +266,7 @@ def promote(candidate): def test_not_preserves_nonbinary_and_shared_boolean_uses(): + """Protect nonbinary values and independent boolean consumers.""" model = _model( [ h.make_node("Cast", ["x"], ["mask"], to=T.BOOL), @@ -268,6 +281,7 @@ def test_not_preserves_nonbinary_and_shared_boolean_uses(): @pytest.mark.parametrize("change", ["alpha", "bias"]) def test_unsupported_gemm_refuses_before_mutation(change): + """Refuse unsupported Gemm layouts without changing the caller.""" attrs = {"alpha": 2.0} if change == "alpha" else {} bias = np.ones((2, 3) if change == "bias" else (3,), dtype=np.float32) model = _model( @@ -284,6 +298,7 @@ def test_unsupported_gemm_refuses_before_mutation(change): @pytest.mark.parametrize("deferred", [False, True]) def test_broadcast_matmul_restores_output_rank(deferred): + """Restore the native broadcast MatMul result layout.""" model = _model( [h.make_node("MatMul", ["x", "w"], ["y"])], [_vi("x")], @@ -301,6 +316,7 @@ def test_broadcast_matmul_restores_output_rank(deferred): @pytest.mark.parametrize("axes", [None, [-1]]) def test_legacy_reduce_restores_scalar_or_reduced_shape(axes): + """Restore scalar and reduced outputs with legacy reduction schemas.""" inputs = ["x"] initializers = [] if axes is not None: @@ -318,6 +334,7 @@ def test_legacy_reduce_restores_scalar_or_reduced_shape(axes): def test_legacy_rank_dependent_op_is_left_intact(): + """Keep unsupported legacy rank-sensitive operators native.""" model = _model( [h.make_node("Einsum", ["x"], ["y"], equation="ij->ji")], [_vi("x")], @@ -329,9 +346,11 @@ def test_legacy_rank_dependent_op_is_left_intact(): def test_dispatcher_keeps_existing_return_and_keyword_contract(monkeypatch): + """Preserve registry dispatch return values and keyword forwarding.""" calls = [] def original(**kwargs): + """Record registry forwarding and return its sentinel value.""" calls.append(kwargs) return {"converted": 3} @@ -351,6 +370,7 @@ def original(**kwargs): def test_control_flow_capture_keeps_initializer(): + """Retain fixed data referenced from a nested control-flow graph.""" branch = h.make_graph( [h.make_node("Identity", ["w"], ["v"])], "branch", [], [_vi("v", shape=[2])] ) @@ -364,6 +384,7 @@ def test_control_flow_capture_keeps_initializer(): def test_rank_converted_topk_chain_preserves_gathered_values(): + """Preserve actual gathered values while removing private index tiling.""" model = _model( [ h.make_node("TopK", ["scores", "k"], ["values", "indices"], axis=3), From bc7f37c0c58eac676fa3d45293631a45549d1852 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 06:30:46 +0000 Subject: [PATCH 05/15] Collapse fixed legacy reshape chains with declared controls Signed-off-by: mgohil --- .../dla_transforms/dla_remove_reshapes.py | 6 +++- .../test_dla_make_compatible.py | 7 ++-- .../test_dla_remove_reshapes.py | 33 +++++++++++++++++++ 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py index 45be9754cd4..e14502e25c2 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py @@ -19,8 +19,10 @@ from onnx import TensorProto, helper from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes +from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _subgraph_captures from .dla_indexing import _operator_nodes +from .dla_rank_conversion import _runtime_input_names __all__ = [] @@ -29,6 +31,7 @@ def _apply_remove_reshapes(model): """Collapse fixed reshape chains while retaining rank-sensitive consumers and runtime controls.""" graph = model.graph cache = GraphCache(graph) + runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) captured = _subgraph_captures(graph) chains = [] @@ -50,7 +53,7 @@ def _apply_remove_reshapes(model): for chain in chains: first, last = chain[0], chain[-1] if any( - node.input[1] not in cache.init_names or node.input[1] in cache.graph_input_names + node.input[1] not in cache.init_names or node.input[1] in runtime_inputs for node in chain ): continue @@ -105,4 +108,5 @@ def _apply_remove_reshapes(model): if id(node) in replacements }, ) + _add_legacy_initializer_inputs(model) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py index 2d8dabe2c82..4599ba5f175 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py @@ -21,10 +21,13 @@ import importlib 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 import run_graph_surgery + ort = pytest.importorskip("onnxruntime", reason="onnxruntime required") _dla_pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible @@ -396,10 +399,6 @@ def test_qdq_dtypes_custom_removes_int8_pair(): def test_qdq_dtypes_threaded_through_run_graph_surgery(tmp_path): """The kwarg is forwarded by the dispatcher to the make-dla-compatible surgery.""" - import onnx - - from modelopt.onnx.graph_surgery import run_graph_surgery - model = _int8_qdq_relu_model() _clamp_ir(model) in_path = tmp_path / "in.onnx" diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py index 3e2889a976d..a762f363dbf 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -34,6 +34,39 @@ def _execute(model, feeds): ) +@pytest.mark.parametrize("ir_version", [3, 9]) +def test_fixed_chain_controls_respect_legacy_initializer_declarations(ir_version): + """Collapse legacy fixed targets, declare new controls, and retain modern runtime defaults.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "first_shape"], ["middle"]), + helper.make_node("Reshape", ["middle", "last_shape"], ["y"]), + ], + "legacy_chain", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("first_shape", TensorProto.INT64, [2]), + helper.make_tensor_value_info("last_shape", TensorProto.INT64, [3]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 1, 2])], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "first_shape"), + numpy_helper.from_array(np.array([3, 1, 2], np.int64), "last_shape"), + ], + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = ir_version + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds) + converted = _apply_remove_reshapes(copy.deepcopy(model)) + assert len(converted.graph.node) == (1 if ir_version == 3 else 2) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("legacy", [False, True]) @pytest.mark.parametrize("cast", [False, True]) def test_native_gather_index_views_are_not_bypassed(legacy, cast): From 5bcc328e6799113c929e22c46a3762e3a9b5c675 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 07:06:24 +0000 Subject: [PATCH 06/15] Verify native Gemm fallback in integration safety checks Signed-off-by: mgohil --- .../dla_transforms/test_dla_rewrite_safety.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py index 499c933bf44..ad2d2a0c899 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -280,8 +280,8 @@ def test_not_preserves_nonbinary_and_shared_boolean_uses(): @pytest.mark.parametrize("change", ["alpha", "bias"]) -def test_unsupported_gemm_refuses_before_mutation(change): - """Refuse unsupported Gemm layouts without changing the caller.""" +def test_ineligible_gemm_retains_executable_native_contract(change): + """Keep valid Gemm attributes and broadcasting unchanged when Conv conversion is ineligible.""" attrs = {"alpha": 2.0} if change == "alpha" else {} bias = np.ones((2, 3) if change == "bias" else (3,), dtype=np.float32) model = _model( @@ -290,10 +290,14 @@ def test_unsupported_gemm_refuses_before_mutation(change): [_vi("y")], [nh.from_array(np.eye(3, dtype=np.float32), "w"), nh.from_array(bias, "b")], ) + onnx.checker.check_model(model, full_check=True) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = ReferenceEvaluator(model).run(None, feeds) before = model.SerializeToString() - with pytest.raises(ValueError): - _apply_matmul_to_transpose_conv_transpose(model) + _apply_matmul_to_transpose_conv_transpose(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("deferred", [False, True]) From 26eb3f22d0c4a03d7f1b4c5ae8dde070ff6d6168 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:22:00 +0000 Subject: [PATCH 07/15] Assert DLA reshape cleanup allocates distinct node names Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_remove_reshapes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py index a762f363dbf..daf3559b942 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -169,6 +169,9 @@ def test_unnamed_public_identity_reshapes_allocate_distinct_node_names(): } expected = _execute(model, feeds) converted = _apply_remove_reshapes(copy.deepcopy(model)) + names = [node.name for node in converted.graph.node] + assert all(names) + assert len(set(names)) == len(names) for actual, reference in zip(_execute(converted, feeds), expected): np.testing.assert_array_equal(actual, reference) From 45eab213fbdd2b00276d3c9479f2e0967b28be62 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:36:17 +0000 Subject: [PATCH 08/15] Preserve initializer input contracts during final DLA cleanup Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_graph_cleanup.py | 8 ++++++++ .../dla_transforms/test_dla_remove_reshapes.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index fda3e1be306..363109e4f23 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -846,7 +846,11 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: for o in graph.output: used.add(o.name) + if model.ir_version >= 4: + used.update(value.name for value in graph.input) + kept = [init for init in graph.initializer if init.name in used] + discarded = {init.name for init in graph.initializer if init.name not in used} removed = len(graph.initializer) - len(kept) if model.ir_version <= 3: removed_names = {init.name for init in graph.initializer if init.name not in used} @@ -855,6 +859,10 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: del graph.input[index] graph.ClearField("initializer") graph.initializer.extend(kept) + if model.ir_version <= 3: + inputs = [value for value in graph.input if value.name not in discarded] + graph.ClearField("input") + graph.input.extend(inputs) logger.debug("remove_unused_initializers: removed %d unused initializer(s).", removed) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py index daf3559b942..94855cb1ca2 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -23,6 +23,7 @@ import pytest from onnx import TensorProto, helper, numpy_helper +from modelopt.onnx.graph_surgery.dla_transforms.dla_graph_cleanup import _remove_unused_initializers from modelopt.onnx.graph_surgery.dla_transforms.dla_remove_reshapes import _apply_remove_reshapes @@ -48,11 +49,13 @@ def test_fixed_chain_controls_respect_legacy_initializer_declarations(ir_version helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), helper.make_tensor_value_info("first_shape", TensorProto.INT64, [2]), helper.make_tensor_value_info("last_shape", TensorProto.INT64, [3]), + helper.make_tensor_value_info("unused_default", TensorProto.INT64, [1]), ], [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 1, 2])], [ numpy_helper.from_array(np.array([2, 3], np.int64), "first_shape"), numpy_helper.from_array(np.array([3, 1, 2], np.int64), "last_shape"), + numpy_helper.from_array(np.array([1], np.int64), "unused_default"), ], value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], ), @@ -63,6 +66,11 @@ def test_fixed_chain_controls_respect_legacy_initializer_declarations(ir_version expected = _execute(model, feeds) converted = _apply_remove_reshapes(copy.deepcopy(model)) assert len(converted.graph.node) == (1 if ir_version == 3 else 2) + converted = _remove_unused_initializers(converted) + input_names = {value.name for value in converted.graph.input} + assert {"first_shape", "last_shape"}.isdisjoint(input_names) == (ir_version == 3) + assert ("unused_default" in input_names) == (ir_version >= 4) + assert {value.name for value in converted.graph.initializer} >= input_names - {"x"} for actual, reference in zip(_execute(converted, feeds), expected): np.testing.assert_array_equal(actual, reference) From 3c952161831089836b23160c6b2466f43dd7968c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:43:38 +0000 Subject: [PATCH 09/15] Declare deferred rank stages in rewrite safety checks Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_rewrite_safety.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py index ad2d2a0c899..0487442a566 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -248,7 +248,7 @@ def promote(candidate): """Declare deferred boundary ranks before applying the operation-specific stage.""" if operation == "Gather": candidate = _apply_gather(candidate) - candidate = _apply_non4d_model_inputs(candidate) + candidate = _apply_non4d_model_inputs(candidate, prepared_input_ops={operation}) candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) if operation == "ArgMax": candidate = _apply_argmax(candidate) From 278da9c9dd947707da5c3f7efa362a19925b123b Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 18:50:30 +0000 Subject: [PATCH 10/15] Collapse interleaved reshape chains by data producer Signed-off-by: mgohil --- .../dla_transforms/dla_remove_reshapes.py | 26 +++++------ .../test_dla_remove_reshapes.py | 45 +++++++++++++++++++ 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py index e14502e25c2..92d71083c78 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py @@ -34,21 +34,21 @@ def _apply_remove_reshapes(model): runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) captured = _subgraph_captures(graph) - chains = [] - chain = [] + chains, tails = [], {} for node in reshapes: - if chain: - consumers = cache.get_consumers(chain[-1].output[0]) - if ( - chain[-1].output[0] in cache.graph_output_names | captured - or len(consumers) != 1 - or consumers[0][0] is not node - ): - chains.append(chain) - chain = [] + producer = cache.get_producer(node.input[0]) + chain = tails.pop(id(producer), None) + if chain is None: + chain = [] + chains.append(chain) chain.append(node) - if chain: - chains.append(chain) + consumers = cache.get_consumers(node.output[0]) + if ( + node.output[0] not in cache.graph_output_names | captured + and len(consumers) == 1 + and consumers[0][1] == 0 + ): + tails[id(node)] = chain replacements = {} for chain in chains: first, last = chain[0], chain[-1] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py index 94855cb1ca2..47674e36ebd 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -245,3 +245,48 @@ def test_padding_view_before_public_unary_output_keeps_its_rank(): expected = _execute(model, feeds)[0] converted = _apply_remove_reshapes(copy.deepcopy(model)) np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("identity", [False, True]) +@pytest.mark.parametrize("public_middle", [False, True]) +def test_interleaved_reshape_chains_follow_data_producers(identity, public_middle): + """Collapse independent interleaved branches without crossing a public intermediate.""" + final_shape = [2, 3] if identity else [6] + outputs = [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, final_shape) for name in ["y", "w"] + ] + if public_middle: + outputs.append(helper.make_tensor_value_info("a", TensorProto.FLOAT, [3, 2])) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "a_shape"], ["a"]), + helper.make_node("Reshape", ["z", "b_shape"], ["b"]), + helper.make_node("Reshape", ["a", "final_shape"], ["y"]), + helper.make_node("Reshape", ["b", "final_shape"], ["w"]), + ], + "interleaved_chains", + [helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) for name in ["x", "z"]], + outputs, + [ + numpy_helper.from_array(np.array(values, np.int64), name) + for name, values in [ + ("a_shape", [3, 2]), + ("b_shape", [1, 6]), + ("final_shape", final_shape), + ] + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(2, 3), + "z": np.arange(10, 16, dtype=np.float32).reshape(2, 3), + } + expected = _execute(model, feeds) + changed = _apply_remove_reshapes(copy.deepcopy(model)) + assert len(changed.graph.node) == (3 if public_middle else 2) + for actual, reference in zip(_execute(changed, feeds), expected): + np.testing.assert_array_equal(actual, reference) From 3433be06a909e11019fbafad355461414779546c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 19:33:25 +0000 Subject: [PATCH 11/15] Preserve existing targets when bypassing interleaved identity reshapes Signed-off-by: mgohil --- .../dla_transforms/dla_remove_reshapes.py | 11 ++++++++++ .../test_dla_remove_reshapes.py | 22 +++++++++++++------ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py index 92d71083c78..ba3b6c4bcb5 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_remove_reshapes.py @@ -15,6 +15,8 @@ """Collapse reshape chains at the final, inferred stage of the DLA pipeline.""" +from itertools import pairwise + import numpy as np from onnx import TensorProto, helper @@ -34,6 +36,7 @@ def _apply_remove_reshapes(model): runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) captured = _subgraph_captures(graph) + positions = {id(node): index for index, node in enumerate(reshapes)} chains, tails = [], {} for node in reshapes: producer = cache.get_producer(node.input[0]) @@ -80,6 +83,14 @@ def _apply_remove_reshapes(model): ) ] elif len(chain) > 1: + if all(cache.get_shape(node.output[0]) == input_shape for node in chain[:-1]) and any( + positions[id(right)] != positions[id(left)] + 1 for left, right in pairwise(chain) + ): + # Interleaved identity prefixes need only bypassing; retain the existing target and attributes. + last.input[0] = first.input[0] + for node in chain[:-1]: + replacements[id(node)] = [] + continue target = add_initializer( graph, cache.unique_name(last.name + "_reshape_shape"), diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py index 47674e36ebd..6a7ec6c94e8 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_remove_reshapes.py @@ -249,21 +249,24 @@ def test_padding_view_before_public_unary_output_keeps_its_rank(): @pytest.mark.parametrize("identity", [False, True]) @pytest.mark.parametrize("public_middle", [False, True]) -def test_interleaved_reshape_chains_follow_data_producers(identity, public_middle): +@pytest.mark.parametrize("identity_prefix", [False, True]) +def test_interleaved_reshape_chains_follow_data_producers(identity, public_middle, identity_prefix): """Collapse independent interleaved branches without crossing a public intermediate.""" final_shape = [2, 3] if identity else [6] + a_shape = [2, 3] if identity_prefix else [3, 2] + b_shape = [2, 3] if identity_prefix else [1, 6] outputs = [ helper.make_tensor_value_info(name, TensorProto.FLOAT, final_shape) for name in ["y", "w"] ] if public_middle: - outputs.append(helper.make_tensor_value_info("a", TensorProto.FLOAT, [3, 2])) + outputs.append(helper.make_tensor_value_info("a", TensorProto.FLOAT, a_shape)) model = helper.make_model( helper.make_graph( [ helper.make_node("Reshape", ["x", "a_shape"], ["a"]), helper.make_node("Reshape", ["z", "b_shape"], ["b"]), - helper.make_node("Reshape", ["a", "final_shape"], ["y"]), - helper.make_node("Reshape", ["b", "final_shape"], ["w"]), + helper.make_node("Reshape", ["a", "final_shape"], ["y"], allowzero=0), + helper.make_node("Reshape", ["b", "final_shape"], ["w"], allowzero=0), ], "interleaved_chains", [helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) for name in ["x", "z"]], @@ -271,9 +274,9 @@ def test_interleaved_reshape_chains_follow_data_producers(identity, public_middl [ numpy_helper.from_array(np.array(values, np.int64), name) for name, values in [ - ("a_shape", [3, 2]), - ("b_shape", [1, 6]), - ("final_shape", final_shape), + ("a_shape", a_shape), + ("b_shape", b_shape), + ("final_shape", final_shape if identity else [-1]), ] ], ), @@ -288,5 +291,10 @@ def test_interleaved_reshape_chains_follow_data_producers(identity, public_middl expected = _execute(model, feeds) changed = _apply_remove_reshapes(copy.deepcopy(model)) assert len(changed.graph.node) == (3 if public_middle else 2) + if identity_prefix and not identity: + for node in changed.graph.node: + if node.output[0] in {"y", "w"}: + assert node.input[1] == "final_shape" + assert next(attr.i for attr in node.attribute if attr.name == "allowzero") == 0 for actual, reference in zip(_execute(changed, feeds), expected): np.testing.assert_array_equal(actual, reference) From f4895bb18e63d316a9964144d72910d34e4a168c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 20:25:54 +0000 Subject: [PATCH 12/15] Remove redundant mock-only dispatcher coverage Signed-off-by: mgohil --- .../dla_transforms/test_dla_rewrite_safety.py | 25 ------------------- 1 file changed, 25 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py index 0487442a566..49acdaa8d20 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rewrite_safety.py @@ -25,7 +25,6 @@ from onnx import numpy_helper as nh from onnx.reference import ReferenceEvaluator -from modelopt.onnx import graph_surgery from modelopt.onnx.graph_surgery.dla_transforms.dla_convert_ops_to_4d import ( _apply_convert_ops_to_4d, ) @@ -349,30 +348,6 @@ def test_legacy_rank_dependent_op_is_left_intact(): ) -def test_dispatcher_keeps_existing_return_and_keyword_contract(monkeypatch): - """Preserve registry dispatch return values and keyword forwarding.""" - calls = [] - - def original(**kwargs): - """Record registry forwarding and return its sentinel value.""" - calls.append(kwargs) - return {"converted": 3} - - monkeypatch.setitem(graph_surgery._SURGERY_REGISTRY, "convert-bf16", original) - result = graph_surgery.run_graph_surgery( - "convert-bf16", "input.onnx", "output.onnx", external_data=False, verbose=False - ) - assert result == {"converted": 3} - assert calls == [ - { - "model_path": "input.onnx", - "output_path": "output.onnx", - "external_data": False, - "verbose": False, - } - ] - - def test_control_flow_capture_keeps_initializer(): """Retain fixed data referenced from a nested control-flow graph.""" branch = h.make_graph( From b994f6b51802dd4a419cf6fe11529ffe074986ee Mon Sep 17 00:00:00 2001 From: mgohil Date: Wed, 28 Oct 2026 03:16:46 +0000 Subject: [PATCH 13/15] Clarify DLA entry points and signed-zero conversion contract Signed-off-by: mgohil --- CHANGELOG.rst | 4 ++-- modelopt/onnx/graph_surgery/__main__.py | 5 ++++- modelopt/onnx/graph_surgery/make_dla_compatible.py | 2 ++ 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a41790cc9c3..3906cc2ff84 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -34,8 +34,8 @@ Changelog *Misc* - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs. -- Add the DLA graph-surgery pipeline with explicitly ordered transforms and a - ``--remove-all-qdq`` option. +- Add the ordered DLA compatibility pipeline through the ``modelopt.onnx.graph_surgery.make_dla_compatible`` + API and ``make-dla-compatible`` CLI command. Use ``--remove-all-qdq`` to remove Q/DQ pairs of every dtype. **Backward Breaking Changes** diff --git a/modelopt/onnx/graph_surgery/__main__.py b/modelopt/onnx/graph_surgery/__main__.py index 50730b38a0c..0c130960de4 100644 --- a/modelopt/onnx/graph_surgery/__main__.py +++ b/modelopt/onnx/graph_surgery/__main__.py @@ -248,7 +248,10 @@ def main(): dla_parser = subparsers.add_parser( "make-dla-compatible", help="Make a model compatible with the NVIDIA DLA accelerator", - description="Apply the ordered DLA graph transforms and validate the resulting ONNX model.", + description=( + "Apply the ordered DLA graph transforms and validate the resulting ONNX model. " + "Where lowering assumes finite FP32 branches and may change the sign of zero." + ), ) dla_parser.add_argument("-i", "--input", required=True, help="Input ONNX model path") dla_parser.add_argument("-o", "--output", required=True, help="Output ONNX model path") diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 0507c7259eb..e33a995da07 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -329,6 +329,8 @@ def dla_make_dla_compatible( Arithmetic lowering of zero-branch ``Where`` assumes finite FP32 branch tensors. NaN and infinity values at those nodes are outside the supported input contract. + Masked multiplication can change the sign of zero even for finite inputs; validate + sign-sensitive downstream operations before using this conversion. Args: model_path: Input file; external weights are loaded from its directory. From 639c06da02195689b87c3cceb541f7e9b3222f6f Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:20:11 +0000 Subject: [PATCH 14/15] Use direct ONNX saves for the DLA pipeline Signed-off-by: mgohil --- .../onnx/graph_surgery/make_dla_compatible.py | 27 +++++--- .../test_dla_make_compatible.py | 61 +++++++++++++++++++ 2 files changed, 80 insertions(+), 8 deletions(-) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index e33a995da07..d9629428a41 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -24,6 +24,7 @@ from __future__ import annotations import copy +import os import time from functools import partial from typing import TYPE_CHECKING @@ -84,7 +85,6 @@ GraphCache, infer_shapes, parse_qdq_quantized_dtype_list, - save_onnx_model, ) if TYPE_CHECKING: @@ -352,11 +352,22 @@ def dla_make_dla_compatible( qdq_quantized_dtypes=qdq_quantized_dtypes, remove_all_qdq=remove_all_qdq, ) - save_onnx_model( - model, - output_path, - use_external_data=use_external_data, - external_data_name=external_data_name, - verbose=verbose, - ) + output_dir = os.path.dirname(output_path) + if output_dir: + os.makedirs(output_dir, exist_ok=True) + if verbose: + logger.info("Saving modified model to: %s", output_path) + if use_external_data: + if external_data_name is None: + external_data_name = os.path.basename(output_path) + "_data" + onnx.save( + model, + output_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=external_data_name, + size_threshold=1024, + ) + else: + onnx.save(model, output_path) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py index 4599ba5f175..4c7d4daf494 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py @@ -413,3 +413,64 @@ def test_qdq_dtypes_threaded_through_run_graph_surgery(tmp_path): ) assert _count_op(result.graph, "QuantizeLinear") == 0 assert _count_op(result.graph, "DequantizeLinear") == 0 + + +@pytest.mark.parametrize("external_name", [False, None, "weights.bin"]) +def test_public_pipeline_saves_with_onnx(tmp_path, external_name): + """Reload and execute inline or external output produced by the public pipeline.""" + shape = [1, 1, 1, 512] + weights = np.arange(512, dtype=np.float32).reshape(shape) + # Include the Windows trailing-0x1A regression pattern in an externalized payload. + weights.view(np.uint8).reshape(-1)[-1] = 0x1A + bias = np.full(shape, 0.5, np.float32) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "weight"], ["sum"]), + helper.make_node("Add", ["sum", "bias"], ["y"]), + ], + "direct_onnx_save", + [_vi("x", TensorProto.FLOAT, shape)], + [_vi("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(weights, "weight"), numpy_helper.from_array(bias, "bias")], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ir_version=_MAX_IR, + ) + source = tmp_path / "source.onnx" + output = tmp_path / "output" / "converted.onnx" + onnx.save(model, source) + _dla_pipeline.dla_make_dla_compatible( + str(source), + str(output), + use_external_data=external_name is not False, + external_data_name=external_name if isinstance(external_name, str) else None, + verbose=False, + ) + onnx.checker.check_model(str(output), full_check=True) + actual = onnx.load(output) + for name, expected in [("weight", weights), ("bias", bias)]: + assert ( + next(t for t in actual.graph.initializer if t.name == name).raw_data + == expected.tobytes() + ) + feeds = {"x": np.ones(shape, np.float32)} + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + result = ort.InferenceSession( + actual.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds)[0] + np.testing.assert_array_equal(result, (feeds["x"] + weights) + bias) + saved = onnx.load(output, load_external_data=False) + external = [ + t for t in saved.graph.initializer if onnx.external_data_helper.uses_external_data(t) + ] + assert bool(external) == (external_name is not False) + if external: + assert len(external) == 2 + expected = external_name or output.name + "_data" + assert (output.parent / expected).is_file() + assert all( + next(e.value for e in t.external_data if e.key == "location") == expected + for t in external + ) From 48677b5e5cb16614194fe4ec207a23658fcc5d47 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 09:16:23 +0000 Subject: [PATCH 15/15] Use native ONNX sidecar naming and remove duplicate cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_graph_cleanup.py | 7 +- .../onnx/graph_surgery/make_dla_compatible.py | 4 +- .../test_dla_make_compatible.py | 70 ++++++++++++++++++- 3 files changed, 71 insertions(+), 10 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py index 363109e4f23..f891c5907ce 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_graph_cleanup.py @@ -853,16 +853,11 @@ def _remove_unused_initializers(model: onnx.ModelProto) -> onnx.ModelProto: discarded = {init.name for init in graph.initializer if init.name not in used} removed = len(graph.initializer) - len(kept) if model.ir_version <= 3: - removed_names = {init.name for init in graph.initializer if init.name not in used} for index in reversed(range(len(graph.input))): - if graph.input[index].name in removed_names: + if graph.input[index].name in discarded: del graph.input[index] graph.ClearField("initializer") graph.initializer.extend(kept) - if model.ir_version <= 3: - inputs = [value for value in graph.input if value.name not in discarded] - graph.ClearField("input") - graph.input.extend(inputs) logger.debug("remove_unused_initializers: removed %d unused initializer(s).", removed) return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index d9629428a41..96958474d61 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -336,7 +336,7 @@ def dla_make_dla_compatible( model_path: Input file; external weights are loaded from its directory. output_path: Destination ONNX file. use_external_data: Store large tensor payloads in a sidecar file. - external_data_name: Sidecar filename; defaults to output filename plus "_data". + external_data_name: Relative sidecar filename; ONNX chooses a unique name by default. verbose: Log each transform stage. qdq_quantized_dtypes: Zero-point dtypes whose QDQ pairs are removed. Defaults to UINT16 and INT16. Accepts ONNX dtype IDs or names. @@ -358,8 +358,6 @@ def dla_make_dla_compatible( if verbose: logger.info("Saving modified model to: %s", output_path) if use_external_data: - if external_data_name is None: - external_data_name = os.path.basename(output_path) + "_data" onnx.save( model, output_path, diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py index 4c7d4daf494..3d9b8f5cdaa 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_make_compatible.py @@ -19,6 +19,7 @@ import copy import importlib +from pathlib import Path import numpy as np import onnx @@ -468,9 +469,76 @@ def test_public_pipeline_saves_with_onnx(tmp_path, external_name): assert bool(external) == (external_name is not False) if external: assert len(external) == 2 - expected = external_name or output.name + "_data" + expected = next(e.value for e in external[0].external_data if e.key == "location") + assert external_name is None or expected == external_name assert (output.parent / expected).is_file() assert all( next(e.value for e in t.external_data if e.key == "location") == expected for t in external ) + + +@pytest.mark.parametrize( + ("output_location", "occupied"), + [("cwd", "cwd"), ("nested", "cwd"), ("nested", "output"), ("nested", "both")], +) +def test_public_pipeline_repeated_default_saves_preserve_sidecars( + tmp_path, monkeypatch, output_location, occupied +): + """Native default naming must preserve existing files and earlier conversion payloads.""" + cwd = tmp_path / "cwd" + cwd.mkdir() + monkeypatch.chdir(cwd) + directory = cwd if output_location == "cwd" else tmp_path / "output" + directory.mkdir(exist_ok=True) + output = directory / "converted.onnx" + source = tmp_path / "source.onnx" + stale = output.name + "_data" + occupied_paths = set() + if occupied in {"cwd", "both"}: + occupied_paths.add(cwd / stale) + if occupied in {"output", "both"}: + occupied_paths.add(directory / stale) + preserved = {} + for path in occupied_paths: + path.write_bytes(b"keep this existing payload") + preserved[path] = path.read_bytes() + locations = [] + shape = [1, 1, 1, 512] + feeds = {"x": np.ones(shape, np.float32)} + for offset in [0, 1]: + weight = np.arange(512, dtype=np.float32).reshape(shape) + offset + model = helper.make_model( + helper.make_graph( + [helper.make_node("Add", ["x", "weight"], ["y"])], + "repeat_native_save", + [_vi("x", TensorProto.FLOAT, shape)], + [_vi("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(weight, "weight")], + ), + opset_imports=[helper.make_opsetid("", _OPSET)], + ir_version=_MAX_IR, + ) + onnx.save(model, source) + _dla_pipeline.dla_make_dla_compatible(str(source), str(output), verbose=False) + onnx.checker.check_model(str(output), full_check=True) + saved = onnx.load(output, load_external_data=False) + tensor = next(t for t in saved.graph.initializer if t.name == "weight") + location = next(e.value for e in tensor.external_data if e.key == "location") + assert location == Path(location).name and location not in locations + locations.append(location) + for path, payload in preserved.items(): + assert path.read_bytes() == payload + sidecar = directory / location + preserved[sidecar] = sidecar.read_bytes() + actual = onnx.load(output) + np.testing.assert_array_equal( + numpy_helper.to_array(next(t for t in actual.graph.initializer if t.name == "weight")), + weight, + ) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + result = ort.InferenceSession( + actual.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds)[0] + np.testing.assert_array_equal(result, feeds["x"] + weight)