From b5c0f3dbc2d861d374fedceef244f93d1bfa7471 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:13 +0000 Subject: [PATCH 01/32] [DLA_transforms_10] [10/17] DLA model boundary promotion Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 259 ++++++++++++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 10 + .../test_dla_model_boundaries.py | 44 +++ 3 files changed, 313 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py new file mode 100644 index 00000000000..7ac8f9af49b --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -0,0 +1,259 @@ +# 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. + +"""Promote model boundaries at their explicit positions in the DLA pipeline.""" + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +from ..utils.graph_utils import GraphCache, add_initializer, pad4d + +__all__ = [] + + +def _apply_non4d_model_inputs(model): + graph = model.graph + cache = GraphCache(graph) + replacements = {} + prepend = [] + promotions = [] + unsqueezed_shapes = {} + for value in graph.input: + shape = cache.get_shape(value.name) + if shape is None or len(shape) not in (2, 3): + continue + consumers = [ + node + for node, _ in cache.get_consumers(value.name) + if (node.input[0] == value.name or node.op_type == "Concat") + and not (node.op_type == "Squeeze" and "Gather" in node.name) + ] + if not consumers: + continue + for consumer in consumers: + if consumer.op_type != "Unsqueeze": + continue + tensor = cache.get_init(consumer.input[1]) + if tensor is None: + raise ValueError(f"Missing Unsqueeze axes for {consumer.name!r}") + output_shape = list(shape) + for axis in numpy_helper.to_array(tensor).reshape(-1).tolist(): + output_shape.insert(axis if axis >= 0 else len(output_shape) + axis + 1, 1) + if ( + pad4d(shape)[-len(output_shape) :] == output_shape + and consumer.output[0] in cache.graph_output_names + ): + raise ValueError(f"Cannot remove graph-output Unsqueeze {consumer.name!r}") + unsqueezed_shapes[id(consumer)] = output_shape + promotions.append((value, shape, consumers)) + for value, shape, consumers in promotions: + axes = list(range(4 - len(shape))) + axes_name = add_initializer( + graph, + value.name + "_unsqueeze_axes", + np.asarray(axes, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = value.name + "_unsqueeze_input" + prepend.insert( + 0, + helper.make_node( + "Unsqueeze", [value.name, axes_name], [output], name=value.name + "_unsqueeze" + ), + ) + promoted_shape = pad4d(shape) + graph.value_info.append( + helper.make_tensor_value_info(output, value.type.tensor_type.elem_type, promoted_shape) + ) + for consumer in consumers: + if consumer.op_type == "Unsqueeze": + output_shape = unsqueezed_shapes[id(consumer)] + if promoted_shape[-len(output_shape) :] == output_shape: + old_output = consumer.output[0] + for downstream, _ in cache.get_consumers(old_output): + for index, name in enumerate(downstream.input): + if name == old_output: + downstream.input[index] = output + replacements[id(consumer)] = [] + else: + target = add_initializer( + graph, + consumer.name + "_reshape_shape", + np.asarray(output_shape, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + replacements[id(consumer)] = [ + helper.make_node( + "Reshape", + [output, target], + list(consumer.output), + name=consumer.name + "_reshape", + ) + ] + else: + for index, name in enumerate(consumer.input): + if name == value.name: + consumer.input[index] = output + nodes = list(graph.node) + del graph.node[:] + graph.node.extend(prepend) + for node in nodes: + graph.node.extend(replacements.get(id(node), [node])) + return model + + +def _apply_non4d_model_outputs(model): + graph = model.graph + cache = GraphCache(graph) + boundary_names = cache.graph_input_names | cache.init_names + append = [] + for value in graph.output: + shape = cache.get_shape(value.name) + if shape is None or len(shape) >= 4: + continue + producer = cache.get_producer(value.name) + if ( + producer is not None + and producer.op_type == "Identity" + and producer.input[0] in boundary_names + ): + continue + if len(shape) == 0: + raise ValueError(f"Unsupported scalar model output {value.name!r}") + axes = [] + existing = producer if producer is not None and producer.op_type == "Squeeze" else None + if existing is not None and len(existing.input) > 1: + initializer = cache.get_init(existing.input[1]) + if initializer is not None: + axes = numpy_helper.to_array(initializer).reshape(-1).tolist() + delta = 4 - len(axes) - len(shape) + axes = list(range(delta)) + [axis + delta for axis in axes] + axes_name = None + if len(shape) > 1: + axes_name = add_initializer( + graph, + value.name + "_squeeze_axes", + np.asarray(axes, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + if existing is not None: + data_input = existing.input[0] + del existing.input[:] + existing.input.extend([data_input] + ([axes_name] if axes_name else [])) + continue + data_input = value.name + "_squeeze_input" + for node in graph.node: + for index, name in enumerate(node.input): + if name == value.name: + node.input[index] = data_input + for index, name in enumerate(node.output): + if name == value.name: + node.output[index] = data_input + append.append( + helper.make_node( + "Squeeze", + [data_input] + ([axes_name] if axes_name else []), + [value.name], + name=value.name + "_squeeze", + ) + ) + graph.node.extend(append) + return model + + +def _apply_remove_intermediary_squeeze_unsqueeze(model): + graph = model.graph + cache = GraphCache(graph) + removed = set() + for node in list(graph.node): + if node.op_type not in {"Squeeze", "Unsqueeze"}: + continue + if ( + node.input[0] in cache.graph_input_names | cache.init_names | cache.graph_output_names + or node.output[0] in cache.graph_output_names + ): + continue + input_shape = cache.get_shape(node.input[0]) + output_shape = cache.get_shape(node.output[0]) + if not input_shape or not output_shape: + continue + short, long = ( + (output_shape, input_shape) + if node.op_type == "Squeeze" + else (input_shape, output_shape) + ) + if len(short) <= len(long) and long[-len(short) :] == short: + old, new = node.input[0], node.output[0] + for other in graph.node: + if other is node: + continue + for index, name in enumerate(other.output): + if name == old: + other.output[index] = new + for index, name in enumerate(other.input): + if name == old: + other.input[index] = new + removed.add(id(node)) + else: + target = add_initializer( + graph, + node.name + "_reshape_shape", + np.asarray(pad4d(output_shape), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.CopyFrom( + helper.make_node( + "Reshape", + [node.input[0], target], + list(node.output), + name=node.name + "_to_reshape", + ) + ) + remaining = [node for node in graph.node if id(node) not in removed] + del graph.node[:] + graph.node.extend(remaining) + return model + + +def _apply_fix_instancenorm_channels(model): + graph = model.graph + cache = GraphCache(graph) + nodes = [] + for node in graph.node: + if node.op_type == "InstanceNormalization": + shape = cache.get_shape(node.input[0]) + if shape is not None and len(shape) == 3: + target = add_initializer( + graph, + node.name + "_rotateleft_shape", + np.asarray([*shape, 1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(node.input[0] + "_rotateleft") + nodes.append( + helper.make_node( + "Reshape", [node.input[0], target], [output], name=node.name + "_rotateleft" + ) + ) + node.input[0] = output + nodes.append(node) + del graph.node[:] + graph.node.extend(nodes) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index bb6168d140d..85cbff385b3 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -37,6 +37,12 @@ from .dla_transforms.dla_matmul_to_transpose_conv_transpose import ( _apply_matmul_to_transpose_conv_transpose, ) +from .dla_transforms.dla_model_boundaries import ( + _apply_fix_instancenorm_channels, + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + _apply_remove_intermediary_squeeze_unsqueeze, +) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes @@ -87,6 +93,10 @@ def _apply_constants_to_initializers(model, **options): "MatMulToTransposeConvTranspose", partial(_apply_matmul_to_transpose_conv_transpose, defer_rank_conversion=True), ), + ("Non4dModelInputs", _apply_non4d_model_inputs), + ("Non4dModelOutputs", _apply_non4d_model_outputs), + ("RemoveIntermediarySqueezeAndUnsqueeze", _apply_remove_intermediary_squeeze_unsqueeze), + ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py new file mode 100644 index 00000000000..8459df24a7d --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -0,0 +1,44 @@ +# 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. + +"""Boundary promotion refuses unsupported exposed adapters before changing the graph.""" + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_non4d_model_inputs, +) + + +def test_input_promotion_refusal_leaves_graph_unchanged(): + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Identity", ["x"], ["first"]), + helper.make_node("Unsqueeze", ["x", "axes"], ["y"], name="exposed"), + ], + "boundary", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3])], + [numpy_helper.from_array(np.array([0, 1], dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before From e206f825bd5a6ceacde156eef96266a5d4f94d1c Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:04:02 +0000 Subject: [PATCH 02/32] Preserve boundary layouts, generated names and captured rank adapters Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 159 +++++-- .../onnx/graph_surgery/make_dla_compatible.py | 5 +- .../test_dla_model_boundaries.py | 398 ++++++++++++++++++ 3 files changed, 518 insertions(+), 44 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 7ac8f9af49b..4dfbfc22681 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -18,14 +18,28 @@ import numpy as np from onnx import TensorProto, helper, numpy_helper -from ..utils.graph_utils import GraphCache, add_initializer, pad4d +from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes, pad4d +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures +from .dla_reorder_nodes import priority_topological_order __all__ = [] +def _require_adapter_opset(model): + """Refuse input-form axes adapters before changing a legacy graph.""" + version = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) + if version < 13: + raise ValueError("DLA boundary rank adapters require ONNX opset 13 or newer") + + def _apply_non4d_model_inputs(model): + """Promote selected boundary uses while preserving explicit adapter results.""" graph = model.graph cache = GraphCache(graph) + captured = _subgraph_captures(graph) replacements = {} prepend = [] promotions = [] @@ -42,36 +56,44 @@ def _apply_non4d_model_inputs(model): ] if not consumers: continue + _require_adapter_opset(model) for consumer in consumers: if consumer.op_type != "Unsqueeze": continue tensor = cache.get_init(consumer.input[1]) - if tensor is None: - raise ValueError(f"Missing Unsqueeze axes for {consumer.name!r}") + if tensor is None or tensor.name in cache.graph_input_names: + raise ValueError(f"Unsqueeze axes must be fixed for {consumer.name!r}") output_shape = list(shape) - for axis in numpy_helper.to_array(tensor).reshape(-1).tolist(): - output_shape.insert(axis if axis >= 0 else len(output_shape) + axis + 1, 1) + axes = numpy_helper.to_array(tensor).reshape(-1).tolist() + output_rank = len(shape) + len(axes) + for axis in sorted(axis % output_rank for axis in axes): + output_shape.insert(axis, 1) if ( pad4d(shape)[-len(output_shape) :] == output_shape - and consumer.output[0] in cache.graph_output_names + and consumer.output[0] in cache.graph_output_names | captured ): - raise ValueError(f"Cannot remove graph-output Unsqueeze {consumer.name!r}") + raise ValueError( + f"Cannot remove graph-output or captured Unsqueeze {consumer.name!r}" + ) unsqueezed_shapes[id(consumer)] = output_shape promotions.append((value, shape, consumers)) for value, shape, consumers in promotions: axes = list(range(4 - len(shape))) axes_name = add_initializer( graph, - value.name + "_unsqueeze_axes", + cache.unique_name(value.name + "_unsqueeze_axes"), np.asarray(axes, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - output = value.name + "_unsqueeze_input" + output = cache.unique_name(value.name + "_unsqueeze_input") prepend.insert( 0, helper.make_node( - "Unsqueeze", [value.name, axes_name], [output], name=value.name + "_unsqueeze" + "Unsqueeze", + [value.name, axes_name], + [output], + name=cache.unique_name(value.name + "_unsqueeze"), ), ) promoted_shape = pad4d(shape) @@ -91,7 +113,7 @@ def _apply_non4d_model_inputs(model): else: target = add_initializer( graph, - consumer.name + "_reshape_shape", + cache.unique_name((consumer.name or consumer.output[0]) + "_reshape_shape"), np.asarray(output_shape, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, @@ -101,29 +123,39 @@ def _apply_non4d_model_inputs(model): "Reshape", [output, target], list(consumer.output), - name=consumer.name + "_reshape", + name=cache.unique_name( + (consumer.name or consumer.output[0]) + "_reshape" + ), ) ] else: for index, name in enumerate(consumer.input): if name == value.name: consumer.input[index] = output - nodes = list(graph.node) - del graph.node[:] - graph.node.extend(prepend) - for node in nodes: - graph.node.extend(replacements.get(id(node), [node])) + batch = { + index: ([node], replacements[id(node)]) + for index, node in enumerate(graph.node) + if id(node) in replacements + } + if prepend: + removed, added = batch.get(0, ([], [])) + batch[0] = (removed, [*prepend, *added]) + batch_replace_nodes(graph, batch) + _add_legacy_initializer_inputs(model) return model def _apply_non4d_model_outputs(model): + """Restore declared output ranks after the ordered rank-promotion sequence.""" graph = model.graph cache = GraphCache(graph) + captured = _subgraph_captures(graph) boundary_names = cache.graph_input_names | cache.init_names append = [] + plans = [] for value in graph.output: shape = cache.get_shape(value.name) - if shape is None or len(shape) >= 4: + if shape is None or len(shape) >= 4 or value.name in boundary_names: continue producer = cache.get_producer(value.name) if ( @@ -136,17 +168,28 @@ def _apply_non4d_model_outputs(model): raise ValueError(f"Unsupported scalar model output {value.name!r}") axes = [] existing = producer if producer is not None and producer.op_type == "Squeeze" else None - if existing is not None and len(existing.input) > 1: - initializer = cache.get_init(existing.input[1]) - if initializer is not None: + if existing is not None: + if len(existing.input) > 1 and existing.input[1]: + initializer = cache.get_init(existing.input[1]) + if initializer is None or initializer.name in cache.graph_input_names: + raise ValueError(f"Output Squeeze axes must be fixed for {value.name!r}") axes = numpy_helper.to_array(initializer).reshape(-1).tolist() + else: + input_shape = cache.get_shape(existing.input[0]) + if input_shape is None: + raise ValueError(f"Output Squeeze input shape must be known for {value.name!r}") + axes = [index for index, dim in enumerate(input_shape) if dim == 1] delta = 4 - len(axes) - len(shape) - axes = list(range(delta)) + [axis + delta for axis in axes] + axes = list(range(delta)) + [axis + delta if axis >= 0 else axis for axis in axes] + plans.append((value, shape, existing, axes)) + if plans: + _require_adapter_opset(model) + for value, shape, existing, axes in plans: axes_name = None - if len(shape) > 1: + if len(shape) > 1 or shape == [1]: axes_name = add_initializer( graph, - value.name + "_squeeze_axes", + cache.unique_name(value.name + "_squeeze_axes"), np.asarray(axes, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, @@ -156,7 +199,7 @@ def _apply_non4d_model_outputs(model): del existing.input[:] existing.input.extend([data_input] + ([axes_name] if axes_name else [])) continue - data_input = value.name + "_squeeze_input" + data_input = cache.unique_name(value.name + "_squeeze_input") for node in graph.node: for index, name in enumerate(node.input): if name == value.name: @@ -169,29 +212,40 @@ def _apply_non4d_model_outputs(model): "Squeeze", [data_input] + ([axes_name] if axes_name else []), [value.name], - name=value.name + "_squeeze", + name=cache.unique_name(value.name + "_squeeze"), ) ) graph.node.extend(append) + if any(value.name in captured for value, _, _, _ in plans): + order, _, _ = priority_topological_order(graph) + nodes = list(graph.node) + ranks = {id(nodes[index]): rank for rank, index in enumerate(order)} + graph.node.sort(key=lambda node: ranks[id(node)]) + _add_legacy_initializer_inputs(model) return model -def _apply_remove_intermediary_squeeze_unsqueeze(model): +def _apply_remove_intermediary_squeeze_unsqueeze(model, *, rank_converted=False): + """Remove no-op adapters, or reconcile adapters within the ordered rank-converted pipeline.""" graph = model.graph cache = GraphCache(graph) + captured = _subgraph_captures(graph) removed = set() for node in list(graph.node): if node.op_type not in {"Squeeze", "Unsqueeze"}: continue if ( node.input[0] in cache.graph_input_names | cache.init_names | cache.graph_output_names - or node.output[0] in cache.graph_output_names + or node.output[0] in cache.graph_output_names | captured + or node.input[0] in captured ): continue input_shape = cache.get_shape(node.input[0]) output_shape = cache.get_shape(node.output[0]) if not input_shape or not output_shape: continue + if not rank_converted and input_shape != output_shape: + continue short, long = ( (output_shape, input_shape) if node.op_type == "Squeeze" @@ -212,7 +266,7 @@ def _apply_remove_intermediary_squeeze_unsqueeze(model): else: target = add_initializer( graph, - node.name + "_reshape_shape", + cache.unique_name((node.name or node.output[0]) + "_reshape_shape"), np.asarray(pad4d(output_shape), dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, @@ -222,38 +276,57 @@ def _apply_remove_intermediary_squeeze_unsqueeze(model): "Reshape", [node.input[0], target], list(node.output), - name=node.name + "_to_reshape", + name=cache.unique_name((node.name or node.output[0]) + "_to_reshape"), ) ) - remaining = [node for node in graph.node if id(node) not in removed] - del graph.node[:] - graph.node.extend(remaining) + batch_replace_nodes( + graph, {index: ([node], []) for index, node in enumerate(graph.node) if id(node) in removed} + ) + _add_legacy_initializer_inputs(model) return model def _apply_fix_instancenorm_channels(model): + """Append spatial singleton dimensions while retaining normalization channels and outputs.""" graph = model.graph cache = GraphCache(graph) - nodes = [] - for node in graph.node: + replacements = {} + for index, node in enumerate(graph.node): if node.op_type == "InstanceNormalization": shape = cache.get_shape(node.input[0]) if shape is not None and len(shape) == 3: target = add_initializer( graph, - node.name + "_rotateleft_shape", + cache.unique_name((node.name or node.output[0]) + "_rotateleft_shape"), np.asarray([*shape, 1], dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) output = cache.unique_name(node.input[0] + "_rotateleft") - nodes.append( - helper.make_node( - "Reshape", [node.input[0], target], [output], name=node.name + "_rotateleft" - ) + reshape = helper.make_node( + "Reshape", + [node.input[0], target], + [output], + name=cache.unique_name((node.name or node.output[0]) + "_rotateleft"), ) node.input[0] = output - nodes.append(node) - del graph.node[:] - graph.node.extend(nodes) + replacements[index] = ([node], [reshape, node]) + for consumer, _ in cache.get_consumers(node.output[0]): + if ( + consumer.op_type != "Squeeze" + or consumer.output[0] not in cache.graph_output_names + or len(consumer.input) < 2 + ): + continue + axes = cache.get_init(consumer.input[1]) + if axes is not None and numpy_helper.to_array(axes).reshape(-1).tolist() == [0]: + consumer.input[1] = add_initializer( + graph, + cache.unique_name(axes.name + "_trailing"), + np.asarray([3], np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + batch_replace_nodes(graph, replacements) + _add_legacy_initializer_inputs(model) return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 85cbff385b3..c047412898c 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -95,7 +95,10 @@ def _apply_constants_to_initializers(model, **options): ), ("Non4dModelInputs", _apply_non4d_model_inputs), ("Non4dModelOutputs", _apply_non4d_model_outputs), - ("RemoveIntermediarySqueezeAndUnsqueeze", _apply_remove_intermediary_squeeze_unsqueeze), + ( + "RemoveIntermediarySqueezeAndUnsqueeze", + partial(_apply_remove_intermediary_squeeze_unsqueeze, rank_converted=True), + ), ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 8459df24a7d..0db7c85c032 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -16,15 +16,22 @@ """Boundary promotion refuses unsupported exposed adapters before changing the graph.""" import numpy as np +import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_fix_instancenorm_channels, _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + _apply_remove_intermediary_squeeze_unsqueeze, ) +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible def test_input_promotion_refusal_leaves_graph_unchanged(): + """Refuse removing a public input adapter before altering any earlier consumer.""" model = helper.make_model( helper.make_graph( [ @@ -42,3 +49,394 @@ def test_input_promotion_refusal_leaves_graph_unchanged(): with pytest.raises(ValueError): _apply_non4d_model_inputs(model) assert model.SerializeToString() == before + + +def _checked_run(model, feeds): + """Check every generated interface and execute through the independent reference runtime.""" + onnx.checker.check_model(model, full_check=True) + return ReferenceEvaluator(model).run(None, feeds) + + +def _make_model(nodes, inputs, outputs, initializers=()): + """Build a small model with explicit public interface shapes.""" + return helper.make_model( + helper.make_graph(nodes, "boundaries", inputs, outputs, initializer=initializers), + opset_imports=[helper.make_opsetid("", 21)], + ) + + +def test_unnamed_input_adapters_keep_distinct_targets(): + """Promote two unnamed trailing-axis adapters without reusing their different targets.""" + inputs = [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 4]), + ] + outputs = [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3, 1]), + helper.make_tensor_value_info("b_out", TensorProto.FLOAT, [3, 4, 1]), + ] + model = _make_model( + [ + helper.make_node("Unsqueeze", ["a", "axes"], ["a_out"]), + helper.make_node("Unsqueeze", ["b", "axes"], ["b_out"]), + ], + inputs, + outputs, + [numpy_helper.from_array(np.asarray([-1], np.int64), "axes")], + ) + feeds = { + "a": np.arange(6, dtype=np.float32).reshape(2, 3), + "b": np.arange(12, dtype=np.float32).reshape(3, 4), + } + expected = _checked_run(model, feeds) + _apply_non4d_model_inputs(model) + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("width", [1, 4]) +@pytest.mark.parametrize("axis", [0, -2]) +def test_ordered_pipeline_preserves_vector_output(width, axis): + """Keep vector outputs, including a singleton vector, through the actual registered sequence.""" + model = _make_model( + [ + helper.make_node("MatMul", ["x", "w"], ["matrix"]), + helper.make_node("Squeeze", ["matrix", "axes"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [width])], + [ + numpy_helper.from_array(np.arange(3 * width, dtype=np.float32).reshape(3, width), "w"), + numpy_helper.from_array(np.asarray([axis], np.int64), "axes"), + ], + ) + feeds = {"x": np.asarray([[1, 2, 3]], np.float32)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + actual = _checked_run(transformed, feeds)[0] + assert actual.shape == (width,) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("axes", [[1], [-2], None]) +def test_ordered_pipeline_preserves_matrix_squeeze_axes(axes): + """Restore matrix outputs with positive, negative or omitted source squeeze axes.""" + controls = [] if axes is None else [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")] + squeeze_inputs = ["matrix"] if axes is None else ["matrix", "axes"] + model = _make_model( + [ + helper.make_node("MatMul", ["x", "w"], ["matrix"]), + helper.make_node("Squeeze", squeeze_inputs, ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 1, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4])], + [numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "w"), *controls], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 1, 3)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) + + +def test_standalone_intermediary_rank_change_preserves_concat_inputs(): + """Retain an internal rank-changing squeeze whose downstream consumer expects rank three.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["before"]), + helper.make_node("Squeeze", ["before", "axes"], ["squeezed"]), + helper.make_node("Concat", ["squeezed", "b"], ["y"], axis=0), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [2, 4, 4]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [5, 4, 4])], + [numpy_helper.from_array(np.asarray([0], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = { + "x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4), + "b": np.ones((2, 4, 4), np.float32), + } + expected = _checked_run(model, feeds)[0] + _apply_remove_intermediary_squeeze_unsqueeze(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +def test_standalone_empty_axes_adapter_can_be_removed(): + """Retain cleanup of a genuine no-op adapter in standalone mode.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["before"]), + helper.make_node("Unsqueeze", ["before", "axes"], ["adapted"]), + helper.make_node("Identity", ["adapted"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4])], + [numpy_helper.from_array(np.asarray([], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + expected = _checked_run(model, feeds)[0] + _apply_remove_intermediary_squeeze_unsqueeze(model) + assert not any(node.op_type == "Unsqueeze" for node in model.graph.node) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_internal_normalization_boundary_restores_trailing_singleton(batch): + """Restore native normalization output shapes after preserving its channel dimension.""" + model = _make_model( + [ + helper.make_node( + "InstanceNormalization", ["x", "scale", "bias"], ["y"], name="normalize" + ) + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [batch, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [batch, 3, 4])], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + feeds = {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_outputs(model) + _apply_fix_instancenorm_channels(model) + np.testing.assert_allclose(_checked_run(model, feeds)[0], expected, rtol=1e-6, atol=1e-6) + + +def test_unnamed_normalizations_keep_distinct_targets(): + """Give unnamed normalization inputs with different shapes independent reshape targets.""" + model = _make_model( + [ + helper.make_node("InstanceNormalization", ["a", "scale", "bias"], ["a_out"]), + helper.make_node("InstanceNormalization", ["b", "scale", "bias"], ["b_out"]), + ], + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 3, 5]), + ], + [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("b_out", TensorProto.FLOAT, [3, 3, 5]), + ], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + feeds = { + "a": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "b": np.arange(45, dtype=np.float32).reshape(3, 3, 5), + } + expected = _checked_run(model, feeds) + _apply_non4d_model_outputs(model) + _apply_fix_instancenorm_channels(model) + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize( + "transform", + [ + _apply_non4d_model_inputs, + _apply_remove_intermediary_squeeze_unsqueeze, + _apply_fix_instancenorm_channels, + ], +) +def test_untouched_constant_payload_stays_in_place(transform): + """Preserve the original protobuf payload object when no node replacement touches it.""" + model = _make_model( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("Identity", ["x"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + ) + constant = model.graph.node[0] + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + expected = _checked_run(model, feeds) + transform(model) + assert model.graph.node[0] is constant + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_input_adapter_runtime_axes_refuses_before_mutation(): + """Refuse baking an overridable axes default before applying earlier promotions.""" + model = _make_model( + [ + helper.make_node("Identity", ["a"], ["a_out"]), + helper.make_node("Unsqueeze", ["x", "axes"], ["y"]), + ], + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("axes", TensorProto.INT64, [1]), + ], + [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["a", "b", "c"]), + ], + [numpy_helper.from_array(np.asarray([-1], np.int64), "axes")], + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="fixed"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("transform", [_apply_non4d_model_inputs, _apply_non4d_model_outputs]) +def test_legacy_boundary_adapter_schema_refuses_before_mutation(transform): + """Refuse input-form axes adapters on a legacy operator schema.""" + model = _make_model( + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model.opset_import[0].version = 11 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="13"): + transform(model) + assert model.SerializeToString() == before + + +def test_legacy_ir_boundary_adapters_register_initializer_inputs(): + """Keep both new boundary-axis initializers legal in an IR3 graph.""" + model = _make_model( + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model.ir_version = 3 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_inputs(model) + _apply_non4d_model_outputs(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +def _capture_branch(name, source, shape): + """Make a branch that reads an outer-scope value.""" + return helper.make_graph( + [helper.make_node("Identity", [source], [name])], + name, + [], + [helper.make_tensor_value_info(name, TensorProto.FLOAT, shape)], + ) + + +def test_input_promotion_captured_adapter_refuses_before_mutation(): + """Retain a captured adapter output instead of leaving its lexical read unbound.""" + model = _make_model( + [ + helper.make_node("Unsqueeze", ["x", "axes"], ["captured"]), + helper.make_node( + "If", + ["condition"], + ["y"], + then_branch=_capture_branch("then", "captured", [1, 1, 2, 3]), + else_branch=_capture_branch("else", "captured", [1, 1, 2, 3]), + ), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3])], + [numpy_helper.from_array(np.asarray([0, 1], np.int64), "axes")], + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="captured"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("rank_converted", [False, True]) +def test_intermediary_cleanup_preserves_captured_aliases(rank_converted): + """Keep the producer name captured by an If when a no-op adapter could be removed.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["captured"]), + helper.make_node("Unsqueeze", ["captured", "axes"], ["adapted"]), + helper.make_node( + "If", + ["condition"], + ["chosen"], + then_branch=_capture_branch("then", "captured", [1, 2, 3, 4]), + else_branch=_capture_branch("else", "captured", [1, 2, 3, 4]), + ), + helper.make_node("Add", ["chosen", "adapted"], ["y"]), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4])], + [numpy_helper.from_array(np.asarray([], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4), "condition": np.bool_(True)} + expected = _checked_run(model, feeds)[0] + if rank_converted: + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _apply_remove_intermediary_squeeze_unsqueeze(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +def test_output_boundary_orders_a_captured_public_value_before_its_reader(): + """Make a restored public value available before an existing lexical consumer.""" + model = _make_model( + [ + helper.make_node("Neg", ["x"], ["y"]), + helper.make_node( + "If", + ["condition"], + ["unused"], + then_branch=_capture_branch("then", "y", [2, 3]), + else_branch=_capture_branch("else", "y", [2, 3]), + ), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3), "condition": np.bool_(True)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_inputs(model) + _apply_non4d_model_outputs(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +@pytest.mark.parametrize( + ("axes", "output_shape"), + [([2, 0], [1, 2, 1, 3]), ([-2, -1], [2, 3, 1, 1]), ([2, -4], [1, 2, 1, 3])], +) +def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_shape): + """Preserve unordered positive, negative and mixed Unsqueeze axis semantics.""" + model = _make_model( + [helper.make_node("Unsqueeze", ["x", "axes"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + transformed = _apply_non4d_model_inputs(model) + np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) From 514c77e1903326f342c67fe3fcea1b884eb92dff Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:23:10 +0000 Subject: [PATCH 03/32] Preserve normalization channels and declare deferred output rank restoration Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 12 ++- .../onnx/graph_surgery/make_dla_compatible.py | 15 +++- .../test_dla_model_boundaries.py | 73 +++++++++++++++++-- 3 files changed, 92 insertions(+), 8 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 4dfbfc22681..c99e421ccda 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -53,6 +53,7 @@ def _apply_non4d_model_inputs(model): for node, _ in cache.get_consumers(value.name) if (node.input[0] == value.name or node.op_type == "Concat") and not (node.op_type == "Squeeze" and "Gather" in node.name) + and node.op_type != "InstanceNormalization" ] if not consumers: continue @@ -145,8 +146,8 @@ def _apply_non4d_model_inputs(model): return model -def _apply_non4d_model_outputs(model): - """Restore declared output ranks after the ordered rank-promotion sequence.""" +def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False): + """Restore four-dimensional outputs, optionally deferring checks to later rank stages.""" graph = model.graph cache = GraphCache(graph) captured = _subgraph_captures(graph) @@ -179,7 +180,14 @@ def _apply_non4d_model_outputs(model): if input_shape is None: raise ValueError(f"Output Squeeze input shape must be known for {value.name!r}") axes = [index for index, dim in enumerate(input_shape) if dim == 1] + source_shape = cache.get_shape(existing.input[0] if existing is not None else value.name) + if not defer_rank_conversion and (source_shape is None or len(source_shape) != 4): + raise ValueError( + f"Output restoration requires a known rank-4 producer for {value.name!r}" + ) delta = 4 - len(axes) - len(shape) + if delta < 0: + raise ValueError(f"Output Squeeze input rank exceeds 4 for {value.name!r}") axes = list(range(delta)) + [axis + delta if axis >= 0 else axis for axis in axes] plans.append((value, shape, existing, axes)) if plans: diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index c047412898c..5ac8cf367c2 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -94,7 +94,7 @@ def _apply_constants_to_initializers(model, **options): partial(_apply_matmul_to_transpose_conv_transpose, defer_rank_conversion=True), ), ("Non4dModelInputs", _apply_non4d_model_inputs), - ("Non4dModelOutputs", _apply_non4d_model_outputs), + ("Non4dModelOutputs", partial(_apply_non4d_model_outputs, defer_rank_conversion=True)), ( "RemoveIntermediarySqueezeAndUnsqueeze", partial(_apply_remove_intermediary_squeeze_unsqueeze, rank_converted=True), @@ -120,6 +120,19 @@ def _transform_make_dla_compatible( name == "RemoveQDQ" for name, _ in DLA_TRANSFORMS ): raise NotImplementedError("QDQ removal is not available in this pipeline") + if not any(name == "StandaloneReduceOperations" for name, _ in DLA_TRANSFORMS): + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + for node in model.graph.node: + if ( + node.op_type in {"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"} + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + ): + raise NotImplementedError( + "Standalone reductions are not available in this pipeline" + ) quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 0db7c85c032..972dc12d5fa 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -21,6 +21,7 @@ from onnx import TensorProto, helper, numpy_helper from onnx.reference import ReferenceEvaluator +import modelopt.onnx.graph_surgery.make_dla_compatible as dla_pipeline from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( _apply_fix_instancenorm_channels, _apply_non4d_model_inputs, @@ -201,7 +202,7 @@ def test_internal_normalization_boundary_restores_trailing_singleton(batch): ) feeds = {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4)} expected = _checked_run(model, feeds)[0] - _apply_non4d_model_outputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) _apply_fix_instancenorm_channels(model) np.testing.assert_allclose(_checked_run(model, feeds)[0], expected, rtol=1e-6, atol=1e-6) @@ -231,7 +232,7 @@ def test_unnamed_normalizations_keep_distinct_targets(): "b": np.arange(45, dtype=np.float32).reshape(3, 3, 5), } expected = _checked_run(model, feeds) - _apply_non4d_model_outputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) _apply_fix_instancenorm_channels(model) for actual, reference in zip(_checked_run(model, feeds), expected): np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=1e-6) @@ -309,7 +310,8 @@ def test_legacy_boundary_adapter_schema_refuses_before_mutation(transform): onnx.checker.check_model(model, full_check=True) before = model.SerializeToString() with pytest.raises(ValueError, match="13"): - transform(model) + options = {"defer_rank_conversion": True} if transform is _apply_non4d_model_outputs else {} + transform(model, **options) assert model.SerializeToString() == before @@ -324,7 +326,7 @@ def test_legacy_ir_boundary_adapters_register_initializer_inputs(): feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} expected = _checked_run(model, feeds)[0] _apply_non4d_model_inputs(model) - _apply_non4d_model_outputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) @@ -420,7 +422,7 @@ def test_output_boundary_orders_a_captured_public_value_before_its_reader(): feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3), "condition": np.bool_(True)} expected = _checked_run(model, feeds)[0] _apply_non4d_model_inputs(model) - _apply_non4d_model_outputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) @@ -440,3 +442,64 @@ def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_ expected = _checked_run(model, feeds)[0] transformed = _apply_non4d_model_inputs(model) np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2, 3]) +def test_registered_normalization_preserves_native_channel_axis(batch): + """Keep direct rank-three normalization inputs on the channel-preserving route.""" + shape = [batch, 3, 4] + model = _make_model( + [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.asarray([1, 2, 3], np.float32), "scale"), + numpy_helper.from_array(np.asarray([-1, 0, 1], np.float32), "bias"), + ], + ) + feeds = {"x": np.random.default_rng(19).standard_normal(shape).astype(np.float32)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(transformed, feeds)[0], expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("entrypoint", ["standalone", "pipeline"]) +def test_native_reduction_output_refuses_before_adding_rank_adapters(entrypoint, monkeypatch): + """Refuse a native producer when no rank-conversion stage can prepare its output.""" + model = _make_model( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[2, 3], keepdims=0)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])], + ) + before = model.SerializeToString() + if entrypoint == "standalone": + with pytest.raises(ValueError, match="rank"): + _apply_non4d_model_outputs(model) + else: + monkeypatch.setattr( + dla_pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in dla_pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + with pytest.raises(NotImplementedError, match="reductions"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("axes", [[0, 2], [0, -3]]) +def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes): + """Avoid remapping axes through an unspecified five-dimensional rank collapse.""" + model = _make_model( + [helper.make_node("Squeeze", ["x", "axes"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 1, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank"): + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before From 2fc9758426ac919ef362a2f7098c2ca470d92c56 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:27:12 +0000 Subject: [PATCH 04/32] Resolve the pipeline module independently of public re-exports Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_model_boundaries.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 972dc12d5fa..2bd1ed5717d 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -15,13 +15,14 @@ """Boundary promotion refuses unsupported exposed adapters before changing the graph.""" +from importlib import import_module + import numpy as np import onnx import pytest from onnx import TensorProto, helper, numpy_helper from onnx.reference import ReferenceEvaluator -import modelopt.onnx.graph_surgery.make_dla_compatible as dla_pipeline from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( _apply_fix_instancenorm_channels, _apply_non4d_model_inputs, @@ -476,6 +477,7 @@ def test_native_reduction_output_refuses_before_adding_rank_adapters(entrypoint, with pytest.raises(ValueError, match="rank"): _apply_non4d_model_outputs(model) else: + dla_pipeline = import_module(_transform_make_dla_compatible.__module__) monkeypatch.setattr( dla_pipeline, "DLA_TRANSFORMS", From 6dc1c5fa388f5d511004d16c88aee920629917bd Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:45:38 +0000 Subject: [PATCH 05/32] Restore normalization output ranks through unary activation paths Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 44 +++++++++++++++---- .../test_dla_model_boundaries.py | 11 ++++- 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index c99e421ccda..9bf92d85ccc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -20,7 +20,7 @@ from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes, pad4d from .dla_constants_to_initializers import _add_legacy_initializer_inputs -from .dla_graph_cleanup import _subgraph_captures +from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_reorder_nodes import priority_topological_order __all__ = [] @@ -299,8 +299,40 @@ def _apply_fix_instancenorm_channels(model): graph = model.graph cache = GraphCache(graph) replacements = {} + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + + def output_squeezes(name): + """Follow layout-preserving unary paths to public output restoration adapters.""" + pending, visited = [name], set() + while pending: + tensor = pending.pop() + if tensor in visited: + continue + visited.add(tensor) + for consumer, input_index in cache.get_consumers(tensor): + if ( + input_index != 0 + or consumer.domain not in {"", "ai.onnx"} + or consumer.overload + or (consumer.domain, consumer.op_type, consumer.overload) in local_functions + ): + continue + if consumer.op_type == "Squeeze" and consumer.output[0] in cache.graph_output_names: + yield consumer + elif ( + consumer.op_type in _RANK_AGNOSTIC_OPS + and len(consumer.input) == 1 + and len(consumer.output) == 1 + ): + pending.append(consumer.output[0]) + for index, node in enumerate(graph.node): - if node.op_type == "InstanceNormalization": + if ( + node.op_type == "InstanceNormalization" + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + ): shape = cache.get_shape(node.input[0]) if shape is not None and len(shape) == 3: target = add_initializer( @@ -319,12 +351,8 @@ def _apply_fix_instancenorm_channels(model): ) node.input[0] = output replacements[index] = ([node], [reshape, node]) - for consumer, _ in cache.get_consumers(node.output[0]): - if ( - consumer.op_type != "Squeeze" - or consumer.output[0] not in cache.graph_output_names - or len(consumer.input) < 2 - ): + for consumer in output_squeezes(node.output[0]): + if len(consumer.input) < 2: continue axes = cache.get_init(consumer.input[1]) if axes is not None and numpy_helper.to_array(axes).reshape(-1).tolist() == [0]: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 2bd1ed5717d..699f74977d6 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -446,11 +446,18 @@ def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_ @pytest.mark.parametrize("batch", [1, 2, 3]) -def test_registered_normalization_preserves_native_channel_axis(batch): +@pytest.mark.parametrize("activations", [(), ("Relu",), ("Relu", "Tanh")]) +def test_registered_normalization_preserves_native_channel_axis(batch, activations): """Keep direct rank-three normalization inputs on the channel-preserving route.""" shape = [batch, 3, 4] + output = "normalized" if activations else "y" + nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], [output])] + for index, operation in enumerate(activations): + following = "y" if index == len(activations) - 1 else f"activation_{index}" + nodes.append(helper.make_node(operation, [output], [following])) + output = following model = _make_model( - [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["y"])], + nodes, [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], [ From 2b96b21ffff660e90680cd251d01d20dd244cd94 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 07:01:49 +0000 Subject: [PATCH 06/32] Validate concat layouts and retain normalization squeeze axes Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 28 +++++++++++- .../test_dla_model_boundaries.py | 43 +++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 9bf92d85ccc..833a8833684 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -78,6 +78,24 @@ def _apply_non4d_model_inputs(model): ) unsqueezed_shapes[id(consumer)] = output_shape promotions.append((value, shape, consumers)) + promoted_shapes = {value.name: shape for value, shape, _ in promotions} + concats = { + id(node): node + for _, _, consumers in promotions + for node in consumers + if node.op_type == "Concat" + } + for node in concats.values(): + if ( + any(name not in promoted_shapes for name in node.input) + or len({len(promoted_shapes[name]) for name in node.input}) != 1 + ): + raise ValueError(f"Cannot promote mixed-layout Concat {node.name!r}") + for node in concats.values(): + for attribute in node.attribute: + if attribute.name == "axis" and attribute.i >= 0: + attribute.i += 4 - len(promoted_shapes[node.input[0]]) + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") for value, shape, consumers in promotions: axes = list(range(4 - len(shape))) axes_name = add_initializer( @@ -355,11 +373,17 @@ def output_squeezes(name): if len(consumer.input) < 2: continue axes = cache.get_init(consumer.input[1]) - if axes is not None and numpy_helper.to_array(axes).reshape(-1).tolist() == [0]: + values = ( + None if axes is None else numpy_helper.to_array(axes).reshape(-1).tolist() + ) + if axes is not None and values and values[0] == 0: + native_axes = [ + value - 1 if value >= 1 else value + 3 for value in values[1:] + ] consumer.input[1] = add_initializer( graph, cache.unique_name(axes.name + "_trailing"), - np.asarray([3], np.int64), + np.asarray([*native_axes, 3], np.int64), dtype=TensorProto.INT64, with_value_info=False, ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 699f74977d6..37dca51f483 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -140,6 +140,45 @@ def test_ordered_pipeline_preserves_matrix_squeeze_axes(axes): np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_registered_concat_preserves_promoted_input_axes(axis): + """Shift Concat axes once after consistently promoting every native operand.""" + shape = [2, 3] + result_shape = shape.copy() + result_shape[axis] *= 2 + model = _make_model( + [helper.make_node("Concat", ["a", "b"], ["y"], axis=axis)], + [helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) for name in ("a", "b")], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, result_shape)], + ) + feeds = {"a": np.arange(6, dtype=np.float32).reshape(shape), "b": np.ones(shape, np.float32)} + expected = _checked_run(model, feeds)[0] + changed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(changed, feeds)[0], expected) + + +@pytest.mark.parametrize("axis", [0, -3]) +def test_normalization_public_squeeze_keeps_original_axes(axis): + """Retarget original batch squeezing together with the generated spatial singleton.""" + model = _make_model( + [ + helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"]), + helper.make_node("Squeeze", ["normalized", "axes"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 4])], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + numpy_helper.from_array(np.array([axis], np.int64), "axes"), + ], + ) + feeds = {"x": np.arange(12, dtype=np.float32).reshape(1, 3, 4)} + expected = _checked_run(model, feeds)[0] + changed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(changed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + def test_standalone_intermediary_rank_change_preserves_concat_inputs(): """Retain an internal rank-changing squeeze whose downstream consumer expects rank three.""" model = _make_model( @@ -163,6 +202,10 @@ def test_standalone_intermediary_rank_change_preserves_concat_inputs(): expected = _checked_run(model, feeds)[0] _apply_remove_intermediary_squeeze_unsqueeze(model) np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Concat"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before def test_standalone_empty_axes_adapter_can_be_removed(): From 2d72135a96dc72f67d51ca4126a243158c8c75b7 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:15:33 +0000 Subject: [PATCH 07/32] Guard DLA boundary rewrites by operator identity and native rank Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 105 ++++++++++++++---- .../test_dla_model_boundaries.py | 91 +++++++++++++++ 2 files changed, 177 insertions(+), 19 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 833a8833684..06f269975f4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -25,6 +25,59 @@ __all__ = [] +_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} + + +def _is_standard_node(node, local_functions, op_types): + """Match builtin operations without interpreting opaque overloads or local definitions.""" + return ( + node.op_type in op_types + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + ) + + +def _output_rank_source(cache, node, local_functions): + """Follow native rank-preserving unary paths before planning public restoration.""" + seen = set() + while ( + node is not None + and _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) + and len(node.input) == 1 + ): + if id(node) in seen: + raise ValueError("Cyclic unary path while planning output restoration") + seen.add(id(node)) + node = cache.get_producer(node.input[0]) + return node + + +def _has_native_symbolic_data(cache, name, local_functions): + """Prove a low-rank symbolic data path was left native by boundary promotion.""" + seen = set() + while name not in seen: + seen.add(name) + info = cache.vi_map.get(name) + rank = ( + len(info.type.tensor_type.shape.dim) + if info is not None and info.type.tensor_type.HasField("shape") + else None + ) + if cache.get_shape(name) is not None or (rank is not None and rank >= 4): + return False + if name in cache.graph_input_names: + return rank is not None and rank < 4 + node = cache.get_producer(name) + if ( + node is None + or not _is_standard_node(node, local_functions, _RANK_FLOW_OPS) + or not node.input + ): + return False + name = node.input[0] + raise ValueError("Cyclic data path while planning native symbolic controls") + def _require_adapter_opset(model): """Refuse input-form axes adapters before changing a legacy graph.""" @@ -44,6 +97,8 @@ def _apply_non4d_model_inputs(model): prepend = [] promotions = [] unsqueezed_shapes = {} + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + adapters = {"Squeeze", "Unsqueeze", "Concat"} for value in graph.input: shape = cache.get_shape(value.name) if shape is None or len(shape) not in (2, 3): @@ -51,15 +106,20 @@ def _apply_non4d_model_inputs(model): consumers = [ node for node, _ in cache.get_consumers(value.name) - if (node.input[0] == value.name or node.op_type == "Concat") - and not (node.op_type == "Squeeze" and "Gather" in node.name) + if (node.op_type not in adapters or _is_standard_node(node, local_functions, adapters)) + and ( + node.input[0] == value.name or _is_standard_node(node, local_functions, {"Concat"}) + ) + and not ( + _is_standard_node(node, local_functions, {"Squeeze"}) and "Gather" in node.name + ) and node.op_type != "InstanceNormalization" ] if not consumers: continue _require_adapter_opset(model) for consumer in consumers: - if consumer.op_type != "Unsqueeze": + if not _is_standard_node(consumer, local_functions, {"Unsqueeze"}): continue tensor = cache.get_init(consumer.input[1]) if tensor is None or tensor.name in cache.graph_input_names: @@ -83,7 +143,7 @@ def _apply_non4d_model_inputs(model): id(node): node for _, _, consumers in promotions for node in consumers - if node.op_type == "Concat" + if _is_standard_node(node, local_functions, {"Concat"}) } for node in concats.values(): if ( @@ -120,7 +180,7 @@ def _apply_non4d_model_inputs(model): helper.make_tensor_value_info(output, value.type.tensor_type.elem_type, promoted_shape) ) for consumer in consumers: - if consumer.op_type == "Unsqueeze": + if _is_standard_node(consumer, local_functions, {"Unsqueeze"}): output_shape = unsqueezed_shapes[id(consumer)] if promoted_shape[-len(output_shape) :] == output_shape: old_output = consumer.output[0] @@ -172,21 +232,35 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False): boundary_names = cache.graph_input_names | cache.init_names append = [] plans = [] + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} for value in graph.output: shape = cache.get_shape(value.name) if shape is None or len(shape) >= 4 or value.name in boundary_names: continue producer = cache.get_producer(value.name) + source = _output_rank_source(cache, producer, local_functions) + if source is not None: + if source.op_type in {"Squeeze", "Unsqueeze", "Concat"} and not _is_standard_node( + source, local_functions, {"Squeeze", "Unsqueeze", "Concat"} + ): + continue + if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): + if _has_native_symbolic_data(cache, source.input[0], local_functions): + continue if ( producer is not None - and producer.op_type == "Identity" + and _is_standard_node(producer, local_functions, {"Identity"}) and producer.input[0] in boundary_names ): continue if len(shape) == 0: raise ValueError(f"Unsupported scalar model output {value.name!r}") axes = [] - existing = producer if producer is not None and producer.op_type == "Squeeze" else None + existing = ( + producer + if producer is not None and _is_standard_node(producer, local_functions, {"Squeeze"}) + else None + ) if existing is not None: if len(existing.input) > 1 and existing.input[1]: initializer = cache.get_init(existing.input[1]) @@ -257,8 +331,9 @@ def _apply_remove_intermediary_squeeze_unsqueeze(model, *, rank_converted=False) cache = GraphCache(graph) captured = _subgraph_captures(graph) removed = set() + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} for node in list(graph.node): - if node.op_type not in {"Squeeze", "Unsqueeze"}: + if not _is_standard_node(node, local_functions, {"Squeeze", "Unsqueeze"}): continue if ( node.input[0] in cache.graph_input_names | cache.init_names | cache.graph_output_names @@ -328,11 +403,8 @@ def output_squeezes(name): continue visited.add(tensor) for consumer, input_index in cache.get_consumers(tensor): - if ( - input_index != 0 - or consumer.domain not in {"", "ai.onnx"} - or consumer.overload - or (consumer.domain, consumer.op_type, consumer.overload) in local_functions + if input_index != 0 or not _is_standard_node( + consumer, local_functions, {consumer.op_type} ): continue if consumer.op_type == "Squeeze" and consumer.output[0] in cache.graph_output_names: @@ -345,12 +417,7 @@ def output_squeezes(name): pending.append(consumer.output[0]) for index, node in enumerate(graph.node): - if ( - node.op_type == "InstanceNormalization" - and node.domain in {"", "ai.onnx"} - and not node.overload - and (node.domain, node.op_type, node.overload) not in local_functions - ): + if _is_standard_node(node, local_functions, {"InstanceNormalization"}): shape = cache.get_shape(node.input[0]) if shape is not None and len(shape) == 3: target = add_initializer( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 37dca51f483..03594c6e063 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -32,6 +32,97 @@ from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +@pytest.mark.parametrize("identity", ["custom", "local", "overload"]) +@pytest.mark.parametrize( + ("operation", "entrypoint"), + [ + ("Squeeze", "pipeline"), + ("Unsqueeze", "pipeline"), + ("Concat", "pipeline"), + ("Squeeze", "cleanup"), + ("Unsqueeze", "cleanup"), + ], +) +def test_operator_lookalikes_keep_function_semantics(identity, operation, entrypoint): + """Execute opaque lookalikes without rewriting their builtin-named calls or controls.""" + domain = "opaque.boundary" if identity == "custom" else "" + overload = "custom" if identity == "overload" else "" + inputs = ["x"] if operation == "Concat" else ["x", "axes"] + function = helper.make_function( + domain, + operation, + inputs, + ["y"], + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_opsetid("", 21)], + attributes=["axis"] if operation == "Concat" else [], + ) + function.overload = overload + opaque = helper.make_node( + operation, + ["middle" if entrypoint == "cleanup" else "x", *inputs[1:]], + ["hidden" if entrypoint == "cleanup" else "y"], + name="opaque", + domain=domain, + **({"axis": 1} if operation == "Concat" else {}), + ) + opaque.overload = overload + nodes = ( + [ + helper.make_node("Identity", ["x"], ["middle"]), + opaque, + helper.make_node("Identity", ["hidden"], ["y"]), + ] + if entrypoint == "cleanup" + else [opaque] + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + [] if operation == "Concat" else [numpy_helper.from_array(np.array([], np.int64), "axes")], + ) + model.functions.append(function) + if domain: + model.opset_import.append(helper.make_opsetid(domain, 1)) + if entrypoint == "cleanup": + model.graph.value_info.extend( + helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) + for name in ("middle", "hidden") + ) + before = next(n for n in model.graph.node if n.name == "opaque").SerializeToString() + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = ( + _transform_make_dla_compatible(model, verbose=False) + if entrypoint == "pipeline" + else _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + ) + assert next(n for n in converted.graph.node if n.name == "opaque").SerializeToString() == before + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_native_controls_do_not_gain_output_adapters(operation, tail): + """Keep static native outputs when symbolic data was not promoted, including unary tails.""" + target = {"Reshape": [3, 2], "Expand": [2, 3], "Tile": [0, 1]}[operation] + output_shape = {"Reshape": [3, 2], "Expand": [2, 3], "Tile": [0, 3]}[operation] + nodes = [helper.make_node(operation, ["x", "control"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.array(target, np.int64), "control")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + def test_input_promotion_refusal_leaves_graph_unchanged(): """Refuse removing a public input adapter before altering any earlier consumer.""" model = helper.make_model( From ba1a3b8ef6cbd41485c84b812f249bf337aea5ae Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:24:33 +0000 Subject: [PATCH 08/32] Reuse shared unary paths during DLA output planning Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_model_boundaries.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 06f269975f4..47f2d8e3843 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -38,7 +38,7 @@ def _is_standard_node(node, local_functions, op_types): ) -def _output_rank_source(cache, node, local_functions): +def _output_rank_source(cache, node, local_functions, sources): """Follow native rank-preserving unary paths before planning public restoration.""" seen = set() while ( @@ -46,10 +46,14 @@ def _output_rank_source(cache, node, local_functions): and _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) and len(node.input) == 1 ): + if id(node) in sources: + node = sources[id(node)] + break if id(node) in seen: raise ValueError("Cyclic unary path while planning output restoration") seen.add(id(node)) node = cache.get_producer(node.input[0]) + sources.update(dict.fromkeys(seen, node)) return node @@ -233,12 +237,13 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False): append = [] plans = [] local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + rank_sources = {} for value in graph.output: shape = cache.get_shape(value.name) if shape is None or len(shape) >= 4 or value.name in boundary_names: continue producer = cache.get_producer(value.name) - source = _output_rank_source(cache, producer, local_functions) + source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is not None: if source.op_type in {"Squeeze", "Unsqueeze", "Concat"} and not _is_standard_node( source, local_functions, {"Squeeze", "Unsqueeze", "Concat"} From 4b7768c92b8f9d4eb15f9cbfa5370033bc9832b0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:49:31 +0000 Subject: [PATCH 09/32] Require registered rank conversion for deferred native reducers Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 21 +++++- .../onnx/graph_surgery/make_dla_compatible.py | 13 +++- .../test_dla_model_boundaries.py | 66 +++++++++++++++++++ 3 files changed, 98 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 47f2d8e3843..ff74177db7f 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -26,6 +26,23 @@ __all__ = [] _RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} +_REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) + + +def _require_output_rank_conversion(node, local_functions, prepared_output_ops): + """Require an explicit later stage before deferring a native rank-reducing operation.""" + if ( + _is_standard_node(node, local_functions, {node.op_type}) + and ( + node.op_type == "Flatten" + or ( + (node.op_type in {"ArgMax", "ArgMin"} or node.op_type.startswith("Reduce")) + and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + ) + ) + and node.op_type not in prepared_output_ops + ): + raise ValueError(f"Output rank conversion is unavailable for {node.op_type}") def _is_standard_node(node, local_functions, op_types): @@ -228,7 +245,7 @@ def _apply_non4d_model_inputs(model): return model -def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False): +def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_output_ops=()): """Restore four-dimensional outputs, optionally deferring checks to later rank stages.""" graph = model.graph cache = GraphCache(graph) @@ -245,6 +262,8 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False): producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is not None: + if defer_rank_conversion: + _require_output_rank_conversion(source, local_functions, prepared_output_ops) if source.op_type in {"Squeeze", "Unsqueeze", "Concat"} and not _is_standard_node( source, local_functions, {"Squeeze", "Unsqueeze", "Concat"} ): diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 5ac8cf367c2..f95a21a920e 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -38,10 +38,12 @@ _apply_matmul_to_transpose_conv_transpose, ) from .dla_transforms.dla_model_boundaries import ( + _REDUCTION_OPS, _apply_fix_instancenorm_channels, _apply_non4d_model_inputs, _apply_non4d_model_outputs, _apply_remove_intermediary_squeeze_unsqueeze, + _require_output_rank_conversion, ) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model @@ -124,7 +126,7 @@ def _transform_make_dla_compatible( local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} for node in model.graph.node: if ( - node.op_type in {"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"} + node.op_type in _REDUCTION_OPS and node.domain in {"", "ai.onnx"} and not node.overload and (node.domain, node.op_type, node.overload) not in local_functions @@ -133,6 +135,13 @@ def _transform_make_dla_compatible( raise NotImplementedError( "Standalone reductions are not available in this pipeline" ) + stages = frozenset(stage for stage, _ in DLA_TRANSFORMS) + prepared_output_ops = stages | ( + _REDUCTION_OPS if "StandaloneReduceOperations" in stages else frozenset() + ) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + for node in model.graph.node: + _require_output_rank_conversion(node, local_functions, prepared_output_ops) quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None @@ -153,6 +162,8 @@ def _transform_make_dla_compatible( } elif name == "ConstantsToInitializers": options = {"verbose": verbose} + elif name == "Non4dModelOutputs": + options = {"prepared_output_ops": prepared_output_ops} result = transform(model, **options) if result is not None: model = result diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 03594c6e063..8cc0bbce89b 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -646,3 +646,69 @@ def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes): with pytest.raises(ValueError, match="rank"): _apply_non4d_model_outputs(model, defer_rank_conversion=True) assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Flatten", "ArgMax", "ArgMin", "ReduceProd", "ReduceL2"]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("entrypoint", ["deferred", "pipeline"]) +def test_native_rank_reducing_outputs_require_registered_conversion( + operation, tail, entrypoint, monkeypatch +): + """Refuse native rank reducers before adding planned adapters to earlier valid outputs.""" + attributes = ( + {"axis": 1} + if operation == "Flatten" + else ( + {"axis": 1, "keepdims": 0} + if operation.startswith("Arg") + else {"axes": [2, 3], "keepdims": 0} + ) + ) + shape = ( + [1, 48] + if operation == "Flatten" + else ([1, 4, 4] if operation.startswith("Arg") else [1, 3]) + ) + nodes = [ + helper.make_node("Relu", ["x"], ["first"]), + helper.make_node(operation, ["z"], ["middle" if tail else "y"], **attributes), + ] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = _make_model( + nodes, + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 3, 4, 4]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info( + "y", TensorProto.INT64 if operation.startswith("Arg") else TensorProto.FLOAT, shape + ), + ], + ) + model.opset_import[0].version = 17 + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(2, 3), + "z": np.linspace(0.1, 0.9, 48, dtype=np.float32).reshape(1, 3, 4, 4), + } + _checked_run(model, feeds) + before = model.SerializeToString() + if entrypoint == "pipeline": + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != operation + ), + ) + with pytest.raises(ValueError, match=r"rank.*conversion"): + _transform_make_dla_compatible(model, verbose=False) + else: + with pytest.raises(ValueError, match=r"rank.*conversion"): + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before From b12d997895ac4062a7004fbd17911203389e4b6c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 09:48:48 +0000 Subject: [PATCH 10/32] Preserve boundary operator identity, axes and normalization branch layouts Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 161 ++++++++++----- .../dla_transforms/dla_reshape_patterns.py | 17 +- .../onnx/graph_surgery/make_dla_compatible.py | 16 +- .../test_dla_model_boundaries.py | 191 +++++++++++++++++- .../test_dla_reshape_patterns.py | 27 +++ 5 files changed, 355 insertions(+), 57 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index ff74177db7f..f44ed65ad12 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -16,9 +16,9 @@ """Promote model boundaries at their explicit positions in the DLA pipeline.""" import numpy as np -from onnx import TensorProto, helper, numpy_helper +from onnx import TensorProto, defs, helper, numpy_helper -from ..utils.graph_utils import GraphCache, add_initializer, batch_replace_nodes, pad4d +from ..utils.graph_utils import GraphCache, _attribute, add_initializer, batch_replace_nodes, pad4d from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_reorder_nodes import priority_topological_order @@ -45,16 +45,28 @@ def _require_output_rank_conversion(node, local_functions, prepared_output_ops): raise ValueError(f"Output rank conversion is unavailable for {node.op_type}") -def _is_standard_node(node, local_functions, op_types): +def _is_standard_node(node, local_functions, op_types, *, contrib_ops=()): """Match builtin operations without interpreting opaque overloads or local definitions.""" return ( node.op_type in op_types - and node.domain in {"", "ai.onnx"} + and ( + node.domain in {"", "ai.onnx"} + or (node.domain == "com.microsoft" and node.op_type in contrib_ops) + ) and not node.overload and (node.domain, node.op_type, node.overload) not in local_functions ) +def _boundary_contrib_ops(model): + """Keep the explicit supported runtime schemas separate from opaque custom functions.""" + return ( + {"GroupQueryAttention", "QuantizeLinear", "DequantizeLinear", "Gelu"} + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import) + else set() + ) + + def _output_rank_source(cache, node, local_functions, sources): """Follow native rank-preserving unary paths before planning public restoration.""" seen = set() @@ -118,8 +130,12 @@ def _apply_non4d_model_inputs(model): prepend = [] promotions = [] unsqueezed_shapes = {} + normalization_axes = {} local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} - adapters = {"Squeeze", "Unsqueeze", "Concat"} + contrib_ops = _boundary_contrib_ops(model) + opset = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) for value in graph.input: shape = cache.get_shape(value.name) if shape is None or len(shape) not in (2, 3): @@ -127,7 +143,7 @@ def _apply_non4d_model_inputs(model): consumers = [ node for node, _ in cache.get_consumers(value.name) - if (node.op_type not in adapters or _is_standard_node(node, local_functions, adapters)) + if _is_standard_node(node, local_functions, {node.op_type}, contrib_ops=contrib_ops) and ( node.input[0] == value.name or _is_standard_node(node, local_functions, {"Concat"}) ) @@ -140,6 +156,17 @@ def _apply_non4d_model_inputs(model): continue _require_adapter_opset(model) for consumer in consumers: + if _is_standard_node(consumer, local_functions, {"Softmax", "LogSoftmax"}): + axis = _attribute(consumer, "axis") + axis_value = ( + axis.i + if axis is not None + else defs.get_schema(consumer.op_type, opset).attributes["axis"].default_value.i + ) + if not -len(shape) <= axis_value < len(shape): + raise ValueError(f"Invalid native normalization axis for {consumer.name!r}") + if axis_value >= 0: + normalization_axes[id(consumer)] = (consumer, axis, axis_value + 4 - len(shape)) if not _is_standard_node(consumer, local_functions, {"Unsqueeze"}): continue tensor = cache.get_init(consumer.input[1]) @@ -177,6 +204,12 @@ def _apply_non4d_model_inputs(model): if attribute.name == "axis" and attribute.i >= 0: attribute.i += 4 - len(promoted_shapes[node.input[0]]) node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") + for node, axis, value in normalization_axes.values(): + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") for value, shape, consumers in promotions: axes = list(range(4 - len(shape))) axes_name = add_initializer( @@ -255,6 +288,7 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o plans = [] local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} rank_sources = {} + contrib_ops = _boundary_contrib_ops(model) for value in graph.output: shape = cache.get_shape(value.name) if shape is None or len(shape) >= 4 or value.name in boundary_names: @@ -264,8 +298,8 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o if source is not None: if defer_rank_conversion: _require_output_rank_conversion(source, local_functions, prepared_output_ops) - if source.op_type in {"Squeeze", "Unsqueeze", "Concat"} and not _is_standard_node( - source, local_functions, {"Squeeze", "Unsqueeze", "Concat"} + if not _is_standard_node( + source, local_functions, {source.op_type}, contrib_ops=contrib_ops ): continue if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): @@ -418,8 +452,8 @@ def _apply_fix_instancenorm_channels(model): replacements = {} local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} - def output_squeezes(name): - """Follow layout-preserving unary paths to public output restoration adapters.""" + def output_consumers(name): + """Find public adapters and binary edges after layout-preserving unary paths.""" pending, visited = [name], set() while pending: tensor = pending.pop() @@ -427,57 +461,84 @@ def output_squeezes(name): continue visited.add(tensor) for consumer, input_index in cache.get_consumers(tensor): - if input_index != 0 or not _is_standard_node( - consumer, local_functions, {consumer.op_type} - ): + if not _is_standard_node(consumer, local_functions, {consumer.op_type}): continue - if consumer.op_type == "Squeeze" and consumer.output[0] in cache.graph_output_names: - yield consumer - elif ( - consumer.op_type in _RANK_AGNOSTIC_OPS - and len(consumer.input) == 1 - and len(consumer.output) == 1 + if ( + input_index == 0 + and consumer.op_type == "Squeeze" + and consumer.output[0] in cache.graph_output_names ): - pending.append(consumer.output[0]) + yield consumer, None, None + elif consumer.op_type in _RANK_AGNOSTIC_OPS and len(consumer.output) == 1: + if len(consumer.input) == 1 and input_index == 0: + pending.append(consumer.output[0]) + elif len(consumer.input) > 1: + yield consumer, tensor, input_index + plans = [] for index, node in enumerate(graph.node): if _is_standard_node(node, local_functions, {"InstanceNormalization"}): shape = cache.get_shape(node.input[0]) if shape is not None and len(shape) == 3: - target = add_initializer( + plans.append((index, node, shape, list(output_consumers(node.output[0])))) + bridges = {} + for index, node, shape, consumers in plans: + target = add_initializer( + graph, + cache.unique_name((node.name or node.output[0]) + "_rotateleft_shape"), + np.asarray([*shape, 1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(node.input[0] + "_rotateleft") + reshape = helper.make_node( + "Reshape", + [node.input[0], target], + [output], + name=cache.unique_name((node.name or node.output[0]) + "_rotateleft"), + ) + node.input[0] = output + replacements[index] = ([node], [reshape, node]) + for consumer, tensor, port in consumers: + if tensor is not None: + if tensor not in bridges: + bridges[tensor] = (shape, []) + bridges[tensor][1].append((consumer, port)) + continue + if len(consumer.input) < 2: + continue + axes = cache.get_init(consumer.input[1]) + values = None if axes is None else numpy_helper.to_array(axes).reshape(-1).tolist() + if axes is not None and values and values[0] == 0: + native_axes = [value - 1 if value >= 1 else value + 3 for value in values[1:]] + consumer.input[1] = add_initializer( graph, - cache.unique_name((node.name or node.output[0]) + "_rotateleft_shape"), - np.asarray([*shape, 1], dtype=np.int64), + cache.unique_name(axes.name + "_trailing"), + np.asarray([*native_axes, 3], np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - output = cache.unique_name(node.input[0] + "_rotateleft") - reshape = helper.make_node( - "Reshape", - [node.input[0], target], - [output], - name=cache.unique_name((node.name or node.output[0]) + "_rotateleft"), - ) - node.input[0] = output - replacements[index] = ([node], [reshape, node]) - for consumer in output_squeezes(node.output[0]): - if len(consumer.input) < 2: - continue - axes = cache.get_init(consumer.input[1]) - values = ( - None if axes is None else numpy_helper.to_array(axes).reshape(-1).tolist() - ) - if axes is not None and values and values[0] == 0: - native_axes = [ - value - 1 if value >= 1 else value + 3 for value in values[1:] - ] - consumer.input[1] = add_initializer( - graph, - cache.unique_name(axes.name + "_trailing"), - np.asarray([*native_axes, 3], np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) + positions = {id(node): index for index, node in enumerate(graph.node)} + prefixes = {} + for tensor, (shape, edges) in bridges.items(): + target = add_initializer( + graph, + cache.unique_name(tensor + "_leading_shape"), + np.asarray(pad4d(shape), np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(tensor + "_leading") + view = helper.make_node( + "Reshape", [tensor, target], [output], name=cache.unique_name(output + "_reshape") + ) + first = min(positions[id(consumer)] for consumer, _ in edges) + prefixes.setdefault(first, []).append(view) + for consumer, port in edges: + consumer.input[port] = output + for index, nodes in prefixes.items(): + removed, added = replacements.get(index, ([graph.node[index]], [graph.node[index]])) + replacements[index] = (removed, [*nodes, *added]) batch_replace_nodes(graph, replacements) _add_legacy_initializer_inputs(model) return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py index 60a260b1cf4..face75ab6e6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -50,7 +50,7 @@ def _follower(cache, output, op_type): ) -def _replace(graph, replacements): +def _replace(graph, replacements, *, validate_only=False): """Apply planned replacements without copying untouched payloads and order late controls.""" if not replacements: return @@ -70,6 +70,8 @@ def _replace(graph, replacements): for added in replacements.get(id(node), []) for name in added.input ) + if validate_only: + return batch_replace_nodes( graph, { @@ -95,7 +97,7 @@ def _apply_reshape_qdq_reduce_sum(model): return _reshape_reduce_sum(model, quantized=True) -def _reshape_reduce_sum(model, *, quantized): +def _reshape_reduce_sum(model, *, quantized, validate_only=False): """Plan disjoint reductions while preserving controls and restoring initializers on refusal.""" graph = model.graph cache = GraphCache(graph) @@ -213,7 +215,16 @@ def _reshape_reduce_sum(model, *, quantized): replacements[id(original)] = [] replacements[id(reshape)] = nodes count += 1 - _replace(graph, replacements) + _replace(graph, replacements, validate_only=validate_only) + if validate_only: + outputs = { + name + for node in graph.node + if node.op_type == "ReduceSum" and id(node) in replacements + for name in node.output + } + del graph.initializer[initializer_count:] + return outputs except Exception: del graph.initializer[initializer_count:] raise diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index f95a21a920e..923bebf7fa9 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -52,6 +52,7 @@ _apply_reshape_qdq_reduce_sum, _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, + _reshape_reduce_sum, ) from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, @@ -122,7 +123,15 @@ def _transform_make_dla_compatible( name == "RemoveQDQ" for name, _ in DLA_TRANSFORMS ): raise NotImplementedError("QDQ removal is not available in this pipeline") - if not any(name == "StandaloneReduceOperations" for name, _ in DLA_TRANSFORMS): + stages = frozenset(stage for stage, _ in DLA_TRANSFORMS) + planned_reduction_outputs: set[str] = set() + if "StandaloneReduceOperations" not in stages: + if any(node.op_type == "ReduceSum" for node in model.graph.node): + for stage, quantized in (("ReshapeReduceSum", False), ("ReshapeQDQReduceSum", True)): + if stage in stages: + planned_reduction_outputs.update( + _reshape_reduce_sum(model, quantized=quantized, validate_only=True) + ) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} for node in model.graph.node: if ( @@ -131,17 +140,18 @@ def _transform_make_dla_compatible( and not node.overload and (node.domain, node.op_type, node.overload) not in local_functions and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + and not planned_reduction_outputs.intersection(node.output) ): raise NotImplementedError( "Standalone reductions are not available in this pipeline" ) - stages = frozenset(stage for stage, _ in DLA_TRANSFORMS) prepared_output_ops = stages | ( _REDUCTION_OPS if "StandaloneReduceOperations" in stages else frozenset() ) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} for node in model.graph.node: - _require_output_rank_conversion(node, local_functions, prepared_output_ops) + if not planned_reduction_outputs.intersection(node.output): + _require_output_rank_conversion(node, local_functions, prepared_output_ops) quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 8cc0bbce89b..bafcad80f24 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -20,7 +20,7 @@ import numpy as np import onnx import pytest -from onnx import TensorProto, helper, numpy_helper +from onnx import TensorProto, helper, inliner, numpy_helper from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( @@ -158,6 +158,59 @@ def _make_model(nodes, inputs, outputs, initializers=()): ) +@pytest.mark.parametrize("identity", ["custom", "local", "overload"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_generic_opaque_transpose_keeps_native_shared_input(identity, tail): + """Preserve opaque function layouts while a separate standard branch promotes the same input.""" + domain = "opaque.boundary" if identity == "custom" else "" + operation = "OpaqueTranspose" if identity == "custom" else "Abs" + overload = "transpose" if identity == "overload" else "" + function = helper.make_function( + domain, + operation, + ["data"], + ["result"], + [helper.make_node("Transpose", ["data"], ["result"], perm=[1, 0])], + [helper.make_opsetid("", 21)], + overload=overload, + ) + opaque = helper.make_node( + operation, ["x"], ["middle" if tail else "y"], domain=domain, overload=overload + ) + nodes = [helper.make_node("Relu", ["x"], ["first"]), opaque] + if tail: + nodes.append(helper.make_node("Relu", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 2])], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 2]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2]), + ], + ) + model.functions.append(function) + if domain: + model.opset_import.append(helper.make_opsetid(domain, 1)) + feeds = {"x": np.arange(-1, 3, dtype=np.float32).reshape(2, 2)} + onnx.checker.check_model(model, full_check=True) + expected = _checked_run(inliner.inline_local_functions(model), feeds) + np.testing.assert_array_equal( + expected[1], np.maximum(feeds["x"].T, 0) if tail else feeds["x"].T + ) + original = opaque.SerializeToString() + converted = _transform_make_dla_compatible(model, verbose=False) + call = next( + node + for node in converted.graph.node + if node.op_type == operation and node.domain == domain and node.overload == overload + ) + assert call.SerializeToString() == original + for actual, reference in zip( + _checked_run(inliner.inline_local_functions(converted), feeds), expected + ): + np.testing.assert_array_equal(actual, reference) + + def test_unnamed_input_adapters_keep_distinct_targets(): """Promote two unnamed trailing-axis adapters without reusing their different targets.""" inputs = [ @@ -187,6 +240,142 @@ def test_unnamed_input_adapters_keep_distinct_targets(): np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"]) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("axis", [0, 1, -1, None]) +@pytest.mark.parametrize("opset", [13, 17]) +def test_boundary_softmax_preserves_native_axes(operation, shape, axis, opset): + """Preserve explicit and schema-default normalization axes across input promotion.""" + model = _make_model( + [helper.make_node(operation, ["x"], ["y"], **({} if axis is None else {"axis": axis}))], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ) + model.opset_import[0].version = opset + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 10} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2, 3]) +@pytest.mark.parametrize("operation", ["Add", "Mul"]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_normalization_binary_branch_preserves_broadcast_layout(batch, operation, tail, reverse): + """Keep shared normalized public values and binary residual inputs in compatible layouts.""" + shape = [batch, 3, 4] + normalized = "tail" if tail else "normalized" + nodes = [ + helper.make_node( + "InstanceNormalization", ["x", "scale", "bias"], ["normalized"], epsilon=1e-4 + ) + ] + if tail: + nodes.append(helper.make_node("Tanh", ["normalized"], ["tail"])) + nodes.append( + helper.make_node(operation, [normalized, "x"] if reverse else ["x", normalized], ["y"]) + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + helper.make_tensor_value_info(normalized, TensorProto.FLOAT, shape), + ], + [ + numpy_helper.from_array(np.asarray([0.5, 1, 1.5], np.float32), "scale"), + numpy_helper.from_array(np.asarray([0.1, -0.2, 0.3], np.float32), "bias"), + ], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 8} + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("width", [8, 12]) +@pytest.mark.parametrize("axis", [-1, 4]) +@pytest.mark.parametrize("quantized", [False, True]) +def test_early_registered_reshape_reduction_is_not_rejected(width, axis, quantized, monkeypatch): + """Allow reductions consumed by actual earlier reshape-pattern stages without a native reducer stage.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + shape = [1, 1, 2, width] + controls = [ + numpy_helper.from_array(np.asarray([1, 1, 2, width // 4, 4], np.int64), "target"), + numpy_helper.from_array(np.asarray([axis], np.int64), "axes"), + ] + nodes = [helper.make_node("Reshape", ["x", "target"], ["expanded"])] + current = "expanded" + if quantized: + controls += [ + numpy_helper.from_array(np.float32(0.125), "scale"), + numpy_helper.from_array(np.int8(0), "zero"), + ] + nodes += [ + helper.make_node("QuantizeLinear", [current, "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["dq"]), + ] + current = "dq" + nodes.append(helper.make_node("ReduceSum", [current, "axes"], ["y"], keepdims=0)) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, width // 4])], + controls, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 8} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False, qdq_quantized_dtypes=[]) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +def test_early_pattern_does_not_allow_an_unmatched_native_reduction(monkeypatch): + """Refuse a native reducer before mutation even when another reducer has a valid early lowering.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + shape = [1, 1, 2, 12] + model = _make_model( + [ + helper.make_node("Reshape", ["x", "target"], ["expanded"]), + helper.make_node("ReduceSum", ["expanded", "axes"], ["y"], keepdims=0), + helper.make_node("ReduceSum", ["x", "axes"], ["native"], keepdims=0), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("native", TensorProto.FLOAT, [1, 1, 2]), + ], + [ + numpy_helper.from_array(np.asarray([1, 1, 2, 3, 4], np.int64), "target"), + numpy_helper.from_array(np.asarray([-1], np.int64), "axes"), + ], + ) + _checked_run(model, {"x": np.ones(shape, np.float32)}) + before = model.SerializeToString() + with pytest.raises(NotImplementedError, match="reductions"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + @pytest.mark.parametrize("width", [1, 4]) @pytest.mark.parametrize("axis", [0, -2]) def test_ordered_pipeline_preserves_vector_output(width, axis): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py index a7fb53af5bc..5f824f2b8b9 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -27,6 +27,7 @@ _apply_reshape_qdq_reduce_sum, _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, + _reshape_reduce_sum, ) from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible @@ -399,3 +400,29 @@ def test_terminal_five_dimensional_target_is_not_an_overlap(chained): np.testing.assert_array_equal( ReferenceEvaluator(transformed).run(None, feeds)[0], expected[0] ) + + +@pytest.mark.parametrize("pattern", ["plain", "qdq", "fallback"]) +def test_reduction_preflight_preserves_native_model_bytes(pattern): + """Return matched output owners without retaining generated controls or changing the native graph.""" + model, feeds = _make_model(pattern) + before = model.SerializeToString() + expected = _execute(model, feeds, "reference") + outputs = _reshape_reduce_sum(model, quantized=pattern == "qdq", validate_only=True) + assert outputs == {"Y"} + assert model.SerializeToString() == before + for actual, reference in zip(_execute(model, feeds, "reference"), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("pattern", ["plain", "qdq"]) +def test_reduction_preflight_capture_refusal_preserves_model_bytes(pattern): + """Validate lexical ownership and roll back all temporary controls during preflight refusal.""" + model, _ = _make_model(pattern) + shape = numpy_helper.to_array(model.graph.initializer[0]).tolist() + _attach_if_capture(model, "reshaped", shape) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="captured"): + _reshape_reduce_sum(model, quantized=pattern == "qdq", validate_only=True) + assert model.SerializeToString() == before From 9c408ad42c77ca93139fdb012ff476ab12165d9e Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:37:04 +0000 Subject: [PATCH 11/32] Validate boundary promotion against prepared pipeline stages Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 22 +++- .../dla_transforms/dla_reshape_patterns.py | 17 +-- .../onnx/graph_surgery/make_dla_compatible.py | 103 +++++++++++------- .../test_dla_model_boundaries.py | 87 ++++++++++++++- .../test_dla_reshape_patterns.py | 27 ----- 5 files changed, 175 insertions(+), 81 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index f44ed65ad12..7512d238b4b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -121,7 +121,7 @@ def _require_adapter_opset(model): raise ValueError("DLA boundary rank adapters require ONNX opset 13 or newer") -def _apply_non4d_model_inputs(model): +def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=False): """Promote selected boundary uses while preserving explicit adapter results.""" graph = model.graph cache = GraphCache(graph) @@ -150,12 +150,26 @@ def _apply_non4d_model_inputs(model): and not ( _is_standard_node(node, local_functions, {"Squeeze"}) and "Gather" in node.name ) - and node.op_type != "InstanceNormalization" + and node.op_type not in {"InstanceNormalization", "Gemm"} ] if not consumers: continue _require_adapter_opset(model) for consumer in consumers: + safe = _RANK_AGNOSTIC_OPS | {"Unsqueeze", "Concat", "Softmax", "LogSoftmax", "Gelu"} + if consumer.op_type == "MatMul": + other = cache.get_shape(consumer.input[1]) + if other is not None and len(other) >= 2: + safe |= {"MatMul"} + elif consumer.op_type == "Conv": + kernel = cache.get_shape(consumer.input[1]) + if kernel is not None and len(kernel) == 4: + safe |= {"Conv"} + elif consumer.op_type in {"QuantizeLinear", "DequantizeLinear"}: + if cache.get_shape(consumer.input[1]) == []: + safe |= {consumer.op_type} + if consumer.op_type not in safe and consumer.op_type not in prepared_input_ops: + raise ValueError(f"Input rank conversion is unavailable for {consumer.op_type}") if _is_standard_node(consumer, local_functions, {"Softmax", "LogSoftmax"}): axis = _attribute(consumer, "axis") axis_value = ( @@ -199,6 +213,8 @@ def _apply_non4d_model_inputs(model): or len({len(promoted_shapes[name]) for name in node.input}) != 1 ): raise ValueError(f"Cannot promote mixed-layout Concat {node.name!r}") + if validate_only: + return model for node in concats.values(): for attribute in node.attribute: if attribute.name == "axis" and attribute.i >= 0: @@ -296,6 +312,8 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is not None: + if _is_standard_node(source, local_functions, {"Gemm"}): + continue if defer_rank_conversion: _require_output_rank_conversion(source, local_functions, prepared_output_ops) if not _is_standard_node( diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py index face75ab6e6..60a260b1cf4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reshape_patterns.py @@ -50,7 +50,7 @@ def _follower(cache, output, op_type): ) -def _replace(graph, replacements, *, validate_only=False): +def _replace(graph, replacements): """Apply planned replacements without copying untouched payloads and order late controls.""" if not replacements: return @@ -70,8 +70,6 @@ def _replace(graph, replacements, *, validate_only=False): for added in replacements.get(id(node), []) for name in added.input ) - if validate_only: - return batch_replace_nodes( graph, { @@ -97,7 +95,7 @@ def _apply_reshape_qdq_reduce_sum(model): return _reshape_reduce_sum(model, quantized=True) -def _reshape_reduce_sum(model, *, quantized, validate_only=False): +def _reshape_reduce_sum(model, *, quantized): """Plan disjoint reductions while preserving controls and restoring initializers on refusal.""" graph = model.graph cache = GraphCache(graph) @@ -215,16 +213,7 @@ def _reshape_reduce_sum(model, *, quantized, validate_only=False): replacements[id(original)] = [] replacements[id(reshape)] = nodes count += 1 - _replace(graph, replacements, validate_only=validate_only) - if validate_only: - outputs = { - name - for node in graph.node - if node.op_type == "ReduceSum" and id(node) in replacements - for name in node.output - } - del graph.initializer[initializer_count:] - return outputs + _replace(graph, replacements) except Exception: del graph.initializer[initializer_count:] raise diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 923bebf7fa9..49f370421f0 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -23,6 +23,7 @@ from __future__ import annotations +import copy import time from functools import partial from typing import TYPE_CHECKING @@ -52,7 +53,6 @@ _apply_reshape_qdq_reduce_sum, _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, - _reshape_reduce_sum, ) from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, @@ -108,6 +108,40 @@ def _apply_constants_to_initializers(model, **options): ) +def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prepared_output_ops): + """Share stage options between the prerequisite probe and the actual ordered pipeline.""" + if name == "RemoveQDQ": + return { + "keep_clip_after_inputs": False, + "max_chained_qdq_pairs": 5, + "quant_types_to_remove": quant_types, + } + if name == "ConstantsToInitializers": + return {"verbose": verbose} + if name == "Non4dModelInputs": + return {"prepared_input_ops": prepared_input_ops} + if name == "Non4dModelOutputs": + return {"prepared_output_ops": prepared_output_ops} + return {} + + +def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): + """Refuse unsupported native layouts before adapting any public boundary.""" + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + for node in model.graph.node: + if ( + node.op_type in _REDUCTION_OPS + and node.op_type not in prepared_output_ops + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + ): + raise NotImplementedError("Standalone reductions are not available in this pipeline") + _require_output_rank_conversion(node, local_functions, prepared_output_ops) + _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) + + def _transform_make_dla_compatible( model: onnx.ModelProto, *, @@ -124,34 +158,22 @@ def _transform_make_dla_compatible( ): raise NotImplementedError("QDQ removal is not available in this pipeline") stages = frozenset(stage for stage, _ in DLA_TRANSFORMS) - planned_reduction_outputs: set[str] = set() - if "StandaloneReduceOperations" not in stages: - if any(node.op_type == "ReduceSum" for node in model.graph.node): - for stage, quantized in (("ReshapeReduceSum", False), ("ReshapeQDQReduceSum", True)): - if stage in stages: - planned_reduction_outputs.update( - _reshape_reduce_sum(model, quantized=quantized, validate_only=True) - ) - local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} - for node in model.graph.node: - if ( - node.op_type in _REDUCTION_OPS - and node.domain in {"", "ai.onnx"} - and not node.overload - and (node.domain, node.op_type, node.overload) not in local_functions - and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) - and not planned_reduction_outputs.intersection(node.output) - ): - raise NotImplementedError( - "Standalone reductions are not available in this pipeline" - ) prepared_output_ops = stages | ( _REDUCTION_OPS if "StandaloneReduceOperations" in stages else frozenset() ) - local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} - for node in model.graph.node: - if not planned_reduction_outputs.intersection(node.output): - _require_output_rank_conversion(node, local_functions, prepared_output_ops) + prepared_input_ops = set(prepared_output_ops) + for stage, operations in { + "Non4DReshape": {"Reshape"}, + "Non4DExpand": {"Expand"}, + "Non4DTile": {"Tile"}, + "Non4DTranspose": {"Transpose"}, + "Non4dSliceAxis": {"Slice"}, + "Non4DLpNorm": {"LpNormalization"}, + "RemoveIntermediarySqueezeAndUnsqueeze": {"Squeeze"}, + "UpdateAxis": {"TopK", "Split", "QuantizeLinear", "DequantizeLinear"}, + }.items(): + if stage in stages: + prepared_input_ops.update(operations) quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None @@ -159,21 +181,28 @@ def _transform_make_dla_compatible( ) if remove_all_qdq: quant_types = frozenset(onnx.TensorProto.DataType.values()) + try: + _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops) + except (ValueError, NotImplementedError): + # Earlier lowering may consume native operators that lack a later rank pass. + probe = infer_shapes(copy.deepcopy(model)) + for name, transform in DLA_TRANSFORMS: + if name == "Non4dModelInputs": + break + options = _boundary_stage_options( + name, quant_types, False, prepared_input_ops, prepared_output_ops + ) + result = transform(probe, **options) + if result is not None: + probe = result + _validate_boundary_contract(probe, prepared_input_ops, prepared_output_ops) model = infer_shapes(model) for name, transform in DLA_TRANSFORMS: if verbose: logger.info("[DLA pipeline] Applying %s", name) - options = {} - if name == "RemoveQDQ": - options = { - "keep_clip_after_inputs": False, - "max_chained_qdq_pairs": 5, - "quant_types_to_remove": quant_types, - } - elif name == "ConstantsToInitializers": - options = {"verbose": verbose} - elif name == "Non4dModelOutputs": - options = {"prepared_output_ops": prepared_output_ops} + options = _boundary_stage_options( + name, quant_types, verbose, prepared_input_ops, prepared_output_ops + ) result = transform(model, **options) if result is not None: model = result diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index bafcad80f24..9d5c522006e 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -298,7 +298,10 @@ def test_normalization_binary_branch_preserves_broadcast_layout(batch, operation @pytest.mark.parametrize("width", [8, 12]) @pytest.mark.parametrize("axis", [-1, 4]) @pytest.mark.parametrize("quantized", [False, True]) -def test_early_registered_reshape_reduction_is_not_rejected(width, axis, quantized, monkeypatch): +@pytest.mark.parametrize("preparation", ["direct", "intermediate", "constant"]) +def test_early_registered_reshape_reduction_is_not_rejected( + width, axis, quantized, preparation, monkeypatch +): """Allow reductions consumed by actual earlier reshape-pattern stages without a native reducer stage.""" pipeline = import_module(_transform_make_dla_compatible.__module__) monkeypatch.setattr( @@ -328,6 +331,14 @@ def test_early_registered_reshape_reduction_is_not_rejected(width, axis, quantiz ] current = "dq" nodes.append(helper.make_node("ReduceSum", [current, "axes"], ["y"], keepdims=0)) + if preparation == "intermediate": + nodes[0].input[0] = "hidden" + nodes.insert(0, helper.make_node("Identity", ["x"], ["hidden"])) + elif preparation == "constant": + nodes = [ + helper.make_node("Constant", [], [tensor.name], value=tensor) for tensor in controls + ] + nodes + controls = [] model = _make_model( nodes, [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], @@ -340,6 +351,80 @@ def test_early_registered_reshape_reduction_is_not_rejected(width, axis, quantiz np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) +@pytest.mark.parametrize( + "operation", ["ReduceSum", "Gather", "Slice", "Split", "ArgMax", "Transpose"] +) +def test_input_promotion_requires_registered_axis_conversion(operation, monkeypatch): + """Reject valid native operators before boundary mutation when their rank stage is absent.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + ( + ("ConstantsToInitializers", pipeline._apply_constants_to_initializers), + ("ShapeInference", pipeline.infer_shapes), + ("Non4dModelInputs", _apply_non4d_model_inputs), + ("Non4dModelOutputs", _apply_non4d_model_outputs), + ), + ) + inputs, attributes, output_shape = ["x"], {}, [2, 1] + controls = [] + dtype = TensorProto.FLOAT + if operation == "Transpose": + attributes, output_shape = {"perm": [1, 0]}, [3, 2] + elif operation == "ArgMax": + attributes, dtype = {"axis": 1, "keepdims": 1}, TensorProto.INT64 + else: + values = {"ReduceSum": [1], "Gather": [0], "Split": [3], "Slice": [0]}[operation] + controls.append(numpy_helper.from_array(np.asarray(values, np.int64), "control")) + inputs.append("control") + if operation == "ReduceSum": + attributes = {"keepdims": 1} + elif operation == "Gather": + attributes, output_shape = {"axis": 0}, [1, 3] + elif operation == "Split": + attributes, output_shape = {"axis": 1}, [2, 3] + else: + controls += [ + numpy_helper.from_array(np.asarray([1], np.int64), "end"), + numpy_helper.from_array(np.asarray([1], np.int64), "axes"), + ] + inputs += ["end", "axes"] + model = _make_model( + [helper.make_node(operation, inputs, ["y"], **attributes)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", dtype, output_shape)], + controls, + ) + _checked_run(model, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Input rank conversion"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("tail", [False, True]) +def test_ineligible_gemm_keeps_native_boundary_and_weights(tail): + """Retain native matrix Gemm semantics when convolution lowering declines its coefficients.""" + nodes = [helper.make_node("Gemm", ["x", "w", "b"], ["hidden" if tail else "y"], alpha=0.5)] + if tail: + nodes.append(helper.make_node("Relu", ["hidden"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4])], + [ + numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "w"), + numpy_helper.from_array(np.ones(4, np.float32), "b"), + ], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + assert [list(tensor.dims) for tensor in converted.graph.initializer] == [[3, 4], [4]] + + def test_early_pattern_does_not_allow_an_unmatched_native_reduction(monkeypatch): """Refuse a native reducer before mutation even when another reducer has a valid early lowering.""" pipeline = import_module(_transform_make_dla_compatible.__module__) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py index 5f824f2b8b9..a7fb53af5bc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reshape_patterns.py @@ -27,7 +27,6 @@ _apply_reshape_qdq_reduce_sum, _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, - _reshape_reduce_sum, ) from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible @@ -400,29 +399,3 @@ def test_terminal_five_dimensional_target_is_not_an_overlap(chained): np.testing.assert_array_equal( ReferenceEvaluator(transformed).run(None, feeds)[0], expected[0] ) - - -@pytest.mark.parametrize("pattern", ["plain", "qdq", "fallback"]) -def test_reduction_preflight_preserves_native_model_bytes(pattern): - """Return matched output owners without retaining generated controls or changing the native graph.""" - model, feeds = _make_model(pattern) - before = model.SerializeToString() - expected = _execute(model, feeds, "reference") - outputs = _reshape_reduce_sum(model, quantized=pattern == "qdq", validate_only=True) - assert outputs == {"Y"} - assert model.SerializeToString() == before - for actual, reference in zip(_execute(model, feeds, "reference"), expected): - np.testing.assert_array_equal(actual, reference) - - -@pytest.mark.parametrize("pattern", ["plain", "qdq"]) -def test_reduction_preflight_capture_refusal_preserves_model_bytes(pattern): - """Validate lexical ownership and roll back all temporary controls during preflight refusal.""" - model, _ = _make_model(pattern) - shape = numpy_helper.to_array(model.graph.initializer[0]).tolist() - _attach_if_capture(model, "reshaped", shape) - onnx.checker.check_model(model, full_check=True) - before = model.SerializeToString() - with pytest.raises(ValueError, match="captured"): - _reshape_reduce_sum(model, quantized=pattern == "qdq", validate_only=True) - assert model.SerializeToString() == before From 87356efeca80369fa51e02f1833ff6cec1731533 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:39:35 +0000 Subject: [PATCH 12/32] Include LayerNormalization in safe boundary promotion Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 13 +++++++++++-- .../dla_transforms/test_dla_model_boundaries.py | 15 +++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 7512d238b4b..09fb6dd1b33 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -156,7 +156,14 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal continue _require_adapter_opset(model) for consumer in consumers: - safe = _RANK_AGNOSTIC_OPS | {"Unsqueeze", "Concat", "Softmax", "LogSoftmax", "Gelu"} + safe = _RANK_AGNOSTIC_OPS | { + "Unsqueeze", + "Concat", + "Softmax", + "LogSoftmax", + "LayerNormalization", + "Gelu", + } if consumer.op_type == "MatMul": other = cache.get_shape(consumer.input[1]) if other is not None and len(other) >= 2: @@ -170,7 +177,9 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal safe |= {consumer.op_type} if consumer.op_type not in safe and consumer.op_type not in prepared_input_ops: raise ValueError(f"Input rank conversion is unavailable for {consumer.op_type}") - if _is_standard_node(consumer, local_functions, {"Softmax", "LogSoftmax"}): + if _is_standard_node( + consumer, local_functions, {"Softmax", "LogSoftmax", "LayerNormalization"} + ): axis = _attribute(consumer, "axis") axis_value = ( axis.i diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 9d5c522006e..d2d7181335b 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -258,6 +258,21 @@ def test_boundary_softmax_preserves_native_axes(operation, shape, axis, opset): np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) +@pytest.mark.parametrize("axis", [1, -1]) +def test_boundary_layer_normalization_preserves_native_axis(axis): + """Preserve LayerNormalization reductions after adding leading singleton dimensions.""" + model = _make_model( + [helper.make_node("LayerNormalization", ["x", "scale"], ["y"], axis=axis)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + [numpy_helper.from_array(np.ones(3, np.float32), "scale")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + @pytest.mark.parametrize("batch", [1, 2, 3]) @pytest.mark.parametrize("operation", ["Add", "Mul"]) @pytest.mark.parametrize("tail", [False, True]) From 0655bc7bc54038aaf4183eb6f196b6769ebc71b6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:42:12 +0000 Subject: [PATCH 13/32] Leave Constant controls to registered prerequisite preparation Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_model_boundaries.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 09fb6dd1b33..8d408814d3e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -193,6 +193,14 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal if not _is_standard_node(consumer, local_functions, {"Unsqueeze"}): continue tensor = cache.get_init(consumer.input[1]) + control = cache.get_producer(consumer.input[1]) + if ( + tensor is None + and validate_only + and control is not None + and _is_standard_node(control, local_functions, {"Constant"}) + ): + continue if tensor is None or tensor.name in cache.graph_input_names: raise ValueError(f"Unsqueeze axes must be fixed for {consumer.name!r}") output_shape = list(shape) From eec512450fd1d3465e6ae5f2aacc8059588fd825 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:46:50 +0000 Subject: [PATCH 14/32] Release prerequisite probe before normal DLA transformation Signed-off-by: mgohil --- modelopt/onnx/graph_surgery/make_dla_compatible.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 49f370421f0..2d12d6c3799 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -196,6 +196,7 @@ def _transform_make_dla_compatible( if result is not None: probe = result _validate_boundary_contract(probe, prepared_input_ops, prepared_output_ops) + del probe model = infer_shapes(model) for name, transform in DLA_TRANSFORMS: if verbose: From 4e3ea92c97b9693035f36b6893112a91f858e5b0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:12:46 +0000 Subject: [PATCH 15/32] Preserve native shape outputs and guard scalar Gather restoration Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 22 ++++- .../onnx/graph_surgery/make_dla_compatible.py | 4 +- .../test_dla_model_boundaries.py | 85 +++++++++++++++++++ 3 files changed, 107 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 8d408814d3e..8fd7b92cdad 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -29,8 +29,24 @@ _REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) -def _require_output_rank_conversion(node, local_functions, prepared_output_ops): +def _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache): """Require an explicit later stage before deferring a native rank-reducing operation.""" + if _is_standard_node(node, local_functions, {"Gather"}) and "Gather" not in prepared_output_ops: + shape = cache.get_shape(node.input[1]) + value = cache.vi_map.get(node.input[1]) + rank = ( + len(shape) + if shape is not None + else ( + len(value.type.tensor_type.shape.dim) + if value is not None and value.type.tensor_type.HasField("shape") + else None + ) + ) + if rank is None or rank == 0: + raise ValueError( + "Output rank conversion is unavailable for scalar or unknown-rank Gather" + ) if ( _is_standard_node(node, local_functions, {node.op_type}) and ( @@ -329,10 +345,10 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is not None: - if _is_standard_node(source, local_functions, {"Gemm"}): + if _is_standard_node(source, local_functions, {"Gemm", "Shape"}): continue if defer_rank_conversion: - _require_output_rank_conversion(source, local_functions, prepared_output_ops) + _require_output_rank_conversion(source, local_functions, prepared_output_ops, cache) if not _is_standard_node( source, local_functions, {source.op_type}, contrib_ops=contrib_ops ): diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 2d12d6c3799..03334da1e99 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -56,6 +56,7 @@ ) from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, + GraphCache, infer_shapes, parse_qdq_quantized_dtype_list, ) @@ -128,6 +129,7 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): """Refuse unsupported native layouts before adapting any public boundary.""" local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + cache = GraphCache(model.graph) for node in model.graph.node: if ( node.op_type in _REDUCTION_OPS @@ -138,7 +140,7 @@ def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) ): raise NotImplementedError("Standalone reductions are not available in this pipeline") - _require_output_rank_conversion(node, local_functions, prepared_output_ops) + _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache) _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index d2d7181335b..45a5204e312 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -440,6 +440,91 @@ def test_ineligible_gemm_keeps_native_boundary_and_weights(tail): assert [list(tensor.dims) for tensor in converted.graph.initializer] == [[3, 4], [4]] +@pytest.mark.parametrize("source", ["initializer", "constant", "runtime"]) +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("tail", [False, True]) +def test_scalar_gather_requires_registered_rank_conversion(source, axis, tail, monkeypatch): + """Refuse unprepared scalar Gather before the actual pipeline changes caller state.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple((name, transform) for name, transform in pipeline.DLA_TRANSFORMS if name != "Gather"), + ) + shape = [2, 3, 4, 4] if axis == 0 else [1, 3, 4, 4] + output_shape = shape[:axis] + shape[axis + 1 :] + nodes = [helper.make_node("Gather", ["x", "indices"], ["hidden" if tail else "y"], axis=axis)] + if tail: + nodes.append(helper.make_node("Identity", ["hidden"], ["y"])) + tensor = numpy_helper.from_array(np.int64(1), "indices") + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + if source == "constant": + nodes.insert(0, helper.make_node("Constant", [], ["indices"], value=tensor)) + elif source == "runtime": + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, [])) + model = _make_model( + nodes, + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [tensor] if source == "initializer" else [], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + if source == "runtime": + feeds["indices"] = np.int64(1) + _checked_run(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"Output rank conversion.*Gather"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("indices_shape", [[2], ["K"]]) +def test_native_vector_gather_keeps_rank_four_output(indices_shape, monkeypatch): + """Keep native rank-four Gather when only its index-vector length is symbolic.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple((name, transform) for name, transform in pipeline.DLA_TRANSFORMS if name != "Gather"), + ) + shape = [1, 3, 4, 4] + model = _make_model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=1)], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, indices_shape[0], 4, 4])], + ) + feeds = { + "x": np.arange(48, dtype=np.float32).reshape(shape), + "indices": np.array([0, 2], np.int64), + } + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("shape", [[3], [1, 3, 4, 4]]) +@pytest.mark.parametrize("tail", [False, True]) +def test_shape_outputs_keep_native_vector_rank(shape, tail): + """Retain native Shape vectors, including a one-element vector through unary tails.""" + nodes = [helper.make_node("Shape", ["x"], ["hidden" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["hidden"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.INT64, [len(shape)])], + ) + feeds = {"x": np.ones(shape, np.float32)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + actual = _checked_run(converted, feeds)[0] + assert actual.shape == (len(shape),) + np.testing.assert_array_equal(actual, expected) + + def test_early_pattern_does_not_allow_an_unmatched_native_reduction(monkeypatch): """Refuse a native reducer before mutation even when another reducer has a valid early lowering.""" pipeline = import_module(_transform_make_dla_compatible.__module__) From 5486c2a9f1ac9947fa1c75cada60687c74090e5a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:48:31 +0000 Subject: [PATCH 16/32] Preserve every widened normalization consumer layout Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 162 +++++++++++++++--- .../onnx/graph_surgery/make_dla_compatible.py | 4 +- .../test_dla_model_boundaries.py | 110 ++++++++++++ 3 files changed, 250 insertions(+), 26 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 8fd7b92cdad..e710398b78a 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -137,6 +137,23 @@ def _require_adapter_opset(model): raise ValueError("DLA boundary rank adapters require ONNX opset 13 or newer") +def _native_conv1d(node, cache): + """Recognize Conv kernels whose native spatial layout is one-dimensional.""" + kernel = ( + cache.get_shape(node.input[1]) if node.op_type == "Conv" and len(node.input) > 1 else None + ) + return kernel is not None and len(kernel) == 3 + + +def _promoted_axis(node, rank, default): + """Map an explicit or schema-default axis after adding leading singleton dimensions.""" + axis = _attribute(node, "axis") + value = axis.i if axis is not None else default + if not -rank <= value < rank: + raise ValueError(f"Invalid native normalization axis for {node.name!r}") + return axis, value + 4 - rank if value >= 0 else value + + def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=False): """Promote selected boundary uses while preserving explicit adapter results.""" graph = model.graph @@ -167,6 +184,7 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal _is_standard_node(node, local_functions, {"Squeeze"}) and "Gather" in node.name ) and node.op_type not in {"InstanceNormalization", "Gemm"} + and not _native_conv1d(node, cache) ] if not consumers: continue @@ -196,16 +214,13 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal if _is_standard_node( consumer, local_functions, {"Softmax", "LogSoftmax", "LayerNormalization"} ): - axis = _attribute(consumer, "axis") - axis_value = ( - axis.i - if axis is not None - else defs.get_schema(consumer.op_type, opset).attributes["axis"].default_value.i + axis, axis_value = _promoted_axis( + consumer, + len(shape), + defs.get_schema(consumer.op_type, opset).attributes["axis"].default_value.i, ) - if not -len(shape) <= axis_value < len(shape): - raise ValueError(f"Invalid native normalization axis for {consumer.name!r}") if axis_value >= 0: - normalization_axes[id(consumer)] = (consumer, axis, axis_value + 4 - len(shape)) + normalization_axes[id(consumer)] = (consumer, axis, axis_value) if not _is_standard_node(consumer, local_functions, {"Unsqueeze"}): continue tensor = cache.get_init(consumer.input[1]) @@ -345,7 +360,9 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is not None: - if _is_standard_node(source, local_functions, {"Gemm", "Shape"}): + if _is_standard_node(source, local_functions, {"Gemm", "Shape"}) or _native_conv1d( + source, cache + ): continue if defer_rank_conversion: _require_output_rank_conversion(source, local_functions, prepared_output_ops, cache) @@ -409,6 +426,9 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o existing.input.extend([data_input] + ([axes_name] if axes_name else [])) continue data_input = cache.unique_name(value.name + "_squeeze_input") + graph.value_info.append( + helper.make_tensor_value_info(data_input, value.type.tensor_type.elem_type, shape) + ) for node in graph.node: for index, name in enumerate(node.input): if name == value.name: @@ -496,15 +516,80 @@ def _apply_remove_intermediary_squeeze_unsqueeze(model, *, rank_converted=False) return model -def _apply_fix_instancenorm_channels(model): +def _apply_fix_instancenorm_channels(model, *, prepared_consumer_ops=()): """Append spatial singleton dimensions while retaining normalization channels and outputs.""" graph = model.graph cache = GraphCache(graph) replacements = {} local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + contrib_ops = _boundary_contrib_ops(model) + opset = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) + axis_plans, perm_plans = {}, {} + axis_ops = { + "Softmax", + "LogSoftmax", + "LayerNormalization", + "Concat", + "Split", + "TopK", + "QuantizeLinear", + "DequantizeLinear", + } + + def consumer_layout(consumer, port, standard): + """Plan native views or axis/permutation preparation before widening any tensor.""" + if ( + not standard + or consumer.op_type in {"Shape", "GroupQueryAttention"} + or _native_conv1d(consumer, cache) + or (port != 0 and consumer.op_type not in {"Concat", "MatMul"}) + ): + return True + if consumer.op_type in axis_ops: + if "axis_updated" in consumer.name: + return False + quantizer = consumer.op_type in {"QuantizeLinear", "DequantizeLinear"} + if quantizer and cache.get_shape(consumer.input[1]) == []: + return False + schema = ( + defs.get_schema(consumer.op_type, opset) + if consumer.domain in {"", "ai.onnx"} + else None + ) + default = schema.attributes.get("axis") if schema is not None else None + if default is not None or (quantizer and consumer.domain == "com.microsoft"): + axis, value = _promoted_axis( + consumer, 3, default.default_value.i if default is not None else 1 + ) + if value >= 0: + axis_plans[id(consumer)] = (consumer, axis, value) + elif consumer.op_type == "Transpose": + perm = _attribute(consumer, "perm") + values = list(perm.ints) if perm is not None else [2, 1, 0] + if len(values) == 3 and sorted(values) == [0, 1, 2]: + perm_plans[id(consumer)] = (consumer, perm, [0, *[value + 1 for value in values]]) + elif len(values) != 4 or sorted(values) != [0, 1, 2, 3]: + raise ValueError(f"Unsupported normalization Transpose {consumer.name!r}") + elif consumer.op_type not in prepared_consumer_ops and consumer.op_type not in { + "Gelu", + "InstanceNormalization", + }: + target = cache.get_shape(consumer.input[1]) if consumer.op_type == "Reshape" else None + operands = ( + [cache.get_shape(name) for name in consumer.input] + if consumer.op_type == "MatMul" + else [] + ) + if target != [4] and not ( + operands and all(shape is not None and len(shape) >= 2 for shape in operands) + ): + raise ValueError(f"Unsupported normalization consumer {consumer.op_type}") + return False def output_consumers(name): - """Find public adapters and binary edges after layout-preserving unary paths.""" + """Plan leading or native views for every edge after layout-preserving unary paths.""" pending, visited = [name], set() while pending: tensor = pending.pop() @@ -512,19 +597,32 @@ def output_consumers(name): continue visited.add(tensor) for consumer, input_index in cache.get_consumers(tensor): - if not _is_standard_node(consumer, local_functions, {consumer.op_type}): - continue + standard = _is_standard_node( + consumer, local_functions, {consumer.op_type}, contrib_ops=contrib_ops + ) if ( - input_index == 0 + standard + and input_index == 0 and consumer.op_type == "Squeeze" and consumer.output[0] in cache.graph_output_names ): - yield consumer, None, None - elif consumer.op_type in _RANK_AGNOSTIC_OPS and len(consumer.output) == 1: + yield consumer, None, None, False + elif ( + standard + and consumer.op_type in _RANK_AGNOSTIC_OPS + and len(consumer.output) == 1 + ): if len(consumer.input) == 1 and input_index == 0: pending.append(consumer.output[0]) elif len(consumer.input) > 1: - yield consumer, tensor, input_index + yield consumer, tensor, input_index, False + else: + yield ( + consumer, + tensor, + input_index, + consumer_layout(consumer, input_index, standard), + ) plans = [] for index, node in enumerate(graph.node): @@ -550,11 +648,12 @@ def output_consumers(name): ) node.input[0] = output replacements[index] = ([node], [reshape, node]) - for consumer, tensor, port in consumers: + for consumer, tensor, port, native in consumers: if tensor is not None: - if tensor not in bridges: - bridges[tensor] = (shape, []) - bridges[tensor][1].append((consumer, port)) + key = (tensor, native) + if key not in bridges: + bridges[key] = (shape, []) + bridges[key][1].append((consumer, port)) continue if len(consumer.input) < 2: continue @@ -571,15 +670,16 @@ def output_consumers(name): ) positions = {id(node): index for index, node in enumerate(graph.node)} prefixes = {} - for tensor, (shape, edges) in bridges.items(): + for (tensor, native), (shape, edges) in bridges.items(): + layout = "native" if native else "leading" target = add_initializer( graph, - cache.unique_name(tensor + "_leading_shape"), - np.asarray(pad4d(shape), np.int64), + cache.unique_name(tensor + "_" + layout + "_shape"), + np.asarray(shape if native else pad4d(shape), np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - output = cache.unique_name(tensor + "_leading") + output = cache.unique_name(tensor + "_" + layout) view = helper.make_node( "Reshape", [tensor, target], [output], name=cache.unique_name(output + "_reshape") ) @@ -590,6 +690,18 @@ def output_consumers(name): for index, nodes in prefixes.items(): removed, added = replacements.get(index, ([graph.node[index]], [graph.node[index]])) replacements[index] = (removed, [*nodes, *added]) + for consumer, axis, value in axis_plans.values(): + if axis is None: + consumer.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + consumer.name = cache.unique_name((consumer.name or consumer.output[0]) + "_axis_updated") + for consumer, perm, values in perm_plans.values(): + if perm is None: + consumer.attribute.append(helper.make_attribute("perm", values)) + else: + del perm.ints[:] + perm.ints.extend(values) batch_replace_nodes(graph, replacements) _add_legacy_initializer_inputs(model) return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 03334da1e99..bbc986dfffa 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -123,6 +123,8 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep return {"prepared_input_ops": prepared_input_ops} if name == "Non4dModelOutputs": return {"prepared_output_ops": prepared_output_ops} + if name == "FixInstanceNormChannelMismatch": + return {"prepared_consumer_ops": prepared_input_ops} return {} @@ -172,7 +174,7 @@ def _transform_make_dla_compatible( "Non4dSliceAxis": {"Slice"}, "Non4DLpNorm": {"LpNormalization"}, "RemoveIntermediarySqueezeAndUnsqueeze": {"Squeeze"}, - "UpdateAxis": {"TopK", "Split", "QuantizeLinear", "DequantizeLinear"}, + "UpdateAxis": {"TopK", "Concat", "Split", "QuantizeLinear", "DequantizeLinear"}, }.items(): if stage in stages: prepared_input_ops.update(operations) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 45a5204e312..c5e4b333ec3 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -273,6 +273,116 @@ def test_boundary_layer_normalization_preserves_native_axis(axis): np.testing.assert_allclose(_checked_run(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) +@pytest.mark.parametrize( + ("operation", "axis"), + [ + ("Softmax", -1), + ("Softmax", 1), + ("LogSoftmax", None), + ("Transpose", None), + ("Transpose", 1), + ("Conv", None), + ("opaque", None), + ("Shape", None), + ], +) +@pytest.mark.parametrize("unary", [False, True]) +def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, axis, unary): + """Execute native and promoted consumers alongside a shared public normalization output.""" + shape = [2, 3, 4] + nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"])] + current = "normalized" + if unary: + nodes.append(helper.make_node("Relu", [current], ["activated"])) + current = "activated" + controls = [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ] + output_shape, dtype, functions = shape, TensorProto.FLOAT, [] + if operation in {"Softmax", "LogSoftmax"}: + nodes.append( + helper.make_node( + operation, [current], ["y"], **({"axis": axis} if axis is not None else {}) + ) + ) + elif operation == "Transpose": + nodes.append( + helper.make_node( + operation, [current], ["y"], **({"perm": [0, 2, 1]} if axis is not None else {}) + ) + ) + output_shape = [2, 4, 3] if axis is not None else [4, 3, 2] + elif operation == "Conv": + controls.append( + numpy_helper.from_array(np.arange(18, dtype=np.float32).reshape(2, 3, 3) / 18, "kernel") + ) + nodes.append(helper.make_node("Conv", [current, "kernel"], ["y"])) + output_shape = [2, 2, 2] + elif operation == "opaque": + functions.append( + helper.make_function( + "opaque.norm", + "NativeLayout", + ["value"], + ["result"], + [helper.make_node("Transpose", ["value"], ["result"], perm=[0, 2, 1])], + [helper.make_opsetid("", 21)], + ) + ) + nodes.append(helper.make_node("NativeLayout", [current], ["y"], domain="opaque.norm")) + output_shape = [2, 4, 3] + else: + nodes.append(helper.make_node("Shape", [current], ["y"])) + output_shape, dtype = [3], TensorProto.INT64 + outputs = [ + helper.make_tensor_value_info("y", dtype, output_shape), + helper.make_tensor_value_info("normalized", TensorProto.FLOAT, shape), + ] + if operation in {"Conv", "opaque", "Shape"}: + nodes.append(helper.make_node("Softmax", [current], ["side"], axis=-1)) + outputs.append(helper.make_tensor_value_info("side", TensorProto.FLOAT, shape)) + model = _make_model( + nodes, [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], outputs, controls + ) + model.functions.extend(functions) + if functions: + model.opset_import.append(helper.make_opsetid("opaque.norm", 1)) + feeds = {"x": np.linspace(-1, 2, 24, dtype=np.float32).reshape(shape)} + expected = _checked_run(inliner.inline_local_functions(model), feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + if operation == "Conv": + assert list( + next(tensor for tensor in converted.graph.initializer if tensor.name == "kernel").dims + ) == [2, 3, 3] + for actual, reference in zip( + _checked_run(inliner.inline_local_functions(converted), feeds), expected + ): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +def test_normalization_refuses_unprepared_consumers_before_mutation(): + """Refuse a rank-sensitive consumer before widening any eligible normalization node.""" + shape = [2, 3, 4] + model = _make_model( + [ + helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"]), + helper.make_node("Hardmax", ["normalized"], ["y"], axis=1), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + _checked_run(model, {"x": np.arange(24, dtype=np.float32).reshape(shape)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Unsupported normalization consumer"): + _apply_fix_instancenorm_channels(model) + assert model.SerializeToString() == before + + @pytest.mark.parametrize("batch", [1, 2, 3]) @pytest.mark.parametrize("operation", ["Add", "Mul"]) @pytest.mark.parametrize("tail", [False, True]) From de1c209ac3d9b1e0105a6d26bd0740fd122b152d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:52:56 +0000 Subject: [PATCH 17/32] Limit native Conv fallback to unpromoted normalization paths Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 21 +++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index e710398b78a..b2d8e3582d1 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -138,11 +138,28 @@ def _require_adapter_opset(model): def _native_conv1d(node, cache): - """Recognize Conv kernels whose native spatial layout is one-dimensional.""" + """Recognize one-dimensional Conv paths whose data stays native before normalization repair.""" kernel = ( cache.get_shape(node.input[1]) if node.op_type == "Conv" and len(node.input) > 1 else None ) - return kernel is not None and len(kernel) == 3 + if kernel is None or len(kernel) != 3: + return False + name, seen = node.input[0], set() + while name not in seen: + seen.add(name) + if name in cache.graph_input_names: + shape = cache.get_shape(name) + return shape is not None and len(shape) == 3 + producer = cache.get_producer(name) + if producer is None: + return False + if producer.op_type == "InstanceNormalization": + shape = cache.get_shape(producer.input[0]) + return shape is not None and len(shape) == 3 + if producer.op_type not in _RANK_AGNOSTIC_OPS or len(producer.input) != 1: + return False + name = producer.input[0] + return False def _promoted_axis(node, rank, default): From e5b25bc5b86baf110ffe21ea7772ee306c8e722a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 12:20:15 +0000 Subject: [PATCH 18/32] Recover native output shapes without conflicting annotations Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 18 ++++++--- .../onnx/graph_surgery/utils/graph_utils.py | 37 ++++++++++++++++++- .../test_dla_model_boundaries.py | 7 ++++ 3 files changed, 54 insertions(+), 8 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index b2d8e3582d1..f2df0beee60 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -18,7 +18,14 @@ import numpy as np from onnx import TensorProto, defs, helper, numpy_helper -from ..utils.graph_utils import GraphCache, _attribute, add_initializer, batch_replace_nodes, pad4d +from ..utils.graph_utils import ( + GraphCache, + _attribute, + _shape_before_boundary_promotion, + add_initializer, + batch_replace_nodes, + pad4d, +) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_reorder_nodes import priority_topological_order @@ -154,7 +161,7 @@ def _native_conv1d(node, cache): if producer is None: return False if producer.op_type == "InstanceNormalization": - shape = cache.get_shape(producer.input[0]) + shape = _shape_before_boundary_promotion(cache, producer.input[0]) return shape is not None and len(shape) == 3 if producer.op_type not in _RANK_AGNOSTIC_OPS or len(producer.input) != 1: return False @@ -443,9 +450,6 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o existing.input.extend([data_input] + ([axes_name] if axes_name else [])) continue data_input = cache.unique_name(value.name + "_squeeze_input") - graph.value_info.append( - helper.make_tensor_value_info(data_input, value.type.tensor_type.elem_type, shape) - ) for node in graph.node: for index, name in enumerate(node.input): if name == value.name: @@ -624,6 +628,8 @@ def output_consumers(name): and consumer.output[0] in cache.graph_output_names ): yield consumer, None, None, False + elif standard and consumer.op_type == "InstanceNormalization" and input_index == 0: + continue elif ( standard and consumer.op_type in _RANK_AGNOSTIC_OPS @@ -644,7 +650,7 @@ def output_consumers(name): plans = [] for index, node in enumerate(graph.node): if _is_standard_node(node, local_functions, {"InstanceNormalization"}): - shape = cache.get_shape(node.input[0]) + shape = _shape_before_boundary_promotion(cache, node.input[0]) if shape is not None and len(shape) == 3: plans.append((index, node, shape, list(output_consumers(node.output[0])))) bridges = {} diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py index b20341fe154..3f744c71d62 100644 --- a/modelopt/onnx/graph_surgery/utils/graph_utils.py +++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py @@ -986,14 +986,47 @@ def _attribute(node, name): return next((attribute for attribute in node.attribute if attribute.name == name), None) +def _allocated_boundary_name(name, hint): + """Match the shared allocator's unchanged hint or numbered collision suffix.""" + return name == hint or (name.startswith(hint + "_") and name[len(hint) + 1 :].isdigit()) + + def _shape_before_boundary_promotion(cache, name): - """Resolve the original shape behind a generated graph-input Unsqueeze.""" + """Recover native shapes at generated input and output boundary views.""" producer = cache.get_producer(name) if producer is not None and producer.op_type == "Unsqueeze": original = producer.input[0] if original in cache.graph_input_names and name == original + "_unsqueeze_input": return cache.get_shape(original) - return cache.get_shape(name) + shape = cache.get_shape(name) + if shape is not None: + return shape + for consumer, port in cache.get_consumers(name): + if ( + port != 0 + or consumer.op_type != "Squeeze" + or consumer.domain not in {"", "ai.onnx"} + or consumer.overload + or consumer.output[0] not in cache.graph_output_names + ): + continue + original = consumer.output[0] + if not _allocated_boundary_name( + name, original + "_squeeze_input" + ) or not _allocated_boundary_name(consumer.name, original + "_squeeze"): + continue + shape = cache.get_shape(original) + axes = cache.get_init_array(consumer.input[1]) if len(consumer.input) > 1 else None + if ( + shape is not None + and 0 < len(shape) < 4 + and ( + (axes is not None and axes.reshape(-1).tolist() == list(range(4 - len(shape)))) + or (axes is None and len(shape) == 1 and shape[0] != 1) + ) + ): + return shape + return None def is_const_dq_input(init_list, input_name, graph): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index c5e4b333ec3..775cc312f20 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -284,6 +284,7 @@ def test_boundary_layer_normalization_preserves_native_axis(axis): ("Conv", None), ("opaque", None), ("Shape", None), + ("InstanceNormalization", None), ], ) @pytest.mark.parametrize("unary", [False, True]) @@ -332,6 +333,8 @@ def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, ) nodes.append(helper.make_node("NativeLayout", [current], ["y"], domain="opaque.norm")) output_shape = [2, 4, 3] + elif operation == "InstanceNormalization": + nodes.append(helper.make_node(operation, [current, "scale", "bias"], ["y"])) else: nodes.append(helper.make_node("Shape", [current], ["y"])) output_shape, dtype = [3], TensorProto.INT64 @@ -339,6 +342,10 @@ def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, helper.make_tensor_value_info("y", dtype, output_shape), helper.make_tensor_value_info("normalized", TensorProto.FLOAT, shape), ] + if operation == "Transpose" and axis is None: + for name in ("normalized_squeeze_input", "normalized_squeeze"): + controls.append(numpy_helper.from_array(np.array(1, np.int64), name)) + outputs.append(helper.make_tensor_value_info(name, TensorProto.INT64, [])) if operation in {"Conv", "opaque", "Shape"}: nodes.append(helper.make_node("Softmax", [current], ["side"], axis=-1)) outputs.append(helper.make_tensor_value_info("side", TensorProto.FLOAT, shape)) From 700e968bf3af8c669fe98ca6f5e62f7af7d01828 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:00:34 +0000 Subject: [PATCH 19/32] Validate deferred public output rank plans before mutation Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 17 +++++- .../onnx/graph_surgery/make_dla_compatible.py | 6 +++ .../test_dla_model_boundaries.py | 53 +++++++++++++++++++ 3 files changed, 75 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index f2df0beee60..6a30b344ba8 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -366,7 +366,9 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal return model -def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_output_ops=()): +def _apply_non4d_model_outputs( + model, *, defer_rank_conversion=False, prepared_output_ops=(), validate_only=False +): """Restore four-dimensional outputs, optionally deferring checks to later rank stages.""" graph = model.graph cache = GraphCache(graph) @@ -383,6 +385,8 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o continue producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) + if source is None and len(shape) == 1: + continue if source is not None: if _is_standard_node(source, local_functions, {"Gemm", "Shape"}) or _native_conv1d( source, cache @@ -397,6 +401,15 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): if _has_native_symbolic_data(cache, source.input[0], local_functions): continue + target = cache.get_shape(source.input[1]) + if ( + source.op_type == "Reshape" + and target is not None + and len(target) == 1 + and target[0] < 4 + and "Non4DReshape" not in prepared_output_ops + ): + raise ValueError("Output rank conversion is unavailable for native Reshape") if ( producer is not None and _is_standard_node(producer, local_functions, {"Identity"}) @@ -434,6 +447,8 @@ def _apply_non4d_model_outputs(model, *, defer_rank_conversion=False, prepared_o plans.append((value, shape, existing, axes)) if plans: _require_adapter_opset(model) + if validate_only: + return model for value, shape, existing, axes in plans: axes_name = None if len(shape) > 1 or shape == [1]: diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index bbc986dfffa..1684e258d1e 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -144,6 +144,12 @@ def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): raise NotImplementedError("Standalone reductions are not available in this pipeline") _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache) _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) + _apply_non4d_model_outputs( + model, + defer_rank_conversion=True, + prepared_output_ops=prepared_output_ops, + validate_only=True, + ) def _transform_make_dla_compatible( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 775cc312f20..9f9f480c25a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1124,6 +1124,59 @@ def test_native_reduction_output_refuses_before_adding_rank_adapters(entrypoint, assert model.SerializeToString() == before +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("entrypoint", ["deferred", "pipeline"]) +def test_native_reshape_output_requires_registered_rank_stage(shape, tail, entrypoint, monkeypatch): + """Refuse low-rank Reshape restoration before changing public boundary state.""" + nodes = [helper.make_node("Reshape", ["x", "target"], ["native" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Relu", ["native"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1] * (4 - len(shape)) + shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(np.asarray(shape, np.int64), "target")], + ) + _checked_run(model, {"x": np.ones([1] * (4 - len(shape)) + shape, np.float32)}) + before = model.SerializeToString() + if entrypoint == "pipeline": + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) + for name, callback in pipeline.DLA_TRANSFORMS + if name != "Non4DReshape" + ), + ) + with pytest.raises(ValueError, match="rank"): + if entrypoint == "pipeline": + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shape", [[1], [3]]) +@pytest.mark.parametrize("tail", [False, True]) +def test_native_vector_unary_outputs_keep_their_rank(shape, tail): + """Leave rank-one unary paths native when boundary inputs are not widened.""" + nodes = [helper.make_node("Relu", ["x"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Neg", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + @pytest.mark.parametrize("axes", [[0, 2], [0, -3]]) def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes): """Avoid remapping axes through an unspecified five-dimensional rank collapse.""" From 45a2197dc6e4095c43fd744481c374c48b07fd3c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:09:23 +0000 Subject: [PATCH 20/32] Limit output target inspection to Reshape producers Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/dla_model_boundaries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 6a30b344ba8..c633c0852eb 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -401,7 +401,7 @@ def _apply_non4d_model_outputs( if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): if _has_native_symbolic_data(cache, source.input[0], local_functions): continue - target = cache.get_shape(source.input[1]) + target = cache.get_shape(source.input[1]) if source.op_type == "Reshape" else None if ( source.op_type == "Reshape" and target is not None From 9ca803a2ccd58ec1586a3e2d6f5e8830aa5d9f58 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:13:50 +0000 Subject: [PATCH 21/32] Avoid prerequisite copies while checking output rank availability Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 4 ++- .../test_dla_model_boundaries.py | 28 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index c633c0852eb..8bf94e4e9cc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -369,7 +369,7 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal def _apply_non4d_model_outputs( model, *, defer_rank_conversion=False, prepared_output_ops=(), validate_only=False ): - """Restore four-dimensional outputs, optionally deferring checks to later rank stages.""" + """Restore outputs or only validate producer rank availability before shape inference.""" graph = model.graph cache = GraphCache(graph) captured = _subgraph_captures(graph) @@ -410,6 +410,8 @@ def _apply_non4d_model_outputs( and "Non4DReshape" not in prepared_output_ops ): raise ValueError("Output rank conversion is unavailable for native Reshape") + if validate_only: + continue if ( producer is not None and _is_standard_node(producer, local_functions, {"Identity"}) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 9f9f480c25a..56c91a9412a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1177,6 +1177,34 @@ def test_native_vector_unary_outputs_keep_their_rank(shape, tail): np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) +def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): + """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" + model = _make_model( + [ + helper.make_node("Relu", ["x"], ["hidden"]), + helper.make_node("Squeeze", ["hidden"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _checked_run(model, feeds)[0] + pipeline = import_module(_transform_make_dla_compatible.__module__) + original = pipeline.copy.deepcopy + copies = [] + + def record_copy(value, *args, **kwargs): + """Observe copies of the caller without changing the real fallback behavior.""" + if value is model: + copies.append(True) + return original(value, *args, **kwargs) + + monkeypatch.setattr(pipeline.copy, "deepcopy", record_copy) + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + assert copies == [] + + @pytest.mark.parametrize("axes", [[0, 2], [0, -3]]) def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes): """Avoid remapping axes through an unspecified five-dimensional rank collapse.""" From f68a294ee4b456ad2c2cdce0f6f68808ece5cd7a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 14:10:29 +0000 Subject: [PATCH 22/32] Require reader rank preparation before bypassing adapters Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 62 +++++++++++++- .../onnx/graph_surgery/make_dla_compatible.py | 17 +++- .../test_dla_model_boundaries.py | 80 +++++++++++++++++++ 3 files changed, 157 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 8bf94e4e9cc..2768b5a25ef 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -492,13 +492,65 @@ def _apply_non4d_model_outputs( return model -def _apply_remove_intermediary_squeeze_unsqueeze(model, *, rank_converted=False): +def _apply_remove_intermediary_squeeze_unsqueeze( + model, *, rank_converted=False, prepared_consumer_ops=(), validate_only=False +): """Remove no-op adapters, or reconcile adapters within the ordered rank-converted pipeline.""" graph = model.graph cache = GraphCache(graph) captured = _subgraph_captures(graph) removed = set() local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + checked_readers = set() + runtime_inputs = ( + cache.graph_input_names - cache.init_names + if model.ir_version <= 3 + else cache.graph_input_names + ) + + def require_prepared_readers(name): + """Follow shared rank-agnostic paths before allowing any adapter rewiring.""" + pending, seen = [name], set() + while pending: + current = pending.pop() + if current in seen or current in checked_readers: + continue + seen.add(current) + if current in captured - cache.graph_output_names: + raise ValueError("Intermediary adapter changes a captured native rank") + for reader, port in cache.get_consumers(current): + standard = _is_standard_node(reader, local_functions, {reader.op_type}) + if ( + standard + and reader.op_type == "Squeeze" + and reader.output[0] in cache.graph_output_names + ): + continue + if standard and reader.op_type in _RANK_AGNOSTIC_OPS | {"Squeeze", "Unsqueeze"}: + pending.extend(reader.output) + continue + target = ( + cache.get_init_array(reader.input[1]) + if standard + and reader.op_type == "Reshape" + and port == 0 + and reader.input[1] not in runtime_inputs + else None + ) + if target is not None and not np.any(target == 0): + continue + if ( + standard + and reader.op_type in prepared_consumer_ops + and (port == 0 or reader.op_type == "Concat") + ): + continue + raise ValueError( + f"Intermediary adapter feeds unprepared rank-sensitive {reader.op_type}" + ) + checked_readers.update(seen) + + plans = [] for node in list(graph.node): if not _is_standard_node(node, local_functions, {"Squeeze", "Unsqueeze"}): continue @@ -511,9 +563,17 @@ def _apply_remove_intermediary_squeeze_unsqueeze(model, *, rank_converted=False) input_shape = cache.get_shape(node.input[0]) output_shape = cache.get_shape(node.output[0]) if not input_shape or not output_shape: + if validate_only and rank_converted: + require_prepared_readers(node.output[0]) continue if not rank_converted and input_shape != output_shape: continue + if rank_converted and len(input_shape) != len(output_shape): + require_prepared_readers(node.output[0]) + plans.append((node, input_shape, output_shape)) + if validate_only: + return model + for node, input_shape, output_shape in plans: short, long = ( (output_shape, input_shape) if node.op_type == "Squeeze" diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 1684e258d1e..afc3bab16e0 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -123,6 +123,8 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep return {"prepared_input_ops": prepared_input_ops} if name == "Non4dModelOutputs": return {"prepared_output_ops": prepared_output_ops} + if name == "RemoveIntermediarySqueezeAndUnsqueeze": + return {"prepared_consumer_ops": prepared_input_ops} if name == "FixInstanceNormChannelMismatch": return {"prepared_consumer_ops": prepared_input_ops} return {} @@ -150,6 +152,10 @@ def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): prepared_output_ops=prepared_output_ops, validate_only=True, ) + if "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_input_ops: + _apply_remove_intermediary_squeeze_unsqueeze( + model, rank_converted=True, prepared_consumer_ops=prepared_input_ops, validate_only=True + ) def _transform_make_dla_compatible( @@ -179,8 +185,17 @@ def _transform_make_dla_compatible( "Non4DTranspose": {"Transpose"}, "Non4dSliceAxis": {"Slice"}, "Non4DLpNorm": {"LpNormalization"}, + "FixInstanceNormChannelMismatch": {"InstanceNormalization"}, "RemoveIntermediarySqueezeAndUnsqueeze": {"Squeeze"}, - "UpdateAxis": {"TopK", "Concat", "Split", "QuantizeLinear", "DequantizeLinear"}, + "UpdateAxis": { + "Softmax", + "LogSoftmax", + "TopK", + "Concat", + "Split", + "QuantizeLinear", + "DequantizeLinear", + }, }.items(): if stage in stages: prepared_input_ops.update(operations) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 56c91a9412a..dbe79805c00 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1177,6 +1177,86 @@ def test_native_vector_unary_outputs_keep_their_rank(shape, tail): np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) +def _rank_sensitive_adapter_model(operation, unary, metadata): + """Build an eligible prefix before a rank-sensitive branch to expose partial rewiring.""" + nodes = [ + helper.make_node("Relu", ["x"], ["a"]), + helper.make_node("Squeeze", ["a", "axes"], ["b"]), + helper.make_node("Relu", ["b"], ["early"]), + helper.make_node("Relu", ["x"], ["r"]), + helper.make_node("Squeeze", ["r", "axes"], ["s"]), + ] + if unary: + nodes.append(helper.make_node("Relu", ["s"], ["native"])) + nodes.append( + helper.make_node( + operation, + ["native" if unary else "s"], + ["y"], + **({"perm": [1, 2, 0]} if operation == "Transpose" else {"axis": 1}), + ) + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [ + helper.make_tensor_value_info("early", TensorProto.FLOAT, [3, 4, 4]), + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [4, 4, 3] if operation == "Transpose" else [3, 4, 4] + ), + ], + [numpy_helper.from_array(np.array([0], np.int64), "axes")], + ) + if metadata: + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.linspace(0.1, 4, 48, dtype=np.float32).reshape(1, 3, 4, 4)} + return model, feeds + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax", "Transpose"]) +@pytest.mark.parametrize("unary", [False, True]) +@pytest.mark.parametrize( + ("entrypoint", "metadata"), [("stage", True), ("pipeline", True), ("pipeline", False)] +) +def test_rank_changing_adapter_requires_actual_consumer_preparation( + operation, unary, entrypoint, metadata, monkeypatch +): + """Refuse unsupported reader ranks before changing an earlier eligible adapter.""" + model, feeds = _rank_sensitive_adapter_model(operation, unary, metadata) + _checked_run(model, feeds) + pipeline = import_module(_transform_make_dla_compatible.__module__) + unavailable = "Non4DTranspose" if operation == "Transpose" else "UpdateAxis" + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) for name, callback in pipeline.DLA_TRANSFORMS if name != unavailable + ), + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank"): + if entrypoint == "stage": + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax", "Transpose"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_changing_adapter_uses_registered_consumer_preparation(operation, unary): + """Execute preserved adapter ownership through the matching later axis stage.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + stage = "Non4DTranspose" if operation == "Transpose" else "UpdateAxis" + if not any(name == stage for name, _ in pipeline.DLA_TRANSFORMS): + pytest.skip("The matching rank stage is introduced by a later series part") + model, feeds = _rank_sensitive_adapter_model(operation, unary, True) + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" model = _make_model( From d175c7a7a9d21cebe619ac360d4c8d91e4ab6e31 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 14:52:48 +0000 Subject: [PATCH 23/32] Preserve safe normalization readers during adapter cleanup Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 22 +++++++++- .../test_dla_model_boundaries.py | 43 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 2768b5a25ef..a94e74d4ffa 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -502,6 +502,7 @@ def _apply_remove_intermediary_squeeze_unsqueeze( removed = set() local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} checked_readers = set() + contrib_ops = _boundary_contrib_ops(model) runtime_inputs = ( cache.graph_input_names - cache.init_names if model.ir_version <= 3 @@ -519,16 +520,33 @@ def require_prepared_readers(name): if current in captured - cache.graph_output_names: raise ValueError("Intermediary adapter changes a captured native rank") for reader, port in cache.get_consumers(current): - standard = _is_standard_node(reader, local_functions, {reader.op_type}) + standard = _is_standard_node( + reader, local_functions, {reader.op_type}, contrib_ops=contrib_ops + ) if ( standard and reader.op_type == "Squeeze" and reader.output[0] in cache.graph_output_names ): continue - if standard and reader.op_type in _RANK_AGNOSTIC_OPS | {"Squeeze", "Unsqueeze"}: + if standard and reader.op_type in _RANK_AGNOSTIC_OPS | { + "Gelu", + "Squeeze", + "Unsqueeze", + }: pending.extend(reader.output) continue + if standard and port == 0: + axis = _attribute(reader, "axis") + if reader.op_type == "LayerNormalization" and (axis is None or axis.i < 0): + pending.extend(reader.output) + continue + if ( + reader.op_type in {"QuantizeLinear", "DequantizeLinear"} + and cache.get_shape(reader.input[1]) == [] + ): + pending.extend(reader.output) + continue target = ( cache.get_init_array(reader.input[1]) if standard diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index dbe79805c00..588323e19dc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -19,6 +19,7 @@ import numpy as np import onnx +import onnxruntime as ort import pytest from onnx import TensorProto, helper, inliner, numpy_helper from onnx.reference import ReferenceEvaluator @@ -1257,6 +1258,48 @@ def test_rank_changing_adapter_uses_registered_consumer_preparation(operation, u np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) +@pytest.mark.parametrize("operation", ["LayerNormalization", "Gelu", "QuantizeLinear"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_changing_adapter_preserves_safe_normalization_readers(operation, unary): + """Execute suffix normalization and scalar quantization across leading singleton promotion.""" + model, feeds = _rank_sensitive_adapter_model("Softmax", unary, True) + model.ir_version = 9 + reader = model.graph.node[-1] + reader.op_type = operation + del reader.attribute[:] + if operation == "Gelu": + reader.domain = "com.microsoft" + model.opset_import.append(helper.make_opsetid("com.microsoft", 1)) + else: + reader.input.append("scale") + model.graph.initializer.append( + numpy_helper.from_array( + np.ones([4, 4], np.float32) + if operation == "LayerNormalization" + else np.asarray(0.5, np.float32), + "scale", + ) + ) + if operation == "LayerNormalization": + reader.attribute.append(helper.make_attribute("axis", -2)) + else: + model.graph.output[-1].type.tensor_type.elem_type = TensorProto.UINT8 + + def execute(value): + """Use the CPU runtime, including its registered contrib Gelu implementation.""" + onnx.checker.check_model(value, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + return ort.InferenceSession( + value.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + expected = execute(model) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(execute(converted), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" model = _make_model( From 3adfa831b3779abe96167b88ea4fbe7cd05d6912 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 14:58:46 +0000 Subject: [PATCH 24/32] Use a supported normalization scale in adapter execution coverage Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_model_boundaries.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 588323e19dc..e79d4d3503f 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1274,14 +1274,14 @@ def test_rank_changing_adapter_preserves_safe_normalization_readers(operation, u reader.input.append("scale") model.graph.initializer.append( numpy_helper.from_array( - np.ones([4, 4], np.float32) + np.ones([4], np.float32) if operation == "LayerNormalization" else np.asarray(0.5, np.float32), "scale", ) ) if operation == "LayerNormalization": - reader.attribute.append(helper.make_attribute("axis", -2)) + reader.attribute.append(helper.make_attribute("axis", -1)) else: model.graph.output[-1].type.tensor_type.elem_type = TensorProto.UINT8 From d304375cb9a8aed9e3fcfc7250129d14b9d9e5bc Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 15:38:23 +0000 Subject: [PATCH 25/32] Preserve matching rank restoration during adapter preflight Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 26 ++++++-- .../test_dla_model_boundaries.py | 59 +++++++++++++++++++ 2 files changed, 80 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index a94e74d4ffa..8d2be7b4cc6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -503,18 +503,25 @@ def _apply_remove_intermediary_squeeze_unsqueeze( local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} checked_readers = set() contrib_ops = _boundary_contrib_ops(model) + protected_inputs = ( + cache.graph_input_names | cache.init_names | cache.graph_output_names | captured + ) + protected_outputs = cache.graph_output_names | captured runtime_inputs = ( cache.graph_input_names - cache.init_names if model.ir_version <= 3 else cache.graph_input_names ) - def require_prepared_readers(name): + def require_prepared_readers(name, output_shape): """Follow shared rank-agnostic paths before allowing any adapter rewiring.""" + layout = tuple(pad4d(output_shape)) if output_shape else None + if layout is not None and len(layout) == 4 and cache.get_shape(name) == list(layout): + return pending, seen = [name], set() while pending: current = pending.pop() - if current in seen or current in checked_readers: + if current in seen or (current, layout) in checked_readers: continue seen.add(current) if current in captured - cache.graph_output_names: @@ -534,6 +541,15 @@ def require_prepared_readers(name): "Squeeze", "Unsqueeze", }: + if ( + reader.op_type in {"Squeeze", "Unsqueeze"} + and layout is not None + and len(layout) == 4 + and cache.get_shape(reader.output[0]) == list(layout) + and reader.input[0] not in protected_inputs + and reader.output[0] not in protected_outputs + ): + continue pending.extend(reader.output) continue if standard and port == 0: @@ -566,7 +582,7 @@ def require_prepared_readers(name): raise ValueError( f"Intermediary adapter feeds unprepared rank-sensitive {reader.op_type}" ) - checked_readers.update(seen) + checked_readers.update((current, layout) for current in seen) plans = [] for node in list(graph.node): @@ -582,12 +598,12 @@ def require_prepared_readers(name): output_shape = cache.get_shape(node.output[0]) if not input_shape or not output_shape: if validate_only and rank_converted: - require_prepared_readers(node.output[0]) + require_prepared_readers(node.output[0], output_shape) continue if not rank_converted and input_shape != output_shape: continue if rank_converted and len(input_shape) != len(output_shape): - require_prepared_readers(node.output[0]) + require_prepared_readers(node.output[0], output_shape) plans.append((node, input_shape, output_shape)) if validate_only: return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index e79d4d3503f..15005f8597a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -31,6 +31,7 @@ _apply_remove_intermediary_squeeze_unsqueeze, ) from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes @pytest.mark.parametrize("identity", ["custom", "local", "overload"]) @@ -1300,6 +1301,64 @@ def execute(value): np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) +@pytest.mark.parametrize("unary", [False, True]) +@pytest.mark.parametrize("same_layout", [False, True]) +@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"]) +def test_adapter_roundtrip_checks_the_restored_layout(unary, same_layout, entrypoint): + """Accept matching rank restoration while retaining refusal for a different bypass layout.""" + nodes = [ + helper.make_node("Relu", ["x"], ["private"]), + helper.make_node("Squeeze", ["private", "squeeze_axes"], ["squeezed"]), + ] + current = "squeezed" + if unary: + nodes.append(helper.make_node("Relu", [current], ["middle"])) + current = "middle" + nodes.extend( + [ + helper.make_node("Unsqueeze", [current, "unsqueeze_axes"], ["restored"]), + helper.make_node("Conv", ["restored", "weight"], ["y"]), + ] + ) + channels = 3 if same_layout else 1 + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1 if same_layout else 3, 2, 4, 4])], + [ + numpy_helper.from_array(np.array([0], np.int64), "squeeze_axes"), + numpy_helper.from_array( + np.array([0 if same_layout else 1], np.int64), "unsqueeze_axes" + ), + numpy_helper.from_array( + np.arange(2 * channels, dtype=np.float32).reshape(2, channels, 1, 1), "weight" + ), + ], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)} + expected = _checked_run(model, feeds) + before = model.SerializeToString() + if not same_layout: + with pytest.raises(ValueError, match="unprepared rank-sensitive Conv"): + if entrypoint == "stage": + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + return + converted = ( + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + if entrypoint == "stage" + else _transform_make_dla_compatible(model, verbose=False) + ) + if entrypoint == "stage": + converted = infer_shapes(converted) + assert not any(node.op_type in {"Squeeze", "Unsqueeze"} for node in converted.graph.node) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" model = _make_model( From aa6781b1a6cfd59d4dd4890d1c5681867881f52c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 15:55:02 +0000 Subject: [PATCH 26/32] Validate the actual planned adapter bypass layout Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 21 ++++++++-------- .../test_dla_model_boundaries.py | 24 +++++++++++++++++++ 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 8d2be7b4cc6..42e4eee924c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -513,9 +513,9 @@ def _apply_remove_intermediary_squeeze_unsqueeze( else cache.graph_input_names ) - def require_prepared_readers(name, output_shape): + def require_prepared_readers(name, planned_shape): """Follow shared rank-agnostic paths before allowing any adapter rewiring.""" - layout = tuple(pad4d(output_shape)) if output_shape else None + layout = tuple(planned_shape) if planned_shape else None if layout is not None and len(layout) == 4 and cache.get_shape(name) == list(layout): return pending, seen = [name], set() @@ -598,22 +598,23 @@ def require_prepared_readers(name, output_shape): output_shape = cache.get_shape(node.output[0]) if not input_shape or not output_shape: if validate_only and rank_converted: - require_prepared_readers(node.output[0], output_shape) + require_prepared_readers(node.output[0], None) continue if not rank_converted and input_shape != output_shape: continue - if rank_converted and len(input_shape) != len(output_shape): - require_prepared_readers(node.output[0], output_shape) - plans.append((node, input_shape, output_shape)) - if validate_only: - return model - for node, input_shape, output_shape in plans: short, long = ( (output_shape, input_shape) if node.op_type == "Squeeze" else (input_shape, output_shape) ) - if len(short) <= len(long) and long[-len(short) :] == short: + bypass = len(short) <= len(long) and long[-len(short) :] == short + if rank_converted and len(input_shape) != len(output_shape): + require_prepared_readers(node.output[0], pad4d(input_shape if bypass else output_shape)) + plans.append((node, output_shape, bypass)) + if validate_only: + return model + for node, output_shape, bypass in plans: + if bypass: old, new = node.input[0], node.output[0] for other in graph.node: if other is node: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 15005f8597a..d2d684df54b 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1359,6 +1359,30 @@ def test_adapter_roundtrip_checks_the_restored_layout(unary, same_layout, entryp np.testing.assert_array_equal(actual, reference) +def test_rank_restoration_requires_an_actual_four_dimensional_bypass(): + """Do not infer a 4D bypass solely from the original adapter output annotation.""" + model = _make_model( + [ + helper.make_node("Reshape", ["x", "target"], ["wide"]), + helper.make_node("Squeeze", ["wide", "axes"], ["restored"]), + helper.make_node("Conv", ["restored", "weight"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 4, 4])], + [ + numpy_helper.from_array(np.array([1, 1, 3, 4, 4], np.int64), "target"), + numpy_helper.from_array(np.array([0], np.int64), "axes"), + numpy_helper.from_array(np.ones([2, 3, 1, 1], np.float32), "weight"), + ], + ) + model = onnx.shape_inference.infer_shapes(model) + _checked_run(model, {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="unprepared rank-sensitive Conv"): + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + assert model.SerializeToString() == before + + def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" model = _make_model( From 915afff95f177364ea52deae38d8ae4ffd8e9912 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:51:35 +0000 Subject: [PATCH 27/32] Validate normalization consumers before changing model layouts Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 12 +++- .../onnx/graph_surgery/make_dla_compatible.py | 4 ++ .../test_dla_model_boundaries.py | 62 ++++++++++++++++++- 3 files changed, 75 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 42e4eee924c..edce3e24dac 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -649,7 +649,7 @@ def require_prepared_readers(name, planned_shape): return model -def _apply_fix_instancenorm_channels(model, *, prepared_consumer_ops=()): +def _apply_fix_instancenorm_channels(model, *, prepared_consumer_ops=(), validate_only=False): """Append spatial singleton dimensions while retaining normalization channels and outputs.""" graph = model.graph cache = GraphCache(graph) @@ -765,6 +765,16 @@ def output_consumers(name): shape = _shape_before_boundary_promotion(cache, node.input[0]) if shape is not None and len(shape) == 3: plans.append((index, node, shape, list(output_consumers(node.output[0])))) + concat_inputs = {} + for _, _, _, consumers in plans: + for consumer, _, port, native in consumers: + if consumer.op_type == "Concat" and not native: + concat_inputs.setdefault(id(consumer), (consumer, set()))[1].add(port) + for consumer, ports in concat_inputs.values(): + if ports != set(range(len(consumer.input))): + raise ValueError("Unsupported normalization Concat with unprepared operands") + if validate_only: + return model bridges = {} for index, node, shape, consumers in plans: target = add_initializer( diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index afc3bab16e0..9324c0451c3 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -145,6 +145,10 @@ def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): ): raise NotImplementedError("Standalone reductions are not available in this pipeline") _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache) + if "FixInstanceNormChannelMismatch" in prepared_input_ops: + _apply_fix_instancenorm_channels( + model, prepared_consumer_ops=prepared_input_ops, validate_only=True + ) _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) _apply_non4d_model_outputs( model, diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index d2d684df54b..a342f3cbfdb 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -370,7 +370,8 @@ def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) -def test_normalization_refuses_unprepared_consumers_before_mutation(): +@pytest.mark.parametrize("through_pipeline", [False, True]) +def test_normalization_refuses_unprepared_consumers_before_mutation(through_pipeline): """Refuse a rank-sensitive consumer before widening any eligible normalization node.""" shape = [2, 3, 4] model = _make_model( @@ -388,7 +389,10 @@ def test_normalization_refuses_unprepared_consumers_before_mutation(): _checked_run(model, {"x": np.arange(24, dtype=np.float32).reshape(shape)}) before = model.SerializeToString() with pytest.raises(ValueError, match="Unsupported normalization consumer"): - _apply_fix_instancenorm_channels(model) + if through_pipeline: + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_fix_instancenorm_channels(model) assert model.SerializeToString() == before @@ -1490,3 +1494,57 @@ def test_native_rank_reducing_outputs_require_registered_conversion( with pytest.raises(ValueError, match=r"rank.*conversion"): _apply_non4d_model_outputs(model, defer_rank_conversion=True) assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("other_normalized", "through_pipeline"), [(False, False), (False, True), (True, True)] +) +@pytest.mark.parametrize("unary", [False, True]) +def test_normalization_concat_requires_prepared_operands(other_normalized, through_pipeline, unary): + """Widen Concat only when every operand receives a normalization bridge.""" + shape = [2, 3, 4] + nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"])] + controls = [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + numpy_helper.from_array(np.ones([2, 5, 4], np.float32), "other"), + ] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + right = "other" + if other_normalized: + controls.extend( + [ + numpy_helper.from_array(np.ones(5, np.float32), "other_scale"), + numpy_helper.from_array(np.zeros(5, np.float32), "other_bias"), + ] + ) + nodes.append( + helper.make_node( + "InstanceNormalization", + ["other", "other_scale", "other_bias"], + ["other_normalized"], + ) + ) + right = "other_normalized" + left = "normalized" + if unary: + nodes.append(helper.make_node("Relu", [left], ["activated"])) + left = "activated" + nodes.append(helper.make_node("Concat", [left, right], ["y"], axis=1)) + model = _make_model( + nodes, inputs, [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 8, 4])], controls + ) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + expected = _checked_run(model, feeds) + before = model.SerializeToString() + if other_normalized: + result = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(result, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + else: + with pytest.raises(ValueError, match="normalization Concat"): + if through_pipeline: + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_fix_instancenorm_channels(model) + assert model.SerializeToString() == before From 6e06a223eff66683d4fde2f5f1d65cee02935aab Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:59:52 +0000 Subject: [PATCH 28/32] Preserve native broadcast output layouts when rank conversion is absent Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 21 ++++++++++++ .../test_dla_model_boundaries.py | 33 +++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index edce3e24dac..516ae248c4c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -377,12 +377,33 @@ def _apply_non4d_model_outputs( append = [] plans = [] local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_broadcasts = set() + if defer_rank_conversion and "Non4dInitializers" not in prepared_output_ops: + live_ranks = { + name: len(shape) + for name in boundary_names + if (shape := cache.get_shape(name)) is not None + } + for node in graph.node: + if ( + _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) + and node.input + and all(name in live_ranks for name in node.input) + ): + rank = max(live_ranks[name] for name in node.input) + live_ranks.update(dict.fromkeys(node.output, rank)) + if rank < 4 and ( + len(node.input) > 1 or any(name in native_broadcasts for name in node.input) + ): + native_broadcasts.update(node.output) rank_sources = {} contrib_ops = _boundary_contrib_ops(model) for value in graph.output: shape = cache.get_shape(value.name) if shape is None or len(shape) >= 4 or value.name in boundary_names: continue + if value.name in native_broadcasts: + continue producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if source is None and len(shape) == 1: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index a342f3cbfdb..d79b0e7416d 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -34,6 +34,39 @@ from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("input_port", [0, 1]) +@pytest.mark.parametrize("unary", [False, True]) +def test_broadcast_output_restoration_requires_a_promoted_source(shape, input_port, unary): + """Leave native broadcasts intact when no operand receives leading dimensions.""" + operands = ["x", "weight"] if input_port == 0 else ["weight", "x"] + nodes = [helper.make_node("Add", operands, ["middle" if unary else "y"])] + if unary: + nodes.append(helper.make_node("Relu", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_broadcast_output", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(np.ones(shape, np.float32), "weight")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + onnx.checker.check_model(model, full_check=True) + expected = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + transformed = _transform_make_dla_compatible(model, verbose=False) + onnx.checker.check_model(transformed, full_check=True) + actual = ort.InferenceSession( + transformed.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + np.testing.assert_array_equal(actual[0], expected[0]) + + @pytest.mark.parametrize("identity", ["custom", "local", "overload"]) @pytest.mark.parametrize( ("operation", "entrypoint"), From 2348f8d560e845cb4036f009a839497cab80aaea Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 09:16:15 +0000 Subject: [PATCH 29/32] Validate output refusals before caller mutation Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 16 +++++-- .../onnx/graph_surgery/make_dla_compatible.py | 3 ++ .../test_dla_model_boundaries.py | 48 +++++++++++++++++++ 3 files changed, 64 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 516ae248c4c..00a405db94a 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -32,6 +32,11 @@ __all__ = [] + +class _BoundaryPreflightError(ValueError): + """Prerequisite lowering cannot freeze an exposed runtime control.""" + + _RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} _REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) @@ -431,8 +436,6 @@ def _apply_non4d_model_outputs( and "Non4DReshape" not in prepared_output_ops ): raise ValueError("Output rank conversion is unavailable for native Reshape") - if validate_only: - continue if ( producer is not None and _is_standard_node(producer, local_functions, {"Identity"}) @@ -451,11 +454,18 @@ def _apply_non4d_model_outputs( if len(existing.input) > 1 and existing.input[1]: initializer = cache.get_init(existing.input[1]) if initializer is None or initializer.name in cache.graph_input_names: - raise ValueError(f"Output Squeeze axes must be fixed for {value.name!r}") + error = ( + _BoundaryPreflightError + if existing.input[1] in cache.graph_input_names + else ValueError + ) + raise error(f"Output Squeeze axes must be fixed for {value.name!r}") axes = numpy_helper.to_array(initializer).reshape(-1).tolist() else: input_shape = cache.get_shape(existing.input[0]) if input_shape is None: + if validate_only: + continue raise ValueError(f"Output Squeeze input shape must be known for {value.name!r}") axes = [index for index, dim in enumerate(input_shape) if dim == 1] source_shape = cache.get_shape(existing.input[0] if existing is not None else value.name) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 9324c0451c3..23923ff6472 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -44,6 +44,7 @@ _apply_non4d_model_inputs, _apply_non4d_model_outputs, _apply_remove_intermediary_squeeze_unsqueeze, + _BoundaryPreflightError, _require_output_rank_conversion, ) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin @@ -212,6 +213,8 @@ def _transform_make_dla_compatible( quant_types = frozenset(onnx.TensorProto.DataType.values()) try: _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops) + except _BoundaryPreflightError: + raise except (ValueError, NotImplementedError): # Earlier lowering may consume native operators that lack a later rank pass. probe = infer_shapes(copy.deepcopy(model)) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index d79b0e7416d..dbeb5b337c6 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -15,6 +15,7 @@ """Boundary promotion refuses unsupported exposed adapters before changing the graph.""" +import copy from importlib import import_module import numpy as np @@ -34,6 +35,53 @@ from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes +@pytest.mark.parametrize("entrypoint", ["validate", "pipeline"]) +@pytest.mark.parametrize("case", ["rank5", "scalar", "runtime_axes"]) +def test_output_refusal_preflight_preserves_caller_state(entrypoint, case, monkeypatch): + """Validate known output refusals before inference or an earlier stage changes the caller.""" + shape, axes, output_shape = { + "rank5": ([1, 2, 1, 3, 4], [0, 2], [2, 3, 4]), + "scalar": ([1], [0], []), + "runtime_axes": ([1, 2, 3, 4], [0], [2, 3, 4]), + }[case] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + if case == "runtime_axes": + inputs.append(helper.make_tensor_value_info("axes", TensorProto.INT64, [len(axes)])) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Squeeze", ["x", "axes"], ["y"])], + "output_refusal", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]).run( + None, {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + ) + before = model.SerializeToString() + original = copy.deepcopy + model_copies = [] + + def track(value, *args, **kwargs): + if isinstance(value, onnx.ModelProto): + model_copies.append(value) + return original(value, *args, **kwargs) + + monkeypatch.setattr(copy, "deepcopy", track) + with pytest.raises(ValueError, match=r"scalar|exceeds 4|must be fixed"): + if entrypoint == "pipeline": + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_non4d_model_outputs(model, defer_rank_conversion=True, validate_only=True) + assert model.SerializeToString() == before + if case == "runtime_axes": + assert not model_copies + + @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) @pytest.mark.parametrize("input_port", [0, 1]) @pytest.mark.parametrize("unary", [False, True]) From fb82039b0e75fcf52a85f8367ef89eb7a87b3301 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 10:50:44 +0000 Subject: [PATCH 30/32] Preflight inferable normalization inputs before caller mutation Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 7 +++++++ .../test_dla_model_boundaries.py | 19 +++++++++++++------ 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 00a405db94a..716040825ef 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -794,6 +794,13 @@ def output_consumers(name): for index, node in enumerate(graph.node): if _is_standard_node(node, local_functions, {"InstanceNormalization"}): shape = _shape_before_boundary_promotion(cache, node.input[0]) + if ( + validate_only + and shape is None + and node.input[0] not in cache.graph_input_names + and cache.get_init(node.input[0]) is None + ): + raise ValueError("Normalization preflight requires inferred intermediate shapes") if shape is not None and len(shape) == 3: plans.append((index, node, shape, list(output_consumers(node.output[0])))) concat_inputs = {} diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index dbeb5b337c6..7d39b6be636 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -451,13 +451,17 @@ def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) -@pytest.mark.parametrize("through_pipeline", [False, True]) -def test_normalization_refuses_unprepared_consumers_before_mutation(through_pipeline): +@pytest.mark.parametrize( + ("through_pipeline", "intermediate"), [(False, False), (True, False), (True, True)] +) +def test_normalization_refuses_unprepared_consumers_before_mutation(through_pipeline, intermediate): """Refuse a rank-sensitive consumer before widening any eligible normalization node.""" shape = [2, 3, 4] + data = "r" if intermediate else "x" model = _make_model( - [ - helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"]), + ([helper.make_node("Relu", ["x"], [data])] if intermediate else []) + + [ + helper.make_node("InstanceNormalization", [data, "scale", "bias"], ["normalized"]), helper.make_node("Hardmax", ["normalized"], ["y"], axis=1), ], [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], @@ -1160,10 +1164,13 @@ def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_ @pytest.mark.parametrize("batch", [1, 2, 3]) @pytest.mark.parametrize("activations", [(), ("Relu",), ("Relu", "Tanh")]) def test_registered_normalization_preserves_native_channel_axis(batch, activations): - """Keep direct rank-three normalization inputs on the channel-preserving route.""" + """Preserve channels with direct inputs or inferable intermediate metadata.""" shape = [batch, 3, 4] output = "normalized" if activations else "y" - nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], [output])] + data = "activated_input" if activations else "x" + nodes = ([helper.make_node("Relu", ["x"], [data])] if activations else []) + [ + helper.make_node("InstanceNormalization", [data, "scale", "bias"], [output]) + ] for index, operation in enumerate(activations): following = "y" if index == len(activations) - 1 else f"activation_{index}" nodes.append(helper.make_node(operation, [output], [following])) From 0109428bc56b3fafd4870745e87b22e65c989b87 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 11:02:28 +0000 Subject: [PATCH 31/32] Retain symbolic normalization shapes during missing-metadata preflight Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_model_boundaries.py | 2 ++ .../dla_transforms/test_dla_model_boundaries.py | 5 +++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 716040825ef..13981288ed7 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -794,9 +794,11 @@ def output_consumers(name): for index, node in enumerate(graph.node): if _is_standard_node(node, local_functions, {"InstanceNormalization"}): shape = _shape_before_boundary_promotion(cache, node.input[0]) + input_info = cache.vi_map.get(node.input[0]) if ( validate_only and shape is None + and (input_info is None or not input_info.type.tensor_type.HasField("shape")) and node.input[0] not in cache.graph_input_names and cache.get_init(node.input[0]) is None ): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index 7d39b6be636..eb50d4ac3fe 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -1161,7 +1161,7 @@ def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_ np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) -@pytest.mark.parametrize("batch", [1, 2, 3]) +@pytest.mark.parametrize("batch", [1, 2, "batch"]) @pytest.mark.parametrize("activations", [(), ("Relu",), ("Relu", "Tanh")]) def test_registered_normalization_preserves_native_channel_axis(batch, activations): """Preserve channels with direct inputs or inferable intermediate metadata.""" @@ -1184,7 +1184,8 @@ def test_registered_normalization_preserves_native_channel_axis(batch, activatio numpy_helper.from_array(np.asarray([-1, 0, 1], np.float32), "bias"), ], ) - feeds = {"x": np.random.default_rng(19).standard_normal(shape).astype(np.float32)} + feed_shape = [2 if isinstance(batch, str) else batch, 3, 4] + feeds = {"x": np.random.default_rng(19).standard_normal(feed_shape).astype(np.float32)} expected = _checked_run(model, feeds)[0] transformed = _transform_make_dla_compatible(model, verbose=False) np.testing.assert_allclose(_checked_run(transformed, feeds)[0], expected, rtol=1e-5, atol=1e-6) From 3daa82b52d4875822939a0327df00cca0ddf8207 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 15:32:43 +0000 Subject: [PATCH 32/32] Limit normalization inference refusal to the initial preflight Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 5 ++- .../onnx/graph_surgery/make_dla_compatible.py | 13 +++++-- .../test_dla_model_boundaries.py | 35 +++++++++++++++++++ 3 files changed, 49 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index 13981288ed7..64a34411bb3 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -680,7 +680,9 @@ def require_prepared_readers(name, planned_shape): return model -def _apply_fix_instancenorm_channels(model, *, prepared_consumer_ops=(), validate_only=False): +def _apply_fix_instancenorm_channels( + model, *, prepared_consumer_ops=(), validate_only=False, require_inferred_shapes=False +): """Append spatial singleton dimensions while retaining normalization channels and outputs.""" graph = model.graph cache = GraphCache(graph) @@ -797,6 +799,7 @@ def output_consumers(name): input_info = cache.vi_map.get(node.input[0]) if ( validate_only + and require_inferred_shapes and shape is None and (input_info is None or not input_info.type.tensor_type.HasField("shape")) and node.input[0] not in cache.graph_input_names diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 23923ff6472..478c818b673 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -131,7 +131,9 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep return {} -def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): +def _validate_boundary_contract( + model, prepared_input_ops, prepared_output_ops, *, require_inferred_shapes=False +): """Refuse unsupported native layouts before adapting any public boundary.""" local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} cache = GraphCache(model.graph) @@ -148,7 +150,10 @@ def _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops): _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache) if "FixInstanceNormChannelMismatch" in prepared_input_ops: _apply_fix_instancenorm_channels( - model, prepared_consumer_ops=prepared_input_ops, validate_only=True + model, + prepared_consumer_ops=prepared_input_ops, + validate_only=True, + require_inferred_shapes=require_inferred_shapes, ) _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) _apply_non4d_model_outputs( @@ -212,7 +217,9 @@ def _transform_make_dla_compatible( if remove_all_qdq: quant_types = frozenset(onnx.TensorProto.DataType.values()) try: - _validate_boundary_contract(model, prepared_input_ops, prepared_output_ops) + _validate_boundary_contract( + model, prepared_input_ops, prepared_output_ops, require_inferred_shapes=True + ) except _BoundaryPreflightError: raise except (ValueError, NotImplementedError): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py index eb50d4ac3fe..1c10edaa960 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -481,6 +481,41 @@ def test_normalization_refuses_unprepared_consumers_before_mutation(through_pipe assert model.SerializeToString() == before +def test_normalization_probe_accepts_unknown_sequence_element_rank(): + """A probe may retain valid native normalization when inference cannot recover its rank.""" + shape = [1, 1, 1, 4] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("SequenceAt", ["sequence", "index"], ["r"]), + helper.make_node("InstanceNormalization", ["r", "scale", "bias"], ["unused"]), + helper.make_node("Identity", ["x"], ["y"]), + ], + "unknown_native_rank", + [ + helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, None), + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.array(0, np.int64), "index"), + numpy_helper.from_array(np.ones(1, np.float32), "scale"), + numpy_helper.from_array(np.zeros(1, np.float32), "bias"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + values = np.arange(4, dtype=np.float32).reshape(shape) + feeds = {"sequence": [values], "x": values} + inferred = infer_shapes(copy.deepcopy(model)) + metadata = next(value for value in inferred.graph.value_info if value.name == "r") + assert not metadata.type.tensor_type.HasField("shape") + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected[0]) + + @pytest.mark.parametrize("batch", [1, 2, 3]) @pytest.mark.parametrize("operation", ["Add", "Mul"]) @pytest.mark.parametrize("tail", [False, True])