diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py new file mode 100644 index 00000000000..64a34411bb3 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -0,0 +1,895 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Promote model boundaries at their explicit positions in the DLA pipeline.""" + +import numpy as np +from onnx import TensorProto, defs, 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_reorder_nodes import priority_topological_order + +__all__ = [] + + +class _BoundaryPreflightError(ValueError): + """Prerequisite lowering cannot freeze an exposed runtime control.""" + + +_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} +_REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) + + +def _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache): + """Require an explicit later stage before deferring a native rank-reducing operation.""" + if _is_standard_node(node, local_functions, {"Gather"}) and "Gather" not in prepared_output_ops: + shape = cache.get_shape(node.input[1]) + value = cache.vi_map.get(node.input[1]) + rank = ( + len(shape) + if shape is not None + else ( + len(value.type.tensor_type.shape.dim) + if value is not None and value.type.tensor_type.HasField("shape") + else None + ) + ) + if rank is None or rank == 0: + raise ValueError( + "Output rank conversion is unavailable for scalar or unknown-rank Gather" + ) + if ( + _is_standard_node(node, local_functions, {node.op_type}) + and ( + node.op_type == "Flatten" + or ( + (node.op_type in {"ArgMax", "ArgMin"} or node.op_type.startswith("Reduce")) + and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + ) + ) + and node.op_type not in prepared_output_ops + ): + raise ValueError(f"Output rank conversion is unavailable for {node.op_type}") + + +def _is_standard_node(node, local_functions, op_types, *, contrib_ops=()): + """Match builtin operations without interpreting opaque overloads or local definitions.""" + return ( + node.op_type in op_types + and ( + node.domain in {"", "ai.onnx"} + or (node.domain == "com.microsoft" and node.op_type in contrib_ops) + ) + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + ) + + +def _boundary_contrib_ops(model): + """Keep the explicit supported runtime schemas separate from opaque custom functions.""" + return ( + {"GroupQueryAttention", "QuantizeLinear", "DequantizeLinear", "Gelu"} + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import) + else set() + ) + + +def _output_rank_source(cache, node, local_functions, sources): + """Follow native rank-preserving unary paths before planning public restoration.""" + seen = set() + while ( + node is not None + and _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) + and len(node.input) == 1 + ): + if id(node) in sources: + node = sources[id(node)] + break + if id(node) in seen: + raise ValueError("Cyclic unary path while planning output restoration") + seen.add(id(node)) + node = cache.get_producer(node.input[0]) + sources.update(dict.fromkeys(seen, node)) + return node + + +def _has_native_symbolic_data(cache, name, local_functions): + """Prove a low-rank symbolic data path was left native by boundary promotion.""" + seen = set() + while name not in seen: + seen.add(name) + info = cache.vi_map.get(name) + rank = ( + len(info.type.tensor_type.shape.dim) + if info is not None and info.type.tensor_type.HasField("shape") + else None + ) + if cache.get_shape(name) is not None or (rank is not None and rank >= 4): + return False + if name in cache.graph_input_names: + return rank is not None and rank < 4 + node = cache.get_producer(name) + if ( + node is None + or not _is_standard_node(node, local_functions, _RANK_FLOW_OPS) + or not node.input + ): + return False + name = node.input[0] + raise ValueError("Cyclic data path while planning native symbolic controls") + + +def _require_adapter_opset(model): + """Refuse input-form axes adapters before changing a legacy graph.""" + version = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) + if version < 13: + raise ValueError("DLA boundary rank adapters require ONNX opset 13 or newer") + + +def _native_conv1d(node, cache): + """Recognize one-dimensional Conv paths whose data stays native before normalization repair.""" + kernel = ( + cache.get_shape(node.input[1]) if node.op_type == "Conv" and len(node.input) > 1 else None + ) + if kernel is None or len(kernel) != 3: + return False + name, seen = node.input[0], set() + while name not in seen: + seen.add(name) + if name in cache.graph_input_names: + shape = cache.get_shape(name) + return shape is not None and len(shape) == 3 + producer = cache.get_producer(name) + if producer is None: + return False + if producer.op_type == "InstanceNormalization": + shape = _shape_before_boundary_promotion(cache, producer.input[0]) + return shape is not None and len(shape) == 3 + if producer.op_type not in _RANK_AGNOSTIC_OPS or len(producer.input) != 1: + return False + name = producer.input[0] + return False + + +def _promoted_axis(node, rank, default): + """Map an explicit or schema-default axis after adding leading singleton dimensions.""" + axis = _attribute(node, "axis") + value = axis.i if axis is not None else default + if not -rank <= value < rank: + raise ValueError(f"Invalid native normalization axis for {node.name!r}") + return axis, value + 4 - rank if value >= 0 else value + + +def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=False): + """Promote selected boundary uses while preserving explicit adapter results.""" + graph = model.graph + cache = GraphCache(graph) + captured = _subgraph_captures(graph) + replacements = {} + prepend = [] + promotions = [] + unsqueezed_shapes = {} + normalization_axes = {} + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + contrib_ops = _boundary_contrib_ops(model) + opset = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) + for value in graph.input: + shape = cache.get_shape(value.name) + if shape is None or len(shape) not in (2, 3): + continue + consumers = [ + node + for node, _ in cache.get_consumers(value.name) + if _is_standard_node(node, local_functions, {node.op_type}, contrib_ops=contrib_ops) + and ( + node.input[0] == value.name or _is_standard_node(node, local_functions, {"Concat"}) + ) + and not ( + _is_standard_node(node, local_functions, {"Squeeze"}) and "Gather" in node.name + ) + and node.op_type not in {"InstanceNormalization", "Gemm"} + and not _native_conv1d(node, cache) + ] + if not consumers: + continue + _require_adapter_opset(model) + for consumer in consumers: + safe = _RANK_AGNOSTIC_OPS | { + "Unsqueeze", + "Concat", + "Softmax", + "LogSoftmax", + "LayerNormalization", + "Gelu", + } + if consumer.op_type == "MatMul": + other = cache.get_shape(consumer.input[1]) + if other is not None and len(other) >= 2: + safe |= {"MatMul"} + elif consumer.op_type == "Conv": + kernel = cache.get_shape(consumer.input[1]) + if kernel is not None and len(kernel) == 4: + safe |= {"Conv"} + elif consumer.op_type in {"QuantizeLinear", "DequantizeLinear"}: + if cache.get_shape(consumer.input[1]) == []: + safe |= {consumer.op_type} + if consumer.op_type not in safe and consumer.op_type not in prepared_input_ops: + raise ValueError(f"Input rank conversion is unavailable for {consumer.op_type}") + if _is_standard_node( + consumer, local_functions, {"Softmax", "LogSoftmax", "LayerNormalization"} + ): + axis, axis_value = _promoted_axis( + consumer, + len(shape), + defs.get_schema(consumer.op_type, opset).attributes["axis"].default_value.i, + ) + if axis_value >= 0: + normalization_axes[id(consumer)] = (consumer, axis, axis_value) + if not _is_standard_node(consumer, local_functions, {"Unsqueeze"}): + continue + tensor = cache.get_init(consumer.input[1]) + control = cache.get_producer(consumer.input[1]) + if ( + tensor is None + and validate_only + and control is not None + and _is_standard_node(control, local_functions, {"Constant"}) + ): + continue + if tensor is None or tensor.name in cache.graph_input_names: + raise ValueError(f"Unsqueeze axes must be fixed for {consumer.name!r}") + output_shape = list(shape) + axes = numpy_helper.to_array(tensor).reshape(-1).tolist() + output_rank = len(shape) + len(axes) + for axis in sorted(axis % output_rank for axis in axes): + output_shape.insert(axis, 1) + if ( + pad4d(shape)[-len(output_shape) :] == output_shape + and consumer.output[0] in cache.graph_output_names | captured + ): + raise ValueError( + f"Cannot remove graph-output or captured Unsqueeze {consumer.name!r}" + ) + unsqueezed_shapes[id(consumer)] = output_shape + promotions.append((value, shape, consumers)) + promoted_shapes = {value.name: shape for value, shape, _ in promotions} + concats = { + id(node): node + for _, _, consumers in promotions + for node in consumers + if _is_standard_node(node, local_functions, {"Concat"}) + } + for node in concats.values(): + if ( + any(name not in promoted_shapes for name in node.input) + or len({len(promoted_shapes[name]) for name in node.input}) != 1 + ): + raise ValueError(f"Cannot promote mixed-layout Concat {node.name!r}") + if validate_only: + return model + for node in concats.values(): + for attribute in node.attribute: + if attribute.name == "axis" and attribute.i >= 0: + attribute.i += 4 - len(promoted_shapes[node.input[0]]) + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") + for node, axis, value in normalization_axes.values(): + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name = cache.unique_name((node.name or node.output[0]) + "_axis_updated") + for value, shape, consumers in promotions: + axes = list(range(4 - len(shape))) + axes_name = add_initializer( + graph, + cache.unique_name(value.name + "_unsqueeze_axes"), + np.asarray(axes, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(value.name + "_unsqueeze_input") + prepend.insert( + 0, + helper.make_node( + "Unsqueeze", + [value.name, axes_name], + [output], + name=cache.unique_name(value.name + "_unsqueeze"), + ), + ) + promoted_shape = pad4d(shape) + graph.value_info.append( + helper.make_tensor_value_info(output, value.type.tensor_type.elem_type, promoted_shape) + ) + for consumer in consumers: + if _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, + ) + 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: + consumer.input[index] = output + batch = { + index: ([node], replacements[id(node)]) + for index, node in enumerate(graph.node) + if id(node) in replacements + } + if prepend: + removed, added = batch.get(0, ([], [])) + batch[0] = (removed, [*prepend, *added]) + batch_replace_nodes(graph, batch) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_non4d_model_outputs( + model, *, defer_rank_conversion=False, prepared_output_ops=(), validate_only=False +): + """Restore outputs or only validate producer rank availability before shape inference.""" + graph = model.graph + cache = GraphCache(graph) + captured = _subgraph_captures(graph) + boundary_names = cache.graph_input_names | cache.init_names + append = [] + plans = [] + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_broadcasts = set() + if defer_rank_conversion and "Non4dInitializers" not in prepared_output_ops: + live_ranks = { + name: len(shape) + for name in boundary_names + if (shape := cache.get_shape(name)) is not None + } + for node in graph.node: + if ( + _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) + and node.input + and all(name in live_ranks for name in node.input) + ): + rank = max(live_ranks[name] for name in node.input) + live_ranks.update(dict.fromkeys(node.output, rank)) + if rank < 4 and ( + len(node.input) > 1 or any(name in native_broadcasts for name in node.input) + ): + native_broadcasts.update(node.output) + rank_sources = {} + contrib_ops = _boundary_contrib_ops(model) + for value in graph.output: + shape = cache.get_shape(value.name) + if shape is None or len(shape) >= 4 or value.name in boundary_names: + continue + if value.name in native_broadcasts: + continue + producer = cache.get_producer(value.name) + source = _output_rank_source(cache, producer, local_functions, rank_sources) + if source is None and len(shape) == 1: + continue + if source is not None: + if _is_standard_node(source, local_functions, {"Gemm", "Shape"}) or _native_conv1d( + source, cache + ): + continue + if defer_rank_conversion: + _require_output_rank_conversion(source, local_functions, prepared_output_ops, cache) + if not _is_standard_node( + source, local_functions, {source.op_type}, contrib_ops=contrib_ops + ): + continue + if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): + if _has_native_symbolic_data(cache, source.input[0], local_functions): + continue + target = cache.get_shape(source.input[1]) if source.op_type == "Reshape" else None + if ( + source.op_type == "Reshape" + and target is not None + and len(target) == 1 + and target[0] < 4 + and "Non4DReshape" not in prepared_output_ops + ): + raise ValueError("Output rank conversion is unavailable for native Reshape") + if ( + producer is not None + and _is_standard_node(producer, local_functions, {"Identity"}) + and producer.input[0] in boundary_names + ): + continue + if len(shape) == 0: + raise ValueError(f"Unsupported scalar model output {value.name!r}") + axes = [] + existing = ( + producer + if producer is not None and _is_standard_node(producer, local_functions, {"Squeeze"}) + else None + ) + if existing is not None: + if len(existing.input) > 1 and existing.input[1]: + initializer = cache.get_init(existing.input[1]) + if initializer is None or initializer.name in cache.graph_input_names: + error = ( + _BoundaryPreflightError + if existing.input[1] in cache.graph_input_names + else ValueError + ) + raise error(f"Output Squeeze axes must be fixed for {value.name!r}") + axes = numpy_helper.to_array(initializer).reshape(-1).tolist() + else: + input_shape = cache.get_shape(existing.input[0]) + if input_shape is None: + if validate_only: + continue + raise ValueError(f"Output Squeeze input shape must be known for {value.name!r}") + axes = [index for index, dim in enumerate(input_shape) if dim == 1] + source_shape = cache.get_shape(existing.input[0] if existing is not None else value.name) + if not defer_rank_conversion and (source_shape is None or len(source_shape) != 4): + raise ValueError( + f"Output restoration requires a known rank-4 producer for {value.name!r}" + ) + delta = 4 - len(axes) - len(shape) + if delta < 0: + raise ValueError(f"Output Squeeze input rank exceeds 4 for {value.name!r}") + axes = list(range(delta)) + [axis + delta if axis >= 0 else axis for axis in axes] + plans.append((value, shape, existing, axes)) + if plans: + _require_adapter_opset(model) + if validate_only: + return model + for value, shape, existing, axes in plans: + axes_name = None + if len(shape) > 1 or shape == [1]: + axes_name = add_initializer( + graph, + cache.unique_name(value.name + "_squeeze_axes"), + np.asarray(axes, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + if existing is not None: + data_input = existing.input[0] + del existing.input[:] + existing.input.extend([data_input] + ([axes_name] if axes_name else [])) + continue + data_input = cache.unique_name(value.name + "_squeeze_input") + for node in graph.node: + for index, name in enumerate(node.input): + if name == value.name: + node.input[index] = data_input + for index, name in enumerate(node.output): + if name == value.name: + node.output[index] = data_input + append.append( + helper.make_node( + "Squeeze", + [data_input] + ([axes_name] if axes_name else []), + [value.name], + name=cache.unique_name(value.name + "_squeeze"), + ) + ) + graph.node.extend(append) + if any(value.name in captured for value, _, _, _ in plans): + order, _, _ = priority_topological_order(graph) + nodes = list(graph.node) + ranks = {id(nodes[index]): rank for rank, index in enumerate(order)} + graph.node.sort(key=lambda node: ranks[id(node)]) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_remove_intermediary_squeeze_unsqueeze( + model, *, rank_converted=False, prepared_consumer_ops=(), validate_only=False +): + """Remove no-op adapters, or reconcile adapters within the ordered rank-converted pipeline.""" + graph = model.graph + cache = GraphCache(graph) + captured = _subgraph_captures(graph) + removed = set() + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + checked_readers = set() + contrib_ops = _boundary_contrib_ops(model) + protected_inputs = ( + cache.graph_input_names | cache.init_names | cache.graph_output_names | captured + ) + protected_outputs = cache.graph_output_names | captured + runtime_inputs = ( + cache.graph_input_names - cache.init_names + if model.ir_version <= 3 + else cache.graph_input_names + ) + + def require_prepared_readers(name, planned_shape): + """Follow shared rank-agnostic paths before allowing any adapter rewiring.""" + layout = tuple(planned_shape) if planned_shape else None + if layout is not None and len(layout) == 4 and cache.get_shape(name) == list(layout): + return + pending, seen = [name], set() + while pending: + current = pending.pop() + if current in seen or (current, layout) in checked_readers: + continue + seen.add(current) + if current in captured - cache.graph_output_names: + raise ValueError("Intermediary adapter changes a captured native rank") + for reader, port in cache.get_consumers(current): + standard = _is_standard_node( + reader, local_functions, {reader.op_type}, contrib_ops=contrib_ops + ) + if ( + standard + and reader.op_type == "Squeeze" + and reader.output[0] in cache.graph_output_names + ): + continue + if standard and reader.op_type in _RANK_AGNOSTIC_OPS | { + "Gelu", + "Squeeze", + "Unsqueeze", + }: + if ( + reader.op_type in {"Squeeze", "Unsqueeze"} + and layout is not None + and len(layout) == 4 + and cache.get_shape(reader.output[0]) == list(layout) + and reader.input[0] not in protected_inputs + and reader.output[0] not in protected_outputs + ): + continue + pending.extend(reader.output) + continue + if standard and port == 0: + axis = _attribute(reader, "axis") + if reader.op_type == "LayerNormalization" and (axis is None or axis.i < 0): + pending.extend(reader.output) + continue + if ( + reader.op_type in {"QuantizeLinear", "DequantizeLinear"} + and cache.get_shape(reader.input[1]) == [] + ): + pending.extend(reader.output) + continue + target = ( + cache.get_init_array(reader.input[1]) + if standard + and reader.op_type == "Reshape" + and port == 0 + and reader.input[1] not in runtime_inputs + else None + ) + if target is not None and not np.any(target == 0): + continue + if ( + standard + and reader.op_type in prepared_consumer_ops + and (port == 0 or reader.op_type == "Concat") + ): + continue + raise ValueError( + f"Intermediary adapter feeds unprepared rank-sensitive {reader.op_type}" + ) + checked_readers.update((current, layout) for current in seen) + + plans = [] + for node in list(graph.node): + if not _is_standard_node(node, local_functions, {"Squeeze", "Unsqueeze"}): + continue + if ( + node.input[0] in cache.graph_input_names | cache.init_names | cache.graph_output_names + or node.output[0] in cache.graph_output_names | captured + or node.input[0] in captured + ): + continue + input_shape = cache.get_shape(node.input[0]) + output_shape = cache.get_shape(node.output[0]) + if not input_shape or not output_shape: + if validate_only and rank_converted: + require_prepared_readers(node.output[0], None) + continue + if not rank_converted and input_shape != output_shape: + continue + short, long = ( + (output_shape, input_shape) + if node.op_type == "Squeeze" + else (input_shape, output_shape) + ) + bypass = len(short) <= len(long) and long[-len(short) :] == short + if rank_converted and len(input_shape) != len(output_shape): + require_prepared_readers(node.output[0], pad4d(input_shape if bypass else output_shape)) + plans.append((node, output_shape, bypass)) + if validate_only: + return model + for node, output_shape, bypass in plans: + if bypass: + old, new = node.input[0], node.output[0] + for other in graph.node: + if other is node: + continue + for index, name in enumerate(other.output): + if name == old: + other.output[index] = new + for index, name in enumerate(other.input): + if name == old: + other.input[index] = new + removed.add(id(node)) + else: + target = add_initializer( + graph, + cache.unique_name((node.name or node.output[0]) + "_reshape_shape"), + np.asarray(pad4d(output_shape), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.CopyFrom( + helper.make_node( + "Reshape", + [node.input[0], target], + list(node.output), + name=cache.unique_name((node.name or node.output[0]) + "_to_reshape"), + ) + ) + batch_replace_nodes( + graph, {index: ([node], []) for index, node in enumerate(graph.node) if id(node) in removed} + ) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_fix_instancenorm_channels( + model, *, prepared_consumer_ops=(), validate_only=False, require_inferred_shapes=False +): + """Append spatial singleton dimensions while retaining normalization channels and outputs.""" + graph = model.graph + cache = GraphCache(graph) + replacements = {} + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + contrib_ops = _boundary_contrib_ops(model) + opset = next( + (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 + ) + axis_plans, perm_plans = {}, {} + axis_ops = { + "Softmax", + "LogSoftmax", + "LayerNormalization", + "Concat", + "Split", + "TopK", + "QuantizeLinear", + "DequantizeLinear", + } + + def consumer_layout(consumer, port, standard): + """Plan native views or axis/permutation preparation before widening any tensor.""" + if ( + not standard + or consumer.op_type in {"Shape", "GroupQueryAttention"} + or _native_conv1d(consumer, cache) + or (port != 0 and consumer.op_type not in {"Concat", "MatMul"}) + ): + return True + if consumer.op_type in axis_ops: + if "axis_updated" in consumer.name: + return False + quantizer = consumer.op_type in {"QuantizeLinear", "DequantizeLinear"} + if quantizer and cache.get_shape(consumer.input[1]) == []: + return False + schema = ( + defs.get_schema(consumer.op_type, opset) + if consumer.domain in {"", "ai.onnx"} + else None + ) + default = schema.attributes.get("axis") if schema is not None else None + if default is not None or (quantizer and consumer.domain == "com.microsoft"): + axis, value = _promoted_axis( + consumer, 3, default.default_value.i if default is not None else 1 + ) + if value >= 0: + axis_plans[id(consumer)] = (consumer, axis, value) + elif consumer.op_type == "Transpose": + perm = _attribute(consumer, "perm") + values = list(perm.ints) if perm is not None else [2, 1, 0] + if len(values) == 3 and sorted(values) == [0, 1, 2]: + perm_plans[id(consumer)] = (consumer, perm, [0, *[value + 1 for value in values]]) + elif len(values) != 4 or sorted(values) != [0, 1, 2, 3]: + raise ValueError(f"Unsupported normalization Transpose {consumer.name!r}") + elif consumer.op_type not in prepared_consumer_ops and consumer.op_type not in { + "Gelu", + "InstanceNormalization", + }: + target = cache.get_shape(consumer.input[1]) if consumer.op_type == "Reshape" else None + operands = ( + [cache.get_shape(name) for name in consumer.input] + if consumer.op_type == "MatMul" + else [] + ) + if target != [4] and not ( + operands and all(shape is not None and len(shape) >= 2 for shape in operands) + ): + raise ValueError(f"Unsupported normalization consumer {consumer.op_type}") + return False + + def output_consumers(name): + """Plan leading or native views for every edge after layout-preserving unary paths.""" + pending, visited = [name], set() + while pending: + tensor = pending.pop() + if tensor in visited: + continue + visited.add(tensor) + for consumer, input_index in cache.get_consumers(tensor): + standard = _is_standard_node( + consumer, local_functions, {consumer.op_type}, contrib_ops=contrib_ops + ) + if ( + standard + and input_index == 0 + and consumer.op_type == "Squeeze" + and consumer.output[0] in cache.graph_output_names + ): + yield consumer, None, None, False + elif standard and consumer.op_type == "InstanceNormalization" and input_index == 0: + continue + elif ( + standard + and consumer.op_type in _RANK_AGNOSTIC_OPS + and len(consumer.output) == 1 + ): + if len(consumer.input) == 1 and input_index == 0: + pending.append(consumer.output[0]) + elif len(consumer.input) > 1: + yield consumer, tensor, input_index, False + else: + yield ( + consumer, + tensor, + input_index, + consumer_layout(consumer, input_index, standard), + ) + + plans = [] + for index, node in enumerate(graph.node): + if _is_standard_node(node, local_functions, {"InstanceNormalization"}): + shape = _shape_before_boundary_promotion(cache, node.input[0]) + input_info = cache.vi_map.get(node.input[0]) + if ( + validate_only + and require_inferred_shapes + and shape is None + and (input_info is None or not input_info.type.tensor_type.HasField("shape")) + and node.input[0] not in cache.graph_input_names + and cache.get_init(node.input[0]) is None + ): + raise ValueError("Normalization preflight requires inferred intermediate shapes") + if shape is not None and len(shape) == 3: + plans.append((index, node, shape, list(output_consumers(node.output[0])))) + concat_inputs = {} + for _, _, _, consumers in plans: + for consumer, _, port, native in consumers: + if consumer.op_type == "Concat" and not native: + concat_inputs.setdefault(id(consumer), (consumer, set()))[1].add(port) + for consumer, ports in concat_inputs.values(): + if ports != set(range(len(consumer.input))): + raise ValueError("Unsupported normalization Concat with unprepared operands") + if validate_only: + return model + bridges = {} + for index, node, shape, consumers in plans: + target = add_initializer( + graph, + cache.unique_name((node.name or node.output[0]) + "_rotateleft_shape"), + np.asarray([*shape, 1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(node.input[0] + "_rotateleft") + reshape = helper.make_node( + "Reshape", + [node.input[0], target], + [output], + name=cache.unique_name((node.name or node.output[0]) + "_rotateleft"), + ) + node.input[0] = output + replacements[index] = ([node], [reshape, node]) + for consumer, tensor, port, native in consumers: + if tensor is not None: + key = (tensor, native) + if key not in bridges: + bridges[key] = (shape, []) + bridges[key][1].append((consumer, port)) + continue + if len(consumer.input) < 2: + continue + axes = cache.get_init(consumer.input[1]) + values = None if axes is None else numpy_helper.to_array(axes).reshape(-1).tolist() + if axes is not None and values and values[0] == 0: + native_axes = [value - 1 if value >= 1 else value + 3 for value in values[1:]] + consumer.input[1] = add_initializer( + graph, + cache.unique_name(axes.name + "_trailing"), + np.asarray([*native_axes, 3], np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + positions = {id(node): index for index, node in enumerate(graph.node)} + prefixes = {} + for (tensor, native), (shape, edges) in bridges.items(): + layout = "native" if native else "leading" + target = add_initializer( + graph, + cache.unique_name(tensor + "_" + layout + "_shape"), + np.asarray(shape if native else pad4d(shape), np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + output = cache.unique_name(tensor + "_" + layout) + view = helper.make_node( + "Reshape", [tensor, target], [output], name=cache.unique_name(output + "_reshape") + ) + first = min(positions[id(consumer)] for consumer, _ in edges) + prefixes.setdefault(first, []).append(view) + for consumer, port in edges: + consumer.input[port] = output + for index, nodes in prefixes.items(): + removed, added = replacements.get(index, ([graph.node[index]], [graph.node[index]])) + replacements[index] = (removed, [*nodes, *added]) + for consumer, axis, value in axis_plans.values(): + if axis is None: + consumer.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + consumer.name = cache.unique_name((consumer.name or consumer.output[0]) + "_axis_updated") + for consumer, perm, values in perm_plans.values(): + if perm is None: + consumer.attribute.append(helper.make_attribute("perm", values)) + else: + del perm.ints[:] + perm.ints.extend(values) + batch_replace_nodes(graph, replacements) + _add_legacy_initializer_inputs(model) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index bb6168d140d..478c818b673 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -23,6 +23,7 @@ from __future__ import annotations +import copy import time from functools import partial from typing import TYPE_CHECKING @@ -37,6 +38,15 @@ from .dla_transforms.dla_matmul_to_transpose_conv_transpose import ( _apply_matmul_to_transpose_conv_transpose, ) +from .dla_transforms.dla_model_boundaries import ( + _REDUCTION_OPS, + _apply_fix_instancenorm_channels, + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + _apply_remove_intermediary_squeeze_unsqueeze, + _BoundaryPreflightError, + _require_output_rank_conversion, +) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes @@ -47,6 +57,7 @@ ) from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, + GraphCache, infer_shapes, parse_qdq_quantized_dtype_list, ) @@ -87,11 +98,76 @@ def _apply_constants_to_initializers(model, **options): "MatMulToTransposeConvTranspose", partial(_apply_matmul_to_transpose_conv_transpose, defer_rank_conversion=True), ), + ("Non4dModelInputs", _apply_non4d_model_inputs), + ("Non4dModelOutputs", partial(_apply_non4d_model_outputs, defer_rank_conversion=True)), + ( + "RemoveIntermediarySqueezeAndUnsqueeze", + partial(_apply_remove_intermediary_squeeze_unsqueeze, rank_converted=True), + ), + ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) +def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prepared_output_ops): + """Share stage options between the prerequisite probe and the actual ordered pipeline.""" + if name == "RemoveQDQ": + return { + "keep_clip_after_inputs": False, + "max_chained_qdq_pairs": 5, + "quant_types_to_remove": quant_types, + } + if name == "ConstantsToInitializers": + return {"verbose": verbose} + if name == "Non4dModelInputs": + return {"prepared_input_ops": prepared_input_ops} + if name == "Non4dModelOutputs": + return {"prepared_output_ops": prepared_output_ops} + if name == "RemoveIntermediarySqueezeAndUnsqueeze": + return {"prepared_consumer_ops": prepared_input_ops} + if name == "FixInstanceNormChannelMismatch": + return {"prepared_consumer_ops": prepared_input_ops} + return {} + + +def _validate_boundary_contract( + model, prepared_input_ops, prepared_output_ops, *, require_inferred_shapes=False +): + """Refuse unsupported native layouts before adapting any public boundary.""" + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + cache = GraphCache(model.graph) + for node in model.graph.node: + if ( + node.op_type in _REDUCTION_OPS + and node.op_type not in prepared_output_ops + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + and any(attr.name == "keepdims" and attr.i == 0 for attr in node.attribute) + ): + raise NotImplementedError("Standalone reductions are not available in this pipeline") + _require_output_rank_conversion(node, local_functions, prepared_output_ops, cache) + if "FixInstanceNormChannelMismatch" in prepared_input_ops: + _apply_fix_instancenorm_channels( + model, + prepared_consumer_ops=prepared_input_ops, + validate_only=True, + require_inferred_shapes=require_inferred_shapes, + ) + _apply_non4d_model_inputs(model, prepared_input_ops=prepared_input_ops, validate_only=True) + _apply_non4d_model_outputs( + model, + defer_rank_conversion=True, + prepared_output_ops=prepared_output_ops, + validate_only=True, + ) + if "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_input_ops: + _apply_remove_intermediary_squeeze_unsqueeze( + model, rank_converted=True, prepared_consumer_ops=prepared_input_ops, validate_only=True + ) + + def _transform_make_dla_compatible( model: onnx.ModelProto, *, @@ -107,6 +183,32 @@ def _transform_make_dla_compatible( name == "RemoveQDQ" for name, _ in DLA_TRANSFORMS ): raise NotImplementedError("QDQ removal is not available in this pipeline") + stages = frozenset(stage for stage, _ in DLA_TRANSFORMS) + prepared_output_ops = stages | ( + _REDUCTION_OPS if "StandaloneReduceOperations" in stages else frozenset() + ) + prepared_input_ops = set(prepared_output_ops) + for stage, operations in { + "Non4DReshape": {"Reshape"}, + "Non4DExpand": {"Expand"}, + "Non4DTile": {"Tile"}, + "Non4DTranspose": {"Transpose"}, + "Non4dSliceAxis": {"Slice"}, + "Non4DLpNorm": {"LpNormalization"}, + "FixInstanceNormChannelMismatch": {"InstanceNormalization"}, + "RemoveIntermediarySqueezeAndUnsqueeze": {"Squeeze"}, + "UpdateAxis": { + "Softmax", + "LogSoftmax", + "TopK", + "Concat", + "Split", + "QuantizeLinear", + "DequantizeLinear", + }, + }.items(): + if stage in stages: + prepared_input_ops.update(operations) quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None @@ -114,19 +216,33 @@ def _transform_make_dla_compatible( ) if remove_all_qdq: quant_types = frozenset(onnx.TensorProto.DataType.values()) + try: + _validate_boundary_contract( + model, prepared_input_ops, prepared_output_ops, require_inferred_shapes=True + ) + except _BoundaryPreflightError: + raise + except (ValueError, NotImplementedError): + # Earlier lowering may consume native operators that lack a later rank pass. + probe = infer_shapes(copy.deepcopy(model)) + for name, transform in DLA_TRANSFORMS: + if name == "Non4dModelInputs": + break + options = _boundary_stage_options( + name, quant_types, False, prepared_input_ops, prepared_output_ops + ) + result = transform(probe, **options) + if result is not None: + probe = result + _validate_boundary_contract(probe, prepared_input_ops, prepared_output_ops) + del probe model = infer_shapes(model) for name, transform in DLA_TRANSFORMS: if verbose: logger.info("[DLA pipeline] Applying %s", name) - options = {} - if name == "RemoveQDQ": - options = { - "keep_clip_after_inputs": False, - "max_chained_qdq_pairs": 5, - "quant_types_to_remove": quant_types, - } - elif name == "ConstantsToInitializers": - options = {"verbose": verbose} + options = _boundary_stage_options( + name, quant_types, verbose, prepared_input_ops, prepared_output_ops + ) result = transform(model, **options) if result is not None: model = result diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py index b20341fe154..3f744c71d62 100644 --- a/modelopt/onnx/graph_surgery/utils/graph_utils.py +++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py @@ -986,14 +986,47 @@ def _attribute(node, name): return next((attribute for attribute in node.attribute if attribute.name == name), None) +def _allocated_boundary_name(name, hint): + """Match the shared allocator's unchanged hint or numbered collision suffix.""" + return name == hint or (name.startswith(hint + "_") and name[len(hint) + 1 :].isdigit()) + + def _shape_before_boundary_promotion(cache, name): - """Resolve the original shape behind a generated graph-input Unsqueeze.""" + """Recover native shapes at generated input and output boundary views.""" producer = cache.get_producer(name) if producer is not None and producer.op_type == "Unsqueeze": original = producer.input[0] if original in cache.graph_input_names and name == original + "_unsqueeze_input": return cache.get_shape(original) - return cache.get_shape(name) + shape = cache.get_shape(name) + if shape is not None: + return shape + for consumer, port in cache.get_consumers(name): + if ( + port != 0 + or consumer.op_type != "Squeeze" + or consumer.domain not in {"", "ai.onnx"} + or consumer.overload + or consumer.output[0] not in cache.graph_output_names + ): + continue + original = consumer.output[0] + if not _allocated_boundary_name( + name, original + "_squeeze_input" + ) or not _allocated_boundary_name(consumer.name, original + "_squeeze"): + continue + shape = cache.get_shape(original) + axes = cache.get_init_array(consumer.input[1]) if len(consumer.input) > 1 else None + if ( + shape is not None + and 0 < len(shape) < 4 + and ( + (axes is not None and axes.reshape(-1).tolist() == list(range(4 - len(shape)))) + or (axes is None and len(shape) == 1 and shape[0] != 1) + ) + ): + return shape + return None def is_const_dq_input(init_list, input_name, graph): diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py new file mode 100644 index 00000000000..1c10edaa960 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_model_boundaries.py @@ -0,0 +1,1674 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Boundary promotion refuses unsupported exposed adapters before changing the graph.""" + +import copy +from importlib import import_module + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, inliner, numpy_helper +from onnx.reference import ReferenceEvaluator + +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_fix_instancenorm_channels, + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + _apply_remove_intermediary_squeeze_unsqueeze, +) +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +from modelopt.onnx.graph_surgery.utils.graph_utils import infer_shapes + + +@pytest.mark.parametrize("entrypoint", ["validate", "pipeline"]) +@pytest.mark.parametrize("case", ["rank5", "scalar", "runtime_axes"]) +def test_output_refusal_preflight_preserves_caller_state(entrypoint, case, monkeypatch): + """Validate known output refusals before inference or an earlier stage changes the caller.""" + shape, axes, output_shape = { + "rank5": ([1, 2, 1, 3, 4], [0, 2], [2, 3, 4]), + "scalar": ([1], [0], []), + "runtime_axes": ([1, 2, 3, 4], [0], [2, 3, 4]), + }[case] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + if case == "runtime_axes": + inputs.append(helper.make_tensor_value_info("axes", TensorProto.INT64, [len(axes)])) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Squeeze", ["x", "axes"], ["y"])], + "output_refusal", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]).run( + None, {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + ) + before = model.SerializeToString() + original = copy.deepcopy + model_copies = [] + + def track(value, *args, **kwargs): + if isinstance(value, onnx.ModelProto): + model_copies.append(value) + return original(value, *args, **kwargs) + + monkeypatch.setattr(copy, "deepcopy", track) + with pytest.raises(ValueError, match=r"scalar|exceeds 4|must be fixed"): + if entrypoint == "pipeline": + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_non4d_model_outputs(model, defer_rank_conversion=True, validate_only=True) + assert model.SerializeToString() == before + if case == "runtime_axes": + assert not model_copies + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("input_port", [0, 1]) +@pytest.mark.parametrize("unary", [False, True]) +def test_broadcast_output_restoration_requires_a_promoted_source(shape, input_port, unary): + """Leave native broadcasts intact when no operand receives leading dimensions.""" + operands = ["x", "weight"] if input_port == 0 else ["weight", "x"] + nodes = [helper.make_node("Add", operands, ["middle" if unary else "y"])] + if unary: + nodes.append(helper.make_node("Relu", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_broadcast_output", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(np.ones(shape, np.float32), "weight")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + onnx.checker.check_model(model, full_check=True) + expected = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + transformed = _transform_make_dla_compatible(model, verbose=False) + onnx.checker.check_model(transformed, full_check=True) + actual = ort.InferenceSession( + transformed.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + np.testing.assert_array_equal(actual[0], expected[0]) + + +@pytest.mark.parametrize("identity", ["custom", "local", "overload"]) +@pytest.mark.parametrize( + ("operation", "entrypoint"), + [ + ("Squeeze", "pipeline"), + ("Unsqueeze", "pipeline"), + ("Concat", "pipeline"), + ("Squeeze", "cleanup"), + ("Unsqueeze", "cleanup"), + ], +) +def test_operator_lookalikes_keep_function_semantics(identity, operation, entrypoint): + """Execute opaque lookalikes without rewriting their builtin-named calls or controls.""" + domain = "opaque.boundary" if identity == "custom" else "" + overload = "custom" if identity == "overload" else "" + inputs = ["x"] if operation == "Concat" else ["x", "axes"] + function = helper.make_function( + domain, + operation, + inputs, + ["y"], + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_opsetid("", 21)], + attributes=["axis"] if operation == "Concat" else [], + ) + function.overload = overload + opaque = helper.make_node( + operation, + ["middle" if entrypoint == "cleanup" else "x", *inputs[1:]], + ["hidden" if entrypoint == "cleanup" else "y"], + name="opaque", + domain=domain, + **({"axis": 1} if operation == "Concat" else {}), + ) + opaque.overload = overload + nodes = ( + [ + helper.make_node("Identity", ["x"], ["middle"]), + opaque, + helper.make_node("Identity", ["hidden"], ["y"]), + ] + if entrypoint == "cleanup" + else [opaque] + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + [] if operation == "Concat" else [numpy_helper.from_array(np.array([], np.int64), "axes")], + ) + model.functions.append(function) + if domain: + model.opset_import.append(helper.make_opsetid(domain, 1)) + if entrypoint == "cleanup": + model.graph.value_info.extend( + helper.make_tensor_value_info(name, TensorProto.FLOAT, [2, 3]) + for name in ("middle", "hidden") + ) + before = next(n for n in model.graph.node if n.name == "opaque").SerializeToString() + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = ( + _transform_make_dla_compatible(model, verbose=False) + if entrypoint == "pipeline" + else _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + ) + assert next(n for n in converted.graph.node if n.name == "opaque").SerializeToString() == before + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_native_controls_do_not_gain_output_adapters(operation, tail): + """Keep static native outputs when symbolic data was not promoted, including unary tails.""" + target = {"Reshape": [3, 2], "Expand": [2, 3], "Tile": [0, 1]}[operation] + output_shape = {"Reshape": [3, 2], "Expand": [2, 3], "Tile": [0, 3]}[operation] + nodes = [helper.make_node(operation, ["x", "control"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.array(target, np.int64), "control")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +def test_input_promotion_refusal_leaves_graph_unchanged(): + """Refuse removing a public input adapter before altering any earlier consumer.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Identity", ["x"], ["first"]), + helper.make_node("Unsqueeze", ["x", "axes"], ["y"], name="exposed"), + ], + "boundary", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3])], + [numpy_helper.from_array(np.array([0, 1], dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 21)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +def _checked_run(model, feeds): + """Check every generated interface and execute through the independent reference runtime.""" + onnx.checker.check_model(model, full_check=True) + return ReferenceEvaluator(model).run(None, feeds) + + +def _make_model(nodes, inputs, outputs, initializers=()): + """Build a small model with explicit public interface shapes.""" + return helper.make_model( + helper.make_graph(nodes, "boundaries", inputs, outputs, initializer=initializers), + opset_imports=[helper.make_opsetid("", 21)], + ) + + +@pytest.mark.parametrize("identity", ["custom", "local", "overload"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_generic_opaque_transpose_keeps_native_shared_input(identity, tail): + """Preserve opaque function layouts while a separate standard branch promotes the same input.""" + domain = "opaque.boundary" if identity == "custom" else "" + operation = "OpaqueTranspose" if identity == "custom" else "Abs" + overload = "transpose" if identity == "overload" else "" + function = helper.make_function( + domain, + operation, + ["data"], + ["result"], + [helper.make_node("Transpose", ["data"], ["result"], perm=[1, 0])], + [helper.make_opsetid("", 21)], + overload=overload, + ) + opaque = helper.make_node( + operation, ["x"], ["middle" if tail else "y"], domain=domain, overload=overload + ) + nodes = [helper.make_node("Relu", ["x"], ["first"]), opaque] + if tail: + nodes.append(helper.make_node("Relu", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 2])], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 2]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2]), + ], + ) + model.functions.append(function) + if domain: + model.opset_import.append(helper.make_opsetid(domain, 1)) + feeds = {"x": np.arange(-1, 3, dtype=np.float32).reshape(2, 2)} + onnx.checker.check_model(model, full_check=True) + expected = _checked_run(inliner.inline_local_functions(model), feeds) + np.testing.assert_array_equal( + expected[1], np.maximum(feeds["x"].T, 0) if tail else feeds["x"].T + ) + original = opaque.SerializeToString() + converted = _transform_make_dla_compatible(model, verbose=False) + call = next( + node + for node in converted.graph.node + if node.op_type == operation and node.domain == domain and node.overload == overload + ) + assert call.SerializeToString() == original + for actual, reference in zip( + _checked_run(inliner.inline_local_functions(converted), feeds), expected + ): + np.testing.assert_array_equal(actual, reference) + + +def test_unnamed_input_adapters_keep_distinct_targets(): + """Promote two unnamed trailing-axis adapters without reusing their different targets.""" + inputs = [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 4]), + ] + outputs = [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3, 1]), + helper.make_tensor_value_info("b_out", TensorProto.FLOAT, [3, 4, 1]), + ] + model = _make_model( + [ + helper.make_node("Unsqueeze", ["a", "axes"], ["a_out"]), + helper.make_node("Unsqueeze", ["b", "axes"], ["b_out"]), + ], + inputs, + outputs, + [numpy_helper.from_array(np.asarray([-1], np.int64), "axes")], + ) + feeds = { + "a": np.arange(6, dtype=np.float32).reshape(2, 3), + "b": np.arange(12, dtype=np.float32).reshape(3, 4), + } + expected = _checked_run(model, feeds) + _apply_non4d_model_inputs(model) + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"]) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("axis", [0, 1, -1, None]) +@pytest.mark.parametrize("opset", [13, 17]) +def test_boundary_softmax_preserves_native_axes(operation, shape, axis, opset): + """Preserve explicit and schema-default normalization axes across input promotion.""" + model = _make_model( + [helper.make_node(operation, ["x"], ["y"], **({} if axis is None else {"axis": axis}))], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ) + model.opset_import[0].version = opset + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 10} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("axis", [1, -1]) +def test_boundary_layer_normalization_preserves_native_axis(axis): + """Preserve LayerNormalization reductions after adding leading singleton dimensions.""" + model = _make_model( + [helper.make_node("LayerNormalization", ["x", "scale"], ["y"], axis=axis)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + [numpy_helper.from_array(np.ones(3, np.float32), "scale")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize( + ("operation", "axis"), + [ + ("Softmax", -1), + ("Softmax", 1), + ("LogSoftmax", None), + ("Transpose", None), + ("Transpose", 1), + ("Conv", None), + ("opaque", None), + ("Shape", None), + ("InstanceNormalization", None), + ], +) +@pytest.mark.parametrize("unary", [False, True]) +def test_normalization_preserves_rank_sensitive_and_opaque_consumers(operation, axis, unary): + """Execute native and promoted consumers alongside a shared public normalization output.""" + shape = [2, 3, 4] + nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"])] + current = "normalized" + if unary: + nodes.append(helper.make_node("Relu", [current], ["activated"])) + current = "activated" + controls = [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ] + output_shape, dtype, functions = shape, TensorProto.FLOAT, [] + if operation in {"Softmax", "LogSoftmax"}: + nodes.append( + helper.make_node( + operation, [current], ["y"], **({"axis": axis} if axis is not None else {}) + ) + ) + elif operation == "Transpose": + nodes.append( + helper.make_node( + operation, [current], ["y"], **({"perm": [0, 2, 1]} if axis is not None else {}) + ) + ) + output_shape = [2, 4, 3] if axis is not None else [4, 3, 2] + elif operation == "Conv": + controls.append( + numpy_helper.from_array(np.arange(18, dtype=np.float32).reshape(2, 3, 3) / 18, "kernel") + ) + nodes.append(helper.make_node("Conv", [current, "kernel"], ["y"])) + output_shape = [2, 2, 2] + elif operation == "opaque": + functions.append( + helper.make_function( + "opaque.norm", + "NativeLayout", + ["value"], + ["result"], + [helper.make_node("Transpose", ["value"], ["result"], perm=[0, 2, 1])], + [helper.make_opsetid("", 21)], + ) + ) + nodes.append(helper.make_node("NativeLayout", [current], ["y"], domain="opaque.norm")) + output_shape = [2, 4, 3] + elif operation == "InstanceNormalization": + nodes.append(helper.make_node(operation, [current, "scale", "bias"], ["y"])) + else: + nodes.append(helper.make_node("Shape", [current], ["y"])) + output_shape, dtype = [3], TensorProto.INT64 + outputs = [ + helper.make_tensor_value_info("y", dtype, output_shape), + helper.make_tensor_value_info("normalized", TensorProto.FLOAT, shape), + ] + if operation == "Transpose" and axis is None: + for name in ("normalized_squeeze_input", "normalized_squeeze"): + controls.append(numpy_helper.from_array(np.array(1, np.int64), name)) + outputs.append(helper.make_tensor_value_info(name, TensorProto.INT64, [])) + if operation in {"Conv", "opaque", "Shape"}: + nodes.append(helper.make_node("Softmax", [current], ["side"], axis=-1)) + outputs.append(helper.make_tensor_value_info("side", TensorProto.FLOAT, shape)) + model = _make_model( + nodes, [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], outputs, controls + ) + model.functions.extend(functions) + if functions: + model.opset_import.append(helper.make_opsetid("opaque.norm", 1)) + feeds = {"x": np.linspace(-1, 2, 24, dtype=np.float32).reshape(shape)} + expected = _checked_run(inliner.inline_local_functions(model), feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + if operation == "Conv": + assert list( + next(tensor for tensor in converted.graph.initializer if tensor.name == "kernel").dims + ) == [2, 3, 3] + for actual, reference in zip( + _checked_run(inliner.inline_local_functions(converted), feeds), expected + ): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + ("through_pipeline", "intermediate"), [(False, False), (True, False), (True, True)] +) +def test_normalization_refuses_unprepared_consumers_before_mutation(through_pipeline, intermediate): + """Refuse a rank-sensitive consumer before widening any eligible normalization node.""" + shape = [2, 3, 4] + data = "r" if intermediate else "x" + model = _make_model( + ([helper.make_node("Relu", ["x"], [data])] if intermediate else []) + + [ + helper.make_node("InstanceNormalization", [data, "scale", "bias"], ["normalized"]), + helper.make_node("Hardmax", ["normalized"], ["y"], axis=1), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + _checked_run(model, {"x": np.arange(24, dtype=np.float32).reshape(shape)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Unsupported normalization consumer"): + if through_pipeline: + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_fix_instancenorm_channels(model) + assert model.SerializeToString() == before + + +def test_normalization_probe_accepts_unknown_sequence_element_rank(): + """A probe may retain valid native normalization when inference cannot recover its rank.""" + shape = [1, 1, 1, 4] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("SequenceAt", ["sequence", "index"], ["r"]), + helper.make_node("InstanceNormalization", ["r", "scale", "bias"], ["unused"]), + helper.make_node("Identity", ["x"], ["y"]), + ], + "unknown_native_rank", + [ + helper.make_tensor_sequence_value_info("sequence", TensorProto.FLOAT, None), + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.array(0, np.int64), "index"), + numpy_helper.from_array(np.ones(1, np.float32), "scale"), + numpy_helper.from_array(np.zeros(1, np.float32), "bias"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + values = np.arange(4, dtype=np.float32).reshape(shape) + feeds = {"sequence": [values], "x": values} + inferred = infer_shapes(copy.deepcopy(model)) + metadata = next(value for value in inferred.graph.value_info if value.name == "r") + assert not metadata.type.tensor_type.HasField("shape") + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("batch", [1, 2, 3]) +@pytest.mark.parametrize("operation", ["Add", "Mul"]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_normalization_binary_branch_preserves_broadcast_layout(batch, operation, tail, reverse): + """Keep shared normalized public values and binary residual inputs in compatible layouts.""" + shape = [batch, 3, 4] + normalized = "tail" if tail else "normalized" + nodes = [ + helper.make_node( + "InstanceNormalization", ["x", "scale", "bias"], ["normalized"], epsilon=1e-4 + ) + ] + if tail: + nodes.append(helper.make_node("Tanh", ["normalized"], ["tail"])) + nodes.append( + helper.make_node(operation, [normalized, "x"] if reverse else ["x", normalized], ["y"]) + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + helper.make_tensor_value_info(normalized, TensorProto.FLOAT, shape), + ], + [ + numpy_helper.from_array(np.asarray([0.5, 1, 1.5], np.float32), "scale"), + numpy_helper.from_array(np.asarray([0.1, -0.2, 0.3], np.float32), "bias"), + ], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 8} + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("width", [8, 12]) +@pytest.mark.parametrize("axis", [-1, 4]) +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("preparation", ["direct", "intermediate", "constant"]) +def test_early_registered_reshape_reduction_is_not_rejected( + width, axis, quantized, preparation, monkeypatch +): + """Allow reductions consumed by actual earlier reshape-pattern stages without a native reducer stage.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + shape = [1, 1, 2, width] + controls = [ + numpy_helper.from_array(np.asarray([1, 1, 2, width // 4, 4], np.int64), "target"), + numpy_helper.from_array(np.asarray([axis], np.int64), "axes"), + ] + nodes = [helper.make_node("Reshape", ["x", "target"], ["expanded"])] + current = "expanded" + if quantized: + controls += [ + numpy_helper.from_array(np.float32(0.125), "scale"), + numpy_helper.from_array(np.int8(0), "zero"), + ] + nodes += [ + helper.make_node("QuantizeLinear", [current, "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["dq"]), + ] + current = "dq" + nodes.append(helper.make_node("ReduceSum", [current, "axes"], ["y"], keepdims=0)) + if preparation == "intermediate": + nodes[0].input[0] = "hidden" + nodes.insert(0, helper.make_node("Identity", ["x"], ["hidden"])) + elif preparation == "constant": + nodes = [ + helper.make_node("Constant", [], [tensor.name], value=tensor) for tensor in controls + ] + nodes + controls = [] + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, width // 4])], + controls, + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 8} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False, qdq_quantized_dtypes=[]) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + "operation", ["ReduceSum", "Gather", "Slice", "Split", "ArgMax", "Transpose"] +) +def test_input_promotion_requires_registered_axis_conversion(operation, monkeypatch): + """Reject valid native operators before boundary mutation when their rank stage is absent.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + ( + ("ConstantsToInitializers", pipeline._apply_constants_to_initializers), + ("ShapeInference", pipeline.infer_shapes), + ("Non4dModelInputs", _apply_non4d_model_inputs), + ("Non4dModelOutputs", _apply_non4d_model_outputs), + ), + ) + inputs, attributes, output_shape = ["x"], {}, [2, 1] + controls = [] + dtype = TensorProto.FLOAT + if operation == "Transpose": + attributes, output_shape = {"perm": [1, 0]}, [3, 2] + elif operation == "ArgMax": + attributes, dtype = {"axis": 1, "keepdims": 1}, TensorProto.INT64 + else: + values = {"ReduceSum": [1], "Gather": [0], "Split": [3], "Slice": [0]}[operation] + controls.append(numpy_helper.from_array(np.asarray(values, np.int64), "control")) + inputs.append("control") + if operation == "ReduceSum": + attributes = {"keepdims": 1} + elif operation == "Gather": + attributes, output_shape = {"axis": 0}, [1, 3] + elif operation == "Split": + attributes, output_shape = {"axis": 1}, [2, 3] + else: + controls += [ + numpy_helper.from_array(np.asarray([1], np.int64), "end"), + numpy_helper.from_array(np.asarray([1], np.int64), "axes"), + ] + inputs += ["end", "axes"] + model = _make_model( + [helper.make_node(operation, inputs, ["y"], **attributes)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", dtype, output_shape)], + controls, + ) + _checked_run(model, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Input rank conversion"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("tail", [False, True]) +def test_ineligible_gemm_keeps_native_boundary_and_weights(tail): + """Retain native matrix Gemm semantics when convolution lowering declines its coefficients.""" + nodes = [helper.make_node("Gemm", ["x", "w", "b"], ["hidden" if tail else "y"], alpha=0.5)] + if tail: + nodes.append(helper.make_node("Relu", ["hidden"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4])], + [ + numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "w"), + numpy_helper.from_array(np.ones(4, np.float32), "b"), + ], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + assert [list(tensor.dims) for tensor in converted.graph.initializer] == [[3, 4], [4]] + + +@pytest.mark.parametrize("source", ["initializer", "constant", "runtime"]) +@pytest.mark.parametrize("axis", [0, 1]) +@pytest.mark.parametrize("tail", [False, True]) +def test_scalar_gather_requires_registered_rank_conversion(source, axis, tail, monkeypatch): + """Refuse unprepared scalar Gather before the actual pipeline changes caller state.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple((name, transform) for name, transform in pipeline.DLA_TRANSFORMS if name != "Gather"), + ) + shape = [2, 3, 4, 4] if axis == 0 else [1, 3, 4, 4] + output_shape = shape[:axis] + shape[axis + 1 :] + nodes = [helper.make_node("Gather", ["x", "indices"], ["hidden" if tail else "y"], axis=axis)] + if tail: + nodes.append(helper.make_node("Identity", ["hidden"], ["y"])) + tensor = numpy_helper.from_array(np.int64(1), "indices") + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + if source == "constant": + nodes.insert(0, helper.make_node("Constant", [], ["indices"], value=tensor)) + elif source == "runtime": + inputs.append(helper.make_tensor_value_info("indices", TensorProto.INT64, [])) + model = _make_model( + nodes, + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [tensor] if source == "initializer" else [], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + if source == "runtime": + feeds["indices"] = np.int64(1) + _checked_run(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"Output rank conversion.*Gather"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("indices_shape", [[2], ["K"]]) +def test_native_vector_gather_keeps_rank_four_output(indices_shape, monkeypatch): + """Keep native rank-four Gather when only its index-vector length is symbolic.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple((name, transform) for name, transform in pipeline.DLA_TRANSFORMS if name != "Gather"), + ) + shape = [1, 3, 4, 4] + model = _make_model( + [helper.make_node("Gather", ["x", "indices"], ["y"], axis=1)], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("indices", TensorProto.INT64, indices_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, indices_shape[0], 4, 4])], + ) + feeds = { + "x": np.arange(48, dtype=np.float32).reshape(shape), + "indices": np.array([0, 2], np.int64), + } + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("shape", [[3], [1, 3, 4, 4]]) +@pytest.mark.parametrize("tail", [False, True]) +def test_shape_outputs_keep_native_vector_rank(shape, tail): + """Retain native Shape vectors, including a one-element vector through unary tails.""" + nodes = [helper.make_node("Shape", ["x"], ["hidden" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["hidden"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.INT64, [len(shape)])], + ) + feeds = {"x": np.ones(shape, np.float32)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + actual = _checked_run(converted, feeds)[0] + assert actual.shape == (len(shape),) + np.testing.assert_array_equal(actual, expected) + + +def test_early_pattern_does_not_allow_an_unmatched_native_reduction(monkeypatch): + """Refuse a native reducer before mutation even when another reducer has a valid early lowering.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + shape = [1, 1, 2, 12] + model = _make_model( + [ + helper.make_node("Reshape", ["x", "target"], ["expanded"]), + helper.make_node("ReduceSum", ["expanded", "axes"], ["y"], keepdims=0), + helper.make_node("ReduceSum", ["x", "axes"], ["native"], keepdims=0), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("native", TensorProto.FLOAT, [1, 1, 2]), + ], + [ + numpy_helper.from_array(np.asarray([1, 1, 2, 3, 4], np.int64), "target"), + numpy_helper.from_array(np.asarray([-1], np.int64), "axes"), + ], + ) + _checked_run(model, {"x": np.ones(shape, np.float32)}) + before = model.SerializeToString() + with pytest.raises(NotImplementedError, match="reductions"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("width", [1, 4]) +@pytest.mark.parametrize("axis", [0, -2]) +def test_ordered_pipeline_preserves_vector_output(width, axis): + """Keep vector outputs, including a singleton vector, through the actual registered sequence.""" + model = _make_model( + [ + helper.make_node("MatMul", ["x", "w"], ["matrix"]), + helper.make_node("Squeeze", ["matrix", "axes"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [width])], + [ + numpy_helper.from_array(np.arange(3 * width, dtype=np.float32).reshape(3, width), "w"), + numpy_helper.from_array(np.asarray([axis], np.int64), "axes"), + ], + ) + feeds = {"x": np.asarray([[1, 2, 3]], np.float32)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + actual = _checked_run(transformed, feeds)[0] + assert actual.shape == (width,) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("axes", [[1], [-2], None]) +def test_ordered_pipeline_preserves_matrix_squeeze_axes(axes): + """Restore matrix outputs with positive, negative or omitted source squeeze axes.""" + controls = [] if axes is None else [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")] + squeeze_inputs = ["matrix"] if axes is None else ["matrix", "axes"] + model = _make_model( + [ + helper.make_node("MatMul", ["x", "w"], ["matrix"]), + helper.make_node("Squeeze", squeeze_inputs, ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 1, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4])], + [numpy_helper.from_array(np.arange(12, dtype=np.float32).reshape(3, 4), "w"), *controls], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 1, 3)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) + + +@pytest.mark.parametrize("axis", [0, 1, -1]) +def test_registered_concat_preserves_promoted_input_axes(axis): + """Shift Concat axes once after consistently promoting every native operand.""" + shape = [2, 3] + result_shape = shape.copy() + result_shape[axis] *= 2 + model = _make_model( + [helper.make_node("Concat", ["a", "b"], ["y"], axis=axis)], + [helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) for name in ("a", "b")], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, result_shape)], + ) + feeds = {"a": np.arange(6, dtype=np.float32).reshape(shape), "b": np.ones(shape, np.float32)} + expected = _checked_run(model, feeds)[0] + changed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(changed, feeds)[0], expected) + + +@pytest.mark.parametrize("axis", [0, -3]) +def test_normalization_public_squeeze_keeps_original_axes(axis): + """Retarget original batch squeezing together with the generated spatial singleton.""" + model = _make_model( + [ + helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"]), + helper.make_node("Squeeze", ["normalized", "axes"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, 4])], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + numpy_helper.from_array(np.array([axis], np.int64), "axes"), + ], + ) + feeds = {"x": np.arange(12, dtype=np.float32).reshape(1, 3, 4)} + expected = _checked_run(model, feeds)[0] + changed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(changed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +def test_standalone_intermediary_rank_change_preserves_concat_inputs(): + """Retain an internal rank-changing squeeze whose downstream consumer expects rank three.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["before"]), + helper.make_node("Squeeze", ["before", "axes"], ["squeezed"]), + helper.make_node("Concat", ["squeezed", "b"], ["y"], axis=0), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [2, 4, 4]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [5, 4, 4])], + [numpy_helper.from_array(np.asarray([0], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = { + "x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4), + "b": np.ones((2, 4, 4), np.float32), + } + expected = _checked_run(model, feeds)[0] + _apply_remove_intermediary_squeeze_unsqueeze(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + before = model.SerializeToString() + with pytest.raises(ValueError, match="Concat"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +def test_standalone_empty_axes_adapter_can_be_removed(): + """Retain cleanup of a genuine no-op adapter in standalone mode.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["before"]), + helper.make_node("Unsqueeze", ["before", "axes"], ["adapted"]), + helper.make_node("Identity", ["adapted"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4])], + [numpy_helper.from_array(np.asarray([], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + expected = _checked_run(model, feeds)[0] + _apply_remove_intermediary_squeeze_unsqueeze(model) + assert not any(node.op_type == "Unsqueeze" for node in model.graph.node) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_internal_normalization_boundary_restores_trailing_singleton(batch): + """Restore native normalization output shapes after preserving its channel dimension.""" + model = _make_model( + [ + helper.make_node( + "InstanceNormalization", ["x", "scale", "bias"], ["y"], name="normalize" + ) + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [batch, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [batch, 3, 4])], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + feeds = {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + _apply_fix_instancenorm_channels(model) + np.testing.assert_allclose(_checked_run(model, feeds)[0], expected, rtol=1e-6, atol=1e-6) + + +def test_unnamed_normalizations_keep_distinct_targets(): + """Give unnamed normalization inputs with different shapes independent reshape targets.""" + model = _make_model( + [ + helper.make_node("InstanceNormalization", ["a", "scale", "bias"], ["a_out"]), + helper.make_node("InstanceNormalization", ["b", "scale", "bias"], ["b_out"]), + ], + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [3, 3, 5]), + ], + [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("b_out", TensorProto.FLOAT, [3, 3, 5]), + ], + [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + ], + ) + feeds = { + "a": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "b": np.arange(45, dtype=np.float32).reshape(3, 3, 5), + } + expected = _checked_run(model, feeds) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + _apply_fix_instancenorm_channels(model) + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize( + "transform", + [ + _apply_non4d_model_inputs, + _apply_remove_intermediary_squeeze_unsqueeze, + _apply_fix_instancenorm_channels, + ], +) +def test_untouched_constant_payload_stays_in_place(transform): + """Preserve the original protobuf payload object when no node replacement touches it.""" + model = _make_model( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("Identity", ["x"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4])], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + ) + constant = model.graph.node[0] + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4)} + expected = _checked_run(model, feeds) + transform(model) + assert model.graph.node[0] is constant + for actual, reference in zip(_checked_run(model, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_input_adapter_runtime_axes_refuses_before_mutation(): + """Refuse baking an overridable axes default before applying earlier promotions.""" + model = _make_model( + [ + helper.make_node("Identity", ["a"], ["a_out"]), + helper.make_node("Unsqueeze", ["x", "axes"], ["y"]), + ], + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("axes", TensorProto.INT64, [1]), + ], + [ + helper.make_tensor_value_info("a_out", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["a", "b", "c"]), + ], + [numpy_helper.from_array(np.asarray([-1], np.int64), "axes")], + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="fixed"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("transform", [_apply_non4d_model_inputs, _apply_non4d_model_outputs]) +def test_legacy_boundary_adapter_schema_refuses_before_mutation(transform): + """Refuse input-form axes adapters on a legacy operator schema.""" + model = _make_model( + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model.opset_import[0].version = 11 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="13"): + options = {"defer_rank_conversion": True} if transform is _apply_non4d_model_outputs else {} + transform(model, **options) + assert model.SerializeToString() == before + + +def test_legacy_ir_boundary_adapters_register_initializer_inputs(): + """Keep both new boundary-axis initializers legal in an IR3 graph.""" + model = _make_model( + [helper.make_node("Neg", ["x"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + model.ir_version = 3 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_inputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +def _capture_branch(name, source, shape): + """Make a branch that reads an outer-scope value.""" + return helper.make_graph( + [helper.make_node("Identity", [source], [name])], + name, + [], + [helper.make_tensor_value_info(name, TensorProto.FLOAT, shape)], + ) + + +def test_input_promotion_captured_adapter_refuses_before_mutation(): + """Retain a captured adapter output instead of leaving its lexical read unbound.""" + model = _make_model( + [ + helper.make_node("Unsqueeze", ["x", "axes"], ["captured"]), + helper.make_node( + "If", + ["condition"], + ["y"], + then_branch=_capture_branch("then", "captured", [1, 1, 2, 3]), + else_branch=_capture_branch("else", "captured", [1, 1, 2, 3]), + ), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 2, 3])], + [numpy_helper.from_array(np.asarray([0, 1], np.int64), "axes")], + ) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="captured"): + _apply_non4d_model_inputs(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("rank_converted", [False, True]) +def test_intermediary_cleanup_preserves_captured_aliases(rank_converted): + """Keep the producer name captured by an If when a no-op adapter could be removed.""" + model = _make_model( + [ + helper.make_node("Identity", ["x"], ["captured"]), + helper.make_node("Unsqueeze", ["captured", "axes"], ["adapted"]), + helper.make_node( + "If", + ["condition"], + ["chosen"], + then_branch=_capture_branch("then", "captured", [1, 2, 3, 4]), + else_branch=_capture_branch("else", "captured", [1, 2, 3, 4]), + ), + helper.make_node("Add", ["chosen", "adapted"], ["y"]), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 3, 4]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 3, 4])], + [numpy_helper.from_array(np.asarray([], np.int64), "axes")], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4), "condition": np.bool_(True)} + expected = _checked_run(model, feeds)[0] + if rank_converted: + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _apply_remove_intermediary_squeeze_unsqueeze(model) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +def test_output_boundary_orders_a_captured_public_value_before_its_reader(): + """Make a restored public value available before an existing lexical consumer.""" + model = _make_model( + [ + helper.make_node("Neg", ["x"], ["y"]), + helper.make_node( + "If", + ["condition"], + ["unused"], + then_branch=_capture_branch("then", "y", [2, 3]), + else_branch=_capture_branch("else", "y", [2, 3]), + ), + ], + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("condition", TensorProto.BOOL, []), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3), "condition": np.bool_(True)} + expected = _checked_run(model, feeds)[0] + _apply_non4d_model_inputs(model) + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + np.testing.assert_array_equal(_checked_run(model, feeds)[0], expected) + + +@pytest.mark.parametrize( + ("axes", "output_shape"), + [([2, 0], [1, 2, 1, 3]), ([-2, -1], [2, 3, 1, 1]), ([2, -4], [1, 2, 1, 3])], +) +def test_input_adapters_resolve_unsqueeze_axes_against_output_rank(axes, output_shape): + """Preserve unordered positive, negative and mixed Unsqueeze axis semantics.""" + model = _make_model( + [helper.make_node("Unsqueeze", ["x", "axes"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _checked_run(model, feeds)[0] + transformed = _apply_non4d_model_inputs(model) + np.testing.assert_array_equal(_checked_run(transformed, feeds)[0], expected) + + +@pytest.mark.parametrize("batch", [1, 2, "batch"]) +@pytest.mark.parametrize("activations", [(), ("Relu",), ("Relu", "Tanh")]) +def test_registered_normalization_preserves_native_channel_axis(batch, activations): + """Preserve channels with direct inputs or inferable intermediate metadata.""" + shape = [batch, 3, 4] + output = "normalized" if activations else "y" + data = "activated_input" if activations else "x" + nodes = ([helper.make_node("Relu", ["x"], [data])] if activations else []) + [ + helper.make_node("InstanceNormalization", [data, "scale", "bias"], [output]) + ] + for index, operation in enumerate(activations): + following = "y" if index == len(activations) - 1 else f"activation_{index}" + nodes.append(helper.make_node(operation, [output], [following])) + output = following + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [ + numpy_helper.from_array(np.asarray([1, 2, 3], np.float32), "scale"), + numpy_helper.from_array(np.asarray([-1, 0, 1], np.float32), "bias"), + ], + ) + feed_shape = [2 if isinstance(batch, str) else batch, 3, 4] + feeds = {"x": np.random.default_rng(19).standard_normal(feed_shape).astype(np.float32)} + expected = _checked_run(model, feeds)[0] + transformed = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_allclose(_checked_run(transformed, feeds)[0], expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("entrypoint", ["standalone", "pipeline"]) +def test_native_reduction_output_refuses_before_adding_rank_adapters(entrypoint, monkeypatch): + """Refuse a native producer when no rank-conversion stage can prepare its output.""" + model = _make_model( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[2, 3], keepdims=0)], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 3])], + ) + before = model.SerializeToString() + if entrypoint == "standalone": + with pytest.raises(ValueError, match="rank"): + _apply_non4d_model_outputs(model) + else: + dla_pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + dla_pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in dla_pipeline.DLA_TRANSFORMS + if name != "StandaloneReduceOperations" + ), + ) + with pytest.raises(NotImplementedError, match="reductions"): + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("entrypoint", ["deferred", "pipeline"]) +def test_native_reshape_output_requires_registered_rank_stage(shape, tail, entrypoint, monkeypatch): + """Refuse low-rank Reshape restoration before changing public boundary state.""" + nodes = [helper.make_node("Reshape", ["x", "target"], ["native" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Relu", ["native"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1] * (4 - len(shape)) + shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(np.asarray(shape, np.int64), "target")], + ) + _checked_run(model, {"x": np.ones([1] * (4 - len(shape)) + shape, np.float32)}) + before = model.SerializeToString() + if entrypoint == "pipeline": + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) + for name, callback in pipeline.DLA_TRANSFORMS + if name != "Non4DReshape" + ), + ) + with pytest.raises(ValueError, match="rank"): + if entrypoint == "pipeline": + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shape", [[1], [3]]) +@pytest.mark.parametrize("tail", [False, True]) +def test_native_vector_unary_outputs_keep_their_rank(shape, tail): + """Leave rank-one unary paths native when boundary inputs are not widened.""" + nodes = [helper.make_node("Relu", ["x"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Neg", ["middle"], ["y"])) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + expected = _checked_run(model, feeds)[0] + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + + +def _rank_sensitive_adapter_model(operation, unary, metadata): + """Build an eligible prefix before a rank-sensitive branch to expose partial rewiring.""" + nodes = [ + helper.make_node("Relu", ["x"], ["a"]), + helper.make_node("Squeeze", ["a", "axes"], ["b"]), + helper.make_node("Relu", ["b"], ["early"]), + helper.make_node("Relu", ["x"], ["r"]), + helper.make_node("Squeeze", ["r", "axes"], ["s"]), + ] + if unary: + nodes.append(helper.make_node("Relu", ["s"], ["native"])) + nodes.append( + helper.make_node( + operation, + ["native" if unary else "s"], + ["y"], + **({"perm": [1, 2, 0]} if operation == "Transpose" else {"axis": 1}), + ) + ) + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [ + helper.make_tensor_value_info("early", TensorProto.FLOAT, [3, 4, 4]), + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [4, 4, 3] if operation == "Transpose" else [3, 4, 4] + ), + ], + [numpy_helper.from_array(np.array([0], np.int64), "axes")], + ) + if metadata: + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.linspace(0.1, 4, 48, dtype=np.float32).reshape(1, 3, 4, 4)} + return model, feeds + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax", "Transpose"]) +@pytest.mark.parametrize("unary", [False, True]) +@pytest.mark.parametrize( + ("entrypoint", "metadata"), [("stage", True), ("pipeline", True), ("pipeline", False)] +) +def test_rank_changing_adapter_requires_actual_consumer_preparation( + operation, unary, entrypoint, metadata, monkeypatch +): + """Refuse unsupported reader ranks before changing an earlier eligible adapter.""" + model, feeds = _rank_sensitive_adapter_model(operation, unary, metadata) + _checked_run(model, feeds) + pipeline = import_module(_transform_make_dla_compatible.__module__) + unavailable = "Non4DTranspose" if operation == "Transpose" else "UpdateAxis" + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) for name, callback in pipeline.DLA_TRANSFORMS if name != unavailable + ), + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank"): + if entrypoint == "stage": + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax", "Transpose"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_changing_adapter_uses_registered_consumer_preparation(operation, unary): + """Execute preserved adapter ownership through the matching later axis stage.""" + pipeline = import_module(_transform_make_dla_compatible.__module__) + stage = "Non4DTranspose" if operation == "Transpose" else "UpdateAxis" + if not any(name == stage for name, _ in pipeline.DLA_TRANSFORMS): + pytest.skip("The matching rank stage is introduced by a later series part") + model, feeds = _rank_sensitive_adapter_model(operation, unary, True) + expected = _checked_run(model, feeds) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("operation", ["LayerNormalization", "Gelu", "QuantizeLinear"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_changing_adapter_preserves_safe_normalization_readers(operation, unary): + """Execute suffix normalization and scalar quantization across leading singleton promotion.""" + model, feeds = _rank_sensitive_adapter_model("Softmax", unary, True) + model.ir_version = 9 + reader = model.graph.node[-1] + reader.op_type = operation + del reader.attribute[:] + if operation == "Gelu": + reader.domain = "com.microsoft" + model.opset_import.append(helper.make_opsetid("com.microsoft", 1)) + else: + reader.input.append("scale") + model.graph.initializer.append( + numpy_helper.from_array( + np.ones([4], np.float32) + if operation == "LayerNormalization" + else np.asarray(0.5, np.float32), + "scale", + ) + ) + if operation == "LayerNormalization": + reader.attribute.append(helper.make_attribute("axis", -1)) + else: + model.graph.output[-1].type.tensor_type.elem_type = TensorProto.UINT8 + + def execute(value): + """Use the CPU runtime, including its registered contrib Gelu implementation.""" + onnx.checker.check_model(value, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + return ort.InferenceSession( + value.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + expected = execute(model) + converted = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(execute(converted), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("unary", [False, True]) +@pytest.mark.parametrize("same_layout", [False, True]) +@pytest.mark.parametrize("entrypoint", ["stage", "pipeline"]) +def test_adapter_roundtrip_checks_the_restored_layout(unary, same_layout, entrypoint): + """Accept matching rank restoration while retaining refusal for a different bypass layout.""" + nodes = [ + helper.make_node("Relu", ["x"], ["private"]), + helper.make_node("Squeeze", ["private", "squeeze_axes"], ["squeezed"]), + ] + current = "squeezed" + if unary: + nodes.append(helper.make_node("Relu", [current], ["middle"])) + current = "middle" + nodes.extend( + [ + helper.make_node("Unsqueeze", [current, "unsqueeze_axes"], ["restored"]), + helper.make_node("Conv", ["restored", "weight"], ["y"]), + ] + ) + channels = 3 if same_layout else 1 + model = _make_model( + nodes, + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1 if same_layout else 3, 2, 4, 4])], + [ + numpy_helper.from_array(np.array([0], np.int64), "squeeze_axes"), + numpy_helper.from_array( + np.array([0 if same_layout else 1], np.int64), "unsqueeze_axes" + ), + numpy_helper.from_array( + np.arange(2 * channels, dtype=np.float32).reshape(2, channels, 1, 1), "weight" + ), + ], + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)} + expected = _checked_run(model, feeds) + before = model.SerializeToString() + if not same_layout: + with pytest.raises(ValueError, match="unprepared rank-sensitive Conv"): + if entrypoint == "stage": + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + else: + _transform_make_dla_compatible(model, verbose=False) + assert model.SerializeToString() == before + return + converted = ( + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + if entrypoint == "stage" + else _transform_make_dla_compatible(model, verbose=False) + ) + if entrypoint == "stage": + converted = infer_shapes(converted) + assert not any(node.op_type in {"Squeeze", "Unsqueeze"} for node in converted.graph.node) + for actual, reference in zip(_checked_run(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_rank_restoration_requires_an_actual_four_dimensional_bypass(): + """Do not infer a 4D bypass solely from the original adapter output annotation.""" + model = _make_model( + [ + helper.make_node("Reshape", ["x", "target"], ["wide"]), + helper.make_node("Squeeze", ["wide", "axes"], ["restored"]), + helper.make_node("Conv", ["restored", "weight"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 2, 4, 4])], + [ + numpy_helper.from_array(np.array([1, 1, 3, 4, 4], np.int64), "target"), + numpy_helper.from_array(np.array([0], np.int64), "axes"), + numpy_helper.from_array(np.ones([2, 3, 1, 1], np.float32), "weight"), + ], + ) + model = onnx.shape_inference.infer_shapes(model) + _checked_run(model, {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)}) + before = model.SerializeToString() + with pytest.raises(ValueError, match="unprepared rank-sensitive Conv"): + _apply_remove_intermediary_squeeze_unsqueeze(model, rank_converted=True) + assert model.SerializeToString() == before + + +def test_output_preflight_does_not_copy_for_missing_intermediate_squeeze_metadata(monkeypatch): + """Avoid a prerequisite model copy when normal shape inference can prepare Squeeze.""" + model = _make_model( + [ + helper.make_node("Relu", ["x"], ["hidden"]), + helper.make_node("Squeeze", ["hidden"], ["y"]), + ], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _checked_run(model, feeds)[0] + pipeline = import_module(_transform_make_dla_compatible.__module__) + original = pipeline.copy.deepcopy + copies = [] + + def record_copy(value, *args, **kwargs): + """Observe copies of the caller without changing the real fallback behavior.""" + if value is model: + copies.append(True) + return original(value, *args, **kwargs) + + monkeypatch.setattr(pipeline.copy, "deepcopy", record_copy) + converted = _transform_make_dla_compatible(model, verbose=False) + np.testing.assert_array_equal(_checked_run(converted, feeds)[0], expected) + assert copies == [] + + +@pytest.mark.parametrize("axes", [[0, 2], [0, -3]]) +def test_output_adapter_refuses_native_rank_five_squeeze_before_mutation(axes): + """Avoid remapping axes through an unspecified five-dimensional rank collapse.""" + model = _make_model( + [helper.make_node("Squeeze", ["x", "axes"], ["y"])], + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 2, 1, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])], + [numpy_helper.from_array(np.asarray(axes, np.int64), "axes")], + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="rank"): + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Flatten", "ArgMax", "ArgMin", "ReduceProd", "ReduceL2"]) +@pytest.mark.parametrize("tail", [False, True]) +@pytest.mark.parametrize("entrypoint", ["deferred", "pipeline"]) +def test_native_rank_reducing_outputs_require_registered_conversion( + operation, tail, entrypoint, monkeypatch +): + """Refuse native rank reducers before adding planned adapters to earlier valid outputs.""" + attributes = ( + {"axis": 1} + if operation == "Flatten" + else ( + {"axis": 1, "keepdims": 0} + if operation.startswith("Arg") + else {"axes": [2, 3], "keepdims": 0} + ) + ) + shape = ( + [1, 48] + if operation == "Flatten" + else ([1, 4, 4] if operation.startswith("Arg") else [1, 3]) + ) + nodes = [ + helper.make_node("Relu", ["x"], ["first"]), + helper.make_node(operation, ["z"], ["middle" if tail else "y"], **attributes), + ] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = _make_model( + nodes, + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 3, 4, 4]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info( + "y", TensorProto.INT64 if operation.startswith("Arg") else TensorProto.FLOAT, shape + ), + ], + ) + model.opset_import[0].version = 17 + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(2, 3), + "z": np.linspace(0.1, 0.9, 48, dtype=np.float32).reshape(1, 3, 4, 4), + } + _checked_run(model, feeds) + before = model.SerializeToString() + if entrypoint == "pipeline": + pipeline = import_module(_transform_make_dla_compatible.__module__) + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple( + (name, transform) + for name, transform in pipeline.DLA_TRANSFORMS + if name != operation + ), + ) + with pytest.raises(ValueError, match=r"rank.*conversion"): + _transform_make_dla_compatible(model, verbose=False) + else: + with pytest.raises(ValueError, match=r"rank.*conversion"): + _apply_non4d_model_outputs(model, defer_rank_conversion=True) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("other_normalized", "through_pipeline"), [(False, False), (False, True), (True, True)] +) +@pytest.mark.parametrize("unary", [False, True]) +def test_normalization_concat_requires_prepared_operands(other_normalized, through_pipeline, unary): + """Widen Concat only when every operand receives a normalization bridge.""" + shape = [2, 3, 4] + nodes = [helper.make_node("InstanceNormalization", ["x", "scale", "bias"], ["normalized"])] + controls = [ + numpy_helper.from_array(np.ones(3, np.float32), "scale"), + numpy_helper.from_array(np.zeros(3, np.float32), "bias"), + numpy_helper.from_array(np.ones([2, 5, 4], np.float32), "other"), + ] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + right = "other" + if other_normalized: + controls.extend( + [ + numpy_helper.from_array(np.ones(5, np.float32), "other_scale"), + numpy_helper.from_array(np.zeros(5, np.float32), "other_bias"), + ] + ) + nodes.append( + helper.make_node( + "InstanceNormalization", + ["other", "other_scale", "other_bias"], + ["other_normalized"], + ) + ) + right = "other_normalized" + left = "normalized" + if unary: + nodes.append(helper.make_node("Relu", [left], ["activated"])) + left = "activated" + nodes.append(helper.make_node("Concat", [left, right], ["y"], axis=1)) + model = _make_model( + nodes, inputs, [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 8, 4])], controls + ) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + expected = _checked_run(model, feeds) + before = model.SerializeToString() + if other_normalized: + result = _transform_make_dla_compatible(model, verbose=False) + for actual, reference in zip(_checked_run(result, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-5, atol=1e-6) + else: + with pytest.raises(ValueError, match="normalization Concat"): + if through_pipeline: + _transform_make_dla_compatible(model, verbose=False) + else: + _apply_fix_instancenorm_channels(model) + assert model.SerializeToString() == before