From 8a70e47c288f0aad60647d359d3d5d81ce827f7e Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:14 +0000 Subject: [PATCH 01/31] [DLA_transforms_11] [11/17] DLA indexing, reductions and attention Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 310 ++++++++++++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 10 + .../dla_transforms/test_dla_indexing.py | 158 +++++++++ 3 files changed, 478 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py new file mode 100644 index 00000000000..3ffbfa2aa0e --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -0,0 +1,310 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DLA indexing, reduction and attention passes with deferred rank promotion.""" + +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_gather(model): + graph = model.graph + cache = GraphCache(graph) + nodes = [] + squeezed_indices = {} + for candidate in cache.nodes_by_op("Gather"): + if cache.get_shape(candidate.input[0]) is None: + raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") + if cache.get_init(candidate.input[1]) is None and cache.get_dtype( + candidate.input[1] + ) not in (TensorProto.INT32, TensorProto.INT64): + producer = cache.get_producer(candidate.input[1]) + if producer is None or producer.op_type != "Cast": + raise ValueError( + f"{candidate.name!r}: dynamic Gather indices require an integer Cast" + ) + for node in graph.node: + if node.op_type != "Gather": + nodes.append(node) + continue + indices = cache.get_init(node.input[1]) + if indices is not None: + values = numpy_helper.to_array(indices).reshape(-1).astype(np.int32) + node.input[1] = add_initializer( + graph, + indices.name + "_gather_indices", + np.asarray(values, dtype=np.int32), + dtype=TensorProto.INT32, + with_value_info=False, + ) + else: + original = node.input[1] + if original in squeezed_indices: + node.input[1] = squeezed_indices[original] + axis = _attribute(node, "axis") + shape = cache.get_shape(node.input[0]) + if axis is None: + node.attribute.append(helper.make_attribute("axis", 0)) + axis = _attribute(node, "axis") + if axis.i >= 0 and shape: + axis.i += {2: 2, 3: 1}.get(len(shape), 0) + node.name += "_axis_updated" + nodes.append(node) + continue + squeezed = cache.unique_name(original + "_squeezed") + squeeze_inputs = [original] + producer = cache.get_producer(original) + if producer is not None and producer.op_type == "Cast": + target = _attribute(producer, "to") + shared = ( + original in cache.graph_output_names or len(cache.get_consumers(original)) > 1 + ) + if shared: + integer = cache.unique_name(original + "_gather_int32") + nodes.append( + helper.make_node( + "Cast", + [producer.input[0]], + [integer], + name=integer, + to=TensorProto.INT32, + ) + ) + squeeze_inputs[0] = integer + else: + target.i = TensorProto.INT32 + for value in graph.value_info: + if value.name == original: + value.type.tensor_type.elem_type = TensorProto.INT32 + axes = add_initializer( + graph, + cache.unique_name(node.name + "_indices_squeeze_axes"), + np.asarray([0, 1, 2], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + squeeze_inputs.append(axes) + elif cache.get_dtype(original) not in (TensorProto.INT32, TensorProto.INT64): + raise ValueError(f"{node.name!r}: dynamic Gather indices require an integer Cast") + nodes.append( + helper.make_node("Squeeze", squeeze_inputs, [squeezed], name=node.name + "_squeeze") + ) + node.input[1] = squeezed + squeezed_indices[original] = squeezed + graph.value_info.append( + helper.make_tensor_value_info( + squeezed, cache.get_dtype(original) or TensorProto.FLOAT, None + ) + ) + if _attribute(node, "axis") is None: + node.attribute.append(helper.make_attribute("axis", 0)) + axis = _attribute(node, "axis") + shape = cache.get_shape(node.input[0]) + if axis is not None: + if axis.i >= 0 and shape: + axis.i += {2: 2, 3: 1}.get(len(shape), 0) + node.name += "_axis_updated" + nodes.append(node) + del graph.node[:] + graph.node.extend(nodes) + return model + + +def _apply_gather_elements(model): + graph = model.graph + cache = GraphCache(graph) + nodes = [] + for node in graph.node: + if node.op_type != "GatherElements": + nodes.append(node) + continue + tensor = cache.get_init(node.input[1]) + if tensor is not None: + array = numpy_helper.to_array(tensor).astype(np.int32) + rank = array.ndim + if rank == 3: + array = array[None] + elif rank == 0: + array = array.reshape(1) + node.input[1] = add_initializer( + graph, + tensor.name + "_gatherelements_indices", + np.asarray(array, dtype=np.int32), + dtype=TensorProto.INT32, + with_value_info=False, + ) + else: + shape = cache.get_shape(node.input[1]) + rank = len(shape) if shape is not None else None + original = node.input[1] + cast = original + "_cast_to_int32" + nodes.append( + helper.make_node("Cast", [original], [cast], name=cast, to=TensorProto.INT32) + ) + node.input[1] = cast + axis = _attribute(node, "axis") + if rank == 3 and axis is not None and axis.i != -1: + axis.i += 1 + node.name += "_axis_updated" + nodes.append(node) + del graph.node[:] + graph.node.extend(nodes) + return model + + +def _apply_standalone_reduce_operations(model): + graph = model.graph + cache = GraphCache(graph) + nodes = [] + count = 0 + plans = {} + for node in graph.node: + if node.op_type not in {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: + continue + shape = _shape_before_boundary_promotion(cache, node.input[0]) + if not shape: + keepdims = _attribute(node, "keepdims") + if keepdims is not None and keepdims.i == 0: + raise ValueError(f"Missing static reduction input shape for {node.name!r}") + continue # Already rank-preserving, including generated 4D reduction branches. + tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None + if len(node.input) > 1 and tensor is None: + raise ValueError(f"Dynamic reduction axes are unsupported for {node.name!r}") + attribute = _attribute(node, "axes") + raw_axes = ( + numpy_helper.to_array(tensor).reshape(-1).tolist() + if tensor is not None + else list(attribute.ints) + if attribute is not None + else list(range(len(shape))) + ) + reduced = {value % len(shape) for value in raw_axes or range(len(shape))} + output_shape = cache.get_shape(node.output[0]) + if output_shape is None and node.output[0].endswith("_squeeze_input"): + output_shape = cache.get_shape(node.output[0].removesuffix("_squeeze_input")) + if output_shape is None: + output_shape = [dim for index, dim in enumerate(shape) if index not in reduced] + plans[id(node)] = (shape, raw_axes, output_shape) + for node in graph.node: + nodes.append(node) + if id(node) not in plans: + continue + shape, raw_axes, output_shape = plans[id(node)] + keepdims = _attribute(node, "keepdims") + reshape = keepdims is not None and keepdims.i == 0 + if keepdims is not None: + keepdims.i = 1 + delta = {2: 2, 3: 1}.get(len(shape), 0) + values = [value + delta if value >= 0 else value for value in raw_axes] + if len(node.input) > 1: + node.input[1] = add_initializer( + graph, + node.name + "_axes", + np.asarray(values, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + else: + attribute = _attribute(node, "axes") + if attribute is not None: + del attribute.ints[:] + attribute.ints.extend(values) + if reshape: + target = pad4d(output_shape) + name = add_initializer( + graph, + f"reshape_shape_{count}", + np.asarray(target, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + original = node.output[0] + node.output[0] = original + "_temp" + nodes.append( + helper.make_node( + "Reshape", [node.output[0], name], [original], name=node.name + "_reshape" + ) + ) + count += 1 + del graph.node[:] + graph.node.extend(nodes) + return model + + +def _apply_group_query_attention(model): + graph = model.graph + cache = GraphCache(graph) + nodes = [] + boundaries = cache.graph_input_names | cache.init_names + squeezed_inputs = {} + for node in graph.node: + if node.op_type != "GroupQueryAttention": + nodes.append(node) + continue + for index, name in enumerate(node.input[:5]): + if name and name not in boundaries: + if name in squeezed_inputs: + node.input[index] = squeezed_inputs[name] + continue + axes = add_initializer( + graph, + name + "_axis", + np.asarray([0], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + squeezed = cache.unique_name(name + "_squeezed") + squeezed_inputs[name] = squeezed + nodes.append( + helper.make_node( + "Squeeze", + [name, axes], + [squeezed], + name=node.name + "_" + name + "_squeeze", + ) + ) + node.input[index] = squeezed + nodes.append(node) + for index, name in enumerate(node.output): + if name and name not in cache.graph_output_names: + axes = add_initializer( + graph, + name + "_axis", + np.asarray([0], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.output[index] = name + "_squeezed" + nodes.append( + helper.make_node( + "Unsqueeze", + [node.output[index], axes], + [name], + name=node.name + "_" + name + "_unsqueeze", + ) + ) + del graph.node[:] + graph.node.extend(nodes) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 478c818b673..3676d0c9bcd 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -35,6 +35,12 @@ from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model from .dla_transforms.dla_decompose_lstm import _apply_decompose_lstm from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers +from .dla_transforms.dla_indexing import ( + _apply_gather, + _apply_gather_elements, + _apply_group_query_attention, + _apply_standalone_reduce_operations, +) from .dla_transforms.dla_matmul_to_transpose_conv_transpose import ( _apply_matmul_to_transpose_conv_transpose, ) @@ -93,6 +99,7 @@ def _apply_constants_to_initializers(model, **options): ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), ("ReshapeReduceSum", _apply_reshape_reduce_sum), + ("Gather", _apply_gather), ("DecomposeLstm", partial(_apply_decompose_lstm, defer_rank_conversion=True)), ( "MatMulToTransposeConvTranspose", @@ -100,10 +107,13 @@ def _apply_constants_to_initializers(model, **options): ), ("Non4dModelInputs", _apply_non4d_model_inputs), ("Non4dModelOutputs", partial(_apply_non4d_model_outputs, defer_rank_conversion=True)), + ("StandaloneReduceOperations", _apply_standalone_reduce_operations), ( "RemoveIntermediarySqueezeAndUnsqueeze", partial(_apply_remove_intermediary_squeeze_unsqueeze, rank_converted=True), ), + ("GroupQueryAttention", _apply_group_query_attention), + ("GatherElements", _apply_gather_elements), ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py new file mode 100644 index 00000000000..7823d182b1c --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -0,0 +1,158 @@ +# 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. + +"""Independent numerical and boundary-contract checks for indexing passes.""" + +import copy + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_indexing import ( + _apply_gather, + _apply_gather_elements, + _apply_group_query_attention, + _apply_standalone_reduce_operations, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, +) + + +@pytest.mark.parametrize("operation", ["Gather", "GatherElements"]) +def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): + if operation == "Gather": + shape, output_shape, axis = [5, 4], [4], 0 + indices = np.int64(2) + transform = _apply_gather + else: + shape, output_shape, axis = [2, 3, 4], [2, 3, 4], 1 + indices = np.ones(shape, dtype=np.int64) + transform = _apply_gather_elements + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x", "indices"], ["y"], name="index", axis=axis)], + "indexing", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializer=[numpy_helper.from_array(np.asarray(indices), "indices")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + candidate = transform(copy.deepcopy(model)) + candidate = _apply_non4d_model_inputs(candidate) + candidate = _apply_non4d_model_outputs(candidate) + 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( + candidate.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data}) + np.testing.assert_array_equal(actual[0], expected[0]) + + +@pytest.mark.parametrize("direct_output", [False, True]) +def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_output): + output = "attention" if direct_output else "y" + nodes = [ + helper.make_node("Unsqueeze", ["x", "axes"], ["q4"]), + helper.make_node( + "GroupQueryAttention", + ["q4", "x", "x", "", ""], + ["attention"], + name="attention", + domain="com.microsoft", + num_heads=1, + kv_num_heads=1, + ), + ] + if not direct_output: + nodes.append(helper.make_node("Identity", ["attention"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "attention", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 2, 4])], + [ + helper.make_tensor_value_info( + output, TensorProto.FLOAT16, [1, 2, 4] if direct_output else [1, 1, 2, 4] + ) + ], + initializer=[numpy_helper.from_array(np.array([0], dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)], + ) + candidate = _apply_group_query_attention(model) + onnx.checker.check_model(candidate) + attention = next(n for n in candidate.graph.node if n.op_type == "GroupQueryAttention") + producer = {name: node for node in candidate.graph.node for name in node.output} + assert producer[attention.input[0]].op_type == "Squeeze" + assert list(attention.input[1:]) == ["x", "x", "", ""] + assert producer["attention"].op_type == ( + "GroupQueryAttention" if direct_output else "Unsqueeze" + ) + + +def test_late_refusal_preserves_the_original_model(): + nodes = [ + helper.make_node("ReduceSum", ["x"], ["first"], keepdims=0), + helper.make_node("ReduceSum", ["x", "dynamic_axes"], ["y"], keepdims=0), + ] + inputs = [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("dynamic_axes", TensorProto.INT64, [1]), + ] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3])] + model = helper.make_model( + helper.make_graph(nodes, "refusal", inputs, outputs), + opset_imports=[helper.make_opsetid("", 21)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError): + _apply_standalone_reduce_operations(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("axes", [[-2], [-1, 0], [0, -1]]) +def test_reduction_preserves_negative_and_mixed_axes(axes): + from onnx.reference import ReferenceEvaluator + + model = helper.make_model( + helper.make_graph( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=int(len(axes) > 1))], + "reduce", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, None)], + [numpy_helper.from_array(np.array(axes, dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + model = onnx.shape_inference.infer_shapes(model) + data = np.arange(6, dtype=np.float32).reshape(2, 3) + expected = ReferenceEvaluator(model).run(None, {"x": data})[0] + candidate = _apply_non4d_model_inputs(copy.deepcopy(model)) + candidate = _apply_non4d_model_outputs(candidate) + candidate = _apply_standalone_reduce_operations(candidate) + onnx.checker.check_model(candidate) + actual = ReferenceEvaluator(candidate).run(None, {"x": data})[0] + np.testing.assert_array_equal(actual, expected) From a93faa15e3c83c2dc1fa3fd3c4ade2664f855b71 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 23:24:40 +0000 Subject: [PATCH 02/31] Declare deferred output ranks in indexing fixtures Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/test_dla_indexing.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 7823d182b1c..2329230420d 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -22,6 +22,7 @@ import onnxruntime as ort import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms.dla_indexing import ( _apply_gather, @@ -58,7 +59,7 @@ def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): model.ir_version = 9 candidate = transform(copy.deepcopy(model)) candidate = _apply_non4d_model_inputs(candidate) - candidate = _apply_non4d_model_outputs(candidate) + candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) options = ort.SessionOptions() options.intra_op_num_threads = 1 data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) @@ -135,8 +136,6 @@ def test_late_refusal_preserves_the_original_model(): @pytest.mark.parametrize("axes", [[-2], [-1, 0], [0, -1]]) def test_reduction_preserves_negative_and_mixed_axes(axes): - from onnx.reference import ReferenceEvaluator - model = helper.make_model( helper.make_graph( [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=int(len(axes) > 1))], @@ -151,7 +150,7 @@ def test_reduction_preserves_negative_and_mixed_axes(axes): data = np.arange(6, dtype=np.float32).reshape(2, 3) expected = ReferenceEvaluator(model).run(None, {"x": data})[0] candidate = _apply_non4d_model_inputs(copy.deepcopy(model)) - candidate = _apply_non4d_model_outputs(candidate) + candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) candidate = _apply_standalone_reduce_operations(candidate) onnx.checker.check_model(candidate) actual = ReferenceEvaluator(candidate).run(None, {"x": data})[0] From 25decf9b93ef8648253c8b490fa198a5c1999913 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:01:08 +0000 Subject: [PATCH 03/31] Preserve indexing ranks, runtime controls and private attention adapters Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 238 ++++++++--- .../dla_transforms/test_dla_indexing.py | 403 +++++++++++++++++- 2 files changed, 568 insertions(+), 73 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 3ffbfa2aa0e..0ec66cd2cef 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -23,18 +23,53 @@ _attribute, _shape_before_boundary_promotion, add_initializer, + batch_replace_nodes, pad4d, ) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures __all__ = [] +def _insert_support_nodes(graph, nodes): + """Insert generated support nodes without copying existing protobuf payloads.""" + positions = {id(node): index for index, node in enumerate(graph.node)} + batch = {} + pending = [] + for node in nodes: + index = positions.get(id(node)) + if index is None: + pending.append(node) + elif pending: + batch[index] = ([], pending) + pending = [] + if pending: + batch[len(graph.node)] = ([], pending) + batch_replace_nodes(graph, batch) + + +def _operator_nodes(model, operators, domains=("", "ai.onnx")): + """Select supported schemas without interpreting opaque model-local functions.""" + functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + return [ + node + for node in model.graph.node + if node.op_type in operators + and node.domain in domains + and not node.overload + and (node.domain, node.op_type, node.overload) not in functions + ] + + def _apply_gather(model): + """Prepare Gather indices for the subsequent ordered data-rank promotion.""" graph = model.graph cache = GraphCache(graph) nodes = [] squeezed_indices = {} - for candidate in cache.nodes_by_op("Gather"): + selected = _operator_nodes(model, {"Gather"}) + for candidate in selected: if cache.get_shape(candidate.input[0]) is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") if cache.get_init(candidate.input[1]) is None and cache.get_dtype( @@ -45,22 +80,35 @@ def _apply_gather(model): raise ValueError( f"{candidate.name!r}: dynamic Gather indices require an integer Cast" ) + selected_ids = {id(node) for node in selected} for node in graph.node: - if node.op_type != "Gather": + if id(node) not in selected_ids: nodes.append(node) continue indices = cache.get_init(node.input[1]) - if indices is not None: - values = numpy_helper.to_array(indices).reshape(-1).astype(np.int32) - node.input[1] = add_initializer( - graph, - indices.name + "_gather_indices", - np.asarray(values, dtype=np.int32), - dtype=TensorProto.INT32, - with_value_info=False, - ) + if indices is not None and node.input[1] not in cache.graph_input_names: + if indices.name not in squeezed_indices: + values = numpy_helper.to_array(indices).reshape(-1).astype(np.int32) + squeezed_indices[indices.name] = add_initializer( + graph, + cache.unique_name(indices.name + "_gather_indices"), + values, + dtype=TensorProto.INT32, + with_value_info=False, + ) + node.input[1] = squeezed_indices[indices.name] else: original = node.input[1] + index_shape = cache.get_shape(original) + if index_shape is not None and len(index_shape) <= 1: + if original not in squeezed_indices: + prepared = cache.unique_name(original + "_gather_int32") + nodes.append( + helper.make_node( + "Cast", [original], [prepared], name=prepared, to=TensorProto.INT32 + ) + ) + squeezed_indices[original] = prepared if original in squeezed_indices: node.input[1] = squeezed_indices[original] axis = _attribute(node, "axis") @@ -70,7 +118,7 @@ def _apply_gather(model): axis = _attribute(node, "axis") if axis.i >= 0 and shape: axis.i += {2: 2, 3: 1}.get(len(shape), 0) - node.name += "_axis_updated" + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") nodes.append(node) continue squeezed = cache.unique_name(original + "_squeezed") @@ -125,64 +173,103 @@ def _apply_gather(model): if axis is not None: if axis.i >= 0 and shape: axis.i += {2: 2, 3: 1}.get(len(shape), 0) - node.name += "_axis_updated" + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") nodes.append(node) - del graph.node[:] - graph.node.extend(nodes) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) return model def _apply_gather_elements(model): + """Prepare matching index ranks and preserve native axis and runtime-value contracts.""" graph = model.graph cache = GraphCache(graph) - nodes = [] - for node in graph.node: - if node.op_type != "GatherElements": - nodes.append(node) + converted_indices = {} + replacements = {} + selected_ids = {id(node) for node in _operator_nodes(model, {"GatherElements"})} + for index, node in enumerate(graph.node): + if id(node) not in selected_ids: continue - tensor = cache.get_init(node.input[1]) - if tensor is not None: + original = node.input[1] + shape = cache.get_shape(original) + tensor = cache.get_init(original) + rank = len(tensor.dims) if tensor is not None else len(shape) if shape is not None else None + added = [] + if original in converted_indices: + prepared, rank = converted_indices[original] + elif tensor is not None and original not in cache.graph_input_names: array = numpy_helper.to_array(tensor).astype(np.int32) - rank = array.ndim - if rank == 3: - array = array[None] + if rank in (2, 3): + array = array.reshape([*([1] * (4 - rank)), *array.shape]) elif rank == 0: array = array.reshape(1) - node.input[1] = add_initializer( + prepared = add_initializer( graph, - tensor.name + "_gatherelements_indices", - np.asarray(array, dtype=np.int32), + cache.unique_name(tensor.name + "_gatherelements_indices"), + array, dtype=TensorProto.INT32, with_value_info=False, ) else: - shape = cache.get_shape(node.input[1]) - rank = len(shape) if shape is not None else None - original = node.input[1] - cast = original + "_cast_to_int32" - nodes.append( - helper.make_node("Cast", [original], [cast], name=cast, to=TensorProto.INT32) + prepared = cache.unique_name(original + "_cast_to_int32") + added.append( + helper.make_node( + "Cast", + [original], + [prepared], + name=cache.unique_name(original + "_cast_to_int32_node"), + to=TensorProto.INT32, + ) ) - node.input[1] = cast - axis = _attribute(node, "axis") - if rank == 3 and axis is not None and axis.i != -1: - axis.i += 1 - node.name += "_axis_updated" - nodes.append(node) - del graph.node[:] - graph.node.extend(nodes) + # Other intermediate producers are promoted by the subsequent rank stages. + if rank in (2, 3) and original in cache.graph_input_names: + axes = add_initializer( + graph, + cache.unique_name(original + "_gatherelements_axes"), + np.arange(4 - rank, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + promoted = cache.unique_name(original + "_gatherelements_indices") + added.append( + helper.make_node( + "Unsqueeze", + [prepared, axes], + [promoted], + name=cache.unique_name(original + "_indices_unsqueeze"), + ) + ) + prepared = promoted + converted_indices[original] = (prepared, rank) + node.input[1] = prepared + if rank in (2, 3): + axis = _attribute(node, "axis") + if axis is None: + node.attribute.append(helper.make_attribute("axis", 0)) + axis = _attribute(node, "axis") + if axis.i >= 0: + axis.i += 4 - rank + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") + replacements[index] = ([node], [*added, node]) + batch_replace_nodes(graph, replacements) + live = {name for node in graph.node for name in node.input} + live.update(cache.graph_input_names | cache.graph_output_names | _subgraph_captures(graph)) + for original in converted_indices.keys() - live: + tensor = cache.get_init(original) + if tensor is not None: + graph.initializer.remove(tensor) + _add_legacy_initializer_inputs(model) return model def _apply_standalone_reduce_operations(model): + """Plan native reduction axes and preserve their output layouts during rank promotion.""" graph = model.graph cache = GraphCache(graph) nodes = [] count = 0 plans = {} - for node in graph.node: - if node.op_type not in {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: - continue + for node in _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}): shape = _shape_before_boundary_promotion(cache, node.input[0]) if not shape: keepdims = _attribute(node, "keepdims") @@ -190,7 +277,7 @@ def _apply_standalone_reduce_operations(model): raise ValueError(f"Missing static reduction input shape for {node.name!r}") continue # Already rank-preserving, including generated 4D reduction branches. tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None - if len(node.input) > 1 and tensor is None: + if len(node.input) > 1 and (tensor is None or tensor.name in cache.graph_input_names): raise ValueError(f"Dynamic reduction axes are unsupported for {node.name!r}") attribute = _attribute(node, "axes") raw_axes = ( @@ -221,7 +308,7 @@ def _apply_standalone_reduce_operations(model): if len(node.input) > 1: node.input[1] = add_initializer( graph, - node.name + "_axes", + cache.unique_name((node.name or node.output[0]) + "_axes"), np.asarray(values, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, @@ -235,32 +322,53 @@ def _apply_standalone_reduce_operations(model): target = pad4d(output_shape) name = add_initializer( graph, - f"reshape_shape_{count}", + cache.unique_name(f"reshape_shape_{count}"), np.asarray(target, dtype=np.int64), dtype=TensorProto.INT64, with_value_info=False, ) original = node.output[0] - node.output[0] = original + "_temp" + node.output[0] = cache.unique_name(original + "_temp") nodes.append( helper.make_node( - "Reshape", [node.output[0], name], [original], name=node.name + "_reshape" + "Reshape", + [node.output[0], name], + [original], + name=cache.unique_name((node.name or original) + "_reshape"), ) ) count += 1 - del graph.node[:] - graph.node.extend(nodes) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) return model def _apply_group_query_attention(model): + """Restore native attention ranks across internal tensors while retaining public boundaries.""" graph = model.graph cache = GraphCache(graph) nodes = [] boundaries = cache.graph_input_names | cache.init_names squeezed_inputs = {} + axes_names = {} + + def axes_for(name): + """Reuse a private axes control for adapters of the same attention tensor.""" + if name not in axes_names: + axes_names[name] = add_initializer( + graph, + cache.unique_name(name + "_axis"), + np.asarray([0], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + return axes_names[name] + + selected_ids = { + id(node) for node in _operator_nodes(model, {"GroupQueryAttention"}, ("com.microsoft",)) + } for node in graph.node: - if node.op_type != "GroupQueryAttention": + if id(node) not in selected_ids: nodes.append(node) continue for index, name in enumerate(node.input[:5]): @@ -268,13 +376,7 @@ def _apply_group_query_attention(model): if name in squeezed_inputs: node.input[index] = squeezed_inputs[name] continue - axes = add_initializer( - graph, - name + "_axis", - np.asarray([0], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) + axes = axes_for(name) squeezed = cache.unique_name(name + "_squeezed") squeezed_inputs[name] = squeezed nodes.append( @@ -282,29 +384,23 @@ def _apply_group_query_attention(model): "Squeeze", [name, axes], [squeezed], - name=node.name + "_" + name + "_squeeze", + name=cache.unique_name(node.name + "_" + name + "_squeeze"), ) ) node.input[index] = squeezed nodes.append(node) for index, name in enumerate(node.output): if name and name not in cache.graph_output_names: - axes = add_initializer( - graph, - name + "_axis", - np.asarray([0], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - node.output[index] = name + "_squeezed" + axes = axes_for(name) + node.output[index] = cache.unique_name(name + "_squeezed") nodes.append( helper.make_node( "Unsqueeze", [node.output[index], axes], [name], - name=node.name + "_" + name + "_unsqueeze", + name=cache.unique_name(node.name + "_" + name + "_unsqueeze"), ) ) - del graph.node[:] - graph.node.extend(nodes) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 2329230420d..2b4a29fb183 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -34,10 +34,12 @@ _apply_non4d_model_inputs, _apply_non4d_model_outputs, ) +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible @pytest.mark.parametrize("operation", ["Gather", "GatherElements"]) def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): + """Index dtype and rank conversion preserves selected values.""" if operation == "Gather": shape, output_shape, axis = [5, 4], [4], 0 indices = np.int64(2) @@ -73,7 +75,9 @@ def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): @pytest.mark.parametrize("direct_output", [False, True]) -def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_output): +@pytest.mark.parametrize("legacy_ir", [False, True]) +def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_output, legacy_ir): + """Attention preserves graph boundaries and wraps internal tensors.""" output = "attention" if direct_output else "y" nodes = [ helper.make_node("Unsqueeze", ["x", "axes"], ["q4"]), @@ -103,6 +107,12 @@ def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_ ), opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)], ) + if legacy_ir: + model.ir_version = 3 + for tensor in model.graph.initializer: + model.graph.input.append( + helper.make_tensor_value_info(tensor.name, tensor.data_type, list(tensor.dims)) + ) candidate = _apply_group_query_attention(model) onnx.checker.check_model(candidate) attention = next(n for n in candidate.graph.node if n.op_type == "GroupQueryAttention") @@ -114,7 +124,9 @@ def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_ ) -def test_late_refusal_preserves_the_original_model(): +@pytest.mark.parametrize("axes_default", [False, True]) +def test_late_refusal_preserves_the_original_model(axes_default): + """Late refusal preserves the original model.""" nodes = [ helper.make_node("ReduceSum", ["x"], ["first"], keepdims=0), helper.make_node("ReduceSum", ["x", "dynamic_axes"], ["y"], keepdims=0), @@ -128,6 +140,10 @@ def test_late_refusal_preserves_the_original_model(): helper.make_graph(nodes, "refusal", inputs, outputs), opset_imports=[helper.make_opsetid("", 21)], ) + if axes_default: + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([0], np.int64), "dynamic_axes") + ) before = model.SerializeToString() with pytest.raises(ValueError): _apply_standalone_reduce_operations(model) @@ -136,6 +152,7 @@ def test_late_refusal_preserves_the_original_model(): @pytest.mark.parametrize("axes", [[-2], [-1, 0], [0, -1]]) def test_reduction_preserves_negative_and_mixed_axes(axes): + """Reduction preserves negative and mixed axes.""" model = helper.make_model( helper.make_graph( [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=int(len(axes) > 1))], @@ -155,3 +172,385 @@ def test_reduction_preserves_negative_and_mixed_axes(axes): onnx.checker.check_model(candidate) actual = ReferenceEvaluator(candidate).run(None, {"x": data})[0] np.testing.assert_array_equal(actual, expected) + + +def _model(nodes, inputs, outputs, initializers=(), extra_opsets=()): + """Build and fully check a small indexing model with typed public interfaces.""" + model = helper.make_model( + helper.make_graph(nodes, "index_contract", inputs, outputs, initializer=initializers), + opset_imports=[helper.make_opsetid("", 21), *extra_opsets], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model + + +def _execute(model, feeds): + """Run CPU graph semantics without backend fusion changing the comparison.""" + 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("axis", [None, 0, 1, 2, -1, -2, -3]) +@pytest.mark.parametrize("kind", ["constant", "dynamic", "overridable"]) +def test_gather_elements_rank_three_preserves_axes_and_runtime_indices(axis, kind): + """Preserve positive, negative and default axes with fixed or runtime indices.""" + shape = [2, 3, 4] + normalized = (axis or 0) % 3 + indices = (np.arange(24).reshape(shape) % shape[normalized]).astype(np.int64) + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + initializers = [] + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + if kind != "constant": + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, shape)) + feeds["indices"] = (indices + 1) % shape[normalized] + if kind != "dynamic": + initializers.append(numpy_helper.from_array(indices, "indices")) + attributes = {} if axis is None else {"axis": axis} + model = _model( + [helper.make_node("GatherElements", ["x", "indices"], ["y"], **attributes)], + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + initializers, + ) + 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) + + +@pytest.mark.parametrize("axis", [None, -2]) +@pytest.mark.parametrize("dynamic", [False, True]) +def test_gather_elements_rank_two_preserves_selected_values(axis, dynamic): + """Keep data and index ranks aligned when two leading dimensions are introduced.""" + shape = [2, 3] + indices = (np.arange(6).reshape(shape) % 2).astype(np.int64) + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + initializers = [] + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + if dynamic: + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, shape)) + feeds["indices"] = indices + else: + initializers.append(numpy_helper.from_array(indices, "indices")) + attrs = {} if axis is None else {"axis": axis} + model = _model( + [helper.make_node("GatherElements", ["x", "indices"], ["y"], **attrs)], + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + initializers, + ) + 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_shared_dynamic_gather_elements_indices_have_one_preparation(): + """Reuse a valid dynamic index view across two GatherElements consumers.""" + shape = [2, 3, 4] + inputs = [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) for name in ["a", "b"] + ] + [helper.make_tensor_value_info("indices", TensorProto.INT64, shape)] + outputs = [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) for name in ["a_out", "b_out"] + ] + model = _model( + [ + helper.make_node("GatherElements", ["a", "indices"], ["a_out"], axis=1), + helper.make_node("GatherElements", ["b", "indices"], ["b_out"], axis=1), + ], + inputs, + outputs, + ) + feeds = { + "a": np.arange(24, dtype=np.float32).reshape(shape), + "b": -np.arange(24, dtype=np.float32).reshape(shape), + "indices": (np.arange(24).reshape(shape) % 3).astype(np.int64), + } + 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("keepdims", [False, True]) +def test_unnamed_reductions_keep_distinct_axes(keepdims): + """Allocate different controls for unnamed reductions with different axes.""" + shape = [2, 3, 4] + data = np.arange(24, dtype=np.float32).reshape(shape) + outputs = [ + helper.make_tensor_value_info( + "a", TensorProto.FLOAT, list(data.sum(axis=0, keepdims=keepdims).shape) + ), + helper.make_tensor_value_info( + "b", TensorProto.FLOAT, list(data.sum(axis=1, keepdims=keepdims).shape) + ), + ] + model = _model( + [ + helper.make_node("ReduceSum", ["x", "axes_a"], ["a"], keepdims=int(keepdims)), + helper.make_node("ReduceSum", ["x", "axes_b"], ["b"], keepdims=int(keepdims)), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + outputs, + [ + numpy_helper.from_array(np.asarray([0], np.int64), "axes_a"), + numpy_helper.from_array(np.asarray([1], np.int64), "axes_b"), + ], + ) + expected = _execute(model, {"x": data}) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, {"x": data}), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_reduction_generated_shape_preserves_existing_public_initializer(): + """Reserve reshape targets outside an unrelated public initializer namespace.""" + model = _model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=0)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [3]), + helper.make_tensor_value_info("reshape_shape_0", TensorProto.INT64, [1]), + ], + [ + numpy_helper.from_array(np.asarray([0], np.int64), "axes"), + numpy_helper.from_array(np.asarray([97], np.int64), "reshape_shape_0"), + ], + ) + 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) + + +@pytest.mark.parametrize("collision", [False, True]) +def test_chained_attention_preserves_native_kernel_inputs_and_values(collision): + """Run two real CPU attention kernels across the internal rank adapters.""" + shape = [1, 2, 8] + nodes = [ + helper.make_node( + "GroupQueryAttention", + ["q", "k", "v", "past_key", "past_value", "lengths", "total"], + ["middle", "pk1", "pv1"], + domain="com.microsoft", + name="first", + num_heads=1, + kv_num_heads=1, + ), + helper.make_node( + "GroupQueryAttention", + ["middle", "k", "v", "past_key", "past_value", "lengths", "total"], + ["y", "pk2", "pv2"], + domain="com.microsoft", + name="second", + num_heads=1, + kv_num_heads=1, + ), + ] + model = _model( + nodes, + [helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) for name in ["q", "k", "v"]], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.asarray([1], np.int32), "lengths"), + numpy_helper.from_array(np.int32(2), "total"), + numpy_helper.from_array(np.zeros((1, 1, 2, 8), np.float32), "past_key"), + numpy_helper.from_array(np.zeros((1, 1, 2, 8), np.float32), "past_value"), + ], + [helper.make_opsetid("com.microsoft", 1)], + ) + rng = np.random.default_rng(15) + feeds = {name: rng.standard_normal(shape).astype(np.float32) for name in ["q", "k", "v"]} + if collision: + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([2, 3], np.int64), "middle_axis") + ) + model.graph.output.append( + helper.make_tensor_value_info("middle_axis", TensorProto.INT64, [2]) + ) + 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-5, atol=1e-6) + + +@pytest.mark.parametrize( + "transform", + [ + _apply_gather, + _apply_gather_elements, + _apply_standalone_reduce_operations, + _apply_group_query_attention, + ], +) +def test_indexing_keeps_untouched_constant_payload_in_place(transform): + """Avoid copying a Constant payload when the stage only changes other operations.""" + model = _model( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("Identity", ["x"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + ) + constant = model.graph.node[0] + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + expected = _execute(model, feeds) + transform(model) + assert model.graph.node[0] is constant + for actual, reference in zip(_execute(model, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("producer", ["Cast", "Identity"]) +def test_intermediate_indices_keep_rank_when_upstream_boundary_is_promoted(producer): + """Preserve index ranks when a shape-preserving producer already receives promoted data.""" + shape = [2, 3, 4] + first = helper.make_node( + producer, + ["raw_indices"], + ["indices"], + **({"to": TensorProto.INT64} if producer == "Cast" else {}), + ) + model = _model( + [first, helper.make_node("GatherElements", ["x", "indices"], ["y"], axis=1)], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("raw_indices", TensorProto.INT64, shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(shape), + "raw_indices": (np.arange(24).reshape(shape) % 3).astype(np.int64), + } + 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) + + +@pytest.mark.parametrize( + ("operator", "transform"), + [ + ("Gather", _apply_gather), + ("GatherElements", _apply_gather_elements), + ("ReduceSum", _apply_standalone_reduce_operations), + ("GroupQueryAttention", _apply_group_query_attention), + ], +) +def test_indexing_preserves_opaque_model_local_operators(operator, transform): + """Leave custom functions intact when their names match indexing-stage operators.""" + shape = [1, 2, 3] + attributes = {"keepdims": 0} if operator == "ReduceSum" else {} + model = _model( + [ + helper.make_node("Identity", ["x"], ["middle"]), + helper.make_node(operator, ["middle"], ["y"], domain="custom", **attributes), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + extra_opsets=[helper.make_opsetid("custom", 1)], + ) + model.functions.append( + helper.make_function( + "custom", + operator, + ["X"], + ["Y"], + [helper.make_node("Neg", ["X"], ["Y"])], + [helper.make_opsetid("", 17)], + attributes=list(attributes), + ) + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + before = model.SerializeToString() + transform(model) + assert model.SerializeToString() == before + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(_execute(model, feeds)[0], expected[0]) + + +@pytest.mark.parametrize( + "indices", [np.int64(1), np.array([1], np.int64), np.array([0, 2], np.int64)] +) +@pytest.mark.parametrize("producer", [False, True]) +def test_gather_preserves_overridable_scalar_and_vector_indices(indices, producer): + """Execute runtime Gather overrides while retaining scalar and singleton-vector ranks.""" + indices = np.asarray(indices) + model = _model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=0)], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 4]), + helper.make_tensor_value_info("indices", TensorProto.INT64, list(indices.shape)), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [*indices.shape, 4])], + [numpy_helper.from_array(indices, "indices")], + ) + if producer: + model.graph.node.insert( + 0, helper.make_node("Cast", ["indices"], ["cast_indices"], to=TensorProto.INT64) + ) + model.graph.node[1].input[1] = "cast_indices" + feeds = { + "x": np.arange(12, dtype=np.float32).reshape(3, 4), + "indices": np.asarray((indices + 1) % 3, dtype=np.int64), + } + 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_reduction_registers_legacy_initializer_inputs(): + """Keep IR3 controls declared while preserving reduction output values and ranks.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[2], keepdims=0)], + "legacy_reduce", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds)[0] + transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) + transformed = _apply_standalone_reduce_operations(transformed) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) + + +def test_gather_preserves_existing_public_initializer_name(): + """Allocate private Gather controls without overwriting a public initializer.""" + model = _model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=0)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 4])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [4]), + helper.make_tensor_value_info("indices_gather_indices", TensorProto.INT64, [1]), + ], + [ + numpy_helper.from_array(np.int64(1), "indices"), + numpy_helper.from_array(np.asarray([17], np.int64), "indices_gather_indices"), + ], + ) + feeds = {"x": np.arange(12, dtype=np.float32).reshape(3, 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) From 30c28d4b91321fd2f9961bc1bd17ae0bcef1b77b Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:27:38 +0000 Subject: [PATCH 04/31] Keep attention caches and public query tensors in their native layouts Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 6 ++-- .../dla_transforms/test_dla_indexing.py | 35 ++++++++++++++++++- 2 files changed, 37 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 0ec66cd2cef..7e03379909b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -348,7 +348,7 @@ def _apply_group_query_attention(model): graph = model.graph cache = GraphCache(graph) nodes = [] - boundaries = cache.graph_input_names | cache.init_names + boundaries = cache.graph_input_names | cache.init_names | cache.graph_output_names squeezed_inputs = {} axes_names = {} @@ -371,7 +371,7 @@ def axes_for(name): if id(node) not in selected_ids: nodes.append(node) continue - for index, name in enumerate(node.input[:5]): + for index, name in enumerate(node.input[:3]): if name and name not in boundaries: if name in squeezed_inputs: node.input[index] = squeezed_inputs[name] @@ -389,7 +389,7 @@ def axes_for(name): ) node.input[index] = squeezed nodes.append(node) - for index, name in enumerate(node.output): + for index, name in enumerate(node.output[:1]): if name and name not in cache.graph_output_names: axes = axes_for(name) node.output[index] = cache.unique_name(name + "_squeezed") diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 2b4a29fb183..4aeb1ee9fc5 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -330,7 +330,8 @@ def test_reduction_generated_shape_preserves_existing_public_initializer(): @pytest.mark.parametrize("collision", [False, True]) -def test_chained_attention_preserves_native_kernel_inputs_and_values(collision): +@pytest.mark.parametrize("cache_views", [False, True]) +def test_chained_attention_preserves_native_kernel_inputs_and_values(collision, cache_views): """Run two real CPU attention kernels across the internal rank adapters.""" shape = [1, 2, 8] nodes = [ @@ -374,7 +375,39 @@ def test_chained_attention_preserves_native_kernel_inputs_and_values(collision): model.graph.output.append( helper.make_tensor_value_info("middle_axis", TensorProto.INT64, [2]) ) + if cache_views: + attention = [node for node in model.graph.node if node.op_type == "GroupQueryAttention"] + for name in ["past_key", "past_value"]: + model.graph.input.append( + helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 1, 2, 8]) + ) + feeds[name] = np.arange(16, dtype=np.float32).reshape(1, 1, 2, 8) / 16 + model.graph.node.insert(0, helper.make_node("Identity", [name], [name + "_internal"])) + for node in attention: + node.input[3] = "past_key_internal" + node.input[4] = "past_value_internal" + next(value for value in model.graph.input if value.name == "q").CopyFrom( + helper.make_tensor_value_info("q", TensorProto.FLOAT, [1, *shape]) + ) + feeds["q"] = feeds["q"].reshape(1, *shape) + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([0], np.int64), "public_q_axes") + ) + model.graph.node.insert( + 0, helper.make_node("Squeeze", ["q", "public_q_axes"], ["public_q"]) + ) + model.graph.output.append( + helper.make_tensor_value_info("public_q", TensorProto.FLOAT, shape) + ) + attention[0].input[0] = "public_q" + for name in ["pk1", "pv1"]: + model.graph.node.append(helper.make_node("Identity", [name], [name + "_public"])) + model.graph.output.append( + helper.make_tensor_value_info(name + "_public", TensorProto.FLOAT, [1, 1, 2, 8]) + ) expected = _execute(model, feeds) + if cache_views: + assert expected[-1].shape == expected[-2].shape == (1, 1, 2, 8) 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-5, atol=1e-6) From fce65a89caba2c40812c888b934afe01a5a54365 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:33:52 +0000 Subject: [PATCH 05/31] Use independent reference execution for opaque operator fixtures Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/test_dla_indexing.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 4aeb1ee9fc5..97988a18c70 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -511,12 +511,13 @@ def test_indexing_preserves_opaque_model_local_operators(operator, transform): ) model = onnx.shape_inference.infer_shapes(model) feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} - expected = _execute(model, feeds) + # Use the independent interpreter for opaque custom-function semantics. + expected = ReferenceEvaluator(model).run(None, feeds) before = model.SerializeToString() transform(model) assert model.SerializeToString() == before onnx.checker.check_model(model, full_check=True) - np.testing.assert_array_equal(_execute(model, feeds)[0], expected[0]) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected[0]) @pytest.mark.parametrize( From dcd3d5f1be1b2ff620ec01d4ac7ab0e8290ca8a8 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 00:46:29 +0000 Subject: [PATCH 06/31] Reject unsupported fixed Gather index ranks before mutation Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 5 +++ .../dla_transforms/test_dla_indexing.py | 32 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 7e03379909b..4cf365971f4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -72,6 +72,11 @@ def _apply_gather(model): for candidate in selected: if cache.get_shape(candidate.input[0]) is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") + fixed = cache.get_init(candidate.input[1]) + if fixed is not None and fixed.name not in cache.graph_input_names and len(fixed.dims) > 1: + raise ValueError( + f"{candidate.name!r}: fixed Gather indices with rank > 1 are unsupported" + ) if cache.get_init(candidate.input[1]) is None and cache.get_dtype( candidate.input[1] ) not in (TensorProto.INT32, TensorProto.INT64): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 97988a18c70..e9b54b45c09 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -588,3 +588,35 @@ def test_gather_preserves_existing_public_initializer_name(): 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("index_shape", [(2, 3), (1, 2, 3)]) +def test_multidimensional_fixed_gather_indices_refuse_before_any_mutation(index_shape): + """Reject unsupported index layouts before changing an earlier valid Gather.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Gather", ["x", "first_indices"], ["first"], axis=0), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=0), + ], + "fixed_index_ranks", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [5, 4])], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [1, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, [*index_shape, 4]), + ], + [ + numpy_helper.from_array(np.asarray([1], np.int64), "first_indices"), + numpy_helper.from_array( + np.arange(6, dtype=np.int64).reshape(index_shape) % 5, "indices" + ), + ], + ), + 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=r"fixed Gather indices.*rank"): + _apply_gather(model) + assert model.SerializeToString() == before From 0de9969783b5912eaf9682e9822b607eab8e9724 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 01:01:23 +0000 Subject: [PATCH 07/31] Reject multidimensional runtime Gather indices before mutation Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_indexing.py | 9 +++++++++ .../dla_transforms/test_dla_indexing.py | 12 ++++++++++-- 2 files changed, 19 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 4cf365971f4..05261fb37db 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -72,6 +72,15 @@ def _apply_gather(model): for candidate in selected: if cache.get_shape(candidate.input[0]) is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") + index_shape = cache.get_shape(candidate.input[1]) + if ( + candidate.input[1] in cache.graph_input_names + and index_shape is not None + and len(index_shape) > 1 + ): + raise ValueError( + f"{candidate.name!r}: graph-input Gather indices with rank > 1 are unsupported" + ) fixed = cache.get_init(candidate.input[1]) if fixed is not None and fixed.name not in cache.graph_input_names and len(fixed.dims) > 1: raise ValueError( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index e9b54b45c09..b9e5d8fa9c6 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -591,7 +591,8 @@ def test_gather_preserves_existing_public_initializer_name(): @pytest.mark.parametrize("index_shape", [(2, 3), (1, 2, 3)]) -def test_multidimensional_fixed_gather_indices_refuse_before_any_mutation(index_shape): +@pytest.mark.parametrize("kind", ["fixed", "input", "default"]) +def test_multidimensional_gather_indices_refuse_before_any_mutation(index_shape, kind): """Reject unsupported index layouts before changing an earlier valid Gather.""" model = helper.make_model( helper.make_graph( @@ -616,7 +617,14 @@ def test_multidimensional_fixed_gather_indices_refuse_before_any_mutation(index_ ) model.ir_version = 9 onnx.checker.check_model(model, full_check=True) + if kind != "fixed": + model.graph.input.append( + helper.make_tensor_value_info("indices", TensorProto.INT64, index_shape) + ) + if kind == "input": + model.graph.initializer.remove(model.graph.initializer[1]) + onnx.checker.check_model(model, full_check=True) before = model.SerializeToString() - with pytest.raises(ValueError, match=r"fixed Gather indices.*rank"): + with pytest.raises(ValueError, match=r"Gather indices.*rank"): _apply_gather(model) assert model.SerializeToString() == before From 03e26a7bb1e95905859cc8eaa54b14ac939db837 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 01:21:17 +0000 Subject: [PATCH 08/31] Validate intermediate Gather index layouts and preserve symbolic vectors Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 94 +++++++-------- .../dla_transforms/test_dla_indexing.py | 108 ++++++++++++++++++ 2 files changed, 147 insertions(+), 55 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 05261fb37db..a1527a9bd2e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -62,6 +62,17 @@ def _operator_nodes(model, operators, domains=("", "ai.onnx")): ] +def _known_rank(cache, name): + """Read tensor rank even when individual dimensions are symbolic.""" + shape = cache.get_shape(name) + if shape is not None: + return len(shape) + value = cache.vi_map.get(name) + if value is not None and value.type.tensor_type.HasField("shape"): + return len(value.type.tensor_type.shape.dim) + return None + + def _apply_gather(model): """Prepare Gather indices for the subsequent ordered data-rank promotion.""" graph = model.graph @@ -69,28 +80,30 @@ def _apply_gather(model): nodes = [] squeezed_indices = {} selected = _operator_nodes(model, {"Gather"}) + casts = {id(node) for node in _operator_nodes(model, {"Cast"})} for candidate in selected: if cache.get_shape(candidate.input[0]) is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") - index_shape = cache.get_shape(candidate.input[1]) - if ( - candidate.input[1] in cache.graph_input_names - and index_shape is not None - and len(index_shape) > 1 + original = candidate.input[1] + index_shape = cache.get_shape(original) + rank = _known_rank(cache, original) + producer = cache.get_producer(original) + fixed = cache.get_init(original) + if rank is None: + raise ValueError(f"{candidate.name!r}: Gather indices require a known rank") + if rank > 1 and ( + original in cache.graph_input_names + or fixed is not None + or (producer is not None and id(producer) in casts) + or index_shape is None + or 1 in index_shape ): - raise ValueError( - f"{candidate.name!r}: graph-input Gather indices with rank > 1 are unsupported" - ) - fixed = cache.get_init(candidate.input[1]) - if fixed is not None and fixed.name not in cache.graph_input_names and len(fixed.dims) > 1: - raise ValueError( - f"{candidate.name!r}: fixed Gather indices with rank > 1 are unsupported" - ) - if cache.get_init(candidate.input[1]) is None and cache.get_dtype( - candidate.input[1] - ) not in (TensorProto.INT32, TensorProto.INT64): - producer = cache.get_producer(candidate.input[1]) - if producer is None or producer.op_type != "Cast": + raise ValueError(f"{candidate.name!r}: Gather indices with rank > 1 are unsupported") + if fixed is None and cache.get_dtype(original) not in ( + TensorProto.INT32, + TensorProto.INT64, + ): + if producer is None or id(producer) not in casts: raise ValueError( f"{candidate.name!r}: dynamic Gather indices require an integer Cast" ) @@ -113,8 +126,8 @@ def _apply_gather(model): node.input[1] = squeezed_indices[indices.name] else: original = node.input[1] - index_shape = cache.get_shape(original) - if index_shape is not None and len(index_shape) <= 1: + rank = _known_rank(cache, original) + if rank is not None and rank <= 1: if original not in squeezed_indices: prepared = cache.unique_name(original + "_gather_int32") nodes.append( @@ -136,42 +149,13 @@ def _apply_gather(model): nodes.append(node) continue squeezed = cache.unique_name(original + "_squeezed") - squeeze_inputs = [original] - producer = cache.get_producer(original) - if producer is not None and producer.op_type == "Cast": - target = _attribute(producer, "to") - shared = ( - original in cache.graph_output_names or len(cache.get_consumers(original)) > 1 - ) - if shared: - integer = cache.unique_name(original + "_gather_int32") - nodes.append( - helper.make_node( - "Cast", - [producer.input[0]], - [integer], - name=integer, - to=TensorProto.INT32, - ) - ) - squeeze_inputs[0] = integer - else: - target.i = TensorProto.INT32 - for value in graph.value_info: - if value.name == original: - value.type.tensor_type.elem_type = TensorProto.INT32 - axes = add_initializer( - graph, - cache.unique_name(node.name + "_indices_squeeze_axes"), - np.asarray([0, 1, 2], dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, - ) - squeeze_inputs.append(axes) - elif cache.get_dtype(original) not in (TensorProto.INT32, TensorProto.INT64): - raise ValueError(f"{node.name!r}: dynamic Gather indices require an integer Cast") nodes.append( - helper.make_node("Squeeze", squeeze_inputs, [squeezed], name=node.name + "_squeeze") + helper.make_node( + "Squeeze", + [original], + [squeezed], + name=cache.unique_name((node.name or node.output[0]) + "_squeeze"), + ) ) node.input[1] = squeezed squeezed_indices[original] = squeezed diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index b9e5d8fa9c6..9625d300496 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -628,3 +628,111 @@ def test_multidimensional_gather_indices_refuse_before_any_mutation(index_shape, with pytest.raises(ValueError, match=r"Gather indices.*rank"): _apply_gather(model) assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("producer", "shape", "refuse"), + [ + ("Cast", [2, 3], True), + ("Cast", [1, 3], True), + ("Cast", [1, 2, 3], True), + ("Cast", [2, 3, 4], True), + ("Cast", [1, 1, 1, 3], True), + ("Identity", [1, 3], True), + ("Identity", [1, 2, 3], True), + ("Identity", [2, 3], False), + ("Identity", [2, 3, 4], False), + ], +) +def test_intermediate_gather_indices_validate_squeeze_layout_before_mutation( + producer, shape, refuse +): + """Reject layout-changing squeezes while preserving higher-rank no-op index views.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Gather", ["x", "first_indices"], ["first"], axis=2), + helper.make_node( + producer, + ["source"], + ["indices"], + **({"to": TensorProto.INT64} if producer == "Cast" else {}), + ), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=2), + ], + "intermediate_indices", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 5, 4]), + helper.make_tensor_value_info("source", TensorProto.INT64, shape), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [1, 1, 1, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, *shape, 4]), + ], + [numpy_helper.from_array(np.asarray([1], np.int64), "first_indices")], + value_info=[helper.make_tensor_value_info("indices", TensorProto.INT64, shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + if refuse: + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"Gather indices.*rank"): + _apply_gather(model) + assert model.SerializeToString() == before + else: + feeds = { + "x": np.arange(20, dtype=np.float32).reshape(1, 1, 5, 4), + "source": np.arange(np.prod(shape), dtype=np.int64).reshape(shape) % 5, + } + expected = _execute(model, feeds) + converted = _apply_gather(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("producer", [None, "Cast", "Identity"]) +@pytest.mark.parametrize("length", [0, 1, 3]) +def test_symbolic_vector_indices_preserve_runtime_length_and_values(producer, length): + """Keep a known vector rank even when its runtime dimension is symbolic.""" + name = "source" if producer is None else "indices" + nodes = ( + [] + if producer is None + else [ + helper.make_node( + producer, + ["source"], + ["indices"], + **({"to": TensorProto.INT64} if producer == "Cast" else {}), + ) + ] + ) + nodes.append(helper.make_node("Gather", ["x", name], ["y"], axis=2)) + model = helper.make_model( + helper.make_graph( + nodes, + "symbolic_indices", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 5, 4]), + helper.make_tensor_value_info("source", TensorProto.INT64, [None]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, None, 4])], + value_info=[] + if producer is None + else [helper.make_tensor_value_info("indices", TensorProto.INT64, [None])], + ), + 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(1, 1, 5, 4), + "source": np.arange(length, dtype=np.int64), + } + expected = _execute(model, feeds)[0] + converted = _apply_gather(copy.deepcopy(model)) + onnx.checker.check_model(converted, full_check=True) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) From 848e0ffac6044844c6c2eeb313ef94b2dcea4f87 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 01:26:36 +0000 Subject: [PATCH 09/31] Recover native reduction ranks behind allocated boundary adapters Signed-off-by: mgohil --- .../onnx/graph_surgery/utils/graph_utils.py | 21 +++++++++++-- .../dla_transforms/test_dla_indexing.py | 30 +++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py index 3f744c71d62..d2a59549c1a 100644 --- a/modelopt/onnx/graph_surgery/utils/graph_utils.py +++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py @@ -994,10 +994,25 @@ def _allocated_boundary_name(name, hint): def _shape_before_boundary_promotion(cache, name): """Recover native shapes at generated input and output boundary views.""" producer = cache.get_producer(name) - if producer is not None and producer.op_type == "Unsqueeze": + if ( + producer is not None + and producer.op_type == "Unsqueeze" + and producer.domain in {"", "ai.onnx"} + and not producer.overload + ): original = producer.input[0] - if original in cache.graph_input_names and name == original + "_unsqueeze_input": - return cache.get_shape(original) + hint = original + "_unsqueeze_input" + allocated = _allocated_boundary_name(name, hint) + shape = cache.get_shape(original) + axes = cache.get_init_array(producer.input[1]) if len(producer.input) > 1 else None + if ( + original in cache.graph_input_names + and allocated + and shape is not None + and axes is not None + and axes.reshape(-1).tolist() == list(range(max(0, 4 - len(shape)))) + ): + return shape shape = cache.get_shape(name) if shape is not None: return shape diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 9625d300496..3279850b0d1 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -736,3 +736,33 @@ def test_symbolic_vector_indices_preserve_runtime_length_and_values(producer, le converted = _apply_gather(copy.deepcopy(model)) onnx.checker.check_model(converted, full_check=True) np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_reduction_recovers_native_axes_behind_an_allocated_input_adapter(axis): + """Recover the native rank when a public tensor occupies the preferred adapter name.""" + shape = [2, 3, 4] + output_shape = shape.copy() + output_shape[axis] = 1 + model = helper.make_model( + helper.make_graph( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=1)], + "allocated_reduction", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape), + helper.make_tensor_value_info("x_unsqueeze_input", TensorProto.FLOAT, [1]), + ], + [ + numpy_helper.from_array(np.asarray([axis], np.int64), "axes"), + numpy_helper.from_array(np.asarray([17], np.float32), "x_unsqueeze_input"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(24, 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) From 95b1fe2a0abcb62cc27ac38754d067e0c954c328 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 02:20:48 +0000 Subject: [PATCH 10/31] Preserve supported native Gather index views across rank stages Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 109 ++++++++++++++-- .../dla_transforms/test_dla_indexing.py | 121 ++++++++++++++++++ 2 files changed, 222 insertions(+), 8 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index a1527a9bd2e..a445a4dce29 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -15,6 +15,8 @@ """DLA indexing, reduction and attention passes with deferred rank promotion.""" +import math + import numpy as np from onnx import TensorProto, helper, numpy_helper @@ -73,6 +75,45 @@ def _known_rank(cache, name): return None +def _gather_flatten_count(cache, node, index_shape): + """Return the flattened count only when leading padding preserves the native output layout.""" + shape = cache.get_shape(node.input[0]) + output = cache.get_shape(node.output[0]) + if ( + index_shape is None + or shape is None + or len(shape) not in (2, 3) + or output is None + or len(output) > 4 + ): + return None + axis = _attribute(node, "axis") + value = axis.i if axis is not None else 0 + if not -len(shape) <= value < len(shape): + return None + value = value % len(shape) + 4 - len(shape) + prepared = pad4d(shape) + count = math.prod(index_shape) + return count if pad4d(output) == [*prepared[:value], count, *prepared[value + 1 :]] else None + + +def _gather_index_needs_view(cache, name, layout_nodes): + """Find native index layouts that later rank stages would otherwise promote.""" + seen = set() + producer = cache.get_producer(name) + while ( + producer is not None + and id(producer) in layout_nodes + and producer.op_type in {"Cast", "Identity"} + ): + if name in seen: + raise ValueError("Cyclic Gather index producers are unsupported") + seen.add(name) + name = producer.input[0] + producer = cache.get_producer(name) + return producer is not None and id(producer) in layout_nodes + + def _apply_gather(model): """Prepare Gather indices for the subsequent ordered data-rank promotion.""" graph = model.graph @@ -80,7 +121,12 @@ def _apply_gather(model): nodes = [] squeezed_indices = {} selected = _operator_nodes(model, {"Gather"}) - casts = {id(node) for node in _operator_nodes(model, {"Cast"})} + layout = _operator_nodes(model, {"Cast", "Identity", "Reshape", "Squeeze", "Unsqueeze"}) + layout_nodes = {id(node) for node in layout} + casts = {id(node) for node in layout if node.op_type == "Cast"} + flattened = set() + index_views = {} + captured = _subgraph_captures(graph) for candidate in selected: if cache.get_shape(candidate.input[0]) is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") @@ -91,6 +137,10 @@ def _apply_gather(model): fixed = cache.get_init(original) if rank is None: raise ValueError(f"{candidate.name!r}: Gather indices require a known rank") + if rank > 1 and original in captured: + raise ValueError( + f"{candidate.name!r}: captured Gather indices with rank > 1 are unsupported" + ) if rank > 1 and ( original in cache.graph_input_names or fixed is not None @@ -98,7 +148,17 @@ def _apply_gather(model): or index_shape is None or 1 in index_shape ): - raise ValueError(f"{candidate.name!r}: Gather indices with rank > 1 are unsupported") + count = _gather_flatten_count(cache, candidate, index_shape) + if count is None: + raise ValueError( + f"{candidate.name!r}: Gather indices with rank > 1 are unsupported" + ) + flattened.add(original) + index_views[original] = 1 + if rank <= 1 and original not in index_views: + index_views[original] = ( + rank if _gather_index_needs_view(cache, original, layout_nodes) else None + ) if fixed is None and cache.get_dtype(original) not in ( TensorProto.INT32, TensorProto.INT64, @@ -127,14 +187,41 @@ def _apply_gather(model): else: original = node.input[1] rank = _known_rank(cache, original) - if rank is not None and rank <= 1: + if rank is not None and (rank <= 1 or original in flattened): if original not in squeezed_indices: - prepared = cache.unique_name(original + "_gather_int32") - nodes.append( - helper.make_node( - "Cast", [original], [prepared], name=prepared, to=TensorProto.INT32 + producer = cache.get_producer(original) + if ( + original in flattened + and producer is not None + and id(producer) in casts + and len(cache.get_consumers(original)) == 1 + and original not in cache.graph_output_names | captured + ): + _attribute(producer, "to").i = TensorProto.INT32 + for value in graph.value_info: + if value.name == original: + value.type.tensor_type.elem_type = TensorProto.INT32 + prepared = original + else: + prepared = cache.unique_name(original + "_gather_int32") + nodes.append( + helper.make_node( + "Cast", [original], [prepared], name=prepared, to=TensorProto.INT32 + ) ) - ) + if index_views.get(original) is not None: + target = add_initializer( + graph, + cache.unique_name(original + "_gather_shape"), + np.asarray([] if index_views[original] == 0 else [-1], np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + view = cache.unique_name(original + "_gather_view") + nodes.append( + helper.make_node("Reshape", [prepared, target], [view], name=view) + ) + prepared = view squeezed_indices[original] = prepared if original in squeezed_indices: node.input[1] = squeezed_indices[original] @@ -174,6 +261,12 @@ def _apply_gather(model): node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") nodes.append(node) _insert_support_nodes(graph, nodes) + live = {name for node in graph.node for name in node.input} + live |= cache.graph_input_names | cache.graph_output_names | captured + for original in squeezed_indices: + tensor = cache.get_init(original) + if tensor is not None and original not in live: + graph.initializer.remove(tensor) _add_legacy_initializer_inputs(model) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 3279850b0d1..01c7e94f2bc 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -766,3 +766,124 @@ def test_reduction_recovers_native_axes_behind_an_allocated_input_adapter(axis): 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) + + +def _leading_gather_index_model(producer, shape, *, captured=False, public=True): + """Build leading-singleton indices with a public or nested owner of their native values.""" + values = np.arange(np.prod(shape), dtype=np.int64).reshape(shape) % 5 + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [5, 4])] + nodes, initializers = [], [] + feeds = {"x": np.arange(20, dtype=np.float32).reshape(5, 4)} + if producer == "fixed": + initializers.append(numpy_helper.from_array(values, "indices")) + elif producer == "input": + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, shape)) + feeds["indices"] = values + else: + dtype = TensorProto.FLOAT if producer == "Cast" else TensorProto.INT64 + inputs.append(helper.make_tensor_value_info("source", dtype, shape)) + feeds["source"] = values.astype(np.float32) if producer == "Cast" else values + nodes.append( + helper.make_node( + producer, + ["source"], + ["indices"], + **({"to": TensorProto.INT64} if producer == "Cast" else {}), + ) + ) + nodes.append(helper.make_node("Gather", ["x", "indices"], ["y"], axis=0)) + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [*shape, 4])] + if captured: + branch = helper.make_graph( + [helper.make_node("Identity", ["indices"], ["out"])], + "branch", + [], + [helper.make_tensor_value_info("out", TensorProto.INT64, shape)], + ) + initializers.append(numpy_helper.from_array(np.asarray(True), "condition")) + nodes.append( + helper.make_node("If", ["condition"], ["view"], then_branch=branch, else_branch=branch) + ) + outputs.append(helper.make_tensor_value_info("view", TensorProto.INT64, shape)) + elif public: + outputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, shape)) + model = helper.make_model( + helper.make_graph( + nodes, + "leading_indices", + inputs, + outputs, + initializers, + value_info=[] + if producer in {"fixed", "input"} + else [helper.make_tensor_value_info("indices", TensorProto.INT64, shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model, feeds + + +@pytest.mark.parametrize( + ("producer", "public"), + [("fixed", True), ("input", True), ("Cast", True), ("Identity", True), ("fixed", False)], +) +@pytest.mark.parametrize("shape", [[1, 1], [1, 3]]) +def test_leading_singleton_indices_preserve_native_gather_and_public_index_layouts( + producer, public, shape +): + """Keep supported native layouts while preparing private stable four-dimensional views.""" + model, feeds = _leading_gather_index_model(producer, shape, public=public) + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_captured_multidimensional_indices_refuse_before_layout_promotion(): + """Leave nested owners unchanged when their index layout cannot be promoted safely.""" + model, _ = _leading_gather_index_model("Cast", [1, 3], captured=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"captured Gather indices.*rank"): + _apply_gather(model) + assert model.SerializeToString() == before + + +def test_leading_runtime_indices_keep_legacy_initializer_declarations(): + """Declare generated shape and axis controls without changing native legacy public values.""" + model, feeds = _leading_gather_index_model("input", [1, 3]) + model.ir_version = 3 + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + onnx.checker.check_model(converted, full_check=True) + assert {t.name for t in converted.graph.initializer} <= {v.name for v in converted.graph.input} + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.timeout(2) +def test_cyclic_index_layout_ancestry_refuses_before_mutation(): + """Bound malformed producer walks instead of hanging during index view planning.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["second"], ["first"], to=TensorProto.INT64), + helper.make_node("Cast", ["first"], ["second"], to=TensorProto.INT64), + helper.make_node("Gather", ["x", "first"], ["y"], axis=2), + ], + "cyclic_indices", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 5, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 4])], + value_info=[ + helper.make_tensor_value_info(n, TensorProto.INT64, [2]) + for n in ["first", "second"] + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Cyclic Gather"): + _apply_gather(model) + assert model.SerializeToString() == before From dffd52b2c9d8ddb519b308cd0bf87b30ee5de673 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 07:01:53 +0000 Subject: [PATCH 11/31] Preserve symbolic GatherElements index and output ranks Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 35 ++++++++++++++-- .../dla_transforms/test_dla_indexing.py | 40 +++++++++++++++++++ 2 files changed, 72 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 a445a4dce29..293cb4c66d5 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -282,9 +282,12 @@ def _apply_gather_elements(model): if id(node) not in selected_ids: continue original = node.input[1] - shape = cache.get_shape(original) tensor = cache.get_init(original) - rank = len(tensor.dims) if tensor is not None else len(shape) if shape is not None else None + rank = len(tensor.dims) if tensor is not None else _known_rank(cache, original) + data_rank = _known_rank(cache, node.input[0]) + if rank in (2, 3) and data_rank != 4 and cache.get_shape(node.input[0]) is None: + # Boundary promotion can leave symbolic data native; retain matching native indices. + continue added = [] if original in converted_indices: prepared, rank = converted_indices[original] @@ -341,7 +344,33 @@ def _apply_gather_elements(model): if axis.i >= 0: axis.i += 4 - rank node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") - replacements[index] = ([node], [*added, node]) + restored = [] + output = node.output[0] + if ( + data_rank == 4 + and output in cache.graph_output_names + and rank in (2, 3) + and _known_rank(cache, output) == rank + ): + # The earlier output stage cannot plan adapters from concrete shapes for symbolic outputs. + internal = cache.unique_name(output + "_gatherelements_4d") + node.output[0] = internal + axes = add_initializer( + graph, + cache.unique_name(output + "_gatherelements_squeeze_axes"), + np.arange(4 - rank, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + restored.append( + helper.make_node( + "Squeeze", + [internal, axes], + [output], + name=cache.unique_name(output + "_restore_rank"), + ) + ) + replacements[index] = ([node], [*added, node, *restored]) batch_replace_nodes(graph, replacements) live = {name for node in graph.node for name in node.input} live.update(cache.graph_input_names | cache.graph_output_names | _subgraph_captures(graph)) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 01c7e94f2bc..73dbb3edcfb 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -249,6 +249,46 @@ def test_gather_elements_rank_two_preserves_selected_values(axis, dynamic): np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) +@pytest.mark.parametrize("length", [1, 3]) +@pytest.mark.parametrize("symbolic_data", [False, True]) +@pytest.mark.parametrize("public_reader", [False, True]) +def test_symbolic_gather_elements_index_rank_preserves_runtime_lengths( + length, symbolic_data, public_reader +): + """Promote known-rank symbolic indices and preserve unequal gather-axis lengths.""" + model = _model( + [helper.make_node("GatherElements", ["x", "indices"], ["y"], axis=1)], + [ + helper.make_tensor_value_info( + "x", TensorProto.FLOAT, [2, "width" if symbolic_data else 4, 3] + ), + helper.make_tensor_value_info("indices", TensorProto.INT64, [2, "count", 3]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, "count", 3])], + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 4, 3), + "indices": (np.arange(6 * length).reshape(2, length, 3) % 4).astype(np.int64), + } + if public_reader: + model.graph.node.append( + helper.make_node("GatherElements", ["y", "reader_indices"], ["read"], axis=1) + ) + model.graph.input.append( + helper.make_tensor_value_info("reader_indices", TensorProto.INT64, [2, "count", 3]) + ) + model.graph.output.append( + helper.make_tensor_value_info("read", TensorProto.FLOAT, [2, "count", 3]) + ) + feeds["reader_indices"] = (np.arange(6 * length).reshape(2, length, 3) % length).astype( + np.int64 + ) + 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_shared_dynamic_gather_elements_indices_have_one_preparation(): """Reuse a valid dynamic index view across two GatherElements consumers.""" shape = [2, 3, 4] From 72931151a9980049cb2aa8b2b6bcb15f4ca6ed7d Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:21:54 +0000 Subject: [PATCH 12/31] Validate intermediate Gather index layouts before DLA promotion Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 4 +- .../dla_transforms/test_dla_indexing.py | 39 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 293cb4c66d5..ef65d2286b9 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -128,7 +128,8 @@ def _apply_gather(model): index_views = {} captured = _subgraph_captures(graph) for candidate in selected: - if cache.get_shape(candidate.input[0]) is None: + data_shape = cache.get_shape(candidate.input[0]) + if data_shape is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") original = candidate.input[1] index_shape = cache.get_shape(original) @@ -147,6 +148,7 @@ def _apply_gather(model): or (producer is not None and id(producer) in casts) or index_shape is None or 1 in index_shape + or len(data_shape) in (2, 3) ): count = _gather_flatten_count(cache, candidate, index_shape) if count is None: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 73dbb3edcfb..439d76ec472 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -904,6 +904,45 @@ def test_leading_runtime_indices_keep_legacy_initializer_declarations(): @pytest.mark.timeout(2) +@pytest.mark.parametrize("shape", [[5, 4], [2, 5, 4]]) +def test_multidimensional_identity_gather_indices_refuse_before_mutation(shape): + """Reject intermediate index layouts that leading data padding cannot preserve.""" + axis = len(shape) - 2 + output = [*shape[:axis], 2, 3, *shape[axis + 1 :]] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Gather", ["x", "first_indices"], ["first"], axis=axis), + helper.make_node("Identity", ["source"], ["indices"]), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=axis), + ], + "native_intermediate_indices", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("source", TensorProto.INT64, [2, 3]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [*shape[:axis], 1, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, output), + ], + [numpy_helper.from_array(np.asarray([1], np.int64), "first_indices")], + value_info=[helper.make_tensor_value_info("indices", TensorProto.INT64, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + feeds = { + "x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape), + "source": np.arange(6, dtype=np.int64).reshape(2, 3) % 5, + } + _execute(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"Gather indices.*rank"): + _apply_gather(model) + assert model.SerializeToString() == before + + def test_cyclic_index_layout_ancestry_refuses_before_mutation(): """Bound malformed producer walks instead of hanging during index view planning.""" model = helper.make_model( From 8e04c47ed0b8498f44b9626773d51cd01511fafd Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:51:56 +0000 Subject: [PATCH 13/31] Declare reduction preparation in focused DLA boundary tests Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/test_dla_indexing.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 439d76ec472..2acfac8f50e 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -167,7 +167,9 @@ def test_reduction_preserves_negative_and_mixed_axes(axes): data = np.arange(6, dtype=np.float32).reshape(2, 3) expected = ReferenceEvaluator(model).run(None, {"x": data})[0] candidate = _apply_non4d_model_inputs(copy.deepcopy(model)) - candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) + candidate = _apply_non4d_model_outputs( + candidate, defer_rank_conversion=True, prepared_output_ops={"ReduceSum"} + ) candidate = _apply_standalone_reduce_operations(candidate) onnx.checker.check_model(candidate) actual = ReferenceEvaluator(candidate).run(None, {"x": data})[0] @@ -604,7 +606,9 @@ def test_reduction_registers_legacy_initializer_inputs(): model.ir_version = 3 feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} expected = _execute(model, feeds)[0] - transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) + transformed = _apply_non4d_model_outputs( + copy.deepcopy(model), defer_rank_conversion=True, prepared_output_ops={"ReduceMean"} + ) transformed = _apply_standalone_reduce_operations(transformed) np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) From 149c66c83b2cb7e34217ff449f11a440551620a5 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 09:00:54 +0000 Subject: [PATCH 14/31] Retain native adapter owners through DLA boundary promotion Signed-off-by: mgohil --- .../dla_transforms/dla_model_boundaries.py | 43 ++++++++----------- .../dla_transforms/test_dla_indexing.py | 38 ++++++++++++++++ 2 files changed, 57 insertions(+), 24 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 64a34411bb3..bcf27d5db2c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -329,31 +329,26 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal for consumer in consumers: if _is_standard_node(consumer, local_functions, {"Unsqueeze"}): output_shape = unsqueezed_shapes[id(consumer)] - if promoted_shape[-len(output_shape) :] == output_shape: - old_output = consumer.output[0] - for downstream, _ in cache.get_consumers(old_output): - for index, name in enumerate(downstream.input): - if name == old_output: - downstream.input[index] = output - replacements[id(consumer)] = [] - else: - target = add_initializer( - graph, - cache.unique_name((consumer.name or consumer.output[0]) + "_reshape_shape"), - np.asarray(output_shape, dtype=np.int64), - dtype=TensorProto.INT64, - with_value_info=False, + target_shape = ( + promoted_shape + if promoted_shape[-len(output_shape) :] == output_shape + else output_shape + ) + target = add_initializer( + graph, + cache.unique_name((consumer.name or consumer.output[0]) + "_reshape_shape"), + np.asarray(target_shape, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + replacements[id(consumer)] = [ + helper.make_node( + "Reshape", + [output, target], + list(consumer.output), + name=cache.unique_name((consumer.name or consumer.output[0]) + "_reshape"), ) - replacements[id(consumer)] = [ - helper.make_node( - "Reshape", - [output, target], - list(consumer.output), - name=cache.unique_name( - (consumer.name or consumer.output[0]) + "_reshape" - ), - ) - ] + ] else: for index, name in enumerate(consumer.input): if name == value.name: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 2acfac8f50e..a1faacf1070 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -592,6 +592,44 @@ def test_gather_preserves_overridable_scalar_and_vector_indices(indices, produce np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) +@pytest.mark.parametrize( + ("shape", "unsqueeze_axes"), [([2, 3], [0]), ([2, 3], [0, 1]), ([2, 3, 4], [0])] +) +@pytest.mark.parametrize("negative", [False, True]) +@pytest.mark.parametrize("keepdims", [0, 1]) +def test_reduction_preserves_rank_of_replaced_boundary_unsqueeze( + shape, unsqueeze_axes, negative, keepdims +): + """Keep the original adapter rank when boundary promotion replaces its data owner.""" + native = [1] * len(unsqueeze_axes) + shape + axis = len(unsqueeze_axes) + output = native.copy() + if keepdims: + output[axis] = 1 + else: + output.pop(axis) + model = _model( + [ + helper.make_node("Unsqueeze", ["x", "adapter_axes"], ["expanded"]), + helper.make_node( + "ReduceMean", ["expanded", "reduction_axes"], ["y"], keepdims=keepdims + ), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output)], + [ + numpy_helper.from_array(np.asarray(unsqueeze_axes, np.int64), "adapter_axes"), + numpy_helper.from_array( + np.asarray([axis - len(native) if negative else axis], np.int64), "reduction_axes" + ), + ], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + 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_reduction_registers_legacy_initializer_inputs(): """Keep IR3 controls declared while preserving reduction output values and ranks.""" model = helper.make_model( From b7cbc09d8c63ae64f8ccf79decd979bcd61b3b04 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 10:39:56 +0000 Subject: [PATCH 15/31] Declare prepared indexing stages in boundary integration tests Signed-off-by: mgohil --- .../onnx/graph_surgery/dla_transforms/test_dla_indexing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index a1faacf1070..381637bae36 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -60,7 +60,7 @@ def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): ) model.ir_version = 9 candidate = transform(copy.deepcopy(model)) - candidate = _apply_non4d_model_inputs(candidate) + candidate = _apply_non4d_model_inputs(candidate, prepared_input_ops={operation}) candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) options = ort.SessionOptions() options.intra_op_num_threads = 1 @@ -166,7 +166,7 @@ def test_reduction_preserves_negative_and_mixed_axes(axes): model = onnx.shape_inference.infer_shapes(model) data = np.arange(6, dtype=np.float32).reshape(2, 3) expected = ReferenceEvaluator(model).run(None, {"x": data})[0] - candidate = _apply_non4d_model_inputs(copy.deepcopy(model)) + candidate = _apply_non4d_model_inputs(copy.deepcopy(model), prepared_input_ops={"ReduceSum"}) candidate = _apply_non4d_model_outputs( candidate, defer_rank_conversion=True, prepared_output_ops={"ReduceSum"} ) From 9f974e32e0f3f93956218ad335142ebed09622c0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 11:48:44 +0000 Subject: [PATCH 16/31] Validate fixed scalar Gather output layouts before mutation Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 11 ++++- .../dla_transforms/test_dla_indexing.py | 46 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index ef65d2286b9..d67bd2fdb40 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -82,7 +82,7 @@ def _gather_flatten_count(cache, node, index_shape): if ( index_shape is None or shape is None - or len(shape) not in (2, 3) + or len(shape) not in (2, 3, 4) or output is None or len(output) > 4 ): @@ -161,6 +161,15 @@ def _apply_gather(model): index_views[original] = ( rank if _gather_index_needs_view(cache, original, layout_nodes) else None ) + if ( + rank == 0 + and fixed is not None + and original not in cache.graph_input_names + and _gather_flatten_count(cache, candidate, []) is None + ): + raise ValueError( + f"{candidate.name!r}: fixed scalar Gather indices change the padded output layout" + ) if fixed is None and cache.get_dtype(original) not in ( TensorProto.INT32, TensorProto.INT64, diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 381637bae36..748dc385c6f 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -712,6 +712,52 @@ def test_multidimensional_gather_indices_refuse_before_any_mutation(index_shape, assert model.SerializeToString() == before +@pytest.mark.parametrize(("shape", "axis"), [([3, 5], 1), ([3, 5, 4], 1), ([3, 5, 4], 2)]) +@pytest.mark.parametrize("negative", [False, True]) +def test_fixed_scalar_gather_validates_padded_layout_before_mutation(shape, axis, negative): + """Refuse a layout-changing scalar flatten before changing an earlier eligible Gather.""" + model = _model( + [ + helper.make_node("Gather", ["first_data", "first_indices"], ["first"], axis=0), + helper.make_node( + "Gather", ["x", "indices"], ["y"], axis=axis - len(shape) if negative else axis + ), + ], + [ + helper.make_tensor_value_info("first_data", TensorProto.FLOAT, [5, 4]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [1, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape[:axis] + shape[axis + 1 :]), + ], + [ + numpy_helper.from_array(np.array([1], np.int64), "first_indices"), + numpy_helper.from_array(np.array(1, np.int64), "indices"), + ], + ) + _execute(model, {"first_data": np.ones((5, 4), np.float32), "x": np.ones(shape, np.float32)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"scalar Gather.*layout"): + _apply_gather(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize(("shape", "axis"), [([3, 4], 0), ([3, 5, 4], 0), ([1, 3, 4, 4], 1)]) +def test_fixed_scalar_gather_keeps_compatible_padded_layouts(shape, axis): + """Retain supported scalar selections for promoted and native four-dimensional data.""" + model = _model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=axis)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape[:axis] + shape[axis + 1 :])], + [numpy_helper.from_array(np.array(1, np.int64), "indices")], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + 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", "shape", "refuse"), [ From d456bccd03877c0c09784b1327250fddeeb0a074 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 12:20:19 +0000 Subject: [PATCH 17/31] Retain fixed reshape repair paths for scalar Gather layouts Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 35 ++++++++++++++++++- .../dla_transforms/test_dla_indexing.py | 33 +++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index d67bd2fdb40..5b4e585bfe4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -29,7 +29,7 @@ pad4d, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs -from .dla_graph_cleanup import _subgraph_captures +from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures __all__ = [] @@ -114,6 +114,34 @@ def _gather_index_needs_view(cache, name, layout_nodes): return producer is not None and id(producer) in layout_nodes +def _gather_layout_reinterpreted(cache, name, supported, captured): + """Prove that fixed four-dimensional reshapes erase a private Gather layout difference.""" + pending, seen = [name], set() + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + if current in cache.graph_output_names | captured: + return False + for consumer, port in cache.get_consumers(current): + if id(consumer) not in supported or port != 0: + return False + if consumer.op_type in _RANK_AGNOSTIC_OPS and len(consumer.input) == 1: + pending.extend(consumer.output) + continue + if ( + consumer.op_type != "Reshape" + or consumer.input[1] in cache.graph_input_names + or cache.get_shape(consumer.input[1]) != [4] + ): + return False + target = cache.get_init_array(consumer.input[1]) + if target is None or np.any((target <= 0) & (target != -1)): + return False + return True + + def _apply_gather(model): """Prepare Gather indices for the subsequent ordered data-rank promotion.""" graph = model.graph @@ -127,6 +155,8 @@ def _apply_gather(model): flattened = set() index_views = {} captured = _subgraph_captures(graph) + shape_readers = _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"}) + shape_reader_ids = {id(node) for node in shape_readers} for candidate in selected: data_shape = cache.get_shape(candidate.input[0]) if data_shape is None: @@ -166,6 +196,9 @@ def _apply_gather(model): and fixed is not None and original not in cache.graph_input_names and _gather_flatten_count(cache, candidate, []) is None + and not _gather_layout_reinterpreted( + cache, candidate.output[0], shape_reader_ids, captured + ) ): raise ValueError( f"{candidate.name!r}: fixed scalar Gather indices change the padded output layout" diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 748dc385c6f..b496398fd81 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -758,6 +758,39 @@ def test_fixed_scalar_gather_keeps_compatible_padded_layouts(shape, axis): np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) +@pytest.mark.parametrize("identity", [False, True]) +@pytest.mark.parametrize("copied_dimension", [False, True]) +def test_fixed_scalar_gather_recognizes_layout_independent_reshape(identity, copied_dimension): + """Keep fixed reshape repair paths while refusing targets that copy native axis positions.""" + nodes = [helper.make_node("Gather", ["x", "indices"], ["selected"], axis=1)] + current = "selected" + if identity: + nodes.append(helper.make_node("Identity", [current], ["alias"])) + current = "alias" + target = [0, 1, 1, 4] if copied_dimension else [1, 1, 3, 4] + output = [3, 1, 1, 4] if copied_dimension else target + nodes.append(helper.make_node("Reshape", [current, "target"], ["y"])) + model = _model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3, 5, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output)], + [ + numpy_helper.from_array(np.array(1, np.int64), "indices"), + numpy_helper.from_array(np.array(target, np.int64), "target"), + ], + ) + feeds = {"x": np.arange(60, dtype=np.float32).reshape(3, 5, 4)} + expected = _execute(model, feeds)[0] + if copied_dimension: + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"scalar Gather.*layout"): + _apply_gather(model) + assert model.SerializeToString() == before + else: + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + @pytest.mark.parametrize( ("producer", "shape", "refuse"), [ From 72f7df5f07c462af63d7234fb75679ddec1707b7 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:03:49 +0000 Subject: [PATCH 18/31] Refuse unprepared internal indexing and reduction ranks Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 36 ++++++++++ .../dla_transforms/test_dla_indexing.py | 68 +++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 5b4e585bfe4..c1d9b3ea8ff 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -315,6 +315,31 @@ def _apply_gather(model): return model +def _native_reshape_data(cache, name, standard, checked): + """Prove that an internal Reshape path still has unprepared rank-two or rank-three data.""" + current, seen = name, set() + while current not in checked: + if current in seen: + raise ValueError("Cyclic data path while checking prepared indexing rank") + seen.add(current) + producer = cache.get_producer(current) + if producer is None or id(producer) not in standard: + native = False + break + if producer.op_type == "Reshape": + target = cache.get_shape(producer.input[1]) + native = target is not None and len(target) == 1 and target[0] in (2, 3) + break + if len(producer.input) != 1: + native = False + break + current = producer.input[0] + else: + native = checked[current] + checked.update(dict.fromkeys(seen, native)) + return native + + def _apply_gather_elements(model): """Prepare matching index ranks and preserve native axis and runtime-value contracts.""" graph = model.graph @@ -322,6 +347,13 @@ def _apply_gather_elements(model): converted_indices = {} replacements = {} selected_ids = {id(node) for node in _operator_nodes(model, {"GatherElements"})} + standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} + checked = {} + for node in graph.node: + if id(node) in selected_ids and _native_reshape_data( + cache, node.input[0], standard, checked + ): + raise ValueError(f"{node.name!r}: GatherElements data requires rank preparation") for index, node in enumerate(graph.node): if id(node) not in selected_ids: continue @@ -433,7 +465,11 @@ def _apply_standalone_reduce_operations(model): nodes = [] count = 0 plans = {} + standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} + checked = {} for node in _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}): + if _native_reshape_data(cache, node.input[0], standard, checked): + raise ValueError(f"{node.name!r}: reduction data requires rank preparation") shape = _shape_before_boundary_promotion(cache, node.input[0]) if not shape: keepdims = _attribute(node, "keepdims") diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index b496398fd81..202f763a687 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -187,6 +187,74 @@ def _model(nodes, inputs, outputs, initializers=(), extra_opsets=()): return model +@pytest.mark.parametrize( + "operation", ["GatherElements", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"] +) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("unary", [False, True]) +def test_internal_native_reshape_refuses_before_index_or_reduction_mutation( + operation, shape, unary +): + """Refuse unprepared internal data before changing an earlier eligible operation.""" + source_shape = [1] * (4 - len(shape)) + shape + axis = len(shape) - 1 if unary else -1 + keepdims = int(unary) + nodes = [helper.make_node("Reshape", ["x", "target"], ["reshaped"])] + if unary: + nodes.append(helper.make_node("Relu", ["reshaped"], ["native"])) + data = "native" if unary else "reshaped" + initializers = [numpy_helper.from_array(np.asarray(shape, np.int64), "target")] + outputs = [] + for name, input_name, input_shape, selected_axis in [ + ("ready", "x", source_shape, 3), + ("native_result", data, shape, axis), + ]: + control = name + "_control" + values = ( + np.zeros(input_shape, np.int64) + if operation == "GatherElements" + else np.array([selected_axis], np.int64) + ) + initializers.append(numpy_helper.from_array(values, control)) + nodes.append( + helper.make_node( + operation, + [input_name, control], + [name], + name=name, + **( + {"axis": selected_axis} + if operation == "GatherElements" + else {"keepdims": keepdims} + ), + ) + ) + output_shape = list(input_shape) + if operation != "GatherElements": + if keepdims: + output_shape[selected_axis] = 1 + else: + output_shape.pop(selected_axis) + outputs.append(helper.make_tensor_value_info(name, TensorProto.FLOAT, output_shape)) + model = _model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + outputs, + initializers, + ) + model = onnx.shape_inference.infer_shapes(model) + _execute(model, {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(source_shape)}) + before = model.SerializeToString() + callback = ( + _apply_gather_elements + if operation == "GatherElements" + else _apply_standalone_reduce_operations + ) + with pytest.raises(ValueError, match="rank"): + callback(model) + assert model.SerializeToString() == before + + def _execute(model, feeds): """Run CPU graph semantics without backend fusion changing the comparison.""" onnx.checker.check_model(model, full_check=True) From a5cfc946540cfb9519f6894c0796fd0429e8e0ae Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:05:43 +0000 Subject: [PATCH 19/31] Retain native symbolic paths during rank preflight Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 10 +++-- .../dla_transforms/test_dla_indexing.py | 39 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index c1d9b3ea8ff..362907080e5 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -350,8 +350,10 @@ def _apply_gather_elements(model): standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} checked = {} for node in graph.node: - if id(node) in selected_ids and _native_reshape_data( - cache, node.input[0], standard, checked + if ( + id(node) in selected_ids + and cache.get_shape(node.input[0]) is not None + and _native_reshape_data(cache, node.input[0], standard, checked) ): raise ValueError(f"{node.name!r}: GatherElements data requires rank preparation") for index, node in enumerate(graph.node): @@ -468,14 +470,14 @@ def _apply_standalone_reduce_operations(model): standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} checked = {} for node in _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}): - if _native_reshape_data(cache, node.input[0], standard, checked): - raise ValueError(f"{node.name!r}: reduction data requires rank preparation") shape = _shape_before_boundary_promotion(cache, node.input[0]) if not shape: keepdims = _attribute(node, "keepdims") if keepdims is not None and keepdims.i == 0: raise ValueError(f"Missing static reduction input shape for {node.name!r}") continue # Already rank-preserving, including generated 4D reduction branches. + if _native_reshape_data(cache, node.input[0], standard, checked): + raise ValueError(f"{node.name!r}: reduction data requires rank preparation") tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None if len(node.input) > 1 and (tensor is None or tensor.name in cache.graph_input_names): raise ValueError(f"Dynamic reduction axes are unsupported for {node.name!r}") diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 202f763a687..0ba9facf325 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -255,6 +255,45 @@ def test_internal_native_reshape_refuses_before_index_or_reduction_mutation( assert model.SerializeToString() == before +@pytest.mark.parametrize("operation", ["GatherElements", "ReduceSum"]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_symbolic_internal_reshape_retains_native_index_and_reduction_layout(operation, batch): + """Keep previously supported symbolic native paths outside concrete rank adaptation.""" + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, "batch", 3])] + control = "indices" if operation == "GatherElements" else "axes" + options = {"axis": -1} if operation == "GatherElements" else {"keepdims": 1} + nodes = [ + helper.make_node("Reshape", ["x", "target"], ["native"]), + helper.make_node(operation, ["native", control], ["y"], **options), + ] + initializers = [numpy_helper.from_array(np.array([-1, 3], np.int64), "target")] + if operation == "GatherElements": + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, ["batch", 3])) + else: + initializers.append(numpy_helper.from_array(np.array([-1], np.int64), "axes")) + output_shape = ["batch", 3 if operation == "GatherElements" else 1] + model = _model( + nodes, + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(batch * 3, dtype=np.float32).reshape(1, 1, batch, 3)} + if operation == "GatherElements": + feeds["indices"] = np.zeros((batch, 3), np.int64) + expected = _execute(model, feeds) + before = model.SerializeToString() + callback = ( + _apply_gather_elements + if operation == "GatherElements" + else _apply_standalone_reduce_operations + ) + converted = callback(model) + assert converted.SerializeToString() == before + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected[0]) + + def _execute(model, feeds): """Run CPU graph semantics without backend fusion changing the comparison.""" onnx.checker.check_model(model, full_check=True) From 02ec7d8690ed5edad5e92e12145ed1c4b4a38298 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 13:40:20 +0000 Subject: [PATCH 20/31] Keep native vector Gather control paths unchanged Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 6 ++- .../dla_transforms/test_dla_indexing.py | 46 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 362907080e5..f61eedfff0d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -148,7 +148,11 @@ def _apply_gather(model): cache = GraphCache(graph) nodes = [] squeezed_indices = {} - selected = _operator_nodes(model, {"Gather"}) + selected = [ + node + for node in _operator_nodes(model, {"Gather"}) + if (shape := cache.get_shape(node.input[0])) is None or len(shape) >= 2 + ] layout = _operator_nodes(model, {"Cast", "Identity", "Reshape", "Squeeze", "Unsqueeze"}) layout_nodes = {id(node) for node in layout} casts = {id(node) for node in layout if node.op_type == "Cast"} diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 0ba9facf325..551875309e5 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -294,6 +294,52 @@ def test_symbolic_internal_reshape_retains_native_index_and_reduction_layout(ope np.testing.assert_array_equal(_execute(converted, feeds)[0], expected[0]) +@pytest.mark.parametrize("source", ["input", "initializer", "shape"]) +@pytest.mark.parametrize("scalar", [False, True]) +@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"]) +def test_rank_one_gather_control_paths_remain_native(source, scalar, entrypoint): + """Preserve scalar/vector selection on native vectors, including Shape control paths.""" + dtype = TensorProto.INT64 if source == "shape" else TensorProto.FLOAT + nodes = [helper.make_node("Shape", ["x"], ["data"])] if source == "shape" else [] + inputs, feeds, initializers = [], {}, [] + if source == "shape": + inputs.append(helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])) + feeds["x"] = np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4) + elif source == "input": + inputs.append(helper.make_tensor_value_info("data", dtype, [4])) + feeds["data"] = np.arange(4, dtype=np.float32) + else: + initializers.append(numpy_helper.from_array(np.arange(4, dtype=np.float32), "data")) + initializers.extend( + [ + numpy_helper.from_array(np.array(1 if scalar else [1], np.int64), "indices"), + numpy_helper.from_array(np.array([0], np.int64), "axes"), + numpy_helper.from_array(np.array([1, 1, 1, -1], np.int64), "output_shape"), + ] + ) + nodes.extend( + [ + helper.make_node("Gather", ["data", "indices"], ["selected"], axis=0), + helper.make_node("Unsqueeze", ["selected", "axes"], ["vector"]), + helper.make_node("Reshape", ["vector", "output_shape"], ["y"]), + ] + ) + model = _model( + nodes, inputs, [helper.make_tensor_value_info("y", dtype, [1, 1, 1, 1])], initializers + ) + model = onnx.shape_inference.infer_shapes(model) + expected = _execute(model, feeds) + before = model.SerializeToString() + converted = ( + _apply_gather(model) + if entrypoint == "stage" + else _transform_make_dla_compatible(model, verbose=False) + ) + if entrypoint == "stage": + assert converted.SerializeToString() == before + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected[0]) + + def _execute(model, feeds): """Run CPU graph semantics without backend fusion changing the comparison.""" onnx.checker.check_model(model, full_check=True) From add251b641448a8cf559f83bcc7a369b897b6e96 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 14:48:44 +0000 Subject: [PATCH 21/31] Validate Gather data rank preparation before rewriting indices Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 153 ++++++++++++++++- .../dla_matmul_to_transpose_conv_transpose.py | 9 +- .../onnx/graph_surgery/make_dla_compatible.py | 2 + .../dla_transforms/test_dla_indexing.py | 161 +++++++++++++++++- 4 files changed, 322 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 f61eedfff0d..233f9467b62 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -30,6 +30,8 @@ ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures +from .dla_matmul_to_transpose_conv_transpose import _apply_matmul_to_transpose_conv_transpose +from .dla_model_boundaries import _has_native_symbolic_data, _native_conv1d __all__ = [] @@ -142,12 +144,153 @@ def _gather_layout_reinterpreted(cache, name, supported, captured): return True -def _apply_gather(model): +def _gather_promoted_data(model, cache, prepared_ops): + """Prove later leading-4D data layouts from 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): + standard.update( + id(node) + for node in _operator_nodes( + model, + {"QuantizeLinear", "DequantizeLinear", "Gelu", "GroupQueryAttention"}, + ("com.microsoft",), + ) + ) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + captured = _subgraph_captures(model.graph) + matrices = set() + if {"MatMulToTransposeConvTranspose", "RemoveIntermediarySqueezeAndUnsqueeze"} <= set( + prepared_ops + ): + _apply_matmul_to_transpose_conv_transpose( + model, defer_rank_conversion=True, validate_only=True, planned_outputs=matrices + ) + expand_ops = {"Add", "Sub", "Mul", "Div", "LSTM"} + + def classify(name): + """Return a proven layout or the rank-preserving parents needed to establish it.""" + shape = cache.get_shape(name) + if name in matrices or (shape is not None and len(shape) == 4): + return True + if name in cache.graph_input_names and shape is not None and len(shape) in (2, 3): + return True + tensor = cache.get_init(name) + if name in cache.graph_input_names and (model.ir_version > 3 or tensor is None): + return False + if tensor is not None: + users = [node for node, _ in cache.get_consumers(name)] + if ( + "Non4dInitializers" not in prepared_ops + or name in captured | cache.graph_output_names + or any( + id(node) not in standard + or node.op_type in {"Gemm", "Unsqueeze", "GroupQueryAttention"} + or _native_conv1d(node, cache) + for node in users + ) + ): + return False + operations = {node.op_type for node in users} + return ( + len(tensor.dims) == 3 + or ( + len(tensor.dims) == 2 + and bool(operations & (expand_ops | {"Gather", "Where", "Split"})) + ) + or (len(tensor.dims) == 1 and bool(operations & expand_ops)) + ) + node = cache.get_producer(name) + if node is None or id(node) not in standard: + return False + if node.op_type == "Reshape": + target = cache.get_shape(node.input[1]) + return target == [4] or ( + "Reshape" in prepared_ops + and target is not None + and len(target) == 1 + and target[0] in (1, 2, 3, 5) + and not _has_native_symbolic_data(cache, node.input[0], local_functions) + ) + if node.op_type in {"Flatten", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: + return node.op_type in prepared_ops + if node.op_type == "InstanceNormalization": + return "FixInstanceNormChannelMismatch" 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 node.op_type in {"Squeeze", "Unsqueeze"}: + if "RemoveIntermediarySqueezeAndUnsqueeze" not in prepared_ops: + return False + source = node.input[0] + if source in captured or name in captured - cache.graph_output_names: + return False + if source in cache.graph_input_names: + original = cache.get_shape(source) + if original is None or len(original) not in (2, 3): + return False + if node.op_type == "Squeeze": + return "Gather" not in node.name + return shape is not None and ( + pad4d(original)[-len(shape) :] == shape or "Reshape" in prepared_ops + ) + return source not in cache.init_names | cache.graph_output_names + if node.op_type in {"QuantizeLinear", "DequantizeLinear"}: + return ( + [node.input[0]] + if cache.get_shape(node.input[1]) == [] or node.op_type in prepared_ops + else False + ) + if node.op_type in {"Softmax", "LogSoftmax", "LayerNormalization"}: + axis = _attribute(node, "axis") + opset = next( + (op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1 + ) + value = axis.i if axis is not None else (-1 if opset >= 13 else 1) + return ( + [node.input[0]] + if value < 0 + or node.op_type in prepared_ops + or node.input[0] in cache.graph_input_names + else False + ) + if node.op_type in _RANK_AGNOSTIC_OPS | {"Gelu"}: + return [value for value in node.input if value] + if node.op_type == "Gather" or node.op_type in prepared_ops: + return [node.input[0]] + return False + + resolved, dependencies, active = {}, {}, set() + pending = [(node.input[0], False) for node in _operator_nodes(model, {"Gather"})] + while pending: + name, expanded = pending.pop() + if name in resolved: + continue + if expanded: + resolved[name] = any(resolved[parent] for parent in dependencies[name]) + active.remove(name) + continue + if name in active: + raise ValueError("Cyclic Gather data producers are unsupported") + result = classify(name) + if isinstance(result, bool): + resolved[name] = result + continue + dependencies[name] = result + active.add(name) + pending.append((name, True)) + pending.extend((parent, False) for parent in result) + return resolved + + +def _apply_gather(model, *, prepared_data_ops=()): """Prepare Gather indices for the subsequent ordered data-rank promotion.""" graph = model.graph cache = GraphCache(graph) nodes = [] squeezed_indices = {} + promoted_data = _gather_promoted_data(model, cache, prepared_data_ops) selected = [ node for node in _operator_nodes(model, {"Gather"}) @@ -159,12 +302,20 @@ def _apply_gather(model): flattened = set() index_views = {} captured = _subgraph_captures(graph) + for candidate in _operator_nodes(model, {"Gather"}): + data_shape = cache.get_shape(candidate.input[0]) + if data_shape is not None and len(data_shape) == 1 and promoted_data[candidate.input[0]]: + raise ValueError( + f"{candidate.name!r}: native vector Gather data has a later rank promotion" + ) shape_readers = _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"}) shape_reader_ids = {id(node) for node in shape_readers} for candidate in selected: data_shape = cache.get_shape(candidate.input[0]) if data_shape is None: raise ValueError(f"{candidate.name!r}: Gather requires a static data shape") + if len(data_shape) in (2, 3) and not promoted_data[candidate.input[0]]: + raise ValueError(f"{candidate.name!r}: Gather data requires rank preparation") original = candidate.input[1] index_shape = cache.get_shape(original) rank = _known_rank(cache, original) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index b75a2acd9fd..7b60d169f4f 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -59,7 +59,9 @@ from .dla_graph_cleanup import _subgraph_captures -def _apply_matmul_to_transpose_conv_transpose(model, *, defer_rank_conversion=False): +def _apply_matmul_to_transpose_conv_transpose( + model, *, defer_rank_conversion=False, validate_only=False, planned_outputs=None +): """Convert eligible matrices and optionally defer fallback rank promotion.""" graph = model.graph cache = GraphCache(graph) @@ -280,6 +282,11 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node ) if needs_adapters and opsets.get("", opsets.get("ai.onnx", 1)) < 13: raise ValueError("DLA matrix rank adapters require ONNX opset 13 or newer") + if planned_outputs is not None: + planned_outputs.update(node.output[0] for node, _, _, _ in eligible_nodes) + planned_outputs.update(node.output[0] for node in active_fallback) + if validate_only: + return model for node in active_fallback: cnt_fb = len(fallback_batch) diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 3676d0c9bcd..1a1c376be48 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -122,6 +122,8 @@ def _apply_constants_to_initializers(model, **options): def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prepared_output_ops): """Share stage options between the prerequisite probe and the actual ordered pipeline.""" + if name == "Gather": + return {"prepared_data_ops": prepared_input_ops} if name == "RemoveQDQ": return { "keep_clip_after_inputs": False, diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 551875309e5..22729e39a1e 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -34,7 +34,10 @@ _apply_non4d_model_inputs, _apply_non4d_model_outputs, ) -from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +from modelopt.onnx.graph_surgery.make_dla_compatible import ( + DLA_TRANSFORMS, + _transform_make_dla_compatible, +) @pytest.mark.parametrize("operation", ["Gather", "GatherElements"]) @@ -340,6 +343,162 @@ def test_rank_one_gather_control_paths_remain_native(source, scalar, entrypoint) np.testing.assert_array_equal(_execute(converted, feeds)[0], expected[0]) +@pytest.mark.parametrize( + ("producer", "shape"), + [(name, shape) for name in ("Reshape", "Squeeze", "opaque") for shape in ([2, 3], [2, 3, 4])] + + [("Gemm", [2, 3])], +) +@pytest.mark.parametrize("unary", [False, True]) +@pytest.mark.parametrize("positive_axis", [False, True]) +def test_gather_refuses_unprepared_internal_data_before_mutation( + producer, shape, unary, positive_axis +): + """Reject native internal data before changing an earlier eligible Gather or control.""" + input_shape = [1] * (4 - len(shape)) + shape + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape)] + initializers = [ + numpy_helper.from_array(np.array([0], np.int64), "indices"), + numpy_helper.from_array(np.array([1, 1, 1, -1], np.int64), "output_shape"), + ] + nodes = [helper.make_node("Gather", ["x", "indices"], ["ready"], axis=3)] + functions, extra_opsets = [], [] + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(input_shape)} + if producer == "opaque": + domain = "opaque.gather" + functions.append( + helper.make_function( + domain, + "NativeData", + ["a"], + ["b"], + [ + helper.make_node( + "Constant", + [], + ["target"], + value=numpy_helper.from_array(np.asarray(shape, np.int64)), + ), + helper.make_node("Reshape", ["a", "target"], ["b"]), + ], + [helper.make_opsetid("", 21)], + ) + ) + extra_opsets.append(helper.make_opsetid(domain, 1)) + nodes.append(helper.make_node("NativeData", ["x"], ["native"], domain=domain)) + elif producer == "Gemm": + inputs.append(helper.make_tensor_value_info("matrix", TensorProto.FLOAT, shape)) + feeds["matrix"] = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + initializers.append(numpy_helper.from_array(np.eye(shape[-1], dtype=np.float32), "weight")) + nodes.append(helper.make_node("Gemm", ["matrix", "weight"], ["native"])) + else: + values = shape if producer == "Reshape" else list(range(4 - len(shape))) + initializers.append(numpy_helper.from_array(np.asarray(values, np.int64), "native_control")) + nodes.append(helper.make_node(producer, ["x", "native_control"], ["native"])) + if unary: + nodes.append(helper.make_node("Relu", ["native"], ["activated"])) + axis = len(shape) - 1 if positive_axis else -1 + nodes.extend( + [ + helper.make_node( + "Gather", ["activated" if unary else "native", "indices"], ["selected"], axis=axis + ), + helper.make_node("Reshape", ["selected", "output_shape"], ["y"]), + ] + ) + ready_shape = input_shape.copy() + ready_shape[-1] = 1 + model = helper.make_model( + helper.make_graph( + nodes, + "native_gather_data", + inputs, + [ + helper.make_tensor_value_info("ready", TensorProto.FLOAT, ready_shape), + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [1, 1, 1, int(np.prod(shape[:-1]))] + ), + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 21), *extra_opsets], + functions=functions, + ir_version=9, + ) + model = onnx.shape_inference.infer_shapes(model) + _execute(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank"): + _apply_gather(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("producer", "shape"), + [ + (op, shape) + for op in ["Relu", "fixed", "MatMul", "Squeeze", "Reshape"] + for shape in [[5, 4], [2, 5, 4]] + ] + + [("Gemm", [5, 4])], +) +def test_gather_accepts_registered_future_data_rank_preparation(producer, shape): + """Execute the complete pipeline when the actual registered stages establish a 4D layout.""" + if producer == "Reshape" and not any(stage == "Non4DReshape" for stage, _ in DLA_TRANSFORMS): + pytest.skip("Reshape rank preparation is introduced by a later slice") + if producer == "fixed" and not any(stage == "Non4dInitializers" for stage, _ in DLA_TRANSFORMS): + pytest.skip("Initializer rank preparation is introduced by a later slice") + values = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + source_shape = [1] * (4 - len(shape)) + shape if producer in {"Squeeze", "Reshape"} else shape + inputs = ( + [] + if producer == "fixed" + else [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)] + ) + initializers = [numpy_helper.from_array(np.asarray([1, 0], np.int64), "indices")] + nodes = [] + feeds = {} if producer == "fixed" else {"x": values.reshape(source_shape)} + if producer == "fixed": + initializers.append(numpy_helper.from_array(values, "data")) + elif producer in {"MatMul", "Gemm"}: + initializers.append(numpy_helper.from_array(np.eye(4, dtype=np.float32), "weight")) + nodes.append(helper.make_node(producer, ["x", "weight"], ["data"])) + elif producer in {"Squeeze", "Reshape"}: + control = ( + np.arange(4 - len(shape), dtype=np.int64) + if producer == "Squeeze" + else np.asarray(shape, np.int64) + ) + initializers.append(numpy_helper.from_array(control, "control")) + nodes.append(helper.make_node("Relu", ["x"], ["internal"])) + nodes.append(helper.make_node(producer, ["internal", "control"], ["data"])) + else: + nodes.append(helper.make_node(producer, ["x"], ["data"])) + nodes.append(helper.make_node("Gather", ["data", "indices"], ["selected"], axis=len(shape) - 2)) + output_shape = [*shape[:-2], 2, 4] + initializers.append( + numpy_helper.from_array( + np.asarray([1] * (4 - len(shape)) + output_shape, np.int64), "target" + ) + ) + nodes.append(helper.make_node("Reshape", ["selected", "target"], ["y"])) + model = _model( + nodes, + inputs, + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [1] * (4 - len(shape)) + output_shape + ) + ], + initializers, + ) + model = onnx.shape_inference.infer_shapes(model) + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + def _execute(model, feeds): """Run CPU graph semantics without backend fusion changing the comparison.""" onnx.checker.check_model(model, full_check=True) From 42e6cd41adeccfaffa2bf8ded3550c3a6da68ad5 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 15:50:40 +0000 Subject: [PATCH 22/31] Require positive indexing rank proofs and preserve native matrix owners Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 97 +++++---- .../dla_matmul_to_transpose_conv_transpose.py | 11 +- .../onnx/graph_surgery/make_dla_compatible.py | 2 +- .../dla_transforms/test_dla_indexing.py | 187 +++++++++++++++++- 4 files changed, 246 insertions(+), 51 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 233f9467b62..897c9afce64 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -144,8 +144,8 @@ def _gather_layout_reinterpreted(cache, name, supported, captured): return True -def _gather_promoted_data(model, cache, prepared_ops): - """Prove later leading-4D data layouts from registered rank stages without mutation.""" +def _prepared_index_data(model, cache, prepared_ops, data_names, *, before_boundary=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): standard.update( @@ -172,7 +172,12 @@ def classify(name): shape = cache.get_shape(name) if name in matrices or (shape is not None and len(shape) == 4): return True - if name in cache.graph_input_names and shape is not None and len(shape) in (2, 3): + if ( + before_boundary + and name in cache.graph_input_names + and shape is not None + and len(shape) in (2, 3) + ): return True tensor = cache.get_init(name) if name in cache.graph_input_names and (model.ir_version > 3 or tensor is None): @@ -211,8 +216,10 @@ 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 in {"Flatten", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: + if node.op_type == "Flatten": return node.op_type in prepared_ops + 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 if node.op_type == "GroupQueryAttention": @@ -220,11 +227,24 @@ def classify(name): if node.op_type == "Conv": kernel = cache.get_shape(node.input[1]) return kernel is not None and len(kernel) == 4 + if node.op_type == "Transpose": + perm = _attribute(node, "perm") + if perm is not None and sorted(perm.ints) == [0, 1, 2, 3]: + return True if node.op_type in {"Squeeze", "Unsqueeze"}: + if node.op_type == "Unsqueeze": + source_rank = _known_rank(cache, node.input[0]) + axes = cache.get_init(node.input[1]) if len(node.input) > 1 else None + if ( + axes is not None + and source_rank is not None + and source_rank + math.prod(axes.dims) == 4 + ): + return True 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 name in captured | cache.graph_output_names: return False if source in cache.graph_input_names: original = cache.get_shape(source) @@ -239,7 +259,9 @@ def classify(name): if node.op_type in {"QuantizeLinear", "DequantizeLinear"}: return ( [node.input[0]] - if cache.get_shape(node.input[1]) == [] or node.op_type in prepared_ops + if not before_boundary + or cache.get_shape(node.input[1]) == [] + or node.op_type in prepared_ops else False ) if node.op_type in {"Softmax", "LogSoftmax", "LayerNormalization"}: @@ -250,7 +272,8 @@ def classify(name): value = axis.i if axis is not None else (-1 if opset >= 13 else 1) return ( [node.input[0]] - if value < 0 + if not before_boundary + or value < 0 or node.op_type in prepared_ops or node.input[0] in cache.graph_input_names else False @@ -262,7 +285,7 @@ def classify(name): return False resolved, dependencies, active = {}, {}, set() - pending = [(node.input[0], False) for node in _operator_nodes(model, {"Gather"})] + pending = [(name, False) for name in data_names] while pending: name, expanded = pending.pop() if name in resolved: @@ -290,7 +313,13 @@ def _apply_gather(model, *, prepared_data_ops=()): cache = GraphCache(graph) nodes = [] squeezed_indices = {} - promoted_data = _gather_promoted_data(model, cache, prepared_data_ops) + promoted_data = _prepared_index_data( + model, + cache, + prepared_data_ops, + [node.input[0] for node in _operator_nodes(model, {"Gather"})], + before_boundary=True, + ) selected = [ node for node in _operator_nodes(model, {"Gather"}) @@ -470,45 +499,25 @@ def _apply_gather(model, *, prepared_data_ops=()): return model -def _native_reshape_data(cache, name, standard, checked): - """Prove that an internal Reshape path still has unprepared rank-two or rank-three data.""" - current, seen = name, set() - while current not in checked: - if current in seen: - raise ValueError("Cyclic data path while checking prepared indexing rank") - seen.add(current) - producer = cache.get_producer(current) - if producer is None or id(producer) not in standard: - native = False - break - if producer.op_type == "Reshape": - target = cache.get_shape(producer.input[1]) - native = target is not None and len(target) == 1 and target[0] in (2, 3) - break - if len(producer.input) != 1: - native = False - break - current = producer.input[0] - else: - native = checked[current] - checked.update(dict.fromkeys(seen, native)) - return native - - -def _apply_gather_elements(model): +def _apply_gather_elements(model, *, prepared_data_ops=()): """Prepare matching index ranks and preserve native axis and runtime-value contracts.""" graph = model.graph cache = GraphCache(graph) converted_indices = {} replacements = {} selected_ids = {id(node) for node in _operator_nodes(model, {"GatherElements"})} - standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} - checked = {} + prepared = _prepared_index_data( + model, + cache, + prepared_data_ops, + [node.input[0] for node in graph.node if id(node) in selected_ids], + ) for node in graph.node: if ( id(node) in selected_ids and cache.get_shape(node.input[0]) is not None - and _native_reshape_data(cache, node.input[0], standard, checked) + and _known_rank(cache, node.input[0]) in (2, 3) + and not prepared[node.input[0]] ): raise ValueError(f"{node.name!r}: GatherElements data requires rank preparation") for index, node in enumerate(graph.node): @@ -615,23 +624,25 @@ def _apply_gather_elements(model): return model -def _apply_standalone_reduce_operations(model): +def _apply_standalone_reduce_operations(model, *, prepared_data_ops=()): """Plan native reduction axes and preserve their output layouts during rank promotion.""" graph = model.graph cache = GraphCache(graph) nodes = [] count = 0 plans = {} - standard = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS | {"Reshape"})} - checked = {} - for node in _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}): + selected = _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) + prepared = _prepared_index_data( + model, cache, prepared_data_ops, [node.input[0] for node in selected] + ) + for node in selected: shape = _shape_before_boundary_promotion(cache, node.input[0]) if not shape: keepdims = _attribute(node, "keepdims") if keepdims is not None and keepdims.i == 0: raise ValueError(f"Missing static reduction input shape for {node.name!r}") continue # Already rank-preserving, including generated 4D reduction branches. - if _native_reshape_data(cache, node.input[0], standard, checked): + if len(shape) in (2, 3) and not prepared[node.input[0]]: raise ValueError(f"{node.name!r}: reduction data requires rank preparation") tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None if len(node.input) > 1 and (tensor is None or tensor.name in cache.graph_input_names): diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index 7b60d169f4f..d6d5f32cb9e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -283,8 +283,15 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node if needs_adapters and opsets.get("", opsets.get("ai.onnx", 1)) < 13: raise ValueError("DLA matrix rank adapters require ONNX opset 13 or newer") if planned_outputs is not None: - planned_outputs.update(node.output[0] for node, _, _, _ in eligible_nodes) - planned_outputs.update(node.output[0] for node in active_fallback) + native_outputs = cache.graph_output_names | captured + planned_outputs.update( + node.output[0] + for node, _, _, _ in eligible_nodes + if node.output[0] not in native_outputs + ) + planned_outputs.update( + node.output[0] for node in active_fallback if node.output[0] not in native_outputs + ) if validate_only: return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 1a1c376be48..cc48631684d 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -122,7 +122,7 @@ def _apply_constants_to_initializers(model, **options): def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prepared_output_ops): """Share stage options between the prerequisite probe and the actual ordered pipeline.""" - if name == "Gather": + if name in {"Gather", "GatherElements", "StandaloneReduceOperations"}: return {"prepared_data_ops": prepared_input_ops} if name == "RemoveQDQ": return { diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 22729e39a1e..31d3f843ed1 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -62,8 +62,12 @@ def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): opset_imports=[helper.make_opsetid("", 17)], ) model.ir_version = 9 - candidate = transform(copy.deepcopy(model)) + candidate = copy.deepcopy(model) + if operation == "Gather": + candidate = transform(candidate) candidate = _apply_non4d_model_inputs(candidate, prepared_input_ops={operation}) + if operation == "GatherElements": + candidate = transform(candidate) candidate = _apply_non4d_model_outputs(candidate, defer_rank_conversion=True) options = ort.SessionOptions() options.intra_op_num_threads = 1 @@ -193,20 +197,62 @@ def _model(nodes, inputs, outputs, initializers=(), extra_opsets=()): @pytest.mark.parametrize( "operation", ["GatherElements", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"] ) -@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize( + ("producer", "shape"), + [ + (producer, shape) + for producer in ["Reshape", "Squeeze", "opaque"] + for shape in [[2, 3], [2, 3, 4]] + ] + + [("Gemm", [2, 3])], +) @pytest.mark.parametrize("unary", [False, True]) def test_internal_native_reshape_refuses_before_index_or_reduction_mutation( - operation, shape, unary + operation, producer, shape, unary ): """Refuse unprepared internal data before changing an earlier eligible operation.""" source_shape = [1] * (4 - len(shape)) + shape axis = len(shape) - 1 if unary else -1 keepdims = int(unary) - nodes = [helper.make_node("Reshape", ["x", "target"], ["reshaped"])] + nodes, extra_opsets, functions = [], [], [] + if producer == "opaque": + domain = "opaque.indexing" + extra_opsets.append(helper.make_opsetid(domain, 1)) + functions.append( + helper.make_function( + domain, + "NativeData", + ["a", "target"], + ["b"], + [helper.make_node("Reshape", ["a", "target"], ["b"])], + [helper.make_opsetid("", 21)], + ) + ) + nodes.append(helper.make_node("NativeData", ["x", "target"], ["reshaped"], domain=domain)) + elif producer == "Squeeze": + nodes.append(helper.make_node("Squeeze", ["x", "target"], ["reshaped"])) + elif producer == "Gemm": + nodes.extend( + [ + helper.make_node("Reshape", ["x", "target"], ["matrix"]), + helper.make_node("Gemm", ["matrix", "weight"], ["reshaped"], alpha=2.0), + ] + ) + else: + nodes.append(helper.make_node("Reshape", ["x", "target"], ["reshaped"])) if unary: nodes.append(helper.make_node("Relu", ["reshaped"], ["native"])) data = "native" if unary else "reshaped" - initializers = [numpy_helper.from_array(np.asarray(shape, np.int64), "target")] + initializers = [ + numpy_helper.from_array( + np.arange(4 - len(shape), dtype=np.int64) + if producer == "Squeeze" + else np.asarray(shape, np.int64), + "target", + ) + ] + if producer == "Gemm": + initializers.append(numpy_helper.from_array(np.eye(shape[-1], dtype=np.float32), "weight")) outputs = [] for name, input_name, input_shape, selected_axis in [ ("ready", "x", source_shape, 3), @@ -244,7 +290,9 @@ def test_internal_native_reshape_refuses_before_index_or_reduction_mutation( [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], outputs, initializers, + extra_opsets, ) + model.functions.extend(functions) model = onnx.shape_inference.infer_shapes(model) _execute(model, {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(source_shape)}) before = model.SerializeToString() @@ -499,6 +547,135 @@ def test_gather_accepts_registered_future_data_rank_preparation(producer, shape) np.testing.assert_array_equal(actual, reference) +@pytest.mark.parametrize("operation", ["GatherElements", "ReduceSum"]) +@pytest.mark.parametrize( + ("producer", "shape"), + [ + (producer, shape) + for producer in ["Relu", "Squeeze", "MatMul", "Reshape"] + for shape in [[5, 4], [2, 5, 4]] + ] + + [("Gemm", [5, 4])], +) +def test_internal_index_data_uses_registered_rank_preparation(operation, producer, shape): + """Execute prepared data paths through the real stage order, including later adapter cleanup.""" + if producer == "Reshape" and not any(name == "Non4DReshape" for name, _ in DLA_TRANSFORMS): + pytest.skip("Reshape rank preparation is introduced by a later slice") + source_shape = [1] * (4 - len(shape)) + shape if producer in {"Squeeze", "Reshape"} else shape + nodes, initializers = [], [] + if producer in {"Squeeze", "Reshape"}: + nodes.append(helper.make_node("Relu", ["x"], ["private"])) + control = ( + np.arange(4 - len(shape), dtype=np.int64) + if producer == "Squeeze" + else np.asarray(shape, np.int64) + ) + initializers.append(numpy_helper.from_array(control, "control")) + nodes.append(helper.make_node(producer, ["private", "control"], ["data"])) + elif producer in {"MatMul", "Gemm"}: + initializers.append(numpy_helper.from_array(np.eye(4, dtype=np.float32), "weight")) + nodes.append(helper.make_node(producer, ["x", "weight"], ["data"])) + else: + nodes.append(helper.make_node("Relu", ["x"], ["data"])) + axis = len(shape) - 2 + output_shape = shape.copy() + if operation == "GatherElements": + initializers.append( + numpy_helper.from_array( + np.arange(np.prod(shape), dtype=np.int64).reshape(shape) % shape[axis], "indices" + ) + ) + nodes.append(helper.make_node(operation, ["data", "indices"], ["selected"], axis=axis)) + else: + initializers.append(numpy_helper.from_array(np.asarray([axis], np.int64), "axes")) + nodes.append(helper.make_node(operation, ["data", "axes"], ["selected"], keepdims=0)) + output_shape.pop(axis) + output_shape = [1] * (4 - len(output_shape)) + output_shape + initializers.append(numpy_helper.from_array(np.asarray(output_shape, np.int64), "target")) + nodes.append(helper.make_node("Reshape", ["selected", "target"], ["y"])) + model = _model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(source_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", "shape"), [("MatMul", [2, 3]), ("MatMul", [2, 2, 3]), ("Gemm", [2, 3])] +) +@pytest.mark.parametrize("owner", ["output", "capture"]) +@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"]) +def test_gather_refuses_native_matrix_outputs_with_other_rank_owners( + operation, shape, owner, entrypoint +): + """Do not promise promotion of matrix results whose native rank has a public or nested owner.""" + native = [*shape[:-1], 4] + selected = [*shape[:-2], 1, 4] + target = [1] * (4 - len(selected)) + selected + nodes = [ + helper.make_node(operation, ["x", "weight"], ["projection"]), + helper.make_node("Gather", ["projection", "indices"], ["selected"], axis=len(shape) - 2), + helper.make_node("Reshape", ["selected", "target"], ["y"]), + ] + initializers = [ + numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "weight"), + numpy_helper.from_array(np.array([1], np.int64), "indices"), + numpy_helper.from_array(np.asarray(target, np.int64), "target"), + ] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, target)] + if owner == "output": + outputs.append(helper.make_tensor_value_info("projection", TensorProto.FLOAT, native)) + else: + branch = helper.make_graph( + [ + helper.make_node("Shape", ["projection"], ["native_shape"]), + helper.make_node( + "Constant", + [], + ["view_shape"], + value=numpy_helper.from_array(np.array([1, 1, 1, len(native)], np.int64)), + ), + helper.make_node("Reshape", ["native_shape", "view_shape"], ["view"]), + ], + "shape_owner", + [], + [helper.make_tensor_value_info("view", TensorProto.INT64, [1, 1, 1, len(native)])], + ) + initializers.append(numpy_helper.from_array(np.asarray(True), "condition")) + nodes.append( + helper.make_node( + "If", ["condition"], ["shape_view"], then_branch=branch, else_branch=branch + ) + ) + outputs.append( + helper.make_tensor_value_info("shape_view", TensorProto.INT64, [1, 1, 1, len(native)]) + ) + model = _model( + nodes, [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], outputs, initializers + ) + model = onnx.shape_inference.infer_shapes(model) + _execute(model, {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank preparation"): + if entrypoint == "stage": + _apply_gather( + model, + prepared_data_ops={ + "MatMulToTransposeConvTranspose", + "RemoveIntermediarySqueezeAndUnsqueeze", + }, + ) + else: + _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert model.SerializeToString() == before + + def _execute(model, feeds): """Run CPU graph semantics without backend fusion changing the comparison.""" onnx.checker.check_model(model, full_check=True) From 27545ec46083959e8168566711ff9c8e146e955b Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 16:07:04 +0000 Subject: [PATCH 23/31] Prove widened batched MatMul outputs from matrix operands Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 38 +++++++++++++++---- 1 file changed, 31 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 897c9afce64..2f9f0da06d0 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -167,6 +167,29 @@ def _prepared_index_data(model, cache, prepared_ops, data_names, *, before_bound ) expand_ops = {"Add", "Sub", "Mul", "Div", "LSTM"} + def unsqueezed_rank(node): + """Read adapter rank from metadata without decoding its fixed axes payload.""" + source_rank = _known_rank(cache, node.input[0]) + axes = cache.get_init(node.input[1]) if len(node.input) > 1 else None + return ( + source_rank + math.prod(axes.dims) + if axes is not None and source_rank is not None + else None + ) + + def matrix_rank(name): + """Prove that an operand is a matrix even when a generated adapter has no value info.""" + rank = _known_rank(cache, name) + producer = cache.get_producer(name) + if ( + rank is None + and producer is not None + and id(producer) in standard + and producer.op_type == "Unsqueeze" + ): + rank = unsqueezed_rank(producer) + return rank + def classify(name): """Return a proven layout or the rank-preserving parents needed to establish it.""" shape = cache.get_shape(name) @@ -227,19 +250,20 @@ def classify(name): 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": + ranks = [matrix_rank(value) for value in node.input[:2]] + return ( + list(node.input[:2]) + if all(rank is not None and rank >= 2 for rank in ranks) + else False + ) if node.op_type == "Transpose": perm = _attribute(node, "perm") if perm is not None and sorted(perm.ints) == [0, 1, 2, 3]: return True if node.op_type in {"Squeeze", "Unsqueeze"}: if node.op_type == "Unsqueeze": - source_rank = _known_rank(cache, node.input[0]) - axes = cache.get_init(node.input[1]) if len(node.input) > 1 else None - if ( - axes is not None - and source_rank is not None - and source_rank + math.prod(axes.dims) == 4 - ): + if unsqueezed_rank(node) == 4: return True if "RemoveIntermediarySqueezeAndUnsqueeze" not in prepared_ops: return False From 75feaf5d367d8fff634f847985f7d5f13ff33027 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 16:08:49 +0000 Subject: [PATCH 24/31] Retain native owners in matrix rank proofs Signed-off-by: mgohil --- modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 2f9f0da06d0..47268f4b103 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -251,6 +251,8 @@ def classify(name): 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: + return False ranks = [matrix_rank(value) for value in node.input[:2]] return ( list(node.input[:2]) From 5c5557e77f1251a22a2ca09d18d5ed62ddc3745d Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:51:38 +0000 Subject: [PATCH 25/31] Preserve native rank-three attention inputs Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 17 ++++++- .../onnx/graph_surgery/make_dla_compatible.py | 2 +- .../dla_transforms/test_dla_indexing.py | 50 +++++++++++++++++++ 3 files changed, 67 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 47268f4b103..6fe414d0dc4 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -737,7 +737,7 @@ def _apply_standalone_reduce_operations(model, *, prepared_data_ops=()): return model -def _apply_group_query_attention(model): +def _apply_group_query_attention(model, *, prepared_data_ops=()): """Restore native attention ranks across internal tensors while retaining public boundaries.""" graph = model.graph cache = GraphCache(graph) @@ -761,12 +761,27 @@ def axes_for(name): selected_ids = { id(node) for node in _operator_nodes(model, {"GroupQueryAttention"}, ("com.microsoft",)) } + prepared = _prepared_index_data( + model, + cache, + set(prepared_data_ops) | {"GroupQueryAttention"}, + [ + name + for node in graph.node + if id(node) in selected_ids + for name in node.input[:3] + if name and name not in boundaries + ], + ) for node in graph.node: if id(node) not in selected_ids: nodes.append(node) continue for index, name in enumerate(node.input[:3]): if name and name not in boundaries: + shape = cache.get_shape(name) + if not prepared[name] or (shape is not None and len(shape) == 4 and shape[0] != 1): + continue if name in squeezed_inputs: node.input[index] = squeezed_inputs[name] continue diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index cc48631684d..d5f8d75c9ee 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -122,7 +122,7 @@ def _apply_constants_to_initializers(model, **options): def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prepared_output_ops): """Share stage options between the prerequisite probe and the actual ordered pipeline.""" - if name in {"Gather", "GatherElements", "StandaloneReduceOperations"}: + if name in {"Gather", "GatherElements", "StandaloneReduceOperations", "GroupQueryAttention"}: return {"prepared_data_ops": prepared_input_ops} if name == "RemoveQDQ": return { diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 31d3f843ed1..2190b2e98d0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -1576,3 +1576,53 @@ def test_cyclic_index_layout_ancestry_refuses_before_mutation(): with pytest.raises(ValueError, match="Cyclic Gather"): _apply_gather(model) assert model.SerializeToString() == before + + +@pytest.mark.parametrize("native_slot", [0, 1, 2]) +def test_attention_keeps_private_native_rank_three_inputs(native_slot): + """Execute private native query, key and value views without another Squeeze.""" + shape = [1, 2, 8] + names = ["q", "k", "v"] + source = names[native_slot] + attention_inputs = names.copy() + attention_inputs[native_slot] = "native" + model = _model( + [ + helper.make_node("Squeeze", [source, "axes"], ["native"]), + helper.make_node( + "GroupQueryAttention", + [*attention_inputs, "past_key", "past_value", "lengths", "total"], + ["y", "present_key", "present_value"], + domain="com.microsoft", + num_heads=1, + kv_num_heads=1, + ), + ], + [ + helper.make_tensor_value_info( + name, TensorProto.FLOAT, [1, *shape] if name == source else shape + ) + for name in names + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.asarray([0], np.int64), "axes"), + numpy_helper.from_array(np.asarray([1], np.int32), "lengths"), + numpy_helper.from_array(np.int32(2), "total"), + numpy_helper.from_array(np.zeros((1, 1, 2, 8), np.float32), "past_key"), + numpy_helper.from_array(np.zeros((1, 1, 2, 8), np.float32), "past_value"), + ], + [helper.make_opsetid("com.microsoft", 1)], + ) + feeds = { + name: (np.arange(16, dtype=np.float32) / 16).reshape( + [1, *shape] if name == source else shape + ) + for name in names + } + expected = _execute(model, feeds) + result = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(result, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + attention = next(n for n in result.graph.node if n.op_type == "GroupQueryAttention") + assert attention.input[native_slot] == "native" From da1a0c1454921a6ffc3c76cf44ac3e9f2512dbcb Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:54:02 +0000 Subject: [PATCH 26/31] Skip attention layout planning when no private inputs need it Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 6fe414d0dc4..62b6ff40945 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -761,17 +761,19 @@ def axes_for(name): selected_ids = { id(node) for node in _operator_nodes(model, {"GroupQueryAttention"}, ("com.microsoft",)) } - prepared = _prepared_index_data( - model, - cache, - set(prepared_data_ops) | {"GroupQueryAttention"}, - [ - name - for node in graph.node - if id(node) in selected_ids - for name in node.input[:3] - if name and name not in boundaries - ], + data_names = [ + name + for node in graph.node + if id(node) in selected_ids + for name in node.input[:3] + if name and name not in boundaries + ] + prepared = ( + _prepared_index_data( + model, cache, set(prepared_data_ops) | {"GroupQueryAttention"}, data_names + ) + if data_names + else {} ) for node in graph.node: if id(node) not in selected_ids: From 0638ba2dc6f43af74f680b066ddf9cf086e7ac72 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:58:21 +0000 Subject: [PATCH 27/31] Validate internal GatherElements index preparation before mutation Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 23 +++++++++++- .../onnx/graph_surgery/make_dla_compatible.py | 2 + .../dla_transforms/test_dla_indexing.py | 37 +++++++++++++++++++ 3 files changed, 61 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 62b6ff40945..87e4b9fff2c 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -525,7 +525,7 @@ def _apply_gather(model, *, prepared_data_ops=()): return model -def _apply_gather_elements(model, *, prepared_data_ops=()): +def _apply_gather_elements(model, *, prepared_data_ops=(), validate_only=False): """Prepare matching index ranks and preserve native axis and runtime-value contracts.""" graph = model.graph cache = GraphCache(graph) @@ -537,7 +537,26 @@ def _apply_gather_elements(model, *, prepared_data_ops=()): cache, prepared_data_ops, [node.input[0] for node in graph.node if id(node) in selected_ids], + before_boundary=validate_only, ) + internal_indices = [ + node.input[1] + for node in graph.node + if id(node) in selected_ids + and prepared[node.input[0]] + and node.input[1] not in cache.graph_input_names | cache.init_names + and _known_rank(cache, node.input[1]) in (2, 3) + ] + index_layouts = ( + _prepared_index_data( + model, cache, prepared_data_ops, internal_indices, before_boundary=validate_only + ) + if internal_indices + else {} + ) + for name in internal_indices: + if not index_layouts[name]: + raise ValueError(f"{name!r}: GatherElements indices require rank preparation") for node in graph.node: if ( id(node) in selected_ids @@ -546,6 +565,8 @@ def _apply_gather_elements(model, *, prepared_data_ops=()): and not prepared[node.input[0]] ): raise ValueError(f"{node.name!r}: GatherElements data requires rank preparation") + if validate_only: + return model for index, node in enumerate(graph.node): if id(node) not in selected_ids: continue diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index d5f8d75c9ee..d3b55fa06aa 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -147,6 +147,8 @@ def _validate_boundary_contract( model, prepared_input_ops, prepared_output_ops, *, require_inferred_shapes=False ): """Refuse unsupported native layouts before adapting any public boundary.""" + if "GatherElements" in prepared_input_ops: + _apply_gather_elements(model, prepared_data_ops=prepared_input_ops, validate_only=True) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} cache = GraphCache(model.graph) for node in model.graph.node: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 2190b2e98d0..637eeff7946 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -40,6 +40,43 @@ ) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("pipeline", [False, True]) +def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline): + """A native internal index view cannot follow data into a prepared four-dimensional layout.""" + source_shape = [1] * (4 - len(shape)) + shape + model = _model( + [ + helper.make_node("Reshape", ["source", "target"], ["indices"]), + helper.make_node("GatherElements", ["x", "indices"], ["y"], axis=-1), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("source", TensorProto.INT64, source_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(np.asarray(shape, np.int64), "target")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = { + "x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape), + "source": np.zeros(source_shape, np.int64), + } + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal( + _execute(model, feeds)[0], np.take(feeds["x"], [0] * shape[-1], axis=-1) + ) + if not pipeline: + _apply_non4d_model_inputs(model, prepared_input_ops={"GatherElements"}) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"indices.*rank preparation"): + if pipeline: + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_gather_elements(model) + assert model.SerializeToString() == before + + @pytest.mark.parametrize("operation", ["Gather", "GatherElements"]) def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): """Index dtype and rank conversion preserves selected values.""" From d2d44d67e2ccb14c0276d11d121e56a7fda7f257 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 08:01:04 +0000 Subject: [PATCH 28/31] Cover native Squeeze indices and scheduled reshape support Signed-off-by: mgohil --- .../dla_transforms/test_dla_indexing.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 637eeff7946..5c7f8020ef0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -42,12 +42,13 @@ @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) @pytest.mark.parametrize("pipeline", [False, True]) -def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline): +@pytest.mark.parametrize("producer", ["Reshape", "Squeeze"]) +def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline, producer): """A native internal index view cannot follow data into a prepared four-dimensional layout.""" source_shape = [1] * (4 - len(shape)) + shape model = _model( [ - helper.make_node("Reshape", ["source", "target"], ["indices"]), + helper.make_node(producer, ["source", "target"], ["indices"]), helper.make_node("GatherElements", ["x", "indices"], ["y"], axis=-1), ], [ @@ -55,7 +56,14 @@ def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline): helper.make_tensor_value_info("source", TensorProto.INT64, source_shape), ], [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], - [numpy_helper.from_array(np.asarray(shape, np.int64), "target")], + [ + numpy_helper.from_array( + np.asarray( + shape if producer == "Reshape" else list(range(4 - len(shape))), np.int64 + ), + "target", + ) + ], ) model = onnx.shape_inference.infer_shapes(model) feeds = { @@ -66,6 +74,16 @@ def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline): np.testing.assert_array_equal( _execute(model, feeds)[0], np.take(feeds["x"], [0] * shape[-1], axis=-1) ) + if ( + pipeline + and producer == "Reshape" + and any(name == "Non4DReshape" for name, _ in DLA_TRANSFORMS) + ): + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal( + _execute(transformed, feeds)[0], np.take(feeds["x"], [0] * shape[-1], axis=-1) + ) + return if not pipeline: _apply_non4d_model_inputs(model, prepared_input_ops={"GatherElements"}) before = model.SerializeToString() From ae368b58d2e4ac062e831f867900902a96b578fb Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 08:08:23 +0000 Subject: [PATCH 29/31] Validate internal index layouts without intermediate metadata Signed-off-by: mgohil --- modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py | 2 +- .../onnx/graph_surgery/dla_transforms/test_dla_indexing.py | 5 ++++- 2 files changed, 5 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 87e4b9fff2c..bf5a6113078 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -545,7 +545,7 @@ def _apply_gather_elements(model, *, prepared_data_ops=(), validate_only=False): if id(node) in selected_ids and prepared[node.input[0]] and node.input[1] not in cache.graph_input_names | cache.init_names - and _known_rank(cache, node.input[1]) in (2, 3) + and _known_rank(cache, node.input[1]) in (None, 2, 3) ] index_layouts = ( _prepared_index_data( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 5c7f8020ef0..fb32768441b 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -43,7 +43,8 @@ @pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) @pytest.mark.parametrize("pipeline", [False, True]) @pytest.mark.parametrize("producer", ["Reshape", "Squeeze"]) -def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline, producer): +@pytest.mark.parametrize("metadata", ["inferred", "omitted"]) +def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline, producer, metadata): """A native internal index view cannot follow data into a prepared four-dimensional layout.""" source_shape = [1] * (4 - len(shape)) + shape model = _model( @@ -74,6 +75,8 @@ def test_unprepared_internal_indices_refuse_before_mutation(shape, pipeline, pro np.testing.assert_array_equal( _execute(model, feeds)[0], np.take(feeds["x"], [0] * shape[-1], axis=-1) ) + if metadata == "omitted": + del model.graph.value_info[:] if ( pipeline and producer == "Reshape" From 55e9f469bcf48ecf9d34f33cb5d05a41a304de3a Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 09:16:19 +0000 Subject: [PATCH 30/31] Validate runtime reduction controls before boundary preparation Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 57 +++++++++++++-- .../onnx/graph_surgery/make_dla_compatible.py | 4 + .../dla_transforms/test_dla_indexing.py | 73 ++++++++++++++++++- 3 files changed, 123 insertions(+), 11 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index bf5a6113078..5dec4a6bf8d 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -31,7 +31,7 @@ from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures from .dla_matmul_to_transpose_conv_transpose import _apply_matmul_to_transpose_conv_transpose -from .dla_model_boundaries import _has_native_symbolic_data, _native_conv1d +from .dla_model_boundaries import _BoundaryPreflightError, _has_native_symbolic_data, _native_conv1d __all__ = [] @@ -671,29 +671,68 @@ def _apply_gather_elements(model, *, prepared_data_ops=(), validate_only=False): return model -def _apply_standalone_reduce_operations(model, *, prepared_data_ops=()): +def _apply_standalone_reduce_operations(model, *, prepared_data_ops=(), validate_only=False): """Plan native reduction axes and preserve their output layouts during rank promotion.""" graph = model.graph + selected = _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) + if not selected: + return model cache = GraphCache(graph) nodes = [] count = 0 plans = {} - selected = _operator_nodes(model, {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}) prepared = _prepared_index_data( - model, cache, prepared_data_ops, [node.input[0] for node in selected] + model, + cache, + prepared_data_ops, + [node.input[0] for node in selected], + before_boundary=validate_only, + ) + shapes = {id(node): _shape_before_boundary_promotion(cache, node.input[0]) for node in selected} + constant_controls = ( + {name for node in _operator_nodes(model, {"Constant"}) for name in node.output} + if validate_only and model.ir_version > 3 and "ConstantsToInitializers" in prepared_data_ops + else set() ) for node in selected: - shape = _shape_before_boundary_promotion(cache, node.input[0]) + tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None + if ( + len(node.input) > 1 + and node.input[1] in constant_controls + and node.input[1] not in cache.graph_input_names + ): + continue + if ( + len(node.input) > 1 + and (tensor is None or tensor.name in cache.graph_input_names) + and ( + (validate_only and prepared[node.input[0]]) + or shapes[id(node)] + or ( + validate_only + and (keepdims := _attribute(node, "keepdims")) is not None + and keepdims.i == 0 + ) + ) + ): + error = ( + _BoundaryPreflightError if node.input[1] in cache.graph_input_names else ValueError + ) + raise error(f"Dynamic reduction axes are unsupported for {node.name!r}") + for node in selected: + shape = shapes[id(node)] + tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None if not shape: + if validate_only: + continue keepdims = _attribute(node, "keepdims") if keepdims is not None and keepdims.i == 0: raise ValueError(f"Missing static reduction input shape for {node.name!r}") continue # Already rank-preserving, including generated 4D reduction branches. if len(shape) in (2, 3) and not prepared[node.input[0]]: raise ValueError(f"{node.name!r}: reduction data requires rank preparation") - tensor = cache.get_init(node.input[1]) if len(node.input) > 1 else None - if len(node.input) > 1 and (tensor is None or tensor.name in cache.graph_input_names): - raise ValueError(f"Dynamic reduction axes are unsupported for {node.name!r}") + if validate_only: + continue attribute = _attribute(node, "axes") raw_axes = ( numpy_helper.to_array(tensor).reshape(-1).tolist() @@ -709,6 +748,8 @@ def _apply_standalone_reduce_operations(model, *, prepared_data_ops=()): if output_shape is None: output_shape = [dim for index, dim in enumerate(shape) if index not in reduced] plans[id(node)] = (shape, raw_axes, output_shape) + if validate_only: + return model for node in graph.node: nodes.append(node) if id(node) not in plans: diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index d3b55fa06aa..e9c26b59172 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -149,6 +149,10 @@ def _validate_boundary_contract( """Refuse unsupported native layouts before adapting any public boundary.""" if "GatherElements" in prepared_input_ops: _apply_gather_elements(model, prepared_data_ops=prepared_input_ops, validate_only=True) + if "StandaloneReduceOperations" in prepared_input_ops: + _apply_standalone_reduce_operations( + model, prepared_data_ops=prepared_input_ops, validate_only=True + ) local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} cache = GraphCache(model.graph) for node in model.graph.node: diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index fb32768441b..9601de9d3e1 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -197,10 +197,10 @@ def test_late_refusal_preserves_the_original_model(axes_default): helper.make_node("ReduceSum", ["x", "dynamic_axes"], ["y"], keepdims=0), ] inputs = [ - helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), helper.make_tensor_value_info("dynamic_axes", TensorProto.INT64, [1]), ] - outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3])] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3])] model = helper.make_model( helper.make_graph(nodes, "refusal", inputs, outputs), opset_imports=[helper.make_opsetid("", 21)], @@ -210,11 +210,78 @@ def test_late_refusal_preserves_the_original_model(axes_default): numpy_helper.from_array(np.asarray([0], np.int64), "dynamic_axes") ) before = model.SerializeToString() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="Dynamic reduction axes"): _apply_standalone_reduce_operations(model) assert model.SerializeToString() == before +@pytest.mark.parametrize("axes_default", [False, True]) +@pytest.mark.parametrize("intermediate", [False, True]) +def test_pipeline_dynamic_reduction_axes_refuse_before_mutation(axes_default, intermediate): + """Reduction-control preflight must cover optional intermediate metadata and input defaults.""" + data = "hidden" if intermediate else "x" + nodes = [helper.make_node("Identity", ["x"], [data])] if intermediate else [] + nodes.extend( + [ + helper.make_node("ReduceSum", [data, "fixed_axes"], ["ready"], keepdims=0), + helper.make_node("ReduceSum", [data, "dynamic_axes"], ["y"], keepdims=0), + ] + ) + initializers = [numpy_helper.from_array(np.array([0], np.int64), "fixed_axes")] + if axes_default: + initializers.append(numpy_helper.from_array(np.array([0], np.int64), "dynamic_axes")) + model = _model( + nodes, + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("dynamic_axes", TensorProto.INT64, [1]), + ], + [helper.make_tensor_value_info(name, TensorProto.FLOAT, [3]) for name in ["ready", "y"]], + initializers, + ) + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(2, 3), + "dynamic_axes": np.array([0], np.int64), + } + _execute(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Dynamic reduction axes"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("control", ["initializer", "Constant"]) +@pytest.mark.parametrize("intermediate", [False, True]) +def test_fixed_reduction_preflight_avoids_model_copy(control, intermediate, monkeypatch): + """Known fixed controls and inferable intermediate metadata need no full-model probe.""" + data = "hidden" if intermediate else "x" + nodes = [helper.make_node("Identity", ["x"], [data])] if intermediate else [] + axes = numpy_helper.from_array(np.array([0], np.int64), "axes") + if control == "Constant": + nodes.append(helper.make_node("Constant", [], ["axes"], value=axes)) + nodes.append(helper.make_node("ReduceSum", [data, "axes"], ["y"], keepdims=0)) + model = _model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3])], + [axes] if control == "initializer" else [], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds) + original = copy.deepcopy + model_copies = [] + + def track(value, *args, **kwargs): + if isinstance(value, onnx.ModelProto): + model_copies.append(value) + return original(value, *args, **kwargs) + + monkeypatch.setattr(copy, "deepcopy", track) + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + assert not model_copies + + @pytest.mark.parametrize("axes", [[-2], [-1, 0], [0, -1]]) def test_reduction_preserves_negative_and_mixed_axes(axes): """Reduction preserves negative and mixed axes.""" From ecfaee432ac503d550460222ef5341180a5fff30 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 15:40:45 +0000 Subject: [PATCH 31/31] Preserve native attention outputs captured by subgraphs Signed-off-by: mgohil --- .../dla_transforms/dla_indexing.py | 3 +- .../dla_transforms/test_dla_indexing.py | 30 +++++++++++++++---- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index 5dec4a6bf8d..dfb9d96ad9e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -804,6 +804,7 @@ def _apply_group_query_attention(model, *, prepared_data_ops=()): graph = model.graph cache = GraphCache(graph) nodes = [] + protected_outputs = cache.graph_output_names | _subgraph_captures(graph) boundaries = cache.graph_input_names | cache.init_names | cache.graph_output_names squeezed_inputs = {} axes_names = {} @@ -863,7 +864,7 @@ def axes_for(name): node.input[index] = squeezed nodes.append(node) for index, name in enumerate(node.output[:1]): - if name and name not in cache.graph_output_names: + if name and name not in protected_outputs: axes = axes_for(name) node.output[index] = cache.unique_name(name + "_squeezed") nodes.append( diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py index 9601de9d3e1..08c977706c0 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -139,10 +139,11 @@ def test_index_dtype_and_rank_conversion_preserves_selected_values(operation): np.testing.assert_array_equal(actual[0], expected[0]) -@pytest.mark.parametrize("direct_output", [False, True]) +@pytest.mark.parametrize("boundary", ["internal", "public", "captured"]) @pytest.mark.parametrize("legacy_ir", [False, True]) -def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_output, legacy_ir): +def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(boundary, legacy_ir): """Attention preserves graph boundaries and wraps internal tensors.""" + direct_output = boundary == "public" output = "attention" if direct_output else "y" nodes = [ helper.make_node("Unsqueeze", ["x", "axes"], ["q4"]), @@ -156,7 +157,22 @@ def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_ kv_num_heads=1, ), ] - if not direct_output: + if boundary == "captured": + branches = [ + helper.make_graph( + [helper.make_node("Identity", ["attention"], ["branch_output"])], + name, + [], + [helper.make_tensor_value_info("branch_output", TensorProto.FLOAT16, [1, 2, 4])], + ) + for name in ("then", "else") + ] + nodes.append( + helper.make_node( + "If", ["condition"], ["y"], then_branch=branches[0], else_branch=branches[1] + ) + ) + elif not direct_output: nodes.append(helper.make_node("Identity", ["attention"], ["y"])) model = helper.make_model( helper.make_graph( @@ -165,13 +181,17 @@ def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_ [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 2, 4])], [ helper.make_tensor_value_info( - output, TensorProto.FLOAT16, [1, 2, 4] if direct_output else [1, 1, 2, 4] + output, + TensorProto.FLOAT16, + [1, 1, 2, 4] if boundary == "internal" else [1, 2, 4], ) ], initializer=[numpy_helper.from_array(np.array([0], dtype=np.int64), "axes")], ), opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)], ) + if boundary == "captured": + model.graph.input.append(helper.make_tensor_value_info("condition", TensorProto.BOOL, [])) if legacy_ir: model.ir_version = 3 for tensor in model.graph.initializer: @@ -185,7 +205,7 @@ def test_attention_preserves_graph_boundaries_and_wraps_internal_tensors(direct_ assert producer[attention.input[0]].op_type == "Squeeze" assert list(attention.input[1:]) == ["x", "x", "", ""] assert producer["attention"].op_type == ( - "GroupQueryAttention" if direct_output else "Unsqueeze" + "Unsqueeze" if boundary == "internal" else "GroupQueryAttention" )