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..dfb9d96ad9e --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -0,0 +1,880 @@ +# 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 math + +import numpy as np +from onnx import TensorProto, helper, numpy_helper + +from ..utils.graph_utils import ( + GraphCache, + _attribute, + _shape_before_boundary_promotion, + add_initializer, + batch_replace_nodes, + pad4d, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures +from .dla_matmul_to_transpose_conv_transpose import _apply_matmul_to_transpose_conv_transpose +from .dla_model_boundaries import _BoundaryPreflightError, _has_native_symbolic_data, _native_conv1d + +__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 _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 _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, 4) + 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 _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 _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( + 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 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) + if name in matrices or (shape is not None and len(shape) == 4): + return True + 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): + 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 == "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": + return "GroupQueryAttention" in prepared_ops + if node.op_type == "Conv": + kernel = cache.get_shape(node.input[1]) + return kernel is not None and len(kernel) == 4 + if not before_boundary and node.op_type == "MatMul": + if name in cache.graph_output_names | captured: + return False + 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": + if unsqueezed_rank(node) == 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: + 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 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"}: + 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 not before_boundary + or 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 = [(name, False) for name in data_names] + 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 = _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"}) + 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"} + 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) + 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 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 + 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: + 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 ( + rank == 0 + 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" + ) + 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" + ) + selected_ids = {id(node) for node in selected} + for node in graph.node: + if id(node) not in selected_ids: + nodes.append(node) + continue + indices = cache.get_init(node.input[1]) + 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] + rank = _known_rank(cache, original) + if rank is not None and (rank <= 1 or original in flattened): + if original not in squeezed_indices: + 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] + 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 = cache.unique_name((node.name or node.output[0]) + "_axis_updated") + nodes.append(node) + continue + squeezed = cache.unique_name(original + "_squeezed") + nodes.append( + 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 + 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 = 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 + + +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) + converted_indices = {} + replacements = {} + selected_ids = {id(node) for node in _operator_nodes(model, {"GatherElements"})} + prepared = _prepared_index_data( + model, + 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 (None, 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 + and cache.get_shape(node.input[0]) is not None + 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") + if validate_only: + return model + for index, node in enumerate(graph.node): + if id(node) not in selected_ids: + continue + original = node.input[1] + tensor = cache.get_init(original) + 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] + elif tensor is not None and original not in cache.graph_input_names: + array = numpy_helper.to_array(tensor).astype(np.int32) + if rank in (2, 3): + array = array.reshape([*([1] * (4 - rank)), *array.shape]) + elif rank == 0: + array = array.reshape(1) + prepared = add_initializer( + graph, + cache.unique_name(tensor.name + "_gatherelements_indices"), + array, + dtype=TensorProto.INT32, + with_value_info=False, + ) + else: + 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, + ) + ) + # 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") + 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)) + 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, *, 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 = {} + prepared = _prepared_index_data( + 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: + 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") + if validate_only: + continue + 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) + if validate_only: + return model + 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, + cache.unique_name((node.name or node.output[0]) + "_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, + 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] = cache.unique_name(original + "_temp") + nodes.append( + helper.make_node( + "Reshape", + [node.output[0], name], + [original], + name=cache.unique_name((node.name or original) + "_reshape"), + ) + ) + count += 1 + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) + return 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) + 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 = {} + + 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",)) + } + 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: + 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 + axes = axes_for(name) + squeezed = cache.unique_name(name + "_squeezed") + squeezed_inputs[name] = squeezed + nodes.append( + helper.make_node( + "Squeeze", + [name, axes], + [squeezed], + name=cache.unique_name(node.name + "_" + name + "_squeeze"), + ) + ) + node.input[index] = squeezed + nodes.append(node) + for index, name in enumerate(node.output[:1]): + if name and name not in protected_outputs: + axes = axes_for(name) + node.output[index] = cache.unique_name(name + "_squeezed") + nodes.append( + helper.make_node( + "Unsqueeze", + [node.output[index], axes], + [name], + name=cache.unique_name(node.name + "_" + name + "_unsqueeze"), + ) + ) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) + return model 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..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 @@ -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,18 @@ 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: + 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 for node in active_fallback: cnt_fb = len(fallback_batch) 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/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 478c818b673..e9c26b59172 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), @@ -112,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 in {"Gather", "GatherElements", "StandaloneReduceOperations", "GroupQueryAttention"}: + return {"prepared_data_ops": prepared_input_ops} if name == "RemoveQDQ": return { "keep_clip_after_inputs": False, @@ -135,6 +147,12 @@ 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) + 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/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 new file mode 100644 index 00000000000..08c977706c0 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_indexing.py @@ -0,0 +1,1773 @@ +# 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 onnx.reference import ReferenceEvaluator + +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, +) +from modelopt.onnx.graph_surgery.make_dla_compatible import ( + DLA_TRANSFORMS, + _transform_make_dla_compatible, +) + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("pipeline", [False, True]) +@pytest.mark.parametrize("producer", ["Reshape", "Squeeze"]) +@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( + [ + helper.make_node(producer, ["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 if producer == "Reshape" else list(range(4 - len(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 metadata == "omitted": + del model.graph.value_info[:] + 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() + 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.""" + 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 = 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 + 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("boundary", ["internal", "public", "captured"]) +@pytest.mark.parametrize("legacy_ir", [False, True]) +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"]), + helper.make_node( + "GroupQueryAttention", + ["q4", "x", "x", "", ""], + ["attention"], + name="attention", + domain="com.microsoft", + num_heads=1, + kv_num_heads=1, + ), + ] + 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( + nodes, + "attention", + [helper.make_tensor_value_info("x", TensorProto.FLOAT16, [1, 2, 4])], + [ + helper.make_tensor_value_info( + 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: + 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") + 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 == ( + "Unsqueeze" if boundary == "internal" else "GroupQueryAttention" + ) + + +@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), + ] + inputs = [ + 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, [1, 2, 3])] + model = helper.make_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, 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.""" + 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), prepared_input_ops={"ReduceSum"}) + 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] + 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 + + +@pytest.mark.parametrize( + "operation", ["GatherElements", "ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"] +) +@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, 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, 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.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), + ("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, + 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() + callback = ( + _apply_gather_elements + if operation == "GatherElements" + else _apply_standalone_reduce_operations + ) + with pytest.raises(ValueError, match="rank"): + callback(model) + 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]) + + +@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]) + + +@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) + + +@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) + 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) + + +@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] + 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]) +@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 = [ + 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]) + ) + 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) + + +@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)} + # 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(ReferenceEvaluator(model).run(None, 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) + + +@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( + 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, prepared_output_ops={"ReduceMean"} + ) + 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) + + +@pytest.mark.parametrize("index_shape", [(2, 3), (1, 2, 3)]) +@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( + [ + 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) + 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"Gather indices.*rank"): + _apply_gather(model) + 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("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"), + [ + ("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) + + +@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) + + +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) +@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( + 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 + + +@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"