From 71edb7bc3852427518fa821bd23c80392e6ac514 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:16 +0000 Subject: [PATCH 01/41] [DLA_transforms_12] [12/17] DLA rank conversion Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 305 ++++++++++++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 22 ++ .../test_dla_rank_conversion.py | 78 +++++ 3 files changed, 405 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py new file mode 100644 index 00000000000..6f87cc0a529 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -0,0 +1,305 @@ +# 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. + +"""Individual rank and axis passes in the logged DLA transform sequence. + +These passes share the shape snapshot produced by the preceding shape-inference +stage. Running inference between passes would change how original axes are +mapped after upstream tensors have been promoted to four dimensions. +""" + +import math + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +from ..utils.graph_utils import ( + GraphCache, + _attribute, + _shape_before_boundary_promotion, + add_initializer, + pad4d, +) + +__all__ = [] + + +def _apply_non4d_slice_axis(model): + cache = GraphCache(model.graph) + for node in model.graph.node: + if node.op_type != "Slice" or len(node.input) < 4: + continue + shape = cache.get_shape(node.input[0]) + axes = cache.get_init(node.input[3]) + if shape is None or len(shape) not in (2, 3) or axes is None: + continue + values = numpy_helper.to_array(axes).reshape(-1) + if values.size and np.any(values >= 0): + node.input[3] = add_initializer( + model.graph, + axes.name + "_axes_only", + np.asarray(np.where(values >= 0, values + 4 - len(shape), values), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + return model + + +def _apply_non4d_reshape(model): + cache = GraphCache(model.graph) + replacements = {} + for node in model.graph.node: + if node.op_type != "Reshape" or len(node.input) < 2: + continue + tensor = cache.get_init(node.input[1]) + if tensor is None: + continue + values = numpy_helper.to_array(tensor).reshape(-1).tolist() + if len(values) < 4: + values = pad4d(values) + output_shape = cache.get_shape(node.output[0]) + if output_shape: + output_shape = pad4d(output_shape) + values = [ + output_shape[i] if value == 0 else value for i, value in enumerate(values) + ] + elif len(values) == 5: + if 1 not in values: + raise ValueError(f"{node.name!r}: 5D Reshape shape has no unary dimension") + values.remove(1) + else: + continue + replacement = numpy_helper.from_array(np.asarray(values, dtype=np.int64), tensor.name) + if tensor.name in replacements and replacements[tensor.name] != replacement: + raise ValueError(f"Shared Reshape target {tensor.name!r} requires different shapes") + replacements[tensor.name] = replacement + for tensor in model.graph.initializer: + if tensor.name in replacements: + tensor.CopyFrom(replacements[tensor.name]) + return model + + +def _apply_non4d_initializers(model): + expand, skip, front, avoid = set(), set(), set(), set() + for node in model.graph.node: + if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: + expand.update(node.input) + if node.op_type == "MatMul": + skip.update(node.input) + if node.op_type in {"Gemm", "Where", "Gather", "Split"}: + front.update(node.input) + if node.op_type in {"Unsqueeze", "GroupQueryAttention"}: + avoid.update(node.input) + updates = [] + expand_or_front = expand | front + for tensor in model.graph.initializer: + if tensor.name in avoid: + continue + rank = len(tensor.dims) + if ( + (rank == 1 and tensor.name in expand) + or (rank == 2 and tensor.name in expand_or_front) + or rank == 3 + ): + shape = pad4d(list(tensor.dims)) + updates.append((tensor, shape)) + elif rank == 2 and tensor.name not in skip: + raise ValueError(f"Initializer {tensor.name!r} has unsupported rank 2 usage") + for tensor, shape in updates: + del tensor.dims[:] + tensor.dims.extend(shape) + return model + + +def _apply_non4d_expand(model): + return _pad_shape_inputs(model, "Expand", "_expand_shape") + + +def _apply_non4d_tile(model): + return _pad_shape_inputs(model, "Tile", "_tile_shape") + + +def _pad_shape_inputs(model, op_type, suffix): + cache = GraphCache(model.graph) + for node in model.graph.node: + if node.op_type != op_type or len(node.input) < 2: + continue + tensor = cache.get_init(node.input[1]) + if tensor is None: + continue + values = numpy_helper.to_array(tensor).reshape(-1).tolist() + if len(values) < 4: + node.input[1] = add_initializer( + model.graph, + tensor.name + suffix, + np.asarray(pad4d(values), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + elif len(values) > 4: + raise ValueError(f"{node.name!r}: {op_type} target has rank {len(values)}") + return model + + +def _apply_non4d_transpose(model): + cache = GraphCache(model.graph) + for node in model.graph.node: + if node.op_type != "Transpose": + continue + perm = _attribute(node, "perm") + if perm is None: + continue + values = list(perm.ints) + if len(values) in (2, 3): + delta = 4 - len(values) + values = list(range(delta)) + [v + delta for v in values] + elif len(values) == 5: + shape = cache.get_shape(node.input[0]) + if shape is None or len(shape) != 5 or 1 not in shape: + raise ValueError(f"{node.name!r}: 5D Transpose requires a known unary dimension") + removed = shape.index(1) + values = [v - int(v > removed) for v in values if v != removed] + else: + continue + del perm.ints[:] + perm.ints.extend(values) + return model + + +def _apply_non4d_lpnorm(model): + cache = GraphCache(model.graph) + for node in model.graph.node: + if node.op_type == "LpNormalization": + axis = _attribute(node, "axis") + shape = _shape_before_boundary_promotion(cache, node.input[0]) + if axis is not None and axis.i >= 0 and shape and len(shape) in (2, 3): + axis.i = -1 if axis.i == len(shape) - 1 else axis.i + 4 - len(shape) + elif axis is None: + node.attribute.append(helper.make_attribute("axis", -1)) + node.name += "_axis_updated" + return model + + +def _apply_flatten(model): + graph = model.graph + cache = GraphCache(graph) + for node in graph.node: + if node.op_type == "Flatten": + shape = _shape_before_boundary_promotion(cache, node.input[0]) + axis = _attribute(node, "axis") + axis = axis.i if axis is not None else 1 + if shape is None or any(d <= 0 for d in shape): + raise ValueError(f"{node.name!r}: Flatten requires a known input shape") + if axis < 0: + axis += len(shape) + batch = math.prod(shape[:axis]) + name = add_initializer( + graph, + "flatten_reshape_shape" if batch == 1 else f"{node.name}_reshape_shape", + np.asarray([1, 1, batch, -1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.CopyFrom( + helper.make_node( + "Reshape", [node.input[0], name], list(node.output), name=node.name + "_reshape" + ) + ) + return model + + +def _apply_argmax(model): + cache = GraphCache(model.graph) + nodes = [] + for node in model.graph.node: + nodes.append(node) + if node.op_type != "ArgMax": + continue + shape = _shape_before_boundary_promotion(cache, node.input[0]) + if not shape: + continue + axis = _attribute(node, "axis") + if axis is None: + node.attribute.append(helper.make_attribute("axis", 0)) + axis = _attribute(node, "axis") + if axis.i >= 0 and len(shape) in (2, 3): + axis.i += 4 - len(shape) + keepdims = _attribute(node, "keepdims") + restore_shape = keepdims is not None and keepdims.i == 0 + if keepdims is not None: + keepdims.i = 1 + node.name += "_axis_updated" + if restore_shape: + original = node.output[0] + output_shape = cache.get_shape(original) + if output_shape is None and original.endswith("_squeeze_input"): + output_shape = cache.get_shape(original.removesuffix("_squeeze_input")) + if output_shape is None: + raise ValueError(f"Unknown original ArgMax output shape for {node.name!r}") + target = add_initializer( + model.graph, + node.name + "_reshape_shape", + np.asarray(pad4d(output_shape), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.output[0] = original + "_keepdims" + nodes.append( + helper.make_node( + "Reshape", [node.output[0], target], [original], name=node.name + "_reshape" + ) + ) + del model.graph.node[:] + model.graph.node.extend(nodes) + return model + + +def _apply_update_axis(model): + update = { + "LogSoftmax", + "Softmax", + "DequantizeLinear", + "TopK", + "Concat", + "Split", + "QuantizeLinear", + } + skip = { + "Gather", + "GatherElements", + "LpNormalization", + "ArgMax", + "ReduceMax", + "ReduceMin", + "ReduceSum", + "ReduceMean", + "Slice", + "Flatten", + } + cache = GraphCache(model.graph) + for node in model.graph.node: + axis = _attribute(node, "axis") + if axis is None or axis.i < 0 or node.op_type in skip: + continue + if node.op_type not in update: + raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") + if "axis_updated" in node.name: + continue + shape = _shape_before_boundary_promotion(cache, node.input[0]) + delta = {2: 2, 3: 1, 5: -1}.get(len(shape) if shape else 0, 0) + if delta: + axis.i += delta + node.name += "_axis_updated" + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index e9c26b59172..1a05a3295df 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -53,6 +53,18 @@ _BoundaryPreflightError, _require_output_rank_conversion, ) +from .dla_transforms.dla_rank_conversion import ( + _apply_argmax, + _apply_flatten, + _apply_non4d_expand, + _apply_non4d_initializers, + _apply_non4d_lpnorm, + _apply_non4d_reshape, + _apply_non4d_slice_axis, + _apply_non4d_tile, + _apply_non4d_transpose, + _apply_update_axis, +) 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 @@ -99,6 +111,7 @@ def _apply_constants_to_initializers(model, **options): ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), ("ReshapeReduceSum", _apply_reshape_reduce_sum), + ("Non4dSliceAxis", _apply_non4d_slice_axis), ("Gather", _apply_gather), ("DecomposeLstm", partial(_apply_decompose_lstm, defer_rank_conversion=True)), ( @@ -107,6 +120,8 @@ def _apply_constants_to_initializers(model, **options): ), ("Non4dModelInputs", _apply_non4d_model_inputs), ("Non4dModelOutputs", partial(_apply_non4d_model_outputs, defer_rank_conversion=True)), + ("Non4DReshape", _apply_non4d_reshape), + ("ArgMax", _apply_argmax), ("StandaloneReduceOperations", _apply_standalone_reduce_operations), ( "RemoveIntermediarySqueezeAndUnsqueeze", @@ -114,7 +129,14 @@ def _apply_constants_to_initializers(model, **options): ), ("GroupQueryAttention", _apply_group_query_attention), ("GatherElements", _apply_gather_elements), + ("Non4dInitializers", _apply_non4d_initializers), + ("Non4DExpand", _apply_non4d_expand), + ("Non4DTile", _apply_non4d_tile), + ("Flatten", _apply_flatten), + ("Non4DTranspose", _apply_non4d_transpose), + ("Non4DLpNorm", _apply_non4d_lpnorm), ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), + ("UpdateAxis", _apply_update_axis), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py new file mode 100644 index 00000000000..5ffc69e0211 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -0,0 +1,78 @@ +# 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. + +"""Numerical checks for independently reviewable rank and boundary passes.""" + +import copy + +import numpy as np +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import ( + _apply_non4d_expand, + _apply_non4d_reshape, + _apply_non4d_tile, + _apply_non4d_transpose, +) + + +@pytest.mark.parametrize("operation", ["Transpose", "Reshape", "Tile", "Expand"]) +def test_rank_conversion_preserves_operator_values(operation): + shape = [1, 3] if operation == "Expand" else [2, 3] + inputs, initializers, attributes = ["x"], [], {} + if operation == "Transpose": + attributes["perm"] = [1, 0] + output_shape = [3, 2] + transform = _apply_non4d_transpose + else: + values = {"Reshape": [3, 2], "Tile": [2, 1], "Expand": [2, 3]}[operation] + output_shape = {"Reshape": [3, 2], "Tile": [4, 3], "Expand": [2, 3]}[operation] + inputs.append("parameters") + initializers.append(numpy_helper.from_array(np.array(values, dtype=np.int64), "parameters")) + transform = { + "Reshape": _apply_non4d_reshape, + "Tile": _apply_non4d_tile, + "Expand": _apply_non4d_expand, + }[operation] + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"], name="operator", **attributes)], + "rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + transformed = copy.deepcopy(model) + for apply in (_apply_non4d_model_inputs, _apply_non4d_model_outputs, transform): + transformed = apply(transformed) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + expected = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data}) + actual = ort.InferenceSession( + transformed.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data}) + np.testing.assert_array_equal(actual[0], expected[0]) From 653073fe8007018f166fb72ef14ae389ea073f42 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:24:45 +0000 Subject: [PATCH 02/41] Declare deferred output ranks in rank-conversion fixtures Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_rank_conversion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 5ffc69e0211..84196e6b5b0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -36,6 +36,7 @@ @pytest.mark.parametrize("operation", ["Transpose", "Reshape", "Tile", "Expand"]) def test_rank_conversion_preserves_operator_values(operation): + """Execute native operators across explicitly deferred boundary and rank stages.""" shape = [1, 3] if operation == "Expand" else [2, 3] inputs, initializers, attributes = ["x"], [], {} if operation == "Transpose": @@ -65,7 +66,8 @@ def test_rank_conversion_preserves_operator_values(operation): model.ir_version = 9 transformed = copy.deepcopy(model) for apply in (_apply_non4d_model_inputs, _apply_non4d_model_outputs, transform): - transformed = apply(transformed) + options = {"defer_rank_conversion": True} if apply is _apply_non4d_model_outputs else {} + transformed = apply(transformed, **options) options = ort.SessionOptions() options.intra_op_num_threads = 1 data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) From 37260570df6fc213b669214ae6222eac20ca38dd Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:00:35 +0000 Subject: [PATCH 03/41] Preserve runtime initializer defaults during deferred rank preparation Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 12 +++++++ .../test_dla_rank_conversion.py | 36 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 6f87cc0a529..f70349e1f42 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -37,6 +37,7 @@ def _apply_non4d_slice_axis(model): + """Remap fixed Slice axes before their data rank is promoted.""" cache = GraphCache(model.graph) for node in model.graph.node: if node.op_type != "Slice" or len(node.input) < 4: @@ -58,6 +59,7 @@ def _apply_non4d_slice_axis(model): def _apply_non4d_reshape(model): + """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) replacements = {} for node in model.graph.node: @@ -92,6 +94,7 @@ def _apply_non4d_reshape(model): def _apply_non4d_initializers(model): + """Promote fixed weight layouts while preserving runtime input defaults.""" expand, skip, front, avoid = set(), set(), set(), set() for node in model.graph.node: if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: @@ -103,6 +106,7 @@ def _apply_non4d_initializers(model): if node.op_type in {"Unsqueeze", "GroupQueryAttention"}: avoid.update(node.input) updates = [] + avoid.update(value.name for value in model.graph.input) expand_or_front = expand | front for tensor in model.graph.initializer: if tensor.name in avoid: @@ -124,14 +128,17 @@ def _apply_non4d_initializers(model): def _apply_non4d_expand(model): + """Prepare fixed Expand targets for four-dimensional data.""" return _pad_shape_inputs(model, "Expand", "_expand_shape") def _apply_non4d_tile(model): + """Prepare fixed Tile repeats for four-dimensional data.""" return _pad_shape_inputs(model, "Tile", "_tile_shape") def _pad_shape_inputs(model, op_type, suffix): + """Pad fixed structural controls for the requested rank stage.""" cache = GraphCache(model.graph) for node in model.graph.node: if node.op_type != op_type or len(node.input) < 2: @@ -154,6 +161,7 @@ def _pad_shape_inputs(model, op_type, suffix): def _apply_non4d_transpose(model): + """Remap transpose permutations around added or removed unary axes.""" cache = GraphCache(model.graph) for node in model.graph.node: if node.op_type != "Transpose": @@ -179,6 +187,7 @@ def _apply_non4d_transpose(model): def _apply_non4d_lpnorm(model): + """Preserve normalization axes across leading-dimension promotion.""" cache = GraphCache(model.graph) for node in model.graph.node: if node.op_type == "LpNormalization": @@ -193,6 +202,7 @@ def _apply_non4d_lpnorm(model): def _apply_flatten(model): + """Replace Flatten with its four-dimensional equivalent shape.""" graph = model.graph cache = GraphCache(graph) for node in graph.node: @@ -221,6 +231,7 @@ def _apply_flatten(model): def _apply_argmax(model): + """Preserve ArgMax axes and restore its declared output layout.""" cache = GraphCache(model.graph) nodes = [] for node in model.graph.node: @@ -267,6 +278,7 @@ def _apply_argmax(model): def _apply_update_axis(model): + """Remap supported operation axes from their native input ranks.""" update = { "LogSoftmax", "Softmax", diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 84196e6b5b0..cc51193a7f9 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -18,6 +18,7 @@ import copy import numpy as np +import onnx import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper @@ -28,6 +29,7 @@ ) from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import ( _apply_non4d_expand, + _apply_non4d_initializers, _apply_non4d_reshape, _apply_non4d_tile, _apply_non4d_transpose, @@ -78,3 +80,37 @@ def test_rank_conversion_preserves_operator_values(operation): transformed.SerializeToString(), options, providers=["CPUExecutionProvider"] ).run(None, {"x": data}) np.testing.assert_array_equal(actual[0], expected[0]) + + +@pytest.mark.parametrize("shape", [[4], [3, 4], [2, 3, 4]]) +def test_initializer_rank_conversion_preserves_overridable_defaults(shape): + """Preserve a runtime input's rank and default payload while preparing fixed weights.""" + default = np.ones(shape, dtype=np.float32) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Add", ["x", "weight"], ["y"])], + "runtime_default", + [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + for name in ("x", "weight") + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(default, "weight")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + transformed = _apply_non4d_initializers(copy.deepcopy(model)) + assert transformed.graph.initializer[0] == model.graph.initializer[0] + onnx.checker.check_model(transformed, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + original = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + converted = ort.InferenceSession( + transformed.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + for feeds in ({"x": data}, {"x": data, "weight": -default}): + np.testing.assert_array_equal(converted.run(None, feeds)[0], original.run(None, feeds)[0]) From c8712ae168f30bae139360c7ebea2b2ffe8ff3bb Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:58:50 +0000 Subject: [PATCH 04/41] Preserve native layout contracts in ordered DLA rank stages Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 319 ++++++--- .../test_dla_rank_conversion.py | 641 ++++++++++++++++++ 2 files changed, 878 insertions(+), 82 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index f70349e1f42..a299816dfd6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -23,7 +23,7 @@ import math 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, @@ -32,6 +32,9 @@ add_initializer, pad4d, ) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures +from .dla_indexing import _insert_support_nodes, _operator_nodes __all__ = [] @@ -39,44 +42,74 @@ def _apply_non4d_slice_axis(model): """Remap fixed Slice axes before their data rank is promoted.""" cache = GraphCache(model.graph) - for node in model.graph.node: - if node.op_type != "Slice" or len(node.input) < 4: + plans = [] + for node in _operator_nodes(model, {"Slice"}): + if len(node.input) < 4: continue - shape = cache.get_shape(node.input[0]) + shape = _shape_before_boundary_promotion(cache, node.input[0]) axes = cache.get_init(node.input[3]) - if shape is None or len(shape) not in (2, 3) or axes is None: + if shape is None or len(shape) not in (2, 3): continue + if axes is None or axes.name in cache.graph_input_names: + raise ValueError(f"{node.name!r}: Slice axes must be fixed") values = numpy_helper.to_array(axes).reshape(-1) if values.size and np.any(values >= 0): - node.input[3] = add_initializer( + plans.append((node, axes.name, np.where(values >= 0, values + 4 - len(shape), values))) + views = {} + for node, name, values in plans: + key = (name, tuple(values)) + if key not in views: + views[key] = add_initializer( model.graph, - axes.name + "_axes_only", - np.asarray(np.where(values >= 0, values + 4 - len(shape), values), dtype=np.int64), + cache.unique_name(name + "_axes_only"), + np.asarray(values, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) + node.input[3] = views[key] + _add_legacy_initializer_inputs(model) return model def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} replacements = {} - for node in model.graph.node: - if node.op_type != "Reshape" or len(node.input) < 2: - continue + users = {} + captured = _subgraph_captures(model.graph) + for node in _operator_nodes(model, {"Reshape"}): tensor = cache.get_init(node.input[1]) - if tensor is None: - continue + if tensor is None or tensor.name in cache.graph_input_names: + raise ValueError(f"{node.name!r}: Reshape target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() if len(values) < 4: + native_rank = len(values) values = pad4d(values) output_shape = cache.get_shape(node.output[0]) + if output_shape is None: + aliases = [ + cache.get_shape(user.output[0]) + for user, index in cache.get_consumers(node.output[0]) + if index == 0 + and user.op_type == "Squeeze" + and user.domain in {"", "ai.onnx"} + and not user.overload + and (user.domain, user.op_type, user.overload) not in local_functions + and user.output[0] in cache.graph_output_names + ] + aliases = [ + shape for shape in aliases if shape is not None and len(shape) == native_rank + ] + if aliases and all(shape == aliases[0] for shape in aliases): + output_shape = aliases[0] if output_shape: output_shape = pad4d(output_shape) values = [ output_shape[i] if value == 0 else value for i, value in enumerate(values) ] + if 0 in values: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") elif len(values) == 5: if 1 not in values: raise ValueError(f"{node.name!r}: 5D Reshape shape has no unary dimension") @@ -87,16 +120,49 @@ def _apply_non4d_reshape(model): if tensor.name in replacements and replacements[tensor.name] != replacement: raise ValueError(f"Shared Reshape target {tensor.name!r} requires different shapes") replacements[tensor.name] = replacement - for tensor in model.graph.initializer: - if tensor.name in replacements: - tensor.CopyFrom(replacements[tensor.name]) + users.setdefault(tensor.name, []).append(node) + for name, replacement in replacements.items(): + private = name in cache.graph_output_names | captured or any( + index != 1 + or user.op_type != "Reshape" + or user.domain not in {"", "ai.onnx"} + or user.overload + or (user.domain, user.op_type, user.overload) in local_functions + for user, index in cache.get_consumers(name) + ) + if private: + target = add_initializer( + model.graph, + cache.unique_name(name + "_reshape_shape"), + numpy_helper.to_array(replacement), + dtype=TensorProto.INT64, + with_value_info=False, + ) + for user in users[name]: + user.input[1] = target + else: + tensor = cache.get_init(name) + assert tensor is not None + tensor.CopyFrom(replacement) + _add_legacy_initializer_inputs(model) return model def _apply_non4d_initializers(model): """Promote fixed weight layouts while preserving runtime input defaults.""" expand, skip, front, avoid = set(), set(), set(), set() - for node in model.graph.node: + supported = _operator_nodes(model, {node.op_type for node in model.graph.node}) + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): + supported += _operator_nodes( + model, {"QuantizeLinear", "DequantizeLinear", "GroupQueryAttention"}, ("com.microsoft",) + ) + selected = {id(node) for node in supported} + protected = _subgraph_captures(model.graph) + protected.update( + name for node in model.graph.node if id(node) not in selected for name in node.input + ) + used = {name for node in supported for name in node.input} if protected else set() + for node in supported: if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: expand.update(node.input) if node.op_type == "MatMul": @@ -112,6 +178,12 @@ def _apply_non4d_initializers(model): if tensor.name in avoid: continue rank = len(tensor.dims) + if tensor.name in protected: + if tensor.name in used and rank in (1, 2, 3): + raise ValueError( + f"Initializer {tensor.name!r} has incompatible opaque or captured usage" + ) + continue if ( (rank == 1 and tensor.name in expand) or (rank == 2 and tensor.name in expand_or_front) @@ -121,7 +193,18 @@ def _apply_non4d_initializers(model): updates.append((tensor, shape)) elif rank == 2 and tensor.name not in skip: raise ValueError(f"Initializer {tensor.name!r} has unsupported rank 2 usage") + value_info = {value.name for value in model.graph.value_info} + native_inputs = { + node.input[0] + for node in model.graph.node + if node.input and node.op_type not in _RANK_AGNOSTIC_OPS + } for tensor, shape in updates: + if tensor.name in native_inputs and tensor.name not in value_info: + # Later axis stages read native ranks from the same deferred metadata snapshot. + model.graph.value_info.append( + helper.make_tensor_value_info(tensor.name, tensor.data_type, list(tensor.dims)) + ) del tensor.dims[:] tensor.dims.extend(shape) return model @@ -138,58 +221,67 @@ def _apply_non4d_tile(model): def _pad_shape_inputs(model, op_type, suffix): - """Pad fixed structural controls for the requested rank stage.""" + """Plan fixed structural views before changing nodes or allocating controls.""" cache = GraphCache(model.graph) - for node in model.graph.node: - if node.op_type != op_type or len(node.input) < 2: - continue + plans = [] + for node in _operator_nodes(model, {op_type}): tensor = cache.get_init(node.input[1]) - if tensor is None: - continue + if tensor is None or tensor.name in cache.graph_input_names: + raise ValueError(f"{node.name!r}: {op_type} target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() if len(values) < 4: - node.input[1] = add_initializer( + plans.append((node, tensor.name, pad4d(values))) + elif len(values) > 4: + raise ValueError(f"{node.name!r}: {op_type} target has rank {len(values)}") + views = {} + for node, name, values in plans: + if name not in views: + views[name] = add_initializer( model.graph, - tensor.name + suffix, - np.asarray(pad4d(values), dtype=np.int64), + cache.unique_name(name + suffix), + np.asarray(values, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - elif len(values) > 4: - raise ValueError(f"{node.name!r}: {op_type} target has rank {len(values)}") + node.input[1] = views[name] + _add_legacy_initializer_inputs(model) return model def _apply_non4d_transpose(model): """Remap transpose permutations around added or removed unary axes.""" cache = GraphCache(model.graph) - for node in model.graph.node: - if node.op_type != "Transpose": - continue + plans = [] + for node in _operator_nodes(model, {"Transpose"}): perm = _attribute(node, "perm") - if perm is None: - continue - values = list(perm.ints) + shape = _shape_before_boundary_promotion(cache, node.input[0]) + if perm is None and shape is None: + raise ValueError(f"{node.name!r}: default Transpose requires a known input rank") + values = list(perm.ints) if perm is not None else list(reversed(range(len(shape)))) if len(values) in (2, 3): delta = 4 - len(values) values = list(range(delta)) + [v + delta for v in values] elif len(values) == 5: - shape = cache.get_shape(node.input[0]) if shape is None or len(shape) != 5 or 1 not in shape: raise ValueError(f"{node.name!r}: 5D Transpose requires a known unary dimension") removed = shape.index(1) values = [v - int(v > removed) for v in values if v != removed] else: continue - del perm.ints[:] - perm.ints.extend(values) + plans.append((node, perm, values)) + for node, perm, values in plans: + if perm is None: + node.attribute.append(helper.make_attribute("perm", values)) + else: + del perm.ints[:] + perm.ints.extend(values) return model def _apply_non4d_lpnorm(model): """Preserve normalization axes across leading-dimension promotion.""" cache = GraphCache(model.graph) - for node in model.graph.node: + for node in _operator_nodes(model, {"LpNormalization"}): if node.op_type == "LpNormalization": axis = _attribute(node, "axis") shape = _shape_before_boundary_promotion(cache, node.input[0]) @@ -205,75 +297,93 @@ def _apply_flatten(model): """Replace Flatten with its four-dimensional equivalent shape.""" graph = model.graph cache = GraphCache(graph) - for node in graph.node: - if node.op_type == "Flatten": - shape = _shape_before_boundary_promotion(cache, node.input[0]) - axis = _attribute(node, "axis") - axis = axis.i if axis is not None else 1 - if shape is None or any(d <= 0 for d in shape): - raise ValueError(f"{node.name!r}: Flatten requires a known input shape") - if axis < 0: - axis += len(shape) - batch = math.prod(shape[:axis]) - name = add_initializer( + plans = [] + for node in _operator_nodes(model, {"Flatten"}): + shape = _shape_before_boundary_promotion(cache, node.input[0]) + axis = _attribute(node, "axis") + axis = axis.i if axis is not None else 1 + if shape is None or any(d <= 0 for d in shape): + raise ValueError(f"{node.name!r}: Flatten requires a known input shape") + if not -len(shape) <= axis <= len(shape): + raise ValueError(f"{node.name!r}: invalid Flatten axis {axis}") + if axis < 0: + axis += len(shape) + plans.append((node, math.prod(shape[:axis]))) + targets = {} + for node, batch in plans: + if batch not in targets: + hint = "flatten_reshape_shape" if batch == 1 else f"{node.name}_reshape_shape" + targets[batch] = add_initializer( graph, - "flatten_reshape_shape" if batch == 1 else f"{node.name}_reshape_shape", + cache.unique_name(hint), np.asarray([1, 1, batch, -1], dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - node.CopyFrom( - helper.make_node( - "Reshape", [node.input[0], name], list(node.output), name=node.name + "_reshape" - ) + node.CopyFrom( + helper.make_node( + "Reshape", + [node.input[0], targets[batch]], + list(node.output), + name=cache.unique_name((node.name or node.output[0]) + "_reshape"), ) + ) + _add_legacy_initializer_inputs(model) return model def _apply_argmax(model): - """Preserve ArgMax axes and restore its declared output layout.""" + """Plan native ArgMax layouts before preparing private output restoration views.""" cache = GraphCache(model.graph) + plans = {} + for node in _operator_nodes(model, {"ArgMax"}): + shape = _shape_before_boundary_promotion(cache, node.input[0]) + if not shape or any(dim <= 0 for dim in shape): + raise ValueError(f"{node.name!r}: ArgMax requires a static input shape") + axis = _attribute(node, "axis") + value = axis.i if axis is not None else 0 + if not -len(shape) <= value < len(shape): + raise ValueError(f"{node.name!r}: invalid ArgMax axis {value}") + output_shape = [dim for index, dim in enumerate(shape) if index != value % len(shape)] + if value >= 0: + value += {2: 2, 3: 1}.get(len(shape), 0) + plans[id(node)] = (value, output_shape) nodes = [] for node in model.graph.node: nodes.append(node) - if node.op_type != "ArgMax": - continue - shape = _shape_before_boundary_promotion(cache, node.input[0]) - if not shape: + if id(node) not in plans: continue + value, output_shape = plans[id(node)] axis = _attribute(node, "axis") if axis is None: - node.attribute.append(helper.make_attribute("axis", 0)) - axis = _attribute(node, "axis") - if axis.i >= 0 and len(shape) in (2, 3): - axis.i += 4 - len(shape) + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value keepdims = _attribute(node, "keepdims") restore_shape = keepdims is not None and keepdims.i == 0 if keepdims is not None: keepdims.i = 1 - node.name += "_axis_updated" + original = node.output[0] + node.name = cache.unique_name((node.name or original) + "_axis_updated") if restore_shape: - original = node.output[0] - output_shape = cache.get_shape(original) - if output_shape is None and original.endswith("_squeeze_input"): - output_shape = cache.get_shape(original.removesuffix("_squeeze_input")) - if output_shape is None: - raise ValueError(f"Unknown original ArgMax output shape for {node.name!r}") target = add_initializer( model.graph, - node.name + "_reshape_shape", + cache.unique_name(node.name + "_reshape_shape"), np.asarray(pad4d(output_shape), dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) - node.output[0] = original + "_keepdims" + node.output[0] = cache.unique_name(original + "_keepdims") nodes.append( helper.make_node( - "Reshape", [node.output[0], target], [original], name=node.name + "_reshape" + "Reshape", + [node.output[0], target], + [original], + name=cache.unique_name(node.name + "_reshape"), ) ) - del model.graph.node[:] - model.graph.node.extend(nodes) + _insert_support_nodes(model.graph, nodes) + _add_legacy_initializer_inputs(model) return model @@ -301,17 +411,62 @@ def _apply_update_axis(model): "Flatten", } cache = GraphCache(model.graph) - for node in model.graph.node: + plans = [] + selected = _operator_nodes(model, {node.op_type for node in model.graph.node}) + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): + selected += _operator_nodes( + model, {"QuantizeLinear", "DequantizeLinear"}, ("com.microsoft",) + ) + versions = {op.domain: op.version for op in model.opset_import} + for node in selected: axis = _attribute(node, "axis") - if axis is None or axis.i < 0 or node.op_type in skip: + shape = _shape_before_boundary_promotion(cache, node.input[0]) if node.input else None + if ( + axis is None + and node.op_type in update + and node.domain in {"", "ai.onnx"} + and shape + and ( + len(shape) == 5 + or ( + len(shape) in (2, 3) + and node.op_type not in {"QuantizeLinear", "DequantizeLinear"} + ) + ) + ): + schema = defs.get_schema( + node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain + ) + default = schema.attributes.get("axis") + if default is not None and default.default_value.type: + value = default.default_value.i + else: + continue + elif axis is None or node.op_type in skip or (axis.i < 0 and node.op_type not in update): continue + else: + value = axis.i if node.op_type not in update: raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") if "axis_updated" in node.name: continue - shape = _shape_before_boundary_promotion(cache, node.input[0]) - delta = {2: 2, 3: 1, 5: -1}.get(len(shape) if shape else 0, 0) - if delta: - axis.i += delta - node.name += "_axis_updated" + rank = len(shape) if shape else 0 + original = value + if shape and rank == 5: + if not -5 <= value < 5 or 1 not in shape or value % 5 == shape.index(1): + raise ValueError(f"{node.name!r}: cannot map 5D axis {value}") + native_axis = value % 5 + value = native_axis - int(native_axis > shape.index(1)) + if original < 0: + value -= 4 + elif value >= 0: + value += {2: 2, 3: 1}.get(rank, 0) + if value != original: + plans.append((node, axis, value)) + for node, axis, value in plans: + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name += "_axis_updated" return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index cc51193a7f9..cf1f69a7f0a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -28,11 +28,20 @@ _apply_non4d_model_outputs, ) from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import ( + _apply_argmax, + _apply_flatten, _apply_non4d_expand, _apply_non4d_initializers, + _apply_non4d_lpnorm, _apply_non4d_reshape, + _apply_non4d_slice_axis, _apply_non4d_tile, _apply_non4d_transpose, + _apply_update_axis, +) +from modelopt.onnx.graph_surgery.make_dla_compatible import ( + DLA_TRANSFORMS, + _transform_make_dla_compatible, ) @@ -114,3 +123,635 @@ def test_initializer_rank_conversion_preserves_overridable_defaults(shape): data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) for feeds in ({"x": data}, {"x": data, "weight": -default}): np.testing.assert_array_equal(converted.run(None, feeds)[0], original.run(None, feeds)[0]) + + +def _execute(model, feeds): + """Fully check and execute the small native-rank contract fixtures.""" + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + +@pytest.mark.parametrize("target", [[0, -1], [0, 0, -1], [0, -1, 0]]) +@pytest.mark.parametrize("collision", [False, True]) +def test_reshape_resolves_copied_dimensions_across_output_aliases(target, collision): + """Resolve copied dimensions even when output adapters allocate a suffixed private name.""" + shape = [2, 3, 4] + output_shape = [2, 12] if len(target) == 2 else shape + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "copied_shape", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(target, np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + if collision: + model.graph.input.append(helper.make_tensor_value_info("other", TensorProto.FLOAT, [1])) + model.graph.node.append(helper.make_node("Identity", ["other"], ["y_squeeze_input"])) + model.graph.output.append( + helper.make_tensor_value_info("y_squeeze_input", TensorProto.FLOAT, [1]) + ) + feeds["other"] = np.asarray([17], np.float32) + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("axis", [0, 1, 2, -1]) +@pytest.mark.parametrize("value_info", [False, True]) +def test_registered_axis_update_reads_native_initializer_rank(axis, value_info): + """Keep Concat axes correct with fixed inputs that lack optional value metadata.""" + shape = [2, 3, 4] + output_shape = shape.copy() + output_shape[axis] *= 2 + values = np.arange(24, dtype=np.float32).reshape(shape) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Concat", ["first", "second"], ["y"], axis=axis)], + "initializer_axis", + [], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(values, "first"), numpy_helper.from_array(-values, "second")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + if value_info: + model.graph.value_info.extend( + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + for name in ["first", "second"] + ) + expected = _execute(model, {})[0] + transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) + for name, transform in DLA_TRANSFORMS: + if name in {"Non4dInitializers", "UpdateAxis"}: + transformed = transform(transformed) + transformed.graph.ClearField("value_info") + np.testing.assert_array_equal(_execute(transformed, {})[0], expected) + + +@pytest.mark.parametrize( + ("shape", "axis"), + [ + ([2, 3, 4, 5, 1], 1), + ([2, 3, 4, 5, 1], -4), + ([2, 3, 1, 4, 5], 3), + ([2, 3, 1, 4, 5], -2), + ([1, 2, 3, 4, 5], 2), + ([1, 2, 3, 4, 5], -3), + ], +) +def test_five_dimensional_axis_tracks_the_removed_unary_position(shape, axis): + """Preserve Softmax values when a later rank stage drops the known unary axis.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Softmax", ["x"], ["y"], axis=axis)], + "unary_axis", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + data = np.random.default_rng(20).standard_normal(shape).astype(np.float32) + expected = _execute(model, {"x": data})[0] + transformed = _apply_update_axis(copy.deepcopy(model)) + collapsed = shape.copy() + collapsed.remove(1) + transformed.graph.input[0].CopyFrom( + helper.make_tensor_value_info("x", TensorProto.FLOAT, collapsed) + ) + transformed.graph.output[0].CopyFrom( + helper.make_tensor_value_info("y", TensorProto.FLOAT, collapsed) + ) + np.testing.assert_allclose( + _execute(transformed, {"x": data.reshape(collapsed)})[0], + expected.reshape(collapsed), + rtol=1e-6, + atol=1e-7, + ) + + +@pytest.mark.parametrize( + ("shape", "axis"), [([1, 2, 3, 4, 5], 0), ([1, 2, 3, 4, 5], -5), ([2, 3, 4, 5, 6], 2)] +) +def test_invalid_five_dimensional_axis_refuses_before_earlier_axis_mutation(shape, axis): + """Validate every axis mapping before changing any previously supported operation.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Softmax", ["first"], ["first_out"], axis=1), + helper.make_node("Softmax", ["second"], ["second_out"], axis=axis), + ], + "axis_refusal", + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, shape), + ], + [ + helper.make_tensor_value_info("first_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("second_out", TensorProto.FLOAT, shape), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="5D axis"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_axis_update_recovers_native_rank_when_boundary_name_is_allocated(axis): + """Preserve axis semantics when an unrelated public tensor occupies the adapter hint.""" + shape = [2, 3, 4] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Softmax", ["x"], ["y"], axis=axis), + helper.make_node("Identity", ["other"], ["x_unsqueeze_input"]), + ], + "allocated_input", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("other", TensorProto.FLOAT, [1]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("x_unsqueeze_input", TensorProto.FLOAT, [1]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.random.default_rng(22).standard_normal(shape).astype(np.float32), + "other": np.asarray([17], np.float32), + } + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=1e-7) + + +def _structural_control_model(operation, *, runtime=False, public=False): + """Create fixed shape controls with optional runtime or public owners.""" + values = [1, 1] if operation == "Tile" else [2, 3] + nodes = [helper.make_node(operation, ["x", "target"], ["y"])] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])] + initializers = [numpy_helper.from_array(np.asarray(values, np.int64), "target")] + if runtime: + inputs.append(helper.make_tensor_value_info("target", TensorProto.INT64, [2])) + nodes.insert(0, helper.make_node(operation, ["x", "earlier"], ["first"])) + initializers.append(numpy_helper.from_array(np.asarray(values, np.int64), "earlier")) + if public: + collision = "target_" + operation.lower() + "_shape" + initializers.append(numpy_helper.from_array(np.asarray([17], np.int64), collision)) + outputs.extend( + [ + helper.make_tensor_value_info("target", TensorProto.INT64, [2]), + helper.make_tensor_value_info(collision, TensorProto.INT64, [1]), + ] + ) + model = helper.make_model( + helper.make_graph(nodes, "structural_control", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +def test_runtime_shape_defaults_refuse_before_any_structural_control_mutation(operation): + """Validate runtime shape controls before changing an earlier fixed-target node.""" + model = _structural_control_model(operation, runtime=True) + before = model.SerializeToString() + transform = { + "Reshape": _apply_non4d_reshape, + "Expand": _apply_non4d_expand, + "Tile": _apply_non4d_tile, + }[operation] + with pytest.raises(ValueError, match="fixed"): + transform(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +def test_private_shape_views_preserve_public_controls_and_existing_names(operation): + """Preserve public shape values when an allocated control hint is already occupied.""" + model = _structural_control_model(operation, public=True) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_argmax_keeps_untouched_constant_payload_in_place(): + """Preserve an unrelated Constant payload while preparing ArgMax.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("ArgMax", ["x"], ["y"], axis=3, keepdims=1), + ], + "argmax_payload", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [ + helper.make_tensor_value_info("y", TensorProto.INT64, [1, 1, 2, 1]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds) + constant = model.graph.node[0] + transformed = _apply_argmax(model) + assert transformed.graph.node[0] is constant + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("collision", [False, True]) +@pytest.mark.parametrize("legacy_ir", [False, True]) +def test_unnamed_argmax_outputs_preserve_distinct_shapes_and_public_aliases(collision, legacy_ir): + """Restore two unnamed reductions while respecting an occupied output-adapter hint.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("ArgMax", ["x"], ["first"], axis=1, keepdims=0), + helper.make_node("ArgMax", ["x"], ["second"], axis=2, keepdims=0), + ], + "argmax_alias", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])], + [ + helper.make_tensor_value_info("first", TensorProto.INT64, [2, 4]), + helper.make_tensor_value_info("second", TensorProto.INT64, [2, 3]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + if collision: + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([17], np.int64), "first_squeeze_input") + ) + model.graph.output.append( + helper.make_tensor_value_info("first_squeeze_input", TensorProto.INT64, [1]) + ) + if legacy_ir: + model.ir_version = 3 + model.graph.input.extend( + helper.make_tensor_value_info(t.name, t.data_type, list(t.dims)) + for t in model.graph.initializer + ) + feeds = {"x": np.random.default_rng(23).standard_normal((2, 3, 4)).astype(np.float32)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_late_unknown_argmax_shape_refuses_without_changing_an_earlier_node(): + """Validate all native ArgMax dimensions before allocating or mutating adapters.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("ArgMax", ["first"], ["first_out"], axis=1, keepdims=0), + helper.make_node("ArgMax", ["second"], ["second_out"], axis=1, keepdims=0), + ], + "argmax_refusal", + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, [None, None, None]), + ], + [ + helper.make_tensor_value_info("first_out", TensorProto.INT64, [2, 4]), + helper.make_tensor_value_info("second_out", TensorProto.INT64, [None, None]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="static"): + _apply_argmax(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +def test_default_transpose_preserves_native_reversed_axes(shape): + """Execute the implicit permutation after leading singleton promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["x"], ["y"])], + "default_transpose", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, list(reversed(shape)))], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +def test_unnamed_flatten_nodes_keep_distinct_batch_shapes_and_public_controls(): + """Allocate private targets for unrelated unnamed Flatten operations.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Flatten", ["x"], ["first"]), + helper.make_node("Flatten", ["z"], ["second"]), + ], + "flatten_names", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, [3, 2, 4]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 12]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, [3, 8]), + helper.make_tensor_value_info("_reshape_shape", TensorProto.INT64, [1]), + ], + [numpy_helper.from_array(np.asarray([17], np.int64), "_reshape_shape")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "z": np.arange(24, dtype=np.float32).reshape(3, 2, 4), + } + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Flatten", "Transpose"]) +def test_late_invalid_layout_refuses_before_earlier_node_mutation(operation): + """Plan every layout before mutating earlier valid operators.""" + shape = [None, 3, 4] if operation == "Flatten" else [2, 3, 4, 5, 6] + first_attrs = {} if operation == "Flatten" else {"perm": [1, 0]} + second_attrs = {} if operation == "Flatten" else {"perm": [4, 3, 2, 1, 0]} + model = helper.make_model( + helper.make_graph( + [ + helper.make_node(operation, ["x"], ["first"], **first_attrs), + helper.make_node(operation, ["z"], ["second"], **second_attrs), + ], + "layout_refusal", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, shape), + ], + [ + helper.make_tensor_value_info( + "first", TensorProto.FLOAT, [2, 3] if operation == "Flatten" else [3, 2] + ), + helper.make_tensor_value_info( + "second", + TensorProto.FLOAT, + [None, 12] if operation == "Flatten" else list(reversed(shape)), + ), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="known"): + {"Flatten": _apply_flatten, "Transpose": _apply_non4d_transpose}[operation](model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"]) +def test_implicit_legacy_axis_tracks_removed_leading_unary_dimension(operation): + """Read the imported standard schema when the native axis attribute is omitted.""" + shape = [1, 2, 3, 4, 5] + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x"], ["y"])], + "implicit_axis", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", 11)], + ) + model.ir_version = 9 + data = np.random.default_rng(25).standard_normal(shape).astype(np.float32) + expected = _execute(model, {"x": data})[0] + transformed = _apply_update_axis(copy.deepcopy(model)) + for value in [transformed.graph.input[0], transformed.graph.output[0]]: + value.CopyFrom(helper.make_tensor_value_info(value.name, TensorProto.FLOAT, shape[1:])) + np.testing.assert_allclose( + _execute(transformed, {"x": data.reshape(shape[1:])})[0].reshape(shape), + expected, + rtol=1e-6, + atol=1e-7, + ) + + +def test_implicit_axis_on_removed_dimension_refuses_without_mutation(): + """Reject an omitted default axis when its unary dimension would disappear.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Softmax", ["x"], ["y"])], + "implicit_refusal", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4, 5, 1])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4, 5, 1])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="5D axis"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("owner", ["function", "capture"]) +@pytest.mark.parametrize("shared", [False, True]) +def test_opaque_initializer_owners_preserve_rank_or_refuse_shared_layout(owner, shared): + """Protect fixed native weights owned by an opaque function or nested scope.""" + shape = [2, 3, 4] + initializers = [numpy_helper.from_array(np.ones(shape, np.float32), "weight")] + functions = [] + opsets = [helper.make_opsetid("", 17)] + if owner == "function": + functions.append( + helper.make_function( + "opaque", + "Layout", + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + ) + ) + nodes = [helper.make_node("Layout", ["weight"], ["y"], domain="opaque")] + opsets.append(helper.make_opsetid("opaque", 1)) + else: + branch = helper.make_graph( + [helper.make_node("Identity", ["weight"], ["out"])], + "branch", + [], + [helper.make_tensor_value_info("out", TensorProto.FLOAT, shape)], + ) + initializers.append(numpy_helper.from_array(np.asarray(True), "condition")) + nodes = [ + helper.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch) + ] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)] + if shared: + nodes.insert(0, helper.make_node("Identity", ["earlier"], ["early"])) + nodes.append(helper.make_node("Add", ["weight", "weight"], ["sum"])) + initializers.insert(0, numpy_helper.from_array(np.ones(shape, np.float32), "earlier")) + outputs.append(helper.make_tensor_value_info("sum", TensorProto.FLOAT, shape)) + model = helper.make_model( + helper.make_graph(nodes, "opaque_weight", [], outputs, initializers), + opset_imports=opsets, + functions=functions, + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + if shared: + with pytest.raises(ValueError, match="opaque or captured"): + _apply_non4d_initializers(model) + else: + _apply_non4d_initializers(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + "operation", ["Flatten", "Transpose", "LpNormalization", "Softmax", "Slice"] +) +def test_rank_passes_leave_model_local_operators_unchanged(operation): + """Skip opaque local functions whose names match standard layout operators.""" + function = helper.make_function( + "opaque", + operation, + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x"], ["y"], domain="opaque")], + "opaque_rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("opaque", 1)], + functions=[function], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + for transform in [ + _apply_flatten, + _apply_non4d_transpose, + _apply_non4d_lpnorm, + _apply_update_axis, + _apply_non4d_slice_axis, + ]: + transform(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("runtime", [False, True]) +def test_shared_slice_axes_preserve_private_views_and_transactional_runtime_refusal(runtime): + """Reuse fixed axis views and reject later runtime axes before changing earlier nodes.""" + nodes = [ + helper.make_node("Slice", ["x", "starts", "ends", "axes"], ["first"]), + helper.make_node( + "Slice", ["x", "starts", "ends", "runtime_axes" if runtime else "axes"], ["second"] + ), + ] + initializers = [ + numpy_helper.from_array(np.asarray(v, np.int64), n) + for n, v in [("starts", [0]), ("ends", [2]), ("axes", [1]), ("axes_axes_only", [17])] + ] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])] + if runtime: + inputs.append(helper.make_tensor_value_info("runtime_axes", TensorProto.INT64, [1])) + initializers.append(numpy_helper.from_array(np.asarray([1], np.int64), "runtime_axes")) + model = helper.make_model( + helper.make_graph( + nodes, + "slice_controls", + inputs, + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [2, 2, 4]) + for n in ["first", "second"] + ] + + [helper.make_tensor_value_info("axes_axes_only", TensorProto.INT64, [1])], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + if runtime: + before = model.SerializeToString() + with pytest.raises(ValueError, match="fixed"): + _apply_non4d_slice_axis(model) + assert model.SerializeToString() == before + else: + feeds = {"x": np.arange(24, dtype=np.float32).reshape(2, 3, 4)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + slices = [n for n in transformed.graph.node if n.op_type == "Slice"] + assert slices[0].input[3] == slices[1].input[3] + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_implicit_split_axis_preserves_rank_three_values(): + """Map the standard default Split axis during leading singleton promotion.""" + shape = [2, 3, 4] + model = helper.make_model( + helper.make_graph( + [helper.make_node("Split", ["x"], ["first", "second"])], + "split_axis", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [1, 3, 4]) + for n in ["first", "second"] + ], + ), + opset_imports=[helper.make_opsetid("", 13)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) From aa4c19225674e62d7d823044778f8403d9c5fbed Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 01:26:25 +0000 Subject: [PATCH 05/41] Validate DLA rank-collapse contexts and prepare implicit Slice axes Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 9 +- .../dla_transforms/dla_rank_conversion.py | 140 ++++++----- .../test_dla_rank_conversion.py | 229 ++++++++++++++---- 3 files changed, 260 insertions(+), 118 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 bcf27d5db2c..eada7ddaca3 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -399,7 +399,12 @@ def _apply_non4d_model_outputs( rank_sources = {} contrib_ops = _boundary_contrib_ops(model) for value in graph.output: - shape = cache.get_shape(value.name) + shape: list | None = cache.get_shape(value.name) + if shape is None and defer_rank_conversion and value.type.tensor_type.HasField("shape"): + shape = [ + dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None + for dim in value.type.tensor_type.shape.dim + ] if shape is None or len(shape) >= 4 or value.name in boundary_names: continue if value.name in native_broadcasts: @@ -479,7 +484,7 @@ def _apply_non4d_model_outputs( return model for value, shape, existing, axes in plans: axes_name = None - if len(shape) > 1 or shape == [1]: + if len(shape) > 1 or shape == [1] or any(not isinstance(dim, int) for dim in shape): axes_name = add_initializer( graph, cache.unique_name(value.name + "_squeeze_axes"), diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index a299816dfd6..85ca5a39353 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -13,14 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Individual rank and axis passes in the logged DLA transform sequence. +"""Ordered rank and axis passes sharing the preceding native shape snapshot. -These passes share the shape snapshot produced by the preceding shape-inference -stage. Running inference between passes would change how original axes are -mapped after upstream tensors have been promoted to four dimensions. +Do not infer shapes between stages: original axes precede upstream rank promotion. """ import math +from functools import partial import numpy as np from onnx import TensorProto, defs, helper, numpy_helper @@ -44,12 +43,19 @@ def _apply_non4d_slice_axis(model): cache = GraphCache(model.graph) plans = [] for node in _operator_nodes(model, {"Slice"}): - if len(node.input) < 4: - continue shape = _shape_before_boundary_promotion(cache, node.input[0]) - axes = cache.get_init(node.input[3]) if shape is None or len(shape) not in (2, 3): continue + if len(node.input) < 4 or not node.input[3]: + starts_shape = cache.get_shape(node.input[1]) + if starts_shape is None: + raise ValueError( + f"{node.name!r}: implicit Slice axes require a known starts length" + ) + values = np.arange(math.prod(starts_shape), dtype=np.int64) + 4 - len(shape) + plans.append((node, node.input[1], values)) + continue + axes = cache.get_init(node.input[3]) if axes is None or axes.name in cache.graph_input_names: raise ValueError(f"{node.name!r}: Slice axes must be fixed") values = numpy_helper.to_array(axes).reshape(-1) @@ -66,7 +72,10 @@ def _apply_non4d_slice_axis(model): dtype=TensorProto.INT64, with_value_info=False, ) - node.input[3] = views[key] + if len(node.input) < 4: + node.input.append(views[key]) + else: + node.input[3] = views[key] _add_legacy_initializer_inputs(model) return model @@ -74,11 +83,13 @@ def _apply_non4d_slice_axis(model): def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) - local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + reshapes = _operator_nodes(model, {"Reshape"}) + reshape_ids = {id(node) for node in reshapes} + squeezes = {id(node) for node in _operator_nodes(model, {"Squeeze"})} replacements = {} users = {} captured = _subgraph_captures(model.graph) - for node in _operator_nodes(model, {"Reshape"}): + for node in reshapes: tensor = cache.get_init(node.input[1]) if tensor is None or tensor.name in cache.graph_input_names: raise ValueError(f"{node.name!r}: Reshape target must be fixed") @@ -92,10 +103,7 @@ def _apply_non4d_reshape(model): cache.get_shape(user.output[0]) for user, index in cache.get_consumers(node.output[0]) if index == 0 - and user.op_type == "Squeeze" - and user.domain in {"", "ai.onnx"} - and not user.overload - and (user.domain, user.op_type, user.overload) not in local_functions + and id(user) in squeezes and user.output[0] in cache.graph_output_names ] aliases = [ @@ -111,6 +119,13 @@ def _apply_non4d_reshape(model): if 0 in values: raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") elif len(values) == 5: + if 0 in values: + output_shape = cache.get_shape(node.output[0]) + if output_shape is None or len(output_shape) != 5: + raise ValueError(f"{node.name!r}: cannot resolve 0 in 5D Reshape target") + values = [ + output_shape[i] if value == 0 else value for i, value in enumerate(values) + ] if 1 not in values: raise ValueError(f"{node.name!r}: 5D Reshape shape has no unary dimension") values.remove(1) @@ -123,12 +138,7 @@ def _apply_non4d_reshape(model): users.setdefault(tensor.name, []).append(node) for name, replacement in replacements.items(): private = name in cache.graph_output_names | captured or any( - index != 1 - or user.op_type != "Reshape" - or user.domain not in {"", "ai.onnx"} - or user.overload - or (user.domain, user.op_type, user.overload) in local_functions - for user, index in cache.get_consumers(name) + index != 1 or id(user) not in reshape_ids for user, index in cache.get_consumers(name) ) if private: target = add_initializer( @@ -210,16 +220,6 @@ def _apply_non4d_initializers(model): return model -def _apply_non4d_expand(model): - """Prepare fixed Expand targets for four-dimensional data.""" - return _pad_shape_inputs(model, "Expand", "_expand_shape") - - -def _apply_non4d_tile(model): - """Prepare fixed Tile repeats for four-dimensional data.""" - return _pad_shape_inputs(model, "Tile", "_tile_shape") - - def _pad_shape_inputs(model, op_type, suffix): """Plan fixed structural views before changing nodes or allocating controls.""" cache = GraphCache(model.graph) @@ -248,10 +248,36 @@ def _pad_shape_inputs(model, op_type, suffix): return model +_apply_non4d_expand = partial(_pad_shape_inputs, op_type="Expand", suffix="_expand_shape") +_apply_non4d_tile = partial(_pad_shape_inputs, op_type="Tile", suffix="_tile_shape") + + +def _collapsed_reshape_axes(cache, name, shape, reshapes): + """Identify singleton axes consistent with an actually collapsed native Reshape.""" + producer = cache.get_producer(name) + target = ( + cache.get_init_array(producer.input[1]) + if producer is not None and id(producer) in reshapes + else None + ) + if target is None or target.shape != (4,): + return [] + values = target.tolist() + return [ + axis + for axis, dim in enumerate(shape) + if dim == 1 + and all( + value in (-1, native) for value, native in zip(values, shape[:axis] + shape[axis + 1 :]) + ) + ] + + def _apply_non4d_transpose(model): """Remap transpose permutations around added or removed unary axes.""" cache = GraphCache(model.graph) plans = [] + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"Transpose"}): perm = _attribute(node, "perm") shape = _shape_before_boundary_promotion(cache, node.input[0]) @@ -264,7 +290,12 @@ def _apply_non4d_transpose(model): elif len(values) == 5: if shape is None or len(shape) != 5 or 1 not in shape: raise ValueError(f"{node.name!r}: 5D Transpose requires a known unary dimension") - removed = shape.index(1) + removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not removed_axes: + continue + if len(removed_axes) != 1: + raise ValueError(f"{node.name!r}: ambiguous collapsed 5D Transpose layout") + removed = removed_axes[0] values = [v - int(v > removed) for v in values if v != removed] else: continue @@ -389,34 +420,16 @@ def _apply_argmax(model): def _apply_update_axis(model): """Remap supported operation axes from their native input ranks.""" - update = { - "LogSoftmax", - "Softmax", - "DequantizeLinear", - "TopK", - "Concat", - "Split", - "QuantizeLinear", - } - skip = { - "Gather", - "GatherElements", - "LpNormalization", - "ArgMax", - "ReduceMax", - "ReduceMin", - "ReduceSum", - "ReduceMean", - "Slice", - "Flatten", - } + quantizers = {"QuantizeLinear", "DequantizeLinear"} + update = {"LogSoftmax", "Softmax", "TopK", "Concat", "Split"} | quantizers + skip = {"Gather", "GatherElements", "LpNormalization", "ArgMax", "Slice", "Flatten"} + skip |= {"ReduceMax", "ReduceMin", "ReduceSum", "ReduceMean"} cache = GraphCache(model.graph) plans = [] + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} selected = _operator_nodes(model, {node.op_type for node in model.graph.node}) if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): - selected += _operator_nodes( - model, {"QuantizeLinear", "DequantizeLinear"}, ("com.microsoft",) - ) + selected += _operator_nodes(model, quantizers, ("com.microsoft",)) versions = {op.domain: op.version for op in model.opset_import} for node in selected: axis = _attribute(node, "axis") @@ -426,13 +439,7 @@ def _apply_update_axis(model): and node.op_type in update and node.domain in {"", "ai.onnx"} and shape - and ( - len(shape) == 5 - or ( - len(shape) in (2, 3) - and node.op_type not in {"QuantizeLinear", "DequantizeLinear"} - ) - ) + and (len(shape) == 5 or (len(shape) in (2, 3) and node.op_type not in quantizers)) ): schema = defs.get_schema( node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain @@ -453,10 +460,13 @@ def _apply_update_axis(model): rank = len(shape) if shape else 0 original = value if shape and rank == 5: - if not -5 <= value < 5 or 1 not in shape or value % 5 == shape.index(1): - raise ValueError(f"{node.name!r}: cannot map 5D axis {value}") + removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not -5 <= value < 5 or not removed_axes or value % 5 in removed_axes: + raise ValueError( + f"{node.name!r}: cannot map 5D axis {value} without an unambiguous collapsed Reshape" + ) native_axis = value % 5 - value = native_axis - int(native_axis > shape.index(1)) + value = native_axis - int(native_axis > removed_axes[0]) if original < 0: value -= 4 elif value >= 0: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index cf1f69a7f0a..c265c2a7dd8 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -200,6 +200,43 @@ def test_registered_axis_update_reads_native_initializer_rank(axis, value_info): np.testing.assert_array_equal(_execute(transformed, {})[0], expected) +def _collapsed_normalization_model( + shape, axis, operation="Softmax", opset=17, target=None, data_shape=None +): + """Build a native 5D normalization between a rank-expanding and rank-restoring Reshape.""" + collapsed = shape.copy() + collapsed.remove(1) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "native_shape"], ["native"]), + helper.make_node( + operation, + ["native"], + ["normalized"], + **({} if axis is None else {"axis": axis}), + ), + helper.make_node("Reshape", ["normalized", "output_shape"], ["y"]), + ], + "collapsed_normalization", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape or collapsed)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, collapsed)], + [ + numpy_helper.from_array(np.asarray(target or shape, np.int64), "native_shape"), + numpy_helper.from_array(np.asarray(collapsed, np.int64), "output_shape"), + ], + value_info=[ + helper.make_tensor_value_info(n, TensorProto.FLOAT, shape) + for n in ["native", "normalized"] + ], + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model + + @pytest.mark.parametrize( ("shape", "axis"), [ @@ -212,34 +249,13 @@ def test_registered_axis_update_reads_native_initializer_rank(axis, value_info): ], ) def test_five_dimensional_axis_tracks_the_removed_unary_position(shape, axis): - """Preserve Softmax values when a later rank stage drops the known unary axis.""" - model = helper.make_model( - helper.make_graph( - [helper.make_node("Softmax", ["x"], ["y"], axis=axis)], - "unary_axis", - [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], - [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], - ), - opset_imports=[helper.make_opsetid("", 17)], - ) - model.ir_version = 9 - data = np.random.default_rng(20).standard_normal(shape).astype(np.float32) - expected = _execute(model, {"x": data})[0] - transformed = _apply_update_axis(copy.deepcopy(model)) - collapsed = shape.copy() - collapsed.remove(1) - transformed.graph.input[0].CopyFrom( - helper.make_tensor_value_info("x", TensorProto.FLOAT, collapsed) - ) - transformed.graph.output[0].CopyFrom( - helper.make_tensor_value_info("y", TensorProto.FLOAT, collapsed) - ) - np.testing.assert_allclose( - _execute(transformed, {"x": data.reshape(collapsed)})[0], - expected.reshape(collapsed), - rtol=1e-6, - atol=1e-7, - ) + """Execute the registered rank collapse and its downstream native axis mapping.""" + model = _collapsed_normalization_model(shape, axis) + input_shape = [d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim] + feeds = {"x": np.random.default_rng(20).standard_normal(input_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) @pytest.mark.parametrize( @@ -553,29 +569,12 @@ def test_late_invalid_layout_refuses_before_earlier_node_mutation(operation): @pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"]) def test_implicit_legacy_axis_tracks_removed_leading_unary_dimension(operation): - """Read the imported standard schema when the native axis attribute is omitted.""" - shape = [1, 2, 3, 4, 5] - model = helper.make_model( - helper.make_graph( - [helper.make_node(operation, ["x"], ["y"])], - "implicit_axis", - [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], - [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], - ), - opset_imports=[helper.make_opsetid("", 11)], - ) - model.ir_version = 9 - data = np.random.default_rng(25).standard_normal(shape).astype(np.float32) - expected = _execute(model, {"x": data})[0] - transformed = _apply_update_axis(copy.deepcopy(model)) - for value in [transformed.graph.input[0], transformed.graph.output[0]]: - value.CopyFrom(helper.make_tensor_value_info(value.name, TensorProto.FLOAT, shape[1:])) - np.testing.assert_allclose( - _execute(transformed, {"x": data.reshape(shape[1:])})[0].reshape(shape), - expected, - rtol=1e-6, - atol=1e-7, - ) + """Read the imported default axis after an actual registered Reshape collapse.""" + model = _collapsed_normalization_model([1, 2, 3, 4, 5], None, operation, 11) + feeds = {"x": np.random.default_rng(25).standard_normal((2, 3, 4, 5)).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) def test_implicit_axis_on_removed_dimension_refuses_without_mutation(): @@ -755,3 +754,131 @@ def test_implicit_split_axis_preserves_rank_three_values(): transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) for actual, reference in zip(_execute(transformed, feeds), expected): np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize( + ("rank", "empty", "runtime", "fixed_data"), + [ + (2, False, False, False), + (3, True, False, False), + (3, False, True, False), + (3, True, True, False), + (3, False, False, True), + ], +) +def test_implicit_slice_axes_preserve_registered_data_layouts(rank, empty, runtime, fixed_data): + """Derive implicit axes from the fixed starts length before promoting data.""" + shape = [2, 3] if rank == 2 else [2, 3, 4] + output_shape = [2, 2] if rank == 2 else [2, 2, 4] + names = ["x", "starts", "ends", ""] if empty else ["x", "starts", "ends"] + initializers = [numpy_helper.from_array(np.asarray([2, 3], np.int64), "ends")] + inputs = [] if fixed_data else [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + feeds = {} if fixed_data else {"x": data} + if fixed_data: + initializers.append(numpy_helper.from_array(data, "x")) + if runtime: + inputs.append(helper.make_tensor_value_info("starts", TensorProto.INT64, [2])) + feeds["starts"] = np.asarray([0, 1], np.int64) + else: + initializers.append(numpy_helper.from_array(np.asarray([0, 1], np.int64), "starts")) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Slice", names, ["y"])], + "implicit_slice", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) + + +def test_unknown_implicit_slice_length_refuses_before_earlier_mutation(): + """Reject an unknown runtime starts length before allocating an earlier fixed axis view.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "fixed", "ends"], ["first"]), + helper.make_node("Slice", ["x", "dynamic", "ends"], ["second"]), + ], + "unknown_starts", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("dynamic", TensorProto.INT64, [None]), + ], + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [None, None]) + for n in ["first", "second"] + ], + [ + numpy_helper.from_array(np.asarray([0, 0], np.int64), "fixed"), + numpy_helper.from_array(np.asarray([2, 3], np.int64), "ends"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + before = model.SerializeToString() + with pytest.raises(ValueError, match="known starts length"): + _apply_non4d_slice_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("producer", ["input", "Unsqueeze"]) +def test_native_five_dimensional_axes_refuse_without_a_collapsed_reshape(producer): + """Keep earlier axes unchanged when a later native input never lost a dimension.""" + shape = [1, 2, 3, 4, 5] + nodes = [helper.make_node("Softmax", ["first"], ["first_out"], axis=1)] + inputs = [helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4])] + initializers = [] + if producer == "input": + inputs.append(helper.make_tensor_value_info("native", TensorProto.FLOAT, shape)) + else: + inputs.append(helper.make_tensor_value_info("x", TensorProto.FLOAT, shape[1:])) + initializers.append(numpy_helper.from_array(np.asarray([0], np.int64), "axes")) + nodes.append(helper.make_node("Unsqueeze", ["x", "axes"], ["native"])) + nodes.append(helper.make_node("Softmax", ["native"], ["y"], axis=2)) + model = helper.make_model( + helper.make_graph( + nodes, + "native_five", + inputs, + [ + helper.make_tensor_value_info("first_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + ], + initializers, + value_info=[helper.make_tensor_value_info("native", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="collapsed Reshape"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("shape", "target", "data_shape", "axis"), + [ + ([1, 2, 1, 3, 4], [-1, 2, 1, 3, 4], [1, 2, 3, 4], 1), + ([1, 2, 4, 3, 5], [1, 2, 0, 3, 5], [2, 3, 4, 5], 2), + ], +) +def test_five_dimensional_axis_tracks_the_actually_removed_target_dimension( + shape, target, data_shape, axis +): + """Respect inferred and copied dimensions when locating the singleton actually removed.""" + model = _collapsed_normalization_model(shape, axis, target=target, data_shape=data_shape) + feeds = {"x": np.random.default_rng(29).standard_normal(data_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) From b4c8d6ae74127d264aa83b24faf30314201477a6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 02:20:51 +0000 Subject: [PATCH 06/41] Preserve native index controls and public fixed layouts in rank conversion Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 40 +++++-- .../test_dla_rank_conversion.py | 106 ++++++++++++++++++ 2 files changed, 137 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 85ca5a39353..8ee22381bb0 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -84,13 +84,23 @@ def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) reshapes = _operator_nodes(model, {"Reshape"}) + gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} reshape_ids = {id(node) for node in reshapes} squeezes = {id(node) for node in _operator_nodes(model, {"Squeeze"})} replacements = {} users = {} captured = _subgraph_captures(model.graph) for node in reshapes: + if any( + index == 1 and id(user) in gathers + for user, index in cache.get_consumers(node.output[0]) + ): + continue tensor = cache.get_init(node.input[1]) + target_shape = cache.get_shape(node.input[1]) + native_input = _shape_before_boundary_promotion(cache, node.input[0]) + if target_shape == [4] and native_input is not None and len(native_input) == 4: + continue if tensor is None or tensor.name in cache.graph_input_names: raise ValueError(f"{node.name!r}: Reshape target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() @@ -129,6 +139,13 @@ def _apply_non4d_reshape(model): if 1 not in values: raise ValueError(f"{node.name!r}: 5D Reshape shape has no unary dimension") values.remove(1) + elif len(values) == 4 and 0 in values: + output_shape = cache.get_shape(node.output[0]) + if output_shape is None or len(output_shape) != 4: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") + values = [output_shape[i] if value == 0 else value for i, value in enumerate(values)] + if 0 in values: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") else: continue replacement = numpy_helper.from_array(np.asarray(values, dtype=np.int64), tensor.name) @@ -167,11 +184,15 @@ def _apply_non4d_initializers(model): model, {"QuantizeLinear", "DequantizeLinear", "GroupQueryAttention"}, ("com.microsoft",) ) selected = {id(node) for node in supported} - protected = _subgraph_captures(model.graph) + protected = _subgraph_captures(model.graph) | {value.name for value in model.graph.output} protected.update( name for node in model.graph.node if id(node) not in selected for name in node.input ) - used = {name for node in supported for name in node.input} if protected else set() + used = ( + {name for node in supported for name in node.input} + if any(tensor.name in protected for tensor in model.graph.initializer) + else set() + ) for node in supported: if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: expand.update(node.input) @@ -188,17 +209,18 @@ def _apply_non4d_initializers(model): if tensor.name in avoid: continue rank = len(tensor.dims) + promote = ( + (rank == 1 and tensor.name in expand) + or (rank == 2 and tensor.name in expand_or_front) + or rank == 3 + ) if tensor.name in protected: - if tensor.name in used and rank in (1, 2, 3): + if tensor.name in used and promote: raise ValueError( - f"Initializer {tensor.name!r} has incompatible opaque or captured usage" + f"Initializer {tensor.name!r} has incompatible public, opaque or captured usage" ) continue - if ( - (rank == 1 and tensor.name in expand) - or (rank == 2 and tensor.name in expand_or_front) - or rank == 3 - ): + if promote: shape = pad4d(list(tensor.dims)) updates.append((tensor, shape)) elif rank == 2 and tensor.name not in skip: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index c265c2a7dd8..2f984faf4cb 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -882,3 +882,109 @@ def test_five_dimensional_axis_tracks_the_actually_removed_target_dimension( expected = _execute(model, feeds)[0] transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +def test_public_initializer_outputs_keep_their_native_dimensions(shape): + """Preserve fixed public values that have no internal layout-changing consumer.""" + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + model = helper.make_model( + helper.make_graph( + [], + "public_weight", + [], + [helper.make_tensor_value_info("weight", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(data, "weight")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + expected = _execute(model, {})[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, {})[0], expected) + + +@pytest.mark.parametrize("target", [[0, 1, 1, -1], [1, 0, 1, -1]]) +def test_four_dimensional_reshape_targets_resolve_native_copied_dimensions(target): + """Resolve copied positions before input promotion shifts native dimensions.""" + output_shape = [2, 1, 1, 3] if target[0] == 0 else [1, 3, 1, 2] + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "copied_four", + [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(target, np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +def test_runtime_four_dimensional_reshape_target_preserves_runtime_overrides(): + """Leave runtime targets native when neither input nor output rank needs promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "runtime_four", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("target", TensorProto.INT64, [4]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, None, None, None])], + [numpy_helper.from_array(np.asarray([1, 1, 2, 3], np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3), + "target": np.asarray([1, 2, 1, 3], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + ("producer", "copied"), + [("Reshape", False), ("Reshape", True), ("Squeeze", False), ("Unsqueeze", False)], +) +def test_native_index_views_survive_rank_preparation(producer, copied): + """Retain scalar and vector control ranks after upstream structural nodes are promoted.""" + source_shape = [] if producer == "Unsqueeze" else ([3, 1] if copied else [1, 3]) + count = 1 if producer == "Unsqueeze" else 3 + values = ( + np.asarray(2, np.int64) + if not source_shape + else np.arange(3, dtype=np.int64).reshape(source_shape) + ) + controls = np.asarray([0] if producer != "Reshape" or copied else [-1], np.int64) + nodes = [ + helper.make_node(producer, ["source", "control"], ["indices"]), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=0), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "native_index_views", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [5, 4]), + helper.make_tensor_value_info("source", TensorProto.INT64, source_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [count, 4])], + [numpy_helper.from_array(controls, "control")], + value_info=[helper.make_tensor_value_info("indices", TensorProto.INT64, [count])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + feeds = {"x": np.arange(20, dtype=np.float32).reshape(5, 4), "source": values} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) From 8c279e2c9ca31edd37288048d0c8254e43a3f3bf Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 02:50:59 +0000 Subject: [PATCH 07/41] Preserve symbolic outputs and legacy rank controls Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 12 +- .../dla_transforms/dla_rank_conversion.py | 32 ++++- .../test_dla_rank_conversion.py | 115 ++++++++++++++++++ 3 files changed, 153 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 eada7ddaca3..49c29b9a11a 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -398,9 +398,19 @@ def _apply_non4d_model_outputs( native_broadcasts.update(node.output) rank_sources = {} contrib_ops = _boundary_contrib_ops(model) + inputs_promotable = all( + cache.get_shape(value.name) is not None + for value in graph.input + if value.name not in cache.init_names + ) for value in graph.output: shape: list | None = cache.get_shape(value.name) - if shape is None and defer_rank_conversion and value.type.tensor_type.HasField("shape"): + if ( + shape is None + and defer_rank_conversion + and inputs_promotable + and value.type.tensor_type.HasField("shape") + ): shape = [ dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None for dim in value.type.tensor_type.shape.dim diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 8ee22381bb0..bed53a5e8be 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -38,9 +38,18 @@ __all__ = [] +def _runtime_input_names(model): + """Exclude fixed initializer declarations required by legacy ONNX IR.""" + names = {value.name for value in model.graph.input} + if model.ir_version <= 3: + names.difference_update(tensor.name for tensor in model.graph.initializer) + return names + + def _apply_non4d_slice_axis(model): """Remap fixed Slice axes before their data rank is promoted.""" cache = GraphCache(model.graph) + runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {"Slice"}): shape = _shape_before_boundary_promotion(cache, node.input[0]) @@ -56,7 +65,7 @@ def _apply_non4d_slice_axis(model): plans.append((node, node.input[1], values)) continue axes = cache.get_init(node.input[3]) - if axes is None or axes.name in cache.graph_input_names: + if axes is None or axes.name in runtime_inputs: raise ValueError(f"{node.name!r}: Slice axes must be fixed") values = numpy_helper.to_array(axes).reshape(-1) if values.size and np.any(values >= 0): @@ -83,6 +92,7 @@ def _apply_non4d_slice_axis(model): def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) + runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} reshape_ids = {id(node) for node in reshapes} @@ -101,7 +111,7 @@ def _apply_non4d_reshape(model): native_input = _shape_before_boundary_promotion(cache, node.input[0]) if target_shape == [4] and native_input is not None and len(native_input) == 4: continue - if tensor is None or tensor.name in cache.graph_input_names: + if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: Reshape target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() if len(values) < 4: @@ -154,7 +164,7 @@ def _apply_non4d_reshape(model): replacements[tensor.name] = replacement users.setdefault(tensor.name, []).append(node) for name, replacement in replacements.items(): - private = name in cache.graph_output_names | captured or any( + private = name in cache.graph_output_names | cache.graph_input_names | captured or any( index != 1 or id(user) not in reshape_ids for user, index in cache.get_consumers(name) ) if private: @@ -203,7 +213,7 @@ def _apply_non4d_initializers(model): if node.op_type in {"Unsqueeze", "GroupQueryAttention"}: avoid.update(node.input) updates = [] - avoid.update(value.name for value in model.graph.input) + avoid.update(_runtime_input_names(model)) expand_or_front = expand | front for tensor in model.graph.initializer: if tensor.name in avoid: @@ -231,6 +241,9 @@ def _apply_non4d_initializers(model): for node in model.graph.node if node.input and node.op_type not in _RANK_AGNOSTIC_OPS } + legacy_inputs = ( + {value.name: value for value in model.graph.input} if model.ir_version <= 3 else {} + ) for tensor, shape in updates: if tensor.name in native_inputs and tensor.name not in value_info: # Later axis stages read native ranks from the same deferred metadata snapshot. @@ -239,16 +252,23 @@ def _apply_non4d_initializers(model): ) del tensor.dims[:] tensor.dims.extend(shape) + if tensor.name in legacy_inputs: + legacy_inputs[tensor.name].type.tensor_type.shape.CopyFrom( + helper.make_tensor_value_info( + tensor.name, tensor.data_type, shape + ).type.tensor_type.shape + ) return model def _pad_shape_inputs(model, op_type, suffix): """Plan fixed structural views before changing nodes or allocating controls.""" cache = GraphCache(model.graph) + runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {op_type}): tensor = cache.get_init(node.input[1]) - if tensor is None or tensor.name in cache.graph_input_names: + if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: {op_type} target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() if len(values) < 4: @@ -476,6 +496,8 @@ def _apply_update_axis(model): else: value = axis.i if node.op_type not in update: + if shape is not None and len(shape) not in (2, 3, 5): + continue raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") if "axis_updated" in node.name: continue diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 2f984faf4cb..a5302608e45 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -988,3 +988,118 @@ def test_native_index_views_survive_rank_preparation(producer, copied): expected = _execute(model, feeds)[0] converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_symbolic_input_output_remains_native(batch): + """Leave symbolic input ranks intact instead of restoring an unpromoted output.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Softmax", ["x"], ["y"], axis=-1)], + "symbolic_native", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["ArgMin", "Hardmax", "LayerNormalization"]) +def test_axis_update_preserves_native_four_dimensional_operators(operation): + """Allow axis-bearing operators that already consume their final rank.""" + shape = [2, 3, 4, 5] + inputs = ["x"] + initializers = [] + output_shape = [2, 1, 4, 5] if operation == "ArgMin" else shape + if operation == "LayerNormalization": + inputs.append("scale") + initializers.append(numpy_helper.from_array(np.ones(shape[1:], np.float32), "scale")) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"], axis=1)], + "native_four_dimensional", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMin" else TensorProto.FLOAT, + output_shape, + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(120, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _apply_update_axis(copy.deepcopy(model)) + assert transformed == model + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile", "Slice", "Add"]) +def test_legacy_initializer_controls_remain_fixed(operation): + """Execute IR3 initializer declarations through boundary and rank conversion.""" + shape = [2, 3] + controls = { + "Reshape": [3, 2], + "Expand": [2, 3], + "Tile": [2, 1], + "Slice": [0], + "Add": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + } + output_shape = {"Reshape": [3, 2], "Tile": [4, 3], "Slice": [1, 3]}.get(operation, shape) + inputs = ["x", "control"] + initializers = [ + numpy_helper.from_array( + np.asarray(controls[operation], np.float32 if operation == "Add" else np.int64), + "control", + ) + ] + if operation == "Slice": + for name, values in (("ends", [1]), ("axes", [0])): + initializers.append(numpy_helper.from_array(np.asarray(values, np.int64), name)) + inputs.append(name) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"])], + "legacy_controls", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + + [ + helper.make_tensor_value_info(t.name, t.data_type, list(t.dims)) + for t in initializers + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = copy.deepcopy(model) + transforms = ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + { + "Reshape": _apply_non4d_reshape, + "Expand": _apply_non4d_expand, + "Tile": _apply_non4d_tile, + "Slice": _apply_non4d_slice_axis, + "Add": _apply_non4d_initializers, + }[operation], + ) + for apply in transforms: + transformed = apply( + transformed, + **({"defer_rank_conversion": True} if apply is _apply_non4d_model_outputs else {}), + ) + if operation == "Add": + assert list(transformed.graph.initializer[0].dims) == [1, 1, 2, 3] + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) From d317d86f7c7167b51ae0c89360b69857ee479212 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:07:38 +0000 Subject: [PATCH 08/41] Map implicit per-channel quantizer axes after promotion Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 11 ++- .../test_dla_rank_conversion.py | 68 +++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index bed53a5e8be..e59166aafe9 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -481,7 +481,16 @@ def _apply_update_axis(model): and node.op_type in update and node.domain in {"", "ai.onnx"} and shape - and (len(shape) == 5 or (len(shape) in (2, 3) and node.op_type not in quantizers)) + and ( + len(shape) == 5 + or ( + len(shape) in (2, 3) + and ( + node.op_type not in quantizers + or cache.get_shape(node.input[1]) not in ([], [1]) + ) + ) + ) ): schema = defs.get_schema( node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index a5302608e45..5dbdec01bbc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1103,3 +1103,71 @@ def test_legacy_initializer_controls_remain_fixed(operation): if operation == "Add": assert list(transformed.graph.initializer[0].dims) == [1, 1, 2, 3] np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("per_channel", [False, True, "singleton"]) +def test_registered_implicit_quantizer_axis_preserves_channel_values(operation, shape, per_channel): + """Map surviving per-channel defaults while leaving scalar quantization attributes native.""" + scales = np.asarray( + [0.25, 0.5, 1.0] if per_channel is True else [0.5] if per_channel else 0.5, np.float32 + ) + zeros = np.asarray([0, 1, 2] if per_channel is True else [0] if per_channel else 0, np.uint8) + input_type = TensorProto.FLOAT if operation == "QuantizeLinear" else TensorProto.UINT8 + output_type = TensorProto.UINT8 if operation == "QuantizeLinear" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x", "scale", "zero"], ["y"])], + "implicit_quantizer", + [helper.make_tensor_value_info("x", input_type, shape)], + [helper.make_tensor_value_info("y", output_type, shape)], + [numpy_helper.from_array(scales, "scale"), numpy_helper.from_array(zeros, "zero")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + data = ( + np.arange(np.prod(shape)) + .reshape(shape) + .astype(np.float32 if input_type == TensorProto.FLOAT else np.uint8) + ) + feeds = {"x": data} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + quantizer = next(node for node in transformed.graph.node if node.op_type == operation) + if per_channel is True: + assert next(attr.i for attr in quantizer.attribute if attr.name == "axis") == 5 - len(shape) + else: + assert all(attr.name != "axis" for attr in quantizer.attribute) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) +def test_legacy_quantizer_schema_keeps_implicit_scalar_axis(operation): + """Keep pre-axis quantizer schemas byte-identical during axis preparation.""" + input_type = TensorProto.FLOAT if operation == "QuantizeLinear" else TensorProto.UINT8 + output_type = TensorProto.UINT8 if operation == "QuantizeLinear" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x", "scale", "zero"], ["y"])], + "legacy_quantizer", + [helper.make_tensor_value_info("x", input_type, [2, 3])], + [helper.make_tensor_value_info("y", output_type, [2, 3])], + [ + numpy_helper.from_array(np.asarray(0.5, np.float32), "scale"), + numpy_helper.from_array(np.asarray(0, np.uint8), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 10)], + ) + model.ir_version = 9 + original = model.SerializeToString() + _apply_update_axis(model) + assert model.SerializeToString() == original + feeds = { + "x": np.arange(6) + .reshape(2, 3) + .astype(np.float32 if input_type == TensorProto.FLOAT else np.uint8) + } + _execute(model, feeds) From 98029bce1d4149881a35dda13c8dcf14a894df9a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:41:51 +0000 Subject: [PATCH 09/41] Preserve single-channel per-axis quantization semantics Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_rank_conversion.py | 5 +---- .../dla_transforms/test_dla_rank_conversion.py | 7 +++++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index e59166aafe9..cc2da23b824 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -485,10 +485,7 @@ def _apply_update_axis(model): len(shape) == 5 or ( len(shape) in (2, 3) - and ( - node.op_type not in quantizers - or cache.get_shape(node.input[1]) not in ([], [1]) - ) + and (node.op_type not in quantizers or cache.get_shape(node.input[1]) != []) ) ) ): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 5dbdec01bbc..3e0e589e224 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1109,7 +1109,10 @@ def test_legacy_initializer_controls_remain_fixed(operation): @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) @pytest.mark.parametrize("per_channel", [False, True, "singleton"]) def test_registered_implicit_quantizer_axis_preserves_channel_values(operation, shape, per_channel): - """Map surviving per-channel defaults while leaving scalar quantization attributes native.""" + """Map every vector scale's channel axis, including a single-channel vector.""" + shape = shape.copy() + if per_channel == "singleton": + shape[1] = 1 scales = np.asarray( [0.25, 0.5, 1.0] if per_channel is True else [0.5] if per_channel else 0.5, np.float32 ) @@ -1136,7 +1139,7 @@ def test_registered_implicit_quantizer_axis_preserves_channel_values(operation, expected = _execute(model, feeds)[0] transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) quantizer = next(node for node in transformed.graph.node if node.op_type == operation) - if per_channel is True: + if per_channel: assert next(attr.i for attr in quantizer.attribute if attr.name == "axis") == 5 - len(shape) else: assert all(attr.name != "axis" for attr in quantizer.attribute) From 88bff6abf6c3c9ed43d95814c4637dc675f9f482 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 04:22:18 +0000 Subject: [PATCH 10/41] Preserve native four-dimensional runtime shape controls Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 8 ++- .../test_dla_rank_conversion.py | 57 +++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index cc2da23b824..598a9c2c0a5 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -33,7 +33,7 @@ ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures -from .dla_indexing import _insert_support_nodes, _operator_nodes +from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes __all__ = [] @@ -262,12 +262,16 @@ def _apply_non4d_initializers(model): def _pad_shape_inputs(model, op_type, suffix): - """Plan fixed structural views before changing nodes or allocating controls.""" + """Preserve native rank-four layouts and plan fixed control padding otherwise.""" cache = GraphCache(model.graph) runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {op_type}): tensor = cache.get_init(node.input[1]) + native = _shape_before_boundary_promotion(cache, node.input[0]) + rank = len(native) if native is not None else _known_rank(cache, node.input[0]) + if rank == 4 and cache.get_shape(node.input[1]) == [4]: + continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: {op_type} target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 3e0e589e224..8fbbe28b6a0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1174,3 +1174,60 @@ def test_legacy_quantizer_schema_keeps_implicit_scalar_axis(operation): .astype(np.float32 if input_type == TensorProto.FLOAT else np.uint8) } _execute(model, feeds) + + +@pytest.mark.parametrize("operation", ["Expand", "Tile"]) +@pytest.mark.parametrize("source", ["input", "default", "shape", "symbolic"]) +def test_native_four_dimensional_runtime_controls_need_no_padding(operation, source): + """Execute native rank-four runtime layouts without freezing their controls.""" + shape = [1, 1, 2, 3] + first = [2, 1, 2, 3] if operation == "Expand" else [1, 1, 2, 1] + second = [3, 1, 2, 3] if operation == "Expand" else [2, 1, 1, 1] + declared = ["batch", 1, 2, 3] if source == "symbolic" else shape + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, declared)] + nodes = [] + initializers = [] + if source == "shape": + nodes.append(helper.make_node("Shape", ["x"], ["control"])) + else: + inputs.append(helper.make_tensor_value_info("control", TensorProto.INT64, [4])) + if source == "default": + initializers.append(numpy_helper.from_array(np.asarray(first, np.int64), "control")) + nodes.append(helper.make_node(operation, ["x", "control"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_runtime_controls", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None] * 4)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + data = np.arange(6, dtype=np.float32).reshape(shape) + feeds = ( + [{"x": data}] + if source in {"shape", "default"} + else [{"x": data, "control": np.asarray(first, np.int64)}] + ) + if source != "shape": + feeds.append({"x": data, "control": np.asarray(second, np.int64)}) + expected = [_execute(model, values)[0] for values in feeds] + if source == "shape": + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + else: + # ORT symbolic inference cannot resolve unconstrained runtime Expand controls. + transformed = dict(DLA_TRANSFORMS)["Non4D" + operation](copy.deepcopy(model)) + assert transformed.SerializeToString() == model.SerializeToString() + if source == "shape": + assert any(node.op_type == "Shape" for node in transformed.graph.node) + else: + assert any(value.name == "control" for value in transformed.graph.input) + if initializers: + assert ( + next(t for t in transformed.graph.initializer if t.name == "control") + == model.graph.initializer[0] + ) + for values, reference in zip(feeds, expected): + np.testing.assert_array_equal(_execute(transformed, values)[0], reference) From 91b59bf7ba02eacb96254e1815367c7ee42bd8d6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 06:30:36 +0000 Subject: [PATCH 11/41] Isolate reshape targets from skipped native owners Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 4 +-- .../test_dla_rank_conversion.py | 36 +++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 598a9c2c0a5..e5997f9feaf 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -95,7 +95,6 @@ def _apply_non4d_reshape(model): runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} - reshape_ids = {id(node) for node in reshapes} squeezes = {id(node) for node in _operator_nodes(model, {"Squeeze"})} replacements = {} users = {} @@ -164,8 +163,9 @@ def _apply_non4d_reshape(model): replacements[tensor.name] = replacement users.setdefault(tensor.name, []).append(node) for name, replacement in replacements.items(): + converted = {id(user) for user in users[name]} private = name in cache.graph_output_names | cache.graph_input_names | captured or any( - index != 1 or id(user) not in reshape_ids for user, index in cache.get_consumers(name) + index != 1 or id(user) not in converted for user, index in cache.get_consumers(name) ) if private: target = add_initializer( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 8fbbe28b6a0..a9c5ebfd479 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -91,6 +91,42 @@ def test_rank_conversion_preserves_operator_values(operation): np.testing.assert_array_equal(actual[0], expected[0]) +def test_converted_reshape_target_preserves_skipped_native_owner(): + """Give the converted Reshape a private target when a native 4D owner is skipped.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["a", "target"], ["ay"]), + helper.make_node("Reshape", ["b", "target"], ["by"]), + ], + "shared_target", + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [5, 6, 1, 1]), + ], + [ + helper.make_tensor_value_info("ay", TensorProto.FLOAT, [2, 1, 1, 3]), + helper.make_tensor_value_info("by", TensorProto.FLOAT, [5, 1, 1, 6]), + ], + [numpy_helper.from_array(np.array([0, 1, 1, -1], np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "a": np.arange(6, dtype=np.float32).reshape(2, 3), + "b": np.arange(30, dtype=np.float32).reshape(5, 6, 1, 1), + } + expected = ort.InferenceSession(model.SerializeToString()).run(None, feeds) + converted = _apply_non4d_reshape(copy.deepcopy(model)) + assert converted.graph.node[0].input[1] != converted.graph.node[1].input[1] + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip( + ort.InferenceSession(converted.SerializeToString()).run(None, feeds), expected + ): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("shape", [[4], [3, 4], [2, 3, 4]]) def test_initializer_rank_conversion_preserves_overridable_defaults(shape): """Preserve a runtime input's rank and default payload while preparing fixed weights.""" From 148c0c4b27e7d87b2c540e846f59ae68e35b2740 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 07:01:56 +0000 Subject: [PATCH 12/41] Remap specialized axes after singleton reshape collapse Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 42 +++++++++++--- .../test_dla_rank_conversion.py | 56 +++++++++++++++++++ 2 files changed, 89 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index e5997f9feaf..591127fe5b7 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -355,18 +355,39 @@ def _apply_non4d_transpose(model): return model +def _specialized_collapsed_axis(cache, node, shape, value, reshapes): + """Map a specialized operator axis only when its native singleton was actually collapsed.""" + removed = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not removed: + return value + if len(removed) != 1 or value % 5 in removed: + raise ValueError( + f"{node.name!r}: cannot map axis {value} through an ambiguous collapsed Reshape" + ) + native = value % 5 + return native - int(native > removed[0]) - (4 if value < 0 else 0) + + def _apply_non4d_lpnorm(model): """Preserve normalization axes across leading-dimension promotion.""" cache = GraphCache(model.graph) + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + plans = [] for node in _operator_nodes(model, {"LpNormalization"}): - if node.op_type == "LpNormalization": - axis = _attribute(node, "axis") - shape = _shape_before_boundary_promotion(cache, node.input[0]) - if axis is not None and axis.i >= 0 and shape and len(shape) in (2, 3): - axis.i = -1 if axis.i == len(shape) - 1 else axis.i + 4 - len(shape) - elif axis is None: - node.attribute.append(helper.make_attribute("axis", -1)) - node.name += "_axis_updated" + axis = _attribute(node, "axis") + shape = _shape_before_boundary_promotion(cache, node.input[0]) + value = axis.i if axis is not None else -1 + if shape and len(shape) == 5: + value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) + elif value >= 0 and shape and len(shape) in (2, 3): + value = -1 if value == len(shape) - 1 else value + 4 - len(shape) + plans.append((node, axis, value)) + for node, axis, value in plans: + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name += "_axis_updated" return model @@ -413,6 +434,7 @@ def _apply_argmax(model): """Plan native ArgMax layouts before preparing private output restoration views.""" cache = GraphCache(model.graph) plans = {} + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"ArgMax"}): shape = _shape_before_boundary_promotion(cache, node.input[0]) if not shape or any(dim <= 0 for dim in shape): @@ -422,7 +444,9 @@ def _apply_argmax(model): if not -len(shape) <= value < len(shape): raise ValueError(f"{node.name!r}: invalid ArgMax axis {value}") output_shape = [dim for index, dim in enumerate(shape) if index != value % len(shape)] - if value >= 0: + if len(shape) == 5: + value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) + elif value >= 0: value += {2: 2, 3: 1}.get(len(shape), 0) plans[id(node)] = (value, output_shape) nodes = [] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index a9c5ebfd479..5b67654c9de 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -920,6 +920,62 @@ def test_five_dimensional_axis_tracks_the_actually_removed_target_dimension( np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) +@pytest.mark.parametrize("operation", ["ArgMax", "LpNormalization"]) +@pytest.mark.parametrize("axis", [2, -3]) +@pytest.mark.parametrize("collapse", [False, True]) +def test_specialized_axes_follow_only_actually_collapsed_reshapes(operation, axis, collapse): + """Remap specialized axes after a singleton collapse while keeping untouched native 5D axes.""" + shape = [1, 2, 3, 4, 5] + data_shape = shape[1:] if collapse else shape + reduced = [1, 2, 4, 5] + output_shape = reduced if operation == "ArgMax" else data_shape + inits = [] + nodes = [] + source = "x" + if collapse: + inits.append(numpy_helper.from_array(np.array(shape, np.int64), "native_shape")) + nodes.append(helper.make_node("Reshape", ["x", "native_shape"], ["native"])) + source = "native" + nodes.append( + helper.make_node( + operation, + [source], + ["result" if collapse else "y"], + axis=axis, + **({"keepdims": 0} if operation == "ArgMax" else {}), + ) + ) + if collapse: + inits.append(numpy_helper.from_array(np.array(output_shape, np.int64), "output_shape")) + nodes.append(helper.make_node("Reshape", ["result", "output_shape"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "specialized_axes", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMax" else TensorProto.FLOAT, + output_shape, + ) + ], + inits, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.random.default_rng(73).standard_normal(data_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + converted = copy.deepcopy(model) + if collapse: + converted = _transform_make_dla_compatible(converted, verbose=False) + else: + converted = (_apply_argmax if operation == "ArgMax" else _apply_non4d_lpnorm)(converted) + np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) def test_public_initializer_outputs_keep_their_native_dimensions(shape): """Preserve fixed public values that have no internal layout-changing consumer.""" From 65c6462d8c864f8c9835be9070427ca900cc38ad Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:17:29 +0000 Subject: [PATCH 13/41] Preserve native symbolic layouts in DLA structural controls Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 14 +++ .../test_dla_rank_conversion.py | 111 ++++++++++++++++++ 2 files changed, 125 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 591127fe5b7..dd91b5220c4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -34,6 +34,7 @@ from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes +from .dla_model_boundaries import _has_native_symbolic_data __all__ = [] @@ -92,6 +93,7 @@ def _apply_non4d_slice_axis(model): def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} @@ -110,6 +112,12 @@ def _apply_non4d_reshape(model): native_input = _shape_before_boundary_promotion(cache, node.input[0]) if target_shape == [4] and native_input is not None and len(native_input) == 4: continue + if native_input is None: + rank = _known_rank(cache, node.input[0]) + if rank is None: + raise ValueError(f"{node.name!r}: Reshape input rank must be known") + if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: Reshape target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() @@ -264,6 +272,7 @@ def _apply_non4d_initializers(model): def _pad_shape_inputs(model, op_type, suffix): """Preserve native rank-four layouts and plan fixed control padding otherwise.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {op_type}): @@ -272,6 +281,11 @@ def _pad_shape_inputs(model, op_type, suffix): rank = len(native) if native is not None else _known_rank(cache, node.input[0]) if rank == 4 and cache.get_shape(node.input[1]) == [4]: continue + if native is None: + if rank is None: + raise ValueError(f"{node.name!r}: {op_type} input rank must be known") + if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: {op_type} target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 5b67654c9de..513f547e8b6 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1082,6 +1082,117 @@ def test_native_index_views_survive_rank_preparation(producer, copied): np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("static_output", [False, True]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_structural_controls_keep_native_data_rank(operation, batch, static_output, tail): + """Retain symbolic native layouts, including fixed public outputs and unary readers.""" + target = { + "Tile": [0, 1] if static_output else [2, 1], + "Expand": [2, 3] if static_output else [1, 3], + "Reshape": [3, 1] if static_output else [3, -1], + }[operation] + if static_output and operation == "Reshape" and batch != 1: + target = [3, 2] + output_shape = ( + {"Tile": [0, 3], "Expand": [2, 3], "Reshape": [3, batch]}[operation] + if static_output + else ([3, "N"] if operation == "Reshape" else ["M", 3]) + ) + nodes = [helper.make_node(operation, ["x", "control"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "symbolic_control", + [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")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 3, dtype=np.float32).reshape(batch, 3)} + expected = _execute(model, feeds)[0] + original_control = model.graph.initializer[0].SerializeToString() + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + assert converted.graph.initializer[0].SerializeToString() == original_control + + +@pytest.mark.parametrize("operation", ["Tile", "Reshape"]) +def test_runtime_slice_extents_retain_promoted_control_layouts(operation): + """Follow already-promoted data through unknown Slice extents before adapting controls.""" + values = [1, 2, 1] if operation == "Tile" else [4, -1] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "starts", "ends"], ["middle"]), + helper.make_node(operation, ["middle", "control"], ["y"]), + ], + "promoted_slice_control", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("starts", TensorProto.INT64, [2]), + ], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, ["A", "B", 4] if operation == "Tile" else ["A", "B"] + ) + ], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "ends"), + numpy_helper.from_array(np.array(values, np.int64), "control"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "starts": np.array([0, 1], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +def test_unknown_structural_data_rank_refuses_before_mutation(operation): + """Refuse unresolved intermediate ranks before allocating any control view.""" + values = [2, 1] if operation == "Tile" else [2, 3] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Identity", ["x"], ["hidden"]), + helper.make_node(operation, ["hidden", "control"], ["y"]), + ], + "unknown_control_rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [4, 3] if operation == "Tile" else [2, 3] + ) + ], + [numpy_helper.from_array(np.array(values, np.int64), "control")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + _execute(model, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}) + before = model.SerializeToString() + callback = { + "Tile": _apply_non4d_tile, + "Expand": _apply_non4d_expand, + "Reshape": _apply_non4d_reshape, + }[operation] + with pytest.raises(ValueError, match="input rank"): + callback(model) + assert model.SerializeToString() == before + + @pytest.mark.parametrize("batch", [1, 2]) def test_symbolic_input_output_remains_native(batch): """Leave symbolic input ranks intact instead of restoring an unpromoted output.""" From f30a5fd607a4e6653632ec2a2205760bf57ae805 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:30:03 +0000 Subject: [PATCH 14/41] Track native symbolic provenance through concrete control-chain metadata Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 4 +-- .../dla_transforms/dla_rank_conversion.py | 20 ++++++------ .../test_dla_rank_conversion.py | 31 +++++++++++++++++++ 3 files changed, 42 insertions(+), 13 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 49c29b9a11a..ad0a0aab9d0 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -125,10 +125,10 @@ def _has_native_symbolic_data(cache, name, local_functions): 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): + if rank is not None and rank >= 4: return False if name in cache.graph_input_names: - return rank is not None and rank < 4 + return cache.get_shape(name) is None and rank is not None and rank < 4 node = cache.get_producer(name) if ( node is None diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index dd91b5220c4..9b7b8414c19 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -112,12 +112,11 @@ def _apply_non4d_reshape(model): native_input = _shape_before_boundary_promotion(cache, node.input[0]) if target_shape == [4] and native_input is not None and len(native_input) == 4: continue - if native_input is None: - rank = _known_rank(cache, node.input[0]) - if rank is None: - raise ValueError(f"{node.name!r}: Reshape input rank must be known") - if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): - continue + rank = len(native_input) if native_input is not None else _known_rank(cache, node.input[0]) + if rank is None: + raise ValueError(f"{node.name!r}: Reshape input rank must be known") + if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: Reshape target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() @@ -281,11 +280,10 @@ def _pad_shape_inputs(model, op_type, suffix): rank = len(native) if native is not None else _known_rank(cache, node.input[0]) if rank == 4 and cache.get_shape(node.input[1]) == [4]: continue - if native is None: - if rank is None: - raise ValueError(f"{node.name!r}: {op_type} input rank must be known") - if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): - continue + if rank is None: + raise ValueError(f"{node.name!r}: {op_type} input rank must be known") + if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: {op_type} target must be fixed") values = numpy_helper.to_array(tensor).reshape(-1).tolist() diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 513f547e8b6..056bee10740 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1122,6 +1122,37 @@ def test_symbolic_structural_controls_keep_native_data_rank(operation, batch, st assert converted.graph.initializer[0].SerializeToString() == original_control +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +def test_symbolic_native_control_chain_preserves_concrete_intermediate_rank(operation): + """Keep native controls after a symbolic input acquires concrete intermediate metadata.""" + control = {"Tile": [1, 2], "Expand": [2, 3], "Reshape": [3, 2]}[operation] + output = {"Tile": [2, 6], "Expand": [2, 3], "Reshape": [3, 2]}[operation] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "first_shape"], ["middle"]), + helper.make_node(operation, ["middle", "control"], ["y"]), + ], + "native_control_chain", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output)], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "first_shape"), + numpy_helper.from_array(np.array(control, np.int64), "control"), + ], + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + for actual, original in zip(converted.graph.initializer, model.graph.initializer): + assert actual.SerializeToString() == original.SerializeToString() + + @pytest.mark.parametrize("operation", ["Tile", "Reshape"]) def test_runtime_slice_extents_retain_promoted_control_layouts(operation): """Follow already-promoted data through unknown Slice extents before adapting controls.""" From 68da33cac30aebbf299b7d3cea8fc2197d319730 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:33:41 +0000 Subject: [PATCH 15/41] Reuse native data provenance across DLA control plans Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 25 ++++++++++++++----- .../dla_transforms/dla_rank_conversion.py | 10 ++++++-- 2 files changed, 27 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 ad0a0aab9d0..7e9c67fe1ff 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -114,10 +114,14 @@ def _output_rank_source(cache, node, local_functions, sources): return node -def _has_native_symbolic_data(cache, name, local_functions): +def _has_native_symbolic_data(cache, name, local_functions, sources=None): """Prove a low-rank symbolic data path was left native by boundary promotion.""" seen = set() + sources = {} if sources is None else sources while name not in seen: + if name in sources: + result = sources[name] + break seen.add(name) info = cache.vi_map.get(name) rank = ( @@ -126,18 +130,24 @@ def _has_native_symbolic_data(cache, name, local_functions): else None ) if rank is not None and rank >= 4: - return False + result = False + break if name in cache.graph_input_names: - return cache.get_shape(name) is None and rank is not None and rank < 4 + result = cache.get_shape(name) is None and rank is not None and rank < 4 + break 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 + result = False + break name = node.input[0] - raise ValueError("Cyclic data path while planning native symbolic controls") + else: + raise ValueError("Cyclic data path while planning native symbolic controls") + sources.update(dict.fromkeys(seen, result)) + return result def _require_adapter_opset(model): @@ -398,6 +408,7 @@ def _apply_non4d_model_outputs( native_broadcasts.update(node.output) rank_sources = {} contrib_ops = _boundary_contrib_ops(model) + native_sources = {} inputs_promotable = all( cache.get_shape(value.name) is not None for value in graph.input @@ -435,7 +446,9 @@ def _apply_non4d_model_outputs( ): continue if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): - if _has_native_symbolic_data(cache, source.input[0], local_functions): + if _has_native_symbolic_data( + cache, source.input[0], local_functions, native_sources + ): continue target = cache.get_shape(source.input[1]) if source.op_type == "Reshape" else None if ( diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 9b7b8414c19..0cb1215bf89 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -94,6 +94,7 @@ def _apply_non4d_reshape(model): """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} runtime_inputs = _runtime_input_names(model) reshapes = _operator_nodes(model, {"Reshape"}) gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} @@ -115,7 +116,9 @@ def _apply_non4d_reshape(model): rank = len(native_input) if native_input is not None else _known_rank(cache, node.input[0]) if rank is None: raise ValueError(f"{node.name!r}: Reshape input rank must be known") - if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + if rank < 4 and _has_native_symbolic_data( + cache, node.input[0], local_functions, native_sources + ): continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: Reshape target must be fixed") @@ -272,6 +275,7 @@ def _pad_shape_inputs(model, op_type, suffix): """Preserve native rank-four layouts and plan fixed control padding otherwise.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {op_type}): @@ -282,7 +286,9 @@ def _pad_shape_inputs(model, op_type, suffix): continue if rank is None: raise ValueError(f"{node.name!r}: {op_type} input rank must be known") - if rank < 4 and _has_native_symbolic_data(cache, node.input[0], local_functions): + if rank < 4 and _has_native_symbolic_data( + cache, node.input[0], local_functions, native_sources + ): continue if tensor is None or tensor.name in runtime_inputs: raise ValueError(f"{node.name!r}: {op_type} target must be fixed") From e053a48fda6c0f58567ba4ac082cfbad550fd9aa Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:51:59 +0000 Subject: [PATCH 16/41] Cover registered native reducers after deferred DLA output planning Signed-off-by: mgohil --- .../test_dla_rank_conversion.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 056bee10740..9df5785af51 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -91,6 +91,35 @@ def test_rank_conversion_preserves_operator_values(operation): np.testing.assert_array_equal(actual[0], expected[0]) +@pytest.mark.parametrize("operation", ["Flatten", "ArgMax"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_registered_rank_reducers_restore_native_public_outputs(operation, tail): + """Allow deferred output restoration when its actual later conversion stage is installed.""" + shape = [1, 48] if operation == "Flatten" else [1, 4, 4] + attributes = {"axis": 1} if operation == "Flatten" else {"axis": 1, "keepdims": 0} + nodes = [helper.make_node(operation, ["x"], ["middle" if tail else "y"], **attributes)] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "registered_rank_reducer", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT if operation == "Flatten" else TensorProto.INT64, shape + ) + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + def test_converted_reshape_target_preserves_skipped_native_owner(): """Give the converted Reshape a private target when a native 4D owner is skipped.""" model = helper.make_model( From 804a967c7cbeb3429c253bed9774639c4f61618e Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 09:00:57 +0000 Subject: [PATCH 17/41] Preserve native symbolic axis layouts across DLA rank stages Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 19 ++++- .../dla_transforms/dla_rank_conversion.py | 46 ++++++++++-- .../test_dla_rank_conversion.py | 75 ++++++++++++++++++- 3 files changed, 131 insertions(+), 9 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 7e9c67fe1ff..3806dbd0091 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -33,11 +33,26 @@ __all__ = [] + class _BoundaryPreflightError(ValueError): """Prerequisite lowering cannot freeze an exposed runtime control.""" -_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} +_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | { + "Slice", + "Transpose", + "Reshape", + "Expand", + "Tile", + "Softmax", + "LogSoftmax", + "LpNormalization", + "ArgMax", + "QuantizeLinear", + "DequantizeLinear", + "TopK", + "Split", +} _REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) @@ -445,7 +460,7 @@ def _apply_non4d_model_outputs( source, local_functions, {source.op_type}, contrib_ops=contrib_ops ): continue - if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): + if _is_standard_node(source, local_functions, _RANK_FLOW_OPS): if _has_native_symbolic_data( cache, source.input[0], local_functions, native_sources ): diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 0cb1215bf89..3548b9a0f30 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -47,13 +47,29 @@ def _runtime_input_names(model): return names +def _promoted_input_shape(cache, name, local_functions, native_sources): + """Distinguish native symbolic paths from inputs prepared for four-dimensional data.""" + shape = _shape_before_boundary_promotion(cache, name) + rank = len(shape) if shape is not None else _known_rank(cache, name) + native = ( + rank is not None + and rank < 4 + and _has_native_symbolic_data(cache, name, local_functions, native_sources) + ) + return shape, native + + def _apply_non4d_slice_axis(model): """Remap fixed Slice axes before their data rank is promoted.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {"Slice"}): - shape = _shape_before_boundary_promotion(cache, node.input[0]) + shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + if native: + continue if shape is None or len(shape) not in (2, 3): continue if len(node.input) < 4 or not node.input[3]: @@ -340,11 +356,15 @@ def _collapsed_reshape_axes(cache, name, shape, reshapes): def _apply_non4d_transpose(model): """Remap transpose permutations around added or removed unary axes.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} plans = [] reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"Transpose"}): perm = _attribute(node, "perm") - shape = _shape_before_boundary_promotion(cache, node.input[0]) + shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + if native: + continue if perm is None and shape is None: raise ValueError(f"{node.name!r}: default Transpose requires a known input rank") values = list(perm.ints) if perm is not None else list(reversed(range(len(shape)))) @@ -389,11 +409,15 @@ def _specialized_collapsed_axis(cache, node, shape, value, reshapes): def _apply_non4d_lpnorm(model): """Preserve normalization axes across leading-dimension promotion.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} plans = [] for node in _operator_nodes(model, {"LpNormalization"}): axis = _attribute(node, "axis") - shape = _shape_before_boundary_promotion(cache, node.input[0]) + shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + if native: + continue value = axis.i if axis is not None else -1 if shape and len(shape) == 5: value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) @@ -451,10 +475,14 @@ def _apply_flatten(model): def _apply_argmax(model): """Plan native ArgMax layouts before preparing private output restoration views.""" cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} plans = {} reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"ArgMax"}): - shape = _shape_before_boundary_promotion(cache, node.input[0]) + shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + if native: + continue if not shape or any(dim <= 0 for dim in shape): raise ValueError(f"{node.name!r}: ArgMax requires a static input shape") axis = _attribute(node, "axis") @@ -513,6 +541,8 @@ def _apply_update_axis(model): skip = {"Gather", "GatherElements", "LpNormalization", "ArgMax", "Slice", "Flatten"} skip |= {"ReduceMax", "ReduceMin", "ReduceSum", "ReduceMean"} cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} plans = [] reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} selected = _operator_nodes(model, {node.op_type for node in model.graph.node}) @@ -521,7 +551,13 @@ def _apply_update_axis(model): versions = {op.domain: op.version for op in model.opset_import} for node in selected: axis = _attribute(node, "axis") - shape = _shape_before_boundary_promotion(cache, node.input[0]) if node.input else None + shape, native = ( + _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + if node.input + else (None, False) + ) + if native: + continue if ( axis is None and node.op_type in update diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 9df5785af51..e5f3aced6ac 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -120,6 +120,75 @@ def test_registered_rank_reducers_restore_native_public_outputs(operation, tail) np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_symbolic_native_transpose_keeps_its_permutation(explicit, batch): + """Keep rank-two permutation semantics when symbolic graph inputs remain native.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["x"], ["y"], **({"perm": [1, 0]} if explicit else {}))], + "native_transpose", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, "N"])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 3, dtype=np.float32).reshape(batch, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + "operation", ["Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Slice", "Transpose"] +) +@pytest.mark.parametrize("negative", [False, True]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_native_axis_chain_preserves_concrete_intermediate_rank(operation, negative, tail): + """Retain native axes after symbolic data acquires concrete Reshape metadata.""" + axis = -1 if negative else 1 + controls = [numpy_helper.from_array(np.asarray([2, 3], np.int64), "target")] + inputs = ["middle"] + attributes = {"axis": axis} + output = [2, 3] + dtype = TensorProto.FLOAT + if operation == "ArgMax": + attributes["keepdims"] = 0 + output, dtype = [2], TensorProto.INT64 + elif operation == "Transpose": + attributes = {} if negative else {"perm": [1, 0]} + output = [3, 2] + elif operation == "Slice": + attributes = {} + output = [2, 2] + for name, values in [("starts", [0]), ("ends", [2]), ("axes", [axis])]: + inputs.append(name) + controls.append(numpy_helper.from_array(np.asarray(values, np.int64), name)) + nodes = [ + helper.make_node("Reshape", ["x", "target"], ["middle"]), + helper.make_node(operation, inputs, ["result" if tail else "y"], **attributes), + ] + if tail: + nodes.append(helper.make_node("Identity", ["result"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_axis_chain", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", dtype, output)], + controls, + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(1, 7, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + def test_converted_reshape_target_preserves_skipped_native_owner(): """Give the converted Reshape a private target when a native 4D owner is skipped.""" model = helper.make_model( @@ -525,11 +594,13 @@ def test_late_unknown_argmax_shape_refuses_without_changing_an_earlier_node(): "argmax_refusal", [ helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4]), - helper.make_tensor_value_info("second", TensorProto.FLOAT, [None, None, None]), + helper.make_tensor_value_info( + "second", TensorProto.FLOAT, [None, None, None, None] + ), ], [ helper.make_tensor_value_info("first_out", TensorProto.INT64, [2, 4]), - helper.make_tensor_value_info("second_out", TensorProto.INT64, [None, None]), + helper.make_tensor_value_info("second_out", TensorProto.INT64, [None, None, None]), ], ), opset_imports=[helper.make_opsetid("", 17)], From 5a41046899515a41da9654820cb375e10385129e Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:36:49 +0000 Subject: [PATCH 18/41] Preserve native Tile ranks and map axes with symbolic extents Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 61 +++++++++++------ .../test_dla_rank_conversion.py | 66 +++++++++++++++++++ 2 files changed, 106 insertions(+), 21 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 3548b9a0f30..6013e1a21d5 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -56,7 +56,7 @@ def _promoted_input_shape(cache, name, local_functions, native_sources): and rank < 4 and _has_native_symbolic_data(cache, name, local_functions, native_sources) ) - return shape, native + return shape, rank, native def _apply_non4d_slice_axis(model): @@ -67,10 +67,14 @@ def _apply_non4d_slice_axis(model): runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {"Slice"}): - shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources + ) if native: continue - if shape is None or len(shape) not in (2, 3): + if rank is None: + raise ValueError(f"{node.name!r}: Slice input rank must be known") + if rank not in (2, 3): continue if len(node.input) < 4 or not node.input[3]: starts_shape = cache.get_shape(node.input[1]) @@ -78,7 +82,7 @@ def _apply_non4d_slice_axis(model): raise ValueError( f"{node.name!r}: implicit Slice axes require a known starts length" ) - values = np.arange(math.prod(starts_shape), dtype=np.int64) + 4 - len(shape) + values = np.arange(math.prod(starts_shape), dtype=np.int64) + 4 - rank plans.append((node, node.input[1], values)) continue axes = cache.get_init(node.input[3]) @@ -86,7 +90,7 @@ def _apply_non4d_slice_axis(model): raise ValueError(f"{node.name!r}: Slice axes must be fixed") values = numpy_helper.to_array(axes).reshape(-1) if values.size and np.any(values >= 0): - plans.append((node, axes.name, np.where(values >= 0, values + 4 - len(shape), values))) + plans.append((node, axes.name, np.where(values >= 0, values + 4 - rank, values))) views = {} for node, name, values in plans: key = (name, tuple(values)) @@ -224,6 +228,7 @@ def _apply_non4d_initializers(model): protected.update( name for node in model.graph.node if id(node) not in selected for name in node.input ) + protected.update(name for node in supported if node.op_type == "Gemm" for name in node.input) used = ( {name for node in supported for name in node.input} if any(tensor.name in protected for tensor in model.graph.initializer) @@ -232,9 +237,9 @@ def _apply_non4d_initializers(model): for node in supported: if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: expand.update(node.input) - if node.op_type == "MatMul": + if node.op_type in {"MatMul", "Gemm"}: skip.update(node.input) - if node.op_type in {"Gemm", "Where", "Gather", "Split"}: + if node.op_type in {"Where", "Gather", "Split"}: front.update(node.input) if node.op_type in {"Unsqueeze", "GroupQueryAttention"}: avoid.update(node.input) @@ -302,6 +307,8 @@ def _pad_shape_inputs(model, op_type, suffix): continue if rank is None: raise ValueError(f"{node.name!r}: {op_type} input rank must be known") + if op_type == "Tile" and rank == 1: + continue if rank < 4 and _has_native_symbolic_data( cache, node.input[0], local_functions, native_sources ): @@ -362,12 +369,14 @@ def _apply_non4d_transpose(model): reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"Transpose"}): perm = _attribute(node, "perm") - shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources + ) if native: continue - if perm is None and shape is None: + if perm is None and rank is None: raise ValueError(f"{node.name!r}: default Transpose requires a known input rank") - values = list(perm.ints) if perm is not None else list(reversed(range(len(shape)))) + values = list(perm.ints) if perm is not None else list(reversed(range(rank))) if len(values) in (2, 3): delta = 4 - len(values) values = list(range(delta)) + [v + delta for v in values] @@ -415,14 +424,18 @@ def _apply_non4d_lpnorm(model): plans = [] for node in _operator_nodes(model, {"LpNormalization"}): axis = _attribute(node, "axis") - shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources + ) if native: continue + if rank is None or (rank == 5 and shape is None): + raise ValueError(f"{node.name!r}: LpNormalization input layout must be known") value = axis.i if axis is not None else -1 if shape and len(shape) == 5: value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) - elif value >= 0 and shape and len(shape) in (2, 3): - value = -1 if value == len(shape) - 1 else value + 4 - len(shape) + elif value >= 0 and rank in (2, 3): + value = -1 if value == rank - 1 else value + 4 - rank plans.append((node, axis, value)) for node, axis, value in plans: if axis is None: @@ -480,7 +493,9 @@ def _apply_argmax(model): plans = {} reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"ArgMax"}): - shape, native = _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources + ) if native: continue if not shape or any(dim <= 0 for dim in shape): @@ -551,22 +566,27 @@ def _apply_update_axis(model): versions = {op.domain: op.version for op in model.opset_import} for node in selected: axis = _attribute(node, "axis") - shape, native = ( + shape, rank, native = ( _promoted_input_shape(cache, node.input[0], local_functions, native_sources) if node.input - else (None, False) + else (None, None, False) ) if native: continue + if ( + node.op_type in update + and "axis_updated" not in node.name + and (rank is None or (rank == 5 and shape is None)) + ): + raise ValueError(f"{node.name!r}: {node.op_type} input layout must be known") if ( axis is None and node.op_type in update and node.domain in {"", "ai.onnx"} - and shape and ( - len(shape) == 5 + rank == 5 or ( - len(shape) in (2, 3) + rank in (2, 3) and (node.op_type not in quantizers or cache.get_shape(node.input[1]) != []) ) ) @@ -589,7 +609,6 @@ def _apply_update_axis(model): raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") if "axis_updated" in node.name: continue - rank = len(shape) if shape else 0 original = value if shape and rank == 5: removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) @@ -601,7 +620,7 @@ def _apply_update_axis(model): value = native_axis - int(native_axis > removed_axes[0]) if original < 0: value -= 4 - elif value >= 0: + elif value >= 0 and rank is not None: value += {2: 2, 3: 1}.get(rank, 0) if value != original: plans.append((node, axis, value)) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index e5f3aced6ac..8c5dc538d8a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1290,6 +1290,72 @@ def test_runtime_slice_extents_retain_promoted_control_layouts(operation): np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) +@pytest.mark.parametrize( + "operation", ["Softmax", "LogSoftmax", "LpNormalization", "Transpose", "Slice"] +) +def test_runtime_slice_extents_retain_promoted_axes(operation): + """Map native axes using known rank when runtime Slice extents remain symbolic.""" + inputs, attributes = ["middle"], {"axis": 1} + if operation == "Transpose": + attributes = {"perm": [1, 2, 0]} + elif operation == "Slice": + inputs += ["starts", "ends", "axes"] + attributes = {} + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "starts", "ends"], ["middle"]), + helper.make_node(operation, inputs, ["y"], **attributes), + ], + "promoted_slice_axes", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("starts", TensorProto.INT64, [2]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["A", "B", "C"])], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "ends"), + numpy_helper.from_array(np.array([1, 2], np.int64), "axes"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 16, + "starts": np.array([0, 1], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize("initializer", [False, True]) +@pytest.mark.parametrize("tail", [False, True]) +def test_rank_one_tile_data_and_repeats_remain_native(initializer, tail): + """Execute rank-one Tile for graph inputs and fixed data without padding repeats.""" + data = np.array([1, 2, 3], np.float32) + nodes = [helper.make_node("Tile", ["x", "repeats"], ["tiled" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Relu", ["tiled"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_tile", + [] if initializer else [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])], + [numpy_helper.from_array(np.array([2], np.int64), "repeats")] + + ([numpy_helper.from_array(data, "x")] if initializer else []), + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {} if initializer else {"x": data} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) def test_unknown_structural_data_rank_refuses_before_mutation(operation): """Refuse unresolved intermediate ranks before allocating any control view.""" From 258b3fb15b018ceee5cb3ad430cb86fc8811bb4b Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:39:59 +0000 Subject: [PATCH 19/41] Declare prepared rank stages in boundary integration tests Signed-off-by: mgohil --- .../dla_transforms/test_dla_rank_conversion.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 8c5dc538d8a..46b3660f73b 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -78,6 +78,8 @@ def test_rank_conversion_preserves_operator_values(operation): transformed = copy.deepcopy(model) for apply in (_apply_non4d_model_inputs, _apply_non4d_model_outputs, transform): options = {"defer_rank_conversion": True} if apply is _apply_non4d_model_outputs else {} + if apply is _apply_non4d_model_inputs: + options = {"prepared_input_ops": {operation}} transformed = apply(transformed, **options) options = ort.SessionOptions() options.intra_op_num_threads = 1 @@ -1498,7 +1500,13 @@ def test_legacy_initializer_controls_remain_fixed(operation): for apply in transforms: transformed = apply( transformed, - **({"defer_rank_conversion": True} if apply is _apply_non4d_model_outputs else {}), + **( + {"defer_rank_conversion": True} + if apply is _apply_non4d_model_outputs + else {"prepared_input_ops": {operation}} + if apply is _apply_non4d_model_inputs + else {} + ), ) if operation == "Add": assert list(transformed.graph.initializer[0].dims) == [1, 1, 2, 3] From c30dc339b87db8b82e50efa1b6c8238287e73de3 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:43:35 +0000 Subject: [PATCH 20/41] Recover generated Slice and quantizer ranks from fixed reshape controls Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 6013e1a21d5..b15d62592d1 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -34,7 +34,7 @@ from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes -from .dla_model_boundaries import _has_native_symbolic_data +from .dla_model_boundaries import _has_native_symbolic_data, _is_standard_node __all__ = [] @@ -51,6 +51,21 @@ def _promoted_input_shape(cache, name, local_functions, native_sources): """Distinguish native symbolic paths from inputs prepared for four-dimensional data.""" shape = _shape_before_boundary_promotion(cache, name) rank = len(shape) if shape is not None else _known_rank(cache, name) + current, seen = name, set() + while rank is None and current not in seen: + seen.add(current) + producer = cache.get_producer(current) + if producer is None or not _is_standard_node( + producer, local_functions, {"Reshape", "Slice", "QuantizeLinear", "DequantizeLinear"} + ): + break + if producer.op_type == "Reshape": + target_shape = cache.get_shape(producer.input[1]) + if target_shape is not None and len(target_shape) == 1: + rank = target_shape[0] + break + current = producer.input[0] + rank = _known_rank(cache, current) native = ( rank is not None and rank < 4 From 167c4e8d265058600d9d2b92b29c63e2113350f6 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:47:29 +0000 Subject: [PATCH 21/41] Respect actual shared initializer promotion in Tile preparation Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 13 +++++++++--- .../test_dla_rank_conversion.py | 21 +++++++++++++------ 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index b15d62592d1..9e2523443fa 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -34,7 +34,7 @@ from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes -from .dla_model_boundaries import _has_native_symbolic_data, _is_standard_node +from .dla_model_boundaries import _REDUCTION_OPS, _has_native_symbolic_data, _is_standard_node __all__ = [] @@ -56,9 +56,14 @@ def _promoted_input_shape(cache, name, local_functions, native_sources): seen.add(current) producer = cache.get_producer(current) if producer is None or not _is_standard_node( - producer, local_functions, {"Reshape", "Slice", "QuantizeLinear", "DequantizeLinear"} + producer, + local_functions, + {"Reshape", "Slice", "QuantizeLinear", "DequantizeLinear"} | _REDUCTION_OPS, ): break + keepdims = _attribute(producer, "keepdims") + if producer.op_type in _REDUCTION_OPS and keepdims is not None and keepdims.i == 0: + break if producer.op_type == "Reshape": target_shape = cache.get_shape(producer.input[1]) if target_shape is not None and len(target_shape) == 1: @@ -323,7 +328,9 @@ def _pad_shape_inputs(model, op_type, suffix): if rank is None: raise ValueError(f"{node.name!r}: {op_type} input rank must be known") if op_type == "Tile" and rank == 1: - continue + data = cache.get_init(node.input[0]) + if data is None or len(data.dims) == 1: + continue if rank < 4 and _has_native_symbolic_data( cache, node.input[0], local_functions, native_sources ): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 46b3660f73b..af039a6314d 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1332,20 +1332,26 @@ def test_runtime_slice_extents_retain_promoted_axes(operation): np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) -@pytest.mark.parametrize("initializer", [False, True]) +@pytest.mark.parametrize(("initializer", "shared"), [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("tail", [False, True]) -def test_rank_one_tile_data_and_repeats_remain_native(initializer, tail): +def test_rank_one_tile_data_and_repeats_remain_native(initializer, shared, tail): """Execute rank-one Tile for graph inputs and fixed data without padding repeats.""" data = np.array([1, 2, 3], np.float32) nodes = [helper.make_node("Tile", ["x", "repeats"], ["tiled" if tail else "y"])] if tail: nodes.append(helper.make_node("Relu", ["tiled"], ["y"])) + inputs = [] if initializer else [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])] + if shared: + nodes.append(helper.make_node("Add", ["other", "x"], ["added"])) + inputs.append(helper.make_tensor_value_info("other", TensorProto.FLOAT, [1, 1, 2, 3])) + outputs.append(helper.make_tensor_value_info("added", TensorProto.FLOAT, [1, 1, 2, 3])) model = helper.make_model( helper.make_graph( nodes, "native_tile", - [] if initializer else [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])], - [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])], + inputs, + outputs, [numpy_helper.from_array(np.array([2], np.int64), "repeats")] + ([numpy_helper.from_array(data, "x")] if initializer else []), ), @@ -1353,9 +1359,12 @@ def test_rank_one_tile_data_and_repeats_remain_native(initializer, tail): ir_version=9, ) feeds = {} if initializer else {"x": data} - expected = _execute(model, feeds)[0] + if shared: + feeds["other"] = np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3) + expected = _execute(model, feeds) converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) - np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) From 7fa19db01a30412d3dfdafa5734011c6fe3ab338 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:50:45 +0000 Subject: [PATCH 22/41] Honor completed boundary axis mapping before generic validation Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 9e2523443fa..7219fef7259 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -593,13 +593,9 @@ def _apply_update_axis(model): if node.input else (None, None, False) ) - if native: + if native or "axis_updated" in node.name: continue - if ( - node.op_type in update - and "axis_updated" not in node.name - and (rank is None or (rank == 5 and shape is None)) - ): + if node.op_type in update and (rank is None or (rank == 5 and shape is None)): raise ValueError(f"{node.name!r}: {node.op_type} input layout must be known") if ( axis is None @@ -629,8 +625,6 @@ def _apply_update_axis(model): if shape is not None and len(shape) not in (2, 3, 5): continue raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") - if "axis_updated" in node.name: - continue original = value if shape and rank == 5: removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) From c846a85ca3ca589c25df2d800b7f3870433bb2c3 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:12:50 +0000 Subject: [PATCH 23/41] Preserve symbolic rank-four axis operators without static extents Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 5 +++- .../test_dla_rank_conversion.py | 30 +++++++++++++++++-- 2 files changed, 31 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 7219fef7259..c879df341e3 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -520,6 +520,9 @@ def _apply_argmax(model): ) if native: continue + keepdims = _attribute(node, "keepdims") + if rank == 4 and (keepdims is None or keepdims.i != 0): + continue if not shape or any(dim <= 0 for dim in shape): raise ValueError(f"{node.name!r}: ArgMax requires a static input shape") axis = _attribute(node, "axis") @@ -622,7 +625,7 @@ def _apply_update_axis(model): else: value = axis.i if node.op_type not in update: - if shape is not None and len(shape) not in (2, 3, 5): + if rank is not None and rank not in (2, 3, 5): continue raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") original = value diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index af039a6314d..b0659eaa0d9 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1421,12 +1421,14 @@ def test_symbolic_input_output_remains_native(batch): @pytest.mark.parametrize("operation", ["ArgMin", "Hardmax", "LayerNormalization"]) -def test_axis_update_preserves_native_four_dimensional_operators(operation): +@pytest.mark.parametrize("symbolic", [False, True]) +def test_axis_update_preserves_native_four_dimensional_operators(operation, symbolic): """Allow axis-bearing operators that already consume their final rank.""" shape = [2, 3, 4, 5] + declared_shape = ["N", *shape[1:]] if symbolic else shape inputs = ["x"] initializers = [] - output_shape = [2, 1, 4, 5] if operation == "ArgMin" else shape + output_shape = [declared_shape[0], 1, 4, 5] if operation == "ArgMin" else declared_shape if operation == "LayerNormalization": inputs.append("scale") initializers.append(numpy_helper.from_array(np.ones(shape[1:], np.float32), "scale")) @@ -1434,7 +1436,7 @@ def test_axis_update_preserves_native_four_dimensional_operators(operation): helper.make_graph( [helper.make_node(operation, inputs, ["y"], axis=1)], "native_four_dimensional", - [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, declared_shape)], [ helper.make_tensor_value_info( "y", @@ -1454,6 +1456,28 @@ def test_axis_update_preserves_native_four_dimensional_operators(operation): np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) +@pytest.mark.parametrize("keepdims", [None, 1]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_registered_argmax_keeps_symbolic_four_dimensional_layout(keepdims, batch): + """Leave rank-four ArgMax native when its symbolic dimensions need no restoration.""" + attributes = {} if keepdims is None else {"keepdims": keepdims} + model = helper.make_model( + helper.make_graph( + [helper.make_node("ArgMax", ["x"], ["y"], axis=1, **attributes)], + "native_argmax", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4, 5])], + [helper.make_tensor_value_info("y", TensorProto.INT64, ["N", 1, 4, 5])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(batch * 60, dtype=np.float32).reshape(batch, 3, 4, 5)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert converted.graph.node[0] == model.graph.node[0] + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + @pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile", "Slice", "Add"]) def test_legacy_initializer_controls_remain_fixed(operation): """Execute IR3 initializer declarations through boundary and rank conversion.""" From 7ad098d5e96d5bc7ac80045e9138a2ae2cc84f96 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:52:54 +0000 Subject: [PATCH 24/41] Retain native one-dimensional Conv initializer layouts Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_rank_conversion.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index c879df341e3..8a56d2cb634 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -34,7 +34,12 @@ from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes -from .dla_model_boundaries import _REDUCTION_OPS, _has_native_symbolic_data, _is_standard_node +from .dla_model_boundaries import ( + _REDUCTION_OPS, + _has_native_symbolic_data, + _is_standard_node, + _native_conv1d, +) __all__ = [] @@ -244,6 +249,8 @@ def _apply_non4d_initializers(model): model, {"QuantizeLinear", "DequantizeLinear", "GroupQueryAttention"}, ("com.microsoft",) ) selected = {id(node) for node in supported} + cache = GraphCache(model.graph) + avoid.update(name for node in supported if _native_conv1d(node, cache) for name in node.input) protected = _subgraph_captures(model.graph) | {value.name for value in model.graph.output} protected.update( name for node in model.graph.node if id(node) not in selected for name in node.input From 5e0dd6b484eb4120f4c3b2c53531cc5797cd7f0d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 12:57:06 +0000 Subject: [PATCH 25/41] Track widened initializer paths in Tile control preparation Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 23 +++++++++++++-- .../test_dla_rank_conversion.py | 29 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 8a56d2cb634..d703b6d724e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -319,12 +319,32 @@ def _apply_non4d_initializers(model): return model +def _rank_one_path_widened(cache, name, rank_preserving): + """Recognize broadcasting from promoted initializers without crossing rank-setting nodes.""" + pending, seen = [name], set() + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + tensor = cache.get_init(current) + if tensor is not None: + if len(tensor.dims) == 4: + return True + continue + producer = cache.get_producer(current) + if producer is not None and id(producer) in rank_preserving: + pending.extend(producer.input) + return False + + def _pad_shape_inputs(model, op_type, suffix): """Preserve native rank-four layouts and plan fixed control padding otherwise.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} runtime_inputs = _runtime_input_names(model) + rank_preserving = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS)} plans = [] for node in _operator_nodes(model, {op_type}): tensor = cache.get_init(node.input[1]) @@ -335,8 +355,7 @@ def _pad_shape_inputs(model, op_type, suffix): if rank is None: raise ValueError(f"{node.name!r}: {op_type} input rank must be known") if op_type == "Tile" and rank == 1: - data = cache.get_init(node.input[0]) - if data is None or len(data.dims) == 1: + if not _rank_one_path_widened(cache, node.input[0], rank_preserving): continue if rank < 4 and _has_native_symbolic_data( cache, node.input[0], local_functions, native_sources diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index b0659eaa0d9..1c5de0a02c1 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1367,6 +1367,35 @@ def test_rank_one_tile_data_and_repeats_remain_native(initializer, shared, tail) np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("operation", ["Add", "Sub", "Mul", "Div"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_one_tile_tracks_widened_elementwise_initializer(operation, unary): + """Keep Tile repeats consistent with rank widened through an elementwise operand.""" + nodes = [helper.make_node(operation, ["x", "weight"], ["combined"])] + if unary: + nodes.append(helper.make_node("Relu", ["combined"], ["activated"])) + nodes.append(helper.make_node("Tile", ["activated" if unary else "combined", "repeats"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "widened_tile", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])], + [ + numpy_helper.from_array(np.array([1, 2, 3], np.float32), "weight"), + numpy_helper.from_array(np.array([2], np.int64), "repeats"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.array([3, 5, 7], np.float32)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) def test_unknown_structural_data_rank_refuses_before_mutation(operation): """Refuse unresolved intermediate ranks before allocating any control view.""" From 231b1a6e84230b98f4a35c38ae90bbed2491f005 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:11:57 +0000 Subject: [PATCH 26/41] Declare deferred Reshape stage in manual rank fixtures Signed-off-by: mgohil --- .../dla_transforms/test_dla_rank_conversion.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 1c5de0a02c1..d55b54e3d64 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -77,7 +77,14 @@ def test_rank_conversion_preserves_operator_values(operation): model.ir_version = 9 transformed = copy.deepcopy(model) for apply in (_apply_non4d_model_inputs, _apply_non4d_model_outputs, transform): - options = {"defer_rank_conversion": True} if apply is _apply_non4d_model_outputs else {} + options = ( + { + "defer_rank_conversion": True, + "prepared_output_ops": {"Non4DReshape"} if operation == "Reshape" else set(), + } + if apply is _apply_non4d_model_outputs + else {} + ) if apply is _apply_non4d_model_inputs: options = {"prepared_input_ops": {operation}} transformed = apply(transformed, **options) @@ -1563,7 +1570,10 @@ def test_legacy_initializer_controls_remain_fixed(operation): transformed = apply( transformed, **( - {"defer_rank_conversion": True} + { + "defer_rank_conversion": True, + "prepared_output_ops": {"Non4DReshape"} if operation == "Reshape" else set(), + } if apply is _apply_non4d_model_outputs else {"prepared_input_ops": {operation}} if apply is _apply_non4d_model_inputs From 42a558ec7392b4de35d279796a109bf754a326f7 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:52:14 +0000 Subject: [PATCH 27/41] Preserve native initializer axes and contrib quantizer defaults Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 22 ++++++---- .../onnx/graph_surgery/utils/graph_utils.py | 12 ++++++ .../test_dla_rank_conversion.py | 42 +++++++++++++------ 3 files changed, 56 insertions(+), 20 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index d703b6d724e..b01e0642d0d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -55,7 +55,10 @@ def _runtime_input_names(model): def _promoted_input_shape(cache, name, local_functions, native_sources): """Distinguish native symbolic paths from inputs prepared for four-dimensional data.""" shape = _shape_before_boundary_promotion(cache, name) + snapshot = cache._initializer_vi_map.get(name) rank = len(shape) if shape is not None else _known_rank(cache, name) + if snapshot is not None and snapshot.type.tensor_type.HasField("shape"): + rank = len(snapshot.type.tensor_type.shape.dim) current, seen = name, set() while rank is None and current not in seen: seen.add(current) @@ -629,7 +632,7 @@ def _apply_update_axis(model): if ( axis is None and node.op_type in update - and node.domain in {"", "ai.onnx"} + and node.domain in {"", "ai.onnx", "com.microsoft"} and ( rank == 5 or ( @@ -638,14 +641,17 @@ def _apply_update_axis(model): ) ) ): - schema = defs.get_schema( - node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain - ) - default = schema.attributes.get("axis") - if default is not None and default.default_value.type: - value = default.default_value.i + if node.domain == "com.microsoft": + value = 1 else: - continue + schema = defs.get_schema( + node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain + ) + default = schema.attributes.get("axis") + if default is not None and default.default_value.type: + value = default.default_value.i + else: + continue elif axis is None or node.op_type in skip or (axis.i < 0 and node.op_type not in update): continue else: diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py index d2a59549c1a..3a4a25a1138 100644 --- a/modelopt/onnx/graph_surgery/utils/graph_utils.py +++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py @@ -744,6 +744,7 @@ class GraphCache: __slots__ = ( "_arr_cache", "_graph", + "_initializer_vi_map", "_names", "consumer_map", "graph_input_names", @@ -772,6 +773,9 @@ def rebuild(self) -> None: self._names.update([node.name, *node.input, *node.output]) self.init_map: dict[str, onnx.TensorProto] = {init.name: init for init in graph.initializer} self.init_names: set[str] = set(self.init_map) + self._initializer_vi_map = { + vi.name: vi for vi in graph.value_info if vi.name in self.init_names + } self.producer_map: dict[str, onnx.NodeProto] = {} for node in graph.node: for out in node.output: @@ -993,6 +997,14 @@ def _allocated_boundary_name(name, hint): def _shape_before_boundary_promotion(cache, name): """Recover native shapes at generated input and output boundary views.""" + snapshot = cache._initializer_vi_map.get(name) + if snapshot is not None and snapshot.type.tensor_type.HasField("shape"): + dimensions = snapshot.type.tensor_type.shape.dim + return ( + [dimension.dim_value for dimension in dimensions] + if all(dimension.HasField("dim_value") for dimension in dimensions) + else None + ) producer = cache.get_producer(name) if ( producer is not None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index d55b54e3d64..e8b6e4c3a50 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -311,8 +311,12 @@ def test_reshape_resolves_copied_dimensions_across_output_aliases(target, collis @pytest.mark.parametrize("axis", [0, 1, 2, -1]) -@pytest.mark.parametrize("value_info", [False, True]) -def test_registered_axis_update_reads_native_initializer_rank(axis, value_info): +@pytest.mark.parametrize("value_info", [False, True, "symbolic"]) +@pytest.mark.parametrize("ir_version", [3, 9]) +@pytest.mark.parametrize("entrypoint", ["stages", "pipeline"]) +def test_registered_axis_update_reads_native_initializer_rank( + axis, value_info, ir_version, entrypoint +): """Keep Concat axes correct with fixed inputs that lack optional value metadata.""" shape = [2, 3, 4] output_shape = shape.copy() @@ -328,18 +332,28 @@ def test_registered_axis_update_reads_native_initializer_rank(axis, value_info): ), opset_imports=[helper.make_opsetid("", 17)], ) - model.ir_version = 9 + model.ir_version = ir_version + if ir_version <= 3: + model.graph.input.extend( + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + for name in ["first", "second"] + ) if value_info: model.graph.value_info.extend( - helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + helper.make_tensor_value_info( + name, TensorProto.FLOAT, ["a", "b", "c"] if value_info == "symbolic" else shape + ) for name in ["first", "second"] ) expected = _execute(model, {})[0] - transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) - for name, transform in DLA_TRANSFORMS: - if name in {"Non4dInitializers", "UpdateAxis"}: - transformed = transform(transformed) - transformed.graph.ClearField("value_info") + if entrypoint == "pipeline": + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + else: + transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) + for name, transform in DLA_TRANSFORMS: + if name in {"Non4dInitializers", "UpdateAxis"}: + transformed = transform(transformed) + transformed.graph.ClearField("value_info") np.testing.assert_array_equal(_execute(transformed, {})[0], expected) @@ -1588,7 +1602,10 @@ def test_legacy_initializer_controls_remain_fixed(operation): @pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) @pytest.mark.parametrize("per_channel", [False, True, "singleton"]) -def test_registered_implicit_quantizer_axis_preserves_channel_values(operation, shape, per_channel): +@pytest.mark.parametrize("domain", ["", "com.microsoft"]) +def test_registered_implicit_quantizer_axis_preserves_channel_values( + operation, shape, per_channel, domain +): """Map every vector scale's channel axis, including a single-channel vector.""" shape = shape.copy() if per_channel == "singleton": @@ -1601,13 +1618,14 @@ def test_registered_implicit_quantizer_axis_preserves_channel_values(operation, output_type = TensorProto.UINT8 if operation == "QuantizeLinear" else TensorProto.FLOAT model = helper.make_model( helper.make_graph( - [helper.make_node(operation, ["x", "scale", "zero"], ["y"])], + [helper.make_node(operation, ["x", "scale", "zero"], ["y"], domain=domain)], "implicit_quantizer", [helper.make_tensor_value_info("x", input_type, shape)], [helper.make_tensor_value_info("y", output_type, shape)], [numpy_helper.from_array(scales, "scale"), numpy_helper.from_array(zeros, "zero")], ), - opset_imports=[helper.make_opsetid("", 17)], + opset_imports=[helper.make_opsetid("", 17)] + + ([helper.make_opsetid(domain, 1)] if domain else []), ) model.ir_version = 9 data = ( From 5bd7c47e2c9ffbf363da9b985f49a21497410e6d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 15:39:21 +0000 Subject: [PATCH 28/41] Pad Tile repeats after registered rank-restoring views Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 11 +++-- .../test_dla_rank_conversion.py | 47 +++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index b01e0642d0d..4e6c9e5b3cd 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -322,8 +322,8 @@ def _apply_non4d_initializers(model): return model -def _rank_one_path_widened(cache, name, rank_preserving): - """Recognize broadcasting from promoted initializers without crossing rank-setting nodes.""" +def _rank_one_path_widened(cache, name, rank_preserving, reshapes): + """Recognize earlier 4D initializer and Reshape promotion through rank-preserving paths.""" pending, seen = [name], set() while pending: current = pending.pop() @@ -336,6 +336,10 @@ def _rank_one_path_widened(cache, name, rank_preserving): return True continue producer = cache.get_producer(current) + if producer is not None and id(producer) in reshapes: + if cache.get_shape(producer.input[1]) == [4]: + return True + continue if producer is not None and id(producer) in rank_preserving: pending.extend(producer.input) return False @@ -348,6 +352,7 @@ def _pad_shape_inputs(model, op_type, suffix): native_sources = {} runtime_inputs = _runtime_input_names(model) rank_preserving = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS)} + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} plans = [] for node in _operator_nodes(model, {op_type}): tensor = cache.get_init(node.input[1]) @@ -358,7 +363,7 @@ def _pad_shape_inputs(model, op_type, suffix): if rank is None: raise ValueError(f"{node.name!r}: {op_type} input rank must be known") if op_type == "Tile" and rank == 1: - if not _rank_one_path_widened(cache, node.input[0], rank_preserving): + if not _rank_one_path_widened(cache, node.input[0], rank_preserving, reshapes): continue if rank < 4 and _has_native_symbolic_data( cache, node.input[0], local_functions, native_sources diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index e8b6e4c3a50..000f90028f5 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1417,6 +1417,53 @@ def test_rank_one_tile_tracks_widened_elementwise_initializer(operation, unary): np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize( + ("producer", "shape"), + [("Reshape", [2, 3]), ("Reshape", [1, 2, 3, 1]), ("ArgMax", [2, 3]), ("ReduceSum", [2, 3])], +) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_one_tile_tracks_registered_rank_restoring_reshape(producer, shape, unary): + """Execute vector Tile after earlier registered stages create a leading 4D view.""" + nodes = [] + initializers = [numpy_helper.from_array(np.array([2], np.int64), "repeats")] + if producer == "Reshape": + initializers.append(numpy_helper.from_array(np.array([-1], np.int64), "target")) + nodes.append(helper.make_node("Reshape", ["x", "target"], ["vector"])) + length = 6 + elif producer == "ArgMax": + nodes.append(helper.make_node("ArgMax", ["x"], ["vector"], axis=0, keepdims=0)) + length = 3 + else: + initializers.append(numpy_helper.from_array(np.array([0], np.int64), "axes")) + nodes.append(helper.make_node("ReduceSum", ["x", "axes"], ["vector"], keepdims=0)) + length = 3 + if unary: + nodes.append(helper.make_node("Identity", ["vector"], ["middle"])) + nodes.append(helper.make_node("Tile", ["middle" if unary else "vector", "repeats"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "restored_tile", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if producer == "ArgMax" else TensorProto.FLOAT, + [2 * length], + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) def test_unknown_structural_data_rank_refuses_before_mutation(operation): """Refuse unresolved intermediate ranks before allocating any control view.""" From fe869f1f3ed7fe07076af840c641a6f0adce0cdb Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 16:08:52 +0000 Subject: [PATCH 29/41] Exercise indexing after unconverted batched matrix rank preparation Signed-off-by: mgohil --- .../test_dla_rank_conversion.py | 45 +++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 000f90028f5..2e5d33fa8cd 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1464,6 +1464,51 @@ def test_rank_one_tile_tracks_registered_rank_restoring_reshape(producer, shape, np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("operation", ["GatherElements", "ReduceSum"]) +@pytest.mark.parametrize("reshape", [False, True]) +def test_rank_preparation_preserves_unconverted_batched_matrix_consumers(operation, reshape): + """Prove runtime matrix rank from a padded view or boundary without requiring Conv eligibility.""" + shape = [2, 5, 4] + source_shape = [1, *shape] if reshape else shape + nodes = [] + initializers = [ + numpy_helper.from_array(np.tile(np.eye(4, dtype=np.float32), (2, 1, 1)), "weight") + ] + if reshape: + initializers.append(numpy_helper.from_array(np.asarray(shape, np.int64), "native_target")) + nodes.append(helper.make_node("Reshape", ["x", "native_target"], ["matrix"])) + nodes.append(helper.make_node("MatMul", ["matrix" if reshape else "x", "weight"], ["data"])) + if operation == "GatherElements": + initializers.append( + numpy_helper.from_array(np.arange(40, dtype=np.int64).reshape(shape) % 5, "control") + ) + nodes.append(helper.make_node(operation, ["data", "control"], ["selected"], axis=1)) + target = [1, *shape] + else: + initializers.append(numpy_helper.from_array(np.asarray([1], np.int64), "control")) + nodes.append(helper.make_node(operation, ["data", "control"], ["selected"], keepdims=0)) + target = [1, 1, 2, 4] + initializers.append(numpy_helper.from_array(np.asarray(target, np.int64), "target")) + nodes.append(helper.make_node("Reshape", ["selected", "target"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "matrix_rank_proof", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, target)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(40, dtype=np.float32).reshape(source_shape)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert any(node.op_type == "MatMul" for node in converted.graph.node) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) def test_unknown_structural_data_rank_refuses_before_mutation(operation): """Refuse unresolved intermediate ranks before allocating any control view.""" From 91095df947fdfc97dd0c753d8834f48c55827084 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 17:15:55 +0000 Subject: [PATCH 30/41] Preserve native symbolic broadcast initializer and boundary ranks Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 11 ++ .../dla_transforms/dla_rank_conversion.py | 38 +++++ .../test_dla_rank_conversion.py | 157 ++++++++++++++++++ 3 files changed, 206 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 3806dbd0091..66be6a9e04a 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -219,6 +219,7 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal unsqueezed_shapes = {} normalization_axes = {} local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} contrib_ops = _boundary_contrib_ops(model) opset = next( (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 @@ -239,6 +240,16 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal ) and node.op_type not in {"InstanceNormalization", "Gemm"} and not _native_conv1d(node, cache) + and not ( + node.op_type in _RANK_AGNOSTIC_OPS - {"Clip"} + and len(node.input) > 1 + and any( + name != value.name + and _has_native_symbolic_data(cache, name, local_functions, native_sources) + for name in node.input + if name + ) + ) ] if not consumers: continue diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 4e6c9e5b3cd..97a495ed48e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -275,6 +275,44 @@ def _apply_non4d_initializers(model): avoid.update(node.input) updates = [] avoid.update(_runtime_input_names(model)) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + broadcasts = _RANK_AGNOSTIC_OPS - {"Clip"} + native_initializers = { + name + for node in supported + if node.op_type in broadcasts + and len(node.input) > 1 + and any( + value not in cache.init_names + and _has_native_symbolic_data(cache, value, local_functions, native_sources) + for value in node.input + if value + ) + for name in node.input + if name in cache.init_names + } + for name in native_initializers: + tensor = cache.get_init(name) + assert tensor is not None + if not 0 < len(tensor.dims) < 4: + continue + for user, port in cache.get_consumers(name): + if ( + id(user) in selected + and (port == 0 or user.op_type in {"Concat", "MatMul"}) + and user.op_type not in {"Gemm", "Unsqueeze", "GroupQueryAttention", "Shape"} + and not _native_conv1d(user, cache) + and not (user.op_type in broadcasts and len(user.input) > 1) + and not ( + user.op_type == "Identity" + and all(output in protected for output in user.output) + ) + ): + raise ValueError( + f"Initializer {name!r} has incompatible native symbolic and rank-prepared usage" + ) + avoid.update(native_initializers) expand_or_front = expand | front for tensor in model.graph.initializer: if tensor.name in avoid: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 2e5d33fa8cd..f5bbe9c140c 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1509,6 +1509,163 @@ def test_rank_preparation_preserves_unconverted_batched_matrix_consumers(operati np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("operation", ["Add", "Sub", "Mul", "Div", "Max"]) +@pytest.mark.parametrize("bias_shape", [[4], [3, 4], [1, 3, 4]]) +def test_symbolic_elementwise_initializers_keep_native_rank(operation, bias_shape): + """Reuse one converted model for multiple batches without widening native symbolic data.""" + nodes = [] + source = "x" + if len(bias_shape) == 2: + nodes.append(helper.make_node("Relu", [source], ["native"])) + source = "native" + operands = [source, "bias"] if len(bias_shape) != 3 else ["bias", source] + nodes.append(helper.make_node(operation, operands, ["y"])) + bias = np.arange(np.prod(bias_shape), dtype=np.float32).reshape(bias_shape) + 1 + model = helper.make_model( + helper.make_graph( + nodes, + "symbolic_bias", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4])], + [numpy_helper.from_array(bias, "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + payload = model.graph.initializer[0].SerializeToString() + cases = [ + {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4) + 1} for batch in [1, 2] + ] + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert ( + next(t for t in changed.graph.initializer if t.name == "bias").SerializeToString() + == payload + ) + for feeds, reference in zip(cases, expected): + np.testing.assert_array_equal(_execute(changed, feeds)[0], reference) + + +@pytest.mark.parametrize("public", [False, True]) +@pytest.mark.parametrize("source", ["constant", "default", "runtime"]) +def test_symbolic_bias_can_stay_native_for_a_shared_static_broadcast(public, source): + """Retain shared constant bytes and both native and already-4D broadcast outputs.""" + outputs = [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_y", TensorProto.FLOAT, [1, 2, 3, 4]), + ] + if public: + outputs.append(helper.make_tensor_value_info("bias", TensorProto.FLOAT, [1, 3, 4])) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "bias"], ["y"]), + helper.make_node("Add", ["static_x", "bias"], ["static_y"]), + ], + "shared_symbolic_bias", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 2, 3, 4]), + ], + outputs, + [numpy_helper.from_array(np.ones([1, 3, 4], np.float32), "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + if source != "constant": + model.graph.input.append( + helper.make_tensor_value_info("bias", TensorProto.FLOAT, [1, 3, 4]) + ) + # Make the boundary planner inspect the bias as the data operand. + model.graph.node[0].input[:] = ["bias", "x"] + if source == "default": + model.ir_version = 3 + else: + del model.graph.initializer[:] + cases = [ + { + "x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4), + "static_x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4), + **({"bias": np.ones([1, 3, 4], np.float32)} if source == "runtime" else {}), + } + for batch in [1, 2] + ] + expected = [_execute(model, feeds) for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, references in zip(cases, expected): + for actual, reference in zip(_execute(changed, feeds), references): + np.testing.assert_array_equal(actual, reference) + + +def test_symbolic_bias_preserves_a_shared_quantization_control(): + """Keep a bias native when another consumer requires its original per-axis scale layout.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "scale"], ["y"]), + helper.make_node( + "QuantizeLinear", ["static_x", "scale", "zero"], ["quantized"], axis=3 + ), + ], + "shared_symbolic_scale", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 1, 3, 4]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("quantized", TensorProto.UINT8, [1, 1, 3, 4]), + ], + [ + numpy_helper.from_array(np.array([1, 2, 3, 4], np.float32), "scale"), + numpy_helper.from_array(np.zeros(4, np.uint8), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "static_x": np.arange(12, dtype=np.float32).reshape(1, 1, 3, 4), + } + expected = _execute(model, feeds) + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(changed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_symbolic_bias_refuses_conflicting_rank_preparation_before_mutation(): + """Refuse a shared native/forced-rank initializer instead of leaving one consumer invalid.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "bias"], ["y"]), + helper.make_node("Concat", ["static_x", "bias"], ["joined"], axis=0), + ], + "mixed_symbolic_rank", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 3, 4]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("joined", TensorProto.FLOAT, [2, 3, 4]), + ], + [numpy_helper.from_array(np.ones([1, 3, 4], np.float32), "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + _execute( + model, {"x": np.ones([2, 3, 4], np.float32), "static_x": np.ones([1, 3, 4], np.float32)} + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"native symbolic.*rank-prepared"): + _apply_non4d_initializers(model) + assert model.SerializeToString() == before + + @pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) def test_unknown_structural_data_rank_refuses_before_mutation(operation): """Refuse unresolved intermediate ranks before allocating any control view.""" From fbd51a85ace5dc641acfb88e9bce2b2992cf0a1c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 17:42:28 +0000 Subject: [PATCH 31/41] Restore symbolic outputs only for promoted rank paths Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 20 +++++++ .../test_dla_rank_conversion.py | 60 +++++++++++++++++++ 2 files changed, 80 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 66be6a9e04a..b4dac29677d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -442,6 +442,7 @@ def _apply_non4d_model_outputs( ) for value in graph.output: shape: list | None = cache.get_shape(value.name) + symbolic_output = shape is None if ( shape is None and defer_rank_conversion @@ -458,6 +459,25 @@ def _apply_non4d_model_outputs( continue producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) + if symbolic_output: + if source is None or not _is_standard_node(source, local_functions, _RANK_FLOW_OPS): + continue + native_shape = _shape_before_boundary_promotion(cache, source.input[0]) + info = cache.vi_map.get(source.input[0]) + rank = ( + len(native_shape) + if native_shape is not None + else ( + len(info.type.tensor_type.shape.dim) + if info is not None and info.type.tensor_type.HasField("shape") + else None + ) + ) + target = cache.get_init(source.input[1]) if source.op_type == "Reshape" else None + if rank not in (2, 3) and not ( + target is not None and len(target.dims) == 1 and target.dims[0] < 4 + ): + continue if source is None and len(shape) == 1: continue if source is not None: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index f5bbe9c140c..0b4d775f9be 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -1353,6 +1353,66 @@ def test_runtime_slice_extents_retain_promoted_axes(operation): np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) +@pytest.mark.parametrize("operation", ["Slice", "NonZero", "Range", "Reshape"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_outputs_restore_only_promoted_producer_paths(operation, tail): + """Reuse runtime-dependent native outputs while restoring a genuinely padded Reshape.""" + output = "middle" if tail else "y" + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [6])] + initializers = [] + if operation in {"Slice", "Reshape"}: + inputs.append(helper.make_tensor_value_info("starts", TensorProto.INT64, [1])) + initializers.append(numpy_helper.from_array(np.array([6], np.int64), "ends")) + nodes = [ + helper.make_node( + "Slice", ["x", "starts", "ends"], ["sliced" if operation == "Reshape" else output] + ) + ] + if operation == "Reshape": + initializers.append(numpy_helper.from_array(np.array([-1], np.int64), "shape")) + nodes.append(helper.make_node("Reshape", ["sliced", "shape"], [output])) + cases = [ + {"x": np.arange(6, dtype=np.float32), "starts": np.array([start], np.int64)} + for start in [0, 2] + ] + elif operation == "NonZero": + nodes = [helper.make_node("NonZero", ["x"], [output])] + cases = [ + {"x": np.asarray(values, np.float32)} + for values in [[0, 0, 1, 0, 2, 0], [1, 2, 3, 4, 5, 0]] + ] + else: + inputs = [helper.make_tensor_value_info("stop", TensorProto.FLOAT, [])] + initializers = [ + numpy_helper.from_array(np.float32(value), name) + for name, value in [("start", 0), ("step", 1)] + ] + nodes = [helper.make_node("Range", ["start", "stop", "step"], [output])] + cases = [{"stop": np.asarray(stop, np.float32)} for stop in [3, 5]] + if tail: + nodes.append(helper.make_node("Identity", [output], ["y"])) + dtype = TensorProto.INT64 if operation == "NonZero" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + nodes, + "native_dynamic_output", + inputs, + [ + helper.make_tensor_value_info( + "y", dtype, [1, "N"] if operation == "NonZero" else ["N"] + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, reference in zip(cases, expected): + np.testing.assert_array_equal(_execute(changed, feeds)[0], reference) + + @pytest.mark.parametrize(("initializer", "shared"), [(False, False), (True, False), (True, True)]) @pytest.mark.parametrize("tail", [False, True]) def test_rank_one_tile_data_and_repeats_remain_native(initializer, shared, tail): From af18121ab26bc33603d5c4986991dc44b4a4633f Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 18:50:27 +0000 Subject: [PATCH 32/41] Require shared positive rank promotion proofs before remapping controls Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 11 +++ .../dla_transforms/dla_model_boundaries.py | 34 ++++----- .../dla_transforms/dla_rank_conversion.py | 76 +++++++++++++++---- .../onnx/graph_surgery/make_dla_compatible.py | 13 +++- .../test_dla_rank_conversion.py | 66 ++++++++++++++++ 5 files changed, 164 insertions(+), 36 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index dfb9d96ad9e..578f56cdaad 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -203,6 +203,8 @@ def classify(name): ): return True tensor = cache.get_init(name) + if tensor is not None and len(tensor.dims) == 4: + return True if name in cache.graph_input_names and (model.ir_version > 3 or tensor is None): return False if tensor is not None: @@ -239,6 +241,15 @@ def classify(name): and target[0] in (1, 2, 3, 5) and not _has_native_symbolic_data(cache, node.input[0], local_functions) ) + if node.op_type == "Expand": + target = cache.get_shape(node.input[1]) + return target == [4] or ( + "Expand" in prepared_ops + and target is not None + and len(target) == 1 + and 0 < target[0] < 4 + and not _has_native_symbolic_data(cache, node.input[0], local_functions) + ) if node.op_type == "Flatten": return node.op_type in prepared_ops if node.op_type in {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: 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 b4dac29677d..47c2abd9c7b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -403,35 +403,27 @@ 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 + model, + *, + defer_rank_conversion=False, + prepared_output_ops=(), + prepared_data_ops=None, + validate_only=False, ): """Restore outputs or only validate producer rank availability before shape inference.""" + # Rank stages import the boundary predicates; defer their shared proof to avoid the cycle. + from .dla_rank_conversion import _rank_preparation_proofs + graph = model.graph cache = GraphCache(graph) + prepared = ( + _rank_preparation_proofs(model, cache, prepared_data_ops) if defer_rank_conversion else {} + ) captured = _subgraph_captures(graph) boundary_names = cache.graph_input_names | cache.init_names 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) native_sources = {} @@ -455,7 +447,7 @@ def _apply_non4d_model_outputs( ] if shape is None or len(shape) >= 4 or value.name in boundary_names: continue - if value.name in native_broadcasts: + if defer_rank_conversion and not prepared[value.name]: continue producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 97a495ed48e..e4fbab020fe 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -33,7 +33,7 @@ ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures -from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes +from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes, _prepared_index_data from .dla_model_boundaries import ( _REDUCTION_OPS, _has_native_symbolic_data, @@ -52,7 +52,44 @@ def _runtime_input_names(model): return names -def _promoted_input_shape(cache, name, local_functions, native_sources): +def _rank_preparation_proofs(model, cache, prepared_data_ops): + """Reuse the cached leading-layout proof for rank and axis stages.""" + if prepared_data_ops is None: + prepared_data_ops = ( + _RANK_AGNOSTIC_OPS + | _REDUCTION_OPS + | { + "Non4dInitializers", + "Reshape", + "Expand", + "Tile", + "Transpose", + "Slice", + "Flatten", + "LpNormalization", + "ArgMax", + "TopK", + "Concat", + "Split", + "QuantizeLinear", + "DequantizeLinear", + "Gather", + "GatherElements", + "MatMulToTransposeConvTranspose", + "RemoveIntermediarySqueezeAndUnsqueeze", + } + ) + return _prepared_index_data( + model, + cache, + prepared_data_ops, + [node.input[0] for node in model.graph.node if node.input] + + [value.name for value in model.graph.output], + before_boundary=True, + ) + + +def _promoted_input_shape(cache, name, local_functions, native_sources, prepared): """Distinguish native symbolic paths from inputs prepared for four-dimensional data.""" shape = _shape_before_boundary_promotion(cache, name) snapshot = cache._initializer_vi_map.get(name) @@ -82,21 +119,25 @@ def _promoted_input_shape(cache, name, local_functions, native_sources): native = ( rank is not None and rank < 4 - and _has_native_symbolic_data(cache, name, local_functions, native_sources) + and ( + not prepared[name] + or _has_native_symbolic_data(cache, name, local_functions, native_sources) + ) ) return shape, rank, native -def _apply_non4d_slice_axis(model): +def _apply_non4d_slice_axis(model, *, prepared_data_ops=None): """Remap fixed Slice axes before their data rank is promoted.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {"Slice"}): shape, rank, native = _promoted_input_shape( - cache, node.input[0], local_functions, native_sources + cache, node.input[0], local_functions, native_sources, prepared ) if native: continue @@ -383,11 +424,12 @@ def _rank_one_path_widened(cache, name, rank_preserving, reshapes): return False -def _pad_shape_inputs(model, op_type, suffix): +def _pad_shape_inputs(model, op_type, suffix, *, prepared_data_ops=None): """Preserve native rank-four layouts and plan fixed control padding otherwise.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) runtime_inputs = _runtime_input_names(model) rank_preserving = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS)} reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} @@ -403,6 +445,8 @@ def _pad_shape_inputs(model, op_type, suffix): if op_type == "Tile" and rank == 1: if not _rank_one_path_widened(cache, node.input[0], rank_preserving, reshapes): continue + if op_type == "Tile" and rank < 4 and not prepared[node.input[0]]: + continue if rank < 4 and _has_native_symbolic_data( cache, node.input[0], local_functions, native_sources ): @@ -454,17 +498,18 @@ def _collapsed_reshape_axes(cache, name, shape, reshapes): ] -def _apply_non4d_transpose(model): +def _apply_non4d_transpose(model, *, prepared_data_ops=None): """Remap transpose permutations around added or removed unary axes.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) plans = [] reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"Transpose"}): perm = _attribute(node, "perm") shape, rank, native = _promoted_input_shape( - cache, node.input[0], local_functions, native_sources + cache, node.input[0], local_functions, native_sources, prepared ) if native: continue @@ -509,17 +554,18 @@ def _specialized_collapsed_axis(cache, node, shape, value, reshapes): return native - int(native > removed[0]) - (4 if value < 0 else 0) -def _apply_non4d_lpnorm(model): +def _apply_non4d_lpnorm(model, *, prepared_data_ops=None): """Preserve normalization axes across leading-dimension promotion.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} plans = [] for node in _operator_nodes(model, {"LpNormalization"}): axis = _attribute(node, "axis") shape, rank, native = _promoted_input_shape( - cache, node.input[0], local_functions, native_sources + cache, node.input[0], local_functions, native_sources, prepared ) if native: continue @@ -579,16 +625,17 @@ def _apply_flatten(model): return model -def _apply_argmax(model): +def _apply_argmax(model, *, prepared_data_ops=None): """Plan native ArgMax layouts before preparing private output restoration views.""" cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) plans = {} reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} for node in _operator_nodes(model, {"ArgMax"}): shape, rank, native = _promoted_input_shape( - cache, node.input[0], local_functions, native_sources + cache, node.input[0], local_functions, native_sources, prepared ) if native: continue @@ -646,7 +693,7 @@ def _apply_argmax(model): return model -def _apply_update_axis(model): +def _apply_update_axis(model, *, prepared_data_ops=None): """Remap supported operation axes from their native input ranks.""" quantizers = {"QuantizeLinear", "DequantizeLinear"} update = {"LogSoftmax", "Softmax", "TopK", "Concat", "Split"} | quantizers @@ -655,6 +702,7 @@ def _apply_update_axis(model): cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) plans = [] reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} selected = _operator_nodes(model, {node.op_type for node in model.graph.node}) @@ -664,7 +712,7 @@ def _apply_update_axis(model): for node in selected: axis = _attribute(node, "axis") shape, rank, native = ( - _promoted_input_shape(cache, node.input[0], local_functions, native_sources) + _promoted_input_shape(cache, node.input[0], local_functions, native_sources, prepared) if node.input else (None, None, False) ) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 1a05a3295df..a91c8985d9d 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -157,9 +157,19 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep if name == "Non4dModelInputs": return {"prepared_input_ops": prepared_input_ops} if name == "Non4dModelOutputs": - return {"prepared_output_ops": prepared_output_ops} + return {"prepared_output_ops": prepared_output_ops, "prepared_data_ops": prepared_input_ops} if name == "RemoveIntermediarySqueezeAndUnsqueeze": return {"prepared_consumer_ops": prepared_input_ops} + if name in { + "Non4dSliceAxis", + "Non4DTranspose", + "Non4DLpNorm", + "ArgMax", + "UpdateAxis", + "Non4DExpand", + "Non4DTile", + }: + return {"prepared_data_ops": prepared_input_ops} if name == "FixInstanceNormChannelMismatch": return {"prepared_consumer_ops": prepared_input_ops} return {} @@ -200,6 +210,7 @@ def _validate_boundary_contract( model, defer_rank_conversion=True, prepared_output_ops=prepared_output_ops, + prepared_data_ops=prepared_input_ops, validate_only=True, ) if "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_input_ops: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 0b4d775f9be..372775cbf74 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -2038,3 +2038,69 @@ def test_native_four_dimensional_runtime_controls_need_no_padding(operation, sou ) for values, reference in zip(feeds, expected): np.testing.assert_array_equal(_execute(transformed, values)[0], reference) + + +@pytest.mark.parametrize( + "operation", + ["Slice", "Transpose", "Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Expand", "Tile"], +) +@pytest.mark.parametrize("rank", [2, 3]) +@pytest.mark.parametrize("source", ["unsqueeze", "unary", "reshape"]) +def test_rank_one_adapters_keep_native_controls_unless_reshape_promotes(operation, rank, source): + """Distinguish native rank-one input adapters from an explicitly padded Reshape.""" + shape = [1] * (rank - 1) + [6] + initializers = [ + numpy_helper.from_array(np.array(list(range(rank - 1)), np.int64), "unsqueeze_axes") + ] + if source == "reshape": + initializers.append(numpy_helper.from_array(np.array(shape, np.int64), "target")) + nodes = [helper.make_node("Reshape", ["x", "target"], ["adapted"])] + else: + nodes = [helper.make_node("Unsqueeze", ["x", "unsqueeze_axes"], ["adapted"])] + data = "adapted" + if source == "unary": + nodes.append(helper.make_node("Relu", [data], ["unary"])) + data = "unary" + operands, attributes = [data], {"axis": rank - 1} + output_shape = shape.copy() + if operation == "Slice": + operands += ["starts", "ends", "slice_axes"] + attributes = {} + for name, values in [("starts", [1]), ("ends", [5]), ("slice_axes", [rank - 1])]: + initializers.append(numpy_helper.from_array(np.array(values, np.int64), name)) + output_shape[-1] = 4 + elif operation == "Transpose": + attributes = {"perm": list(reversed(range(rank)))} + output_shape.reverse() + elif operation == "ArgMax": + attributes["keepdims"] = 1 + output_shape[-1] = 1 + elif operation in {"Expand", "Tile"}: + output_shape[0] = 2 + values = output_shape if operation == "Expand" else [2] + [1] * (rank - 1) + initializers.append(numpy_helper.from_array(np.array(values, np.int64), "control")) + operands.append("control") + attributes = {} + nodes.append(helper.make_node(operation, operands, ["y"], **attributes)) + model = helper.make_model( + helper.make_graph( + nodes, + "rank_one_adapter_contract", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [6])], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMax" else TensorProto.FLOAT, + output_shape, + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + cases = [{"x": np.arange(1, 7, dtype=np.float32)}, {"x": np.arange(6, 0, -1, dtype=np.float32)}] + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, reference in zip(cases, expected): + np.testing.assert_allclose(_execute(changed, feeds)[0], reference, rtol=1e-6, atol=1e-7) From 52531a788543a8d31857eb562fc5461c1711e13c Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 19:05:43 +0000 Subject: [PATCH 33/41] Preserve native owner distinctions during output rank restoration Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 19 ++++++++++++++----- .../dla_transforms/dla_model_boundaries.py | 8 +++++--- .../dla_transforms/dla_rank_conversion.py | 3 ++- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 578f56cdaad..2a15fdebb13 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -144,7 +144,9 @@ def _gather_layout_reinterpreted(cache, name, supported, captured): return True -def _prepared_index_data(model, cache, prepared_ops, data_names, *, before_boundary=False): +def _prepared_index_data( + model, cache, prepared_ops, data_names, *, before_boundary=False, restoring_outputs=False +): """Prove leading-4D data from actual layouts and registered rank stages without mutation.""" standard = {id(node) for node in _operator_nodes(model, {n.op_type for n in model.graph.node})} if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): @@ -255,14 +257,17 @@ def classify(name): if node.op_type in {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: return [node.input[0]] if node.op_type in prepared_ops else False if node.op_type == "InstanceNormalization": - return "FixInstanceNormChannelMismatch" in prepared_ops + return ( + "FixInstanceNormChannelMismatch" in prepared_ops + or "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_ops + ) if node.op_type == "GroupQueryAttention": return "GroupQueryAttention" in prepared_ops if node.op_type == "Conv": kernel = cache.get_shape(node.input[1]) return kernel is not None and len(kernel) == 4 - if not before_boundary and node.op_type == "MatMul": - if name in cache.graph_output_names | captured: + if (not before_boundary or restoring_outputs) and node.op_type == "MatMul": + if not restoring_outputs and name in cache.graph_output_names | captured: return False ranks = [matrix_rank(value) for value in node.input[:2]] return ( @@ -281,7 +286,9 @@ def classify(name): if "RemoveIntermediarySqueezeAndUnsqueeze" not in prepared_ops: return False source = node.input[0] - if source in captured or name in captured | cache.graph_output_names: + if source in captured or ( + not restoring_outputs and name in captured | cache.graph_output_names + ): return False if source in cache.graph_input_names: original = cache.get_shape(source) @@ -311,6 +318,8 @@ def classify(name): [node.input[0]] if not before_boundary or value < 0 + or "axis_updated" in node.name + or "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_ops or node.op_type in prepared_ops or node.input[0] in cache.graph_input_names else False 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 47c2abd9c7b..90fbeeeadc8 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -417,7 +417,9 @@ def _apply_non4d_model_outputs( graph = model.graph cache = GraphCache(graph) prepared = ( - _rank_preparation_proofs(model, cache, prepared_data_ops) if defer_rank_conversion else {} + _rank_preparation_proofs(model, cache, prepared_data_ops, restoring_outputs=True) + if defer_rank_conversion + else {} ) captured = _subgraph_captures(graph) boundary_names = cache.graph_input_names | cache.init_names @@ -447,8 +449,6 @@ def _apply_non4d_model_outputs( ] if shape is None or len(shape) >= 4 or value.name in boundary_names: continue - if defer_rank_conversion and not prepared[value.name]: - continue producer = cache.get_producer(value.name) source = _output_rank_source(cache, producer, local_functions, rank_sources) if symbolic_output: @@ -538,6 +538,8 @@ def _apply_non4d_model_outputs( 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] + if defer_rank_conversion and not prepared[value.name]: + continue plans.append((value, shape, existing, axes)) if plans: _require_adapter_opset(model) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index e4fbab020fe..715adeb4566 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -52,7 +52,7 @@ def _runtime_input_names(model): return names -def _rank_preparation_proofs(model, cache, prepared_data_ops): +def _rank_preparation_proofs(model, cache, prepared_data_ops, *, restoring_outputs=False): """Reuse the cached leading-layout proof for rank and axis stages.""" if prepared_data_ops is None: prepared_data_ops = ( @@ -86,6 +86,7 @@ def _rank_preparation_proofs(model, cache, prepared_data_ops): [node.input[0] for node in model.graph.node if node.input] + [value.name for value in model.graph.output], before_boundary=True, + restoring_outputs=restoring_outputs, ) From a9d766c96aafb913e060728b546031ea5ee34ec2 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 19:32:47 +0000 Subject: [PATCH 34/41] Use actual layouts when proving registered rank-stage inputs Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/dla_indexing.py | 9 +++++++-- .../graph_surgery/dla_transforms/dla_rank_conversion.py | 4 +++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 2a15fdebb13..489486b62c7 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -264,8 +264,13 @@ def classify(name): if node.op_type == "GroupQueryAttention": return "GroupQueryAttention" in prepared_ops if node.op_type == "Conv": - kernel = cache.get_shape(node.input[1]) - return kernel is not None and len(kernel) == 4 + weight = cache.get_init(node.input[1]) + if weight is not None: + return len(weight.dims) == 4 + kernel_rank = _known_rank(cache, node.input[1]) + if kernel_rank is not None: + return kernel_rank == 4 + return [node.input[1]] if (not before_boundary or restoring_outputs) and node.op_type == "MatMul": if not restoring_outputs and name in cache.graph_output_names | captured: return False diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 715adeb4566..b144f5af54b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -54,6 +54,8 @@ def _runtime_input_names(model): def _rank_preparation_proofs(model, cache, prepared_data_ops, *, restoring_outputs=False): """Reuse the cached leading-layout proof for rank and axis stages.""" + # Standalone stages retain deferred inputs; registered calls inspect current layouts. + before_boundary = prepared_data_ops is None if prepared_data_ops is None: prepared_data_ops = ( _RANK_AGNOSTIC_OPS @@ -85,7 +87,7 @@ def _rank_preparation_proofs(model, cache, prepared_data_ops, *, restoring_outpu prepared_data_ops, [node.input[0] for node in model.graph.node if node.input] + [value.name for value in model.graph.output], - before_boundary=True, + before_boundary=before_boundary, restoring_outputs=restoring_outputs, ) From 9331a320c006e70a0120dd8abf756046c461434e Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 19:49:30 +0000 Subject: [PATCH 35/41] Respect the early Slice stage when proving future input promotion Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 14 ++++++-- .../test_dla_rank_conversion.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index b144f5af54b..9187af0fccf 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -52,10 +52,13 @@ def _runtime_input_names(model): return names -def _rank_preparation_proofs(model, cache, prepared_data_ops, *, restoring_outputs=False): +def _rank_preparation_proofs( + model, cache, prepared_data_ops, *, restoring_outputs=False, before_boundary=None +): """Reuse the cached leading-layout proof for rank and axis stages.""" # Standalone stages retain deferred inputs; registered calls inspect current layouts. - before_boundary = prepared_data_ops is None + if before_boundary is None: + before_boundary = prepared_data_ops is None if prepared_data_ops is None: prepared_data_ops = ( _RANK_AGNOSTIC_OPS @@ -135,7 +138,12 @@ def _apply_non4d_slice_axis(model, *, prepared_data_ops=None): cache = GraphCache(model.graph) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} - prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + prepared = _rank_preparation_proofs( + model, + cache, + prepared_data_ops, + before_boundary=prepared_data_ops is None or "Non4dModelInputs" in prepared_data_ops, + ) runtime_inputs = _runtime_input_names(model) plans = [] for node in _operator_nodes(model, {"Slice"}): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 372775cbf74..a0d881cba1e 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -23,6 +23,7 @@ import pytest from onnx import TensorProto, helper, numpy_helper +from modelopt.onnx.graph_surgery import make_dla_compatible as dla_pipeline from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( _apply_non4d_model_inputs, _apply_non4d_model_outputs, @@ -2104,3 +2105,34 @@ def test_rank_one_adapters_keep_native_controls_unless_reshape_promotes(operatio changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) for feeds, reference in zip(cases, expected): np.testing.assert_allclose(_execute(changed, feeds)[0], reference, rtol=1e-6, atol=1e-7) + + +def test_registered_slice_without_input_promotion_retains_native_axes(monkeypatch): + """Keep native Slice controls when the registered sequence omits input promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Slice", ["x", "starts", "ends", "axes"], ["y"])], + "slice_without_input_promotion", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2])], + [ + numpy_helper.from_array(np.array(values, np.int64), name) + for name, values in [("starts", [0]), ("ends", [2]), ("axes", [1])] + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + monkeypatch.setattr( + dla_pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) + for name, callback in dla_pipeline.DLA_TRANSFORMS + if name != "Non4dModelInputs" + ), + ) + changed = dla_pipeline._transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(changed, feeds)[0], expected) From edcabbd5859e35b89095828e098f716e042a853b Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 20:05:14 +0000 Subject: [PATCH 36/41] Preserve prospective matrix proofs across the early Slice stage Signed-off-by: mgohil --- modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py | 7 ++++++- .../dla_transforms/test_dla_rank_conversion.py | 4 +++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 489486b62c7..0c48ff96fcc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -271,7 +271,12 @@ def classify(name): if kernel_rank is not None: return kernel_rank == 4 return [node.input[1]] - if (not before_boundary or restoring_outputs) and node.op_type == "MatMul": + if node.op_type == "MatMul" and ( + not before_boundary + or restoring_outputs + or {"MatMulToTransposeConvTranspose", "RemoveIntermediarySqueezeAndUnsqueeze"} + <= set(prepared_ops) + ): if not restoring_outputs and name in cache.graph_output_names | captured: return False ranks = [matrix_rank(value) for value in node.input[:2]] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index a0d881cba1e..3e9e1861cdc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -16,6 +16,7 @@ """Numerical checks for independently reviewable rank and boundary passes.""" import copy +import importlib import numpy as np import onnx @@ -23,7 +24,6 @@ import pytest from onnx import TensorProto, helper, numpy_helper -from modelopt.onnx.graph_surgery import make_dla_compatible as dla_pipeline from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( _apply_non4d_model_inputs, _apply_non4d_model_outputs, @@ -45,6 +45,8 @@ _transform_make_dla_compatible, ) +dla_pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") + @pytest.mark.parametrize("operation", ["Transpose", "Reshape", "Tile", "Expand"]) def test_rank_conversion_preserves_operator_values(operation): From f000870f812033c46cdbbfb3d44f2be9d3b823be Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 20:32:41 +0000 Subject: [PATCH 37/41] Refuse Flatten conversion before native Gemm consumers Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 2 ++ .../test_dla_rank_conversion.py | 31 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 9187af0fccf..4157f6be88c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -603,6 +603,8 @@ def _apply_flatten(model): cache = GraphCache(graph) plans = [] for node in _operator_nodes(model, {"Flatten"}): + if any(reader.op_type == "Gemm" for reader, _ in cache.get_consumers(node.output[0])): + raise ValueError(f"{node.name!r}: Flatten output still feeds an unconverted Gemm") shape = _shape_before_boundary_promotion(cache, node.input[0]) axis = _attribute(node, "axis") axis = axis.i if axis is not None else 1 diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index 3e9e1861cdc..f3490fe0a5f 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -2138,3 +2138,34 @@ def test_registered_slice_without_input_promotion_retains_native_axes(monkeypatc ) changed = dla_pipeline._transform_make_dla_compatible(copy.deepcopy(model), verbose=False) np.testing.assert_array_equal(_execute(changed, feeds)[0], expected) + + +@pytest.mark.parametrize("bad_first", [False, True]) +def test_flatten_refuses_native_gemm_before_mutating_other_branches(bad_first): + """Keep every branch unchanged when a native Gemm still needs a matrix input.""" + flatten = helper.make_node("Flatten", ["x"], ["flat"], name="gemm_flatten") + peer = helper.make_node("Flatten", ["x"], ["peer"], name="peer_flatten") + nodes = [flatten, peer] if bad_first else [peer, flatten] + nodes.append(helper.make_node("Gemm", ["flat", "weight"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_gemm", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("weight", TensorProto.FLOAT, [12, 5]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 5]), + helper.make_tensor_value_info("peer", TensorProto.FLOAT, [2, 12]), + ], + value_info=[helper.make_tensor_value_info("flat", TensorProto.FLOAT, [2, 12])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + original = model.SerializeToString() + with pytest.raises(ValueError, match="Gemm"): + _apply_flatten(model) + assert model.SerializeToString() == original From 975bafa82e4ac54486181976bcd71de864e7b126 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:54:04 +0000 Subject: [PATCH 38/41] Preserve native internal Flatten ranks in the registered pipeline Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 11 ++++-- .../onnx/graph_surgery/make_dla_compatible.py | 1 + .../test_dla_rank_conversion.py | 34 +++++++++++++++++++ 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 4157f6be88c..228ca6b878a 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -597,15 +597,22 @@ def _apply_non4d_lpnorm(model, *, prepared_data_ops=None): return model -def _apply_flatten(model): +def _apply_flatten(model, *, prepared_data_ops=None): """Replace Flatten with its four-dimensional equivalent shape.""" graph = model.graph cache = GraphCache(graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) plans = [] for node in _operator_nodes(model, {"Flatten"}): + shape, _, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native and prepared_data_ops is not None: + continue if any(reader.op_type == "Gemm" for reader, _ in cache.get_consumers(node.output[0])): raise ValueError(f"{node.name!r}: Flatten output still feeds an unconverted Gemm") - shape = _shape_before_boundary_promotion(cache, node.input[0]) axis = _attribute(node, "axis") axis = axis.i if axis is not None else 1 if shape is None or any(d <= 0 for d in shape): diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index a91c8985d9d..b14fe04c876 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -165,6 +165,7 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep "Non4DTranspose", "Non4DLpNorm", "ArgMax", + "Flatten", "UpdateAxis", "Non4DExpand", "Non4DTile", diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index f3490fe0a5f..c68ed8657c9 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -2169,3 +2169,37 @@ def test_flatten_refuses_native_gemm_before_mutating_other_branches(bad_first): with pytest.raises(ValueError, match="Gemm"): _apply_flatten(model) assert model.SerializeToString() == original + + +@pytest.mark.parametrize("native_shape", [[2, 12], [2, 3, 4]]) +@pytest.mark.parametrize("unary", [False, True]) +def test_internal_native_flatten_keeps_its_output_rank(native_shape, unary): + """Keep a native internal Flatten rank visible to a downstream Shape reader.""" + source_shape = [1] * (4 - len(native_shape)) + native_shape + nodes = [helper.make_node("Squeeze", ["x", "axes"], ["native"])] + data = "native" + if unary: + nodes.append(helper.make_node("Relu", [data], ["activated"])) + data = "activated" + nodes.extend( + [ + helper.make_node("Flatten", [data], ["flat"], axis=1), + helper.make_node("Shape", ["flat"], ["y"]), + ] + ) + model = helper.make_model( + helper.make_graph( + nodes, + "native_internal_flatten", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2])], + [numpy_helper.from_array(np.arange(4 - len(native_shape), dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(source_shape)} + expected = _execute(model, feeds)[0] + result = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(result, feeds)[0], expected) + assert any(node.op_type == "Flatten" for node in result.graph.node) From 84a558a63621581675a1783c8ed57877e0517bc5 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:58:51 +0000 Subject: [PATCH 39/41] Preserve shared controls through private initializer views Signed-off-by: mgohil --- .../dla_transforms/dla_rank_conversion.py | 35 ++++++++ .../test_dla_rank_conversion.py | 79 +++++++++++++++++++ 2 files changed, 114 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py index 228ca6b878a..4a771d8b35e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -29,6 +29,7 @@ _attribute, _shape_before_boundary_promotion, add_initializer, + add_value_info, pad4d, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs @@ -330,6 +331,18 @@ def _apply_non4d_initializers(model): local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} native_sources = {} broadcasts = _RANK_AGNOSTIC_OPS - {"Clip"} + arithmetic = {"Add", "Sub", "Mul", "Div"} + private_users = {} + for tensor in model.graph.initializer: + consumers = cache.get_consumers(tensor.name) + if any( + user.op_type == "Shape" + or (port > 0 and user.op_type not in broadcasts | {"Concat", "MatMul", "LSTM"}) + for user, port in consumers + ): + readers = [(user, port) for user, port in consumers if user.op_type in arithmetic] + if readers: + private_users[tensor.name] = readers native_initializers = { name for node in supported @@ -395,7 +408,26 @@ def _apply_non4d_initializers(model): legacy_inputs = ( {value.name: value for value in model.graph.input} if model.ir_version <= 3 else {} ) + views = [] for tensor, shape in updates: + if tensor.name in private_users: + axes = add_initializer( + model.graph, + cache.unique_name(tensor.name + "_4d_axes"), + np.arange(4 - len(tensor.dims), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + view = cache.unique_name(tensor.name + "_4d") + views.append( + helper.make_node( + "Unsqueeze", [tensor.name, axes], [view], name=cache.unique_name(view + "_view") + ) + ) + add_value_info(model.graph, view, tensor.data_type, shape) + for user, port in private_users[tensor.name]: + user.input[port] = view + continue if tensor.name in native_inputs and tensor.name not in value_info: # Later axis stages read native ranks from the same deferred metadata snapshot. model.graph.value_info.append( @@ -409,6 +441,9 @@ def _apply_non4d_initializers(model): tensor.name, tensor.data_type, shape ).type.tensor_type.shape ) + if views: + _insert_support_nodes(model.graph, [*views, *model.graph.node]) + _add_legacy_initializer_inputs(model) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index c68ed8657c9..c51fda0ef6a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -48,6 +48,85 @@ dla_pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") +@pytest.mark.parametrize("existing_metadata", [False, True]) +def test_final_initializer_metadata_matches_materialized_layout(existing_metadata): + """Deferred native initializer ranks expire after axis conversion, before final inference.""" + weight = numpy_helper.from_array(np.arange(24, dtype=np.float32).reshape(2, 3, 4), "weight") + metadata = ( + [helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2, 3, 4])] + if existing_metadata + else [] + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["weight"], ["y"], perm=[0, 2, 1])], + "initializer_metadata", + [], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4, 3])], + [weight], + value_info=metadata, + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + expected = _execute(model, {}) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + tensors = {tensor.name: tensor for tensor in transformed.graph.initializer} + for info in transformed.graph.value_info: + if info.name in tensors: + assert [dim.dim_value for dim in info.type.tensor_type.shape.dim] == list( + tensors[info.name].dims + ) + for actual, reference in zip(_execute(transformed, {}), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Slice", "Reshape", "Tile"]) +def test_initializer_promotion_preserves_shared_structural_controls(operation): + """Arithmetic readers receive a private view while structural controls keep their native rank.""" + shape = [1, 1, 2, 4] if operation == "Tile" else [1, 1, 2, 3] + control = {"Slice": [1], "Reshape": [6], "Tile": [1, 1, 1, 2]}[operation] + initializers = [numpy_helper.from_array(np.asarray(control, np.int64), "control")] + inputs = ["x", "control"] + if operation == "Slice": + initializers.extend( + [ + numpy_helper.from_array(np.array([2], np.int64), "ends"), + numpy_helper.from_array(np.array([2], np.int64), "axes"), + ] + ) + inputs.extend(["ends", "axes"]) + result_shape = {"Slice": [1, 1, 1, 3], "Reshape": [6], "Tile": [1, 1, 2, 8]}[operation] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "control"], ["sum"]), + helper.make_node(operation, inputs, ["result"]), + ], + "shared_structural_control", + [helper.make_tensor_value_info("x", TensorProto.INT64, shape)], + [ + helper.make_tensor_value_info("sum", TensorProto.INT64, shape), + helper.make_tensor_value_info("result", TensorProto.INT64, result_shape), + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.int64).reshape(shape)} + expected = _execute(model, feeds) + tensor = model.graph.initializer[0] + before = tensor.SerializeToString() + transformed = _apply_non4d_initializers(model) + assert tensor is transformed.graph.initializer[0] + assert tensor.SerializeToString() == before + add = next(node for node in transformed.graph.node if node.op_type == "Add") + assert add.input[1] != "control" + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("operation", ["Transpose", "Reshape", "Tile", "Expand"]) def test_rank_conversion_preserves_operator_values(operation): """Execute native operators across explicitly deferred boundary and rank stages.""" From ed0c336a819016026542396d0eca45f8cde4c4de Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 08:03:49 +0000 Subject: [PATCH 40/41] Verify shared initializer views with legacy input declarations Signed-off-by: mgohil --- .../dla_transforms/test_dla_rank_conversion.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py index c51fda0ef6a..1f138505cd5 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -82,7 +82,8 @@ def test_final_initializer_metadata_matches_materialized_layout(existing_metadat @pytest.mark.parametrize("operation", ["Slice", "Reshape", "Tile"]) -def test_initializer_promotion_preserves_shared_structural_controls(operation): +@pytest.mark.parametrize("ir_version", [3, 9]) +def test_initializer_promotion_preserves_shared_structural_controls(operation, ir_version): """Arithmetic readers receive a private view while structural controls keep their native rank.""" shape = [1, 1, 2, 4] if operation == "Tile" else [1, 1, 2, 3] control = {"Slice": [1], "Reshape": [6], "Tile": [1, 1, 1, 2]}[operation] @@ -112,8 +113,13 @@ def test_initializer_promotion_preserves_shared_structural_controls(operation): initializers, ), opset_imports=[helper.make_opsetid("", 21)], - ir_version=9, + ir_version=ir_version, ) + if ir_version <= 3: + model.graph.input.extend( + helper.make_tensor_value_info(tensor.name, tensor.data_type, list(tensor.dims)) + for tensor in model.graph.initializer + ) feeds = {"x": np.arange(np.prod(shape), dtype=np.int64).reshape(shape)} expected = _execute(model, feeds) tensor = model.graph.initializer[0] From d6f46eaecd73cc7c48acd169bbd9777e96f433ea Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 09:24:50 +0000 Subject: [PATCH 41/41] Normalize boundary preflight formatting after restack Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/dla_model_boundaries.py | 1 - 1 file changed, 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 90fbeeeadc8..37391c2f6d0 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -33,7 +33,6 @@ __all__ = [] - class _BoundaryPreflightError(ValueError): """Prerequisite lowering cannot freeze an exposed runtime control."""