diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py index dfb9d96ad9e..0c48ff96fcc 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_indexing.py @@ -144,7 +144,9 @@ def _gather_layout_reinterpreted(cache, name, supported, captured): return True -def _prepared_index_data(model, cache, prepared_ops, data_names, *, before_boundary=False): +def _prepared_index_data( + model, cache, prepared_ops, data_names, *, before_boundary=False, restoring_outputs=False +): """Prove leading-4D data from actual layouts and registered rank stages without mutation.""" standard = {id(node) for node in _operator_nodes(model, {n.op_type for n in model.graph.node})} if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): @@ -203,6 +205,8 @@ def classify(name): ): return True tensor = cache.get_init(name) + if tensor is not None and len(tensor.dims) == 4: + return True if name in cache.graph_input_names and (model.ir_version > 3 or tensor is None): return False if tensor is not None: @@ -239,19 +243,41 @@ def classify(name): and target[0] in (1, 2, 3, 5) and not _has_native_symbolic_data(cache, node.input[0], local_functions) ) + if node.op_type == "Expand": + target = cache.get_shape(node.input[1]) + return target == [4] or ( + "Expand" in prepared_ops + and target is not None + and len(target) == 1 + and 0 < target[0] < 4 + and not _has_native_symbolic_data(cache, node.input[0], local_functions) + ) if node.op_type == "Flatten": return node.op_type in prepared_ops if node.op_type in {"ReduceSum", "ReduceMean", "ReduceMax", "ReduceMin"}: return [node.input[0]] if node.op_type in prepared_ops else False if node.op_type == "InstanceNormalization": - return "FixInstanceNormChannelMismatch" in prepared_ops + return ( + "FixInstanceNormChannelMismatch" in prepared_ops + or "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_ops + ) if node.op_type == "GroupQueryAttention": return "GroupQueryAttention" in prepared_ops if node.op_type == "Conv": - kernel = cache.get_shape(node.input[1]) - return kernel is not None and len(kernel) == 4 - if not before_boundary and node.op_type == "MatMul": - if name in cache.graph_output_names | captured: + weight = cache.get_init(node.input[1]) + if weight is not None: + return len(weight.dims) == 4 + kernel_rank = _known_rank(cache, node.input[1]) + if kernel_rank is not None: + return kernel_rank == 4 + return [node.input[1]] + if node.op_type == "MatMul" and ( + not before_boundary + or restoring_outputs + or {"MatMulToTransposeConvTranspose", "RemoveIntermediarySqueezeAndUnsqueeze"} + <= set(prepared_ops) + ): + if not restoring_outputs and name in cache.graph_output_names | captured: return False ranks = [matrix_rank(value) for value in node.input[:2]] return ( @@ -270,7 +296,9 @@ def classify(name): if "RemoveIntermediarySqueezeAndUnsqueeze" not in prepared_ops: return False source = node.input[0] - if source in captured or name in captured | cache.graph_output_names: + if source in captured or ( + not restoring_outputs and name in captured | cache.graph_output_names + ): return False if source in cache.graph_input_names: original = cache.get_shape(source) @@ -300,6 +328,8 @@ def classify(name): [node.input[0]] if not before_boundary or value < 0 + or "axis_updated" in node.name + or "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_ops or node.op_type in prepared_ops or node.input[0] in cache.graph_input_names else False diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py index bcf27d5db2c..37391c2f6d0 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_model_boundaries.py @@ -37,7 +37,21 @@ class _BoundaryPreflightError(ValueError): """Prerequisite lowering cannot freeze an exposed runtime control.""" -_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | {"Slice", "Transpose", "Reshape", "Expand", "Tile"} +_RANK_FLOW_OPS = _RANK_AGNOSTIC_OPS | { + "Slice", + "Transpose", + "Reshape", + "Expand", + "Tile", + "Softmax", + "LogSoftmax", + "LpNormalization", + "ArgMax", + "QuantizeLinear", + "DequantizeLinear", + "TopK", + "Split", +} _REDUCTION_OPS = frozenset({"ReduceSum", "ReduceMean", "ReduceMin", "ReduceMax"}) @@ -114,10 +128,14 @@ def _output_rank_source(cache, node, local_functions, sources): return node -def _has_native_symbolic_data(cache, name, local_functions): +def _has_native_symbolic_data(cache, name, local_functions, sources=None): """Prove a low-rank symbolic data path was left native by boundary promotion.""" seen = set() + sources = {} if sources is None else sources while name not in seen: + if name in sources: + result = sources[name] + break seen.add(name) info = cache.vi_map.get(name) rank = ( @@ -125,19 +143,25 @@ def _has_native_symbolic_data(cache, name, local_functions): if info is not None and info.type.tensor_type.HasField("shape") else None ) - if cache.get_shape(name) is not None or (rank is not None and rank >= 4): - return False + if rank is not None and rank >= 4: + result = False + break if name in cache.graph_input_names: - return rank is not None and rank < 4 + result = cache.get_shape(name) is None and rank is not None and rank < 4 + break node = cache.get_producer(name) if ( node is None or not _is_standard_node(node, local_functions, _RANK_FLOW_OPS) or not node.input ): - return False + result = False + break name = node.input[0] - raise ValueError("Cyclic data path while planning native symbolic controls") + else: + raise ValueError("Cyclic data path while planning native symbolic controls") + sources.update(dict.fromkeys(seen, result)) + return result def _require_adapter_opset(model): @@ -194,6 +218,7 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal unsqueezed_shapes = {} normalization_axes = {} local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} contrib_ops = _boundary_contrib_ops(model) opset = next( (entry.version for entry in model.opset_import if entry.domain in {"", "ai.onnx"}), 1 @@ -214,6 +239,16 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal ) and node.op_type not in {"InstanceNormalization", "Gemm"} and not _native_conv1d(node, cache) + and not ( + node.op_type in _RANK_AGNOSTIC_OPS - {"Clip"} + and len(node.input) > 1 + and any( + name != value.name + and _has_native_symbolic_data(cache, name, local_functions, native_sources) + for name in node.input + if name + ) + ) ] if not consumers: continue @@ -367,45 +402,73 @@ def _apply_non4d_model_inputs(model, *, prepared_input_ops=(), validate_only=Fal def _apply_non4d_model_outputs( - model, *, defer_rank_conversion=False, prepared_output_ops=(), validate_only=False + model, + *, + defer_rank_conversion=False, + prepared_output_ops=(), + prepared_data_ops=None, + validate_only=False, ): """Restore outputs or only validate producer rank availability before shape inference.""" + # Rank stages import the boundary predicates; defer their shared proof to avoid the cycle. + from .dla_rank_conversion import _rank_preparation_proofs + graph = model.graph cache = GraphCache(graph) + prepared = ( + _rank_preparation_proofs(model, cache, prepared_data_ops, restoring_outputs=True) + if defer_rank_conversion + else {} + ) captured = _subgraph_captures(graph) boundary_names = cache.graph_input_names | cache.init_names append = [] plans = [] local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} - native_broadcasts = set() - if defer_rank_conversion and "Non4dInitializers" not in prepared_output_ops: - live_ranks = { - name: len(shape) - for name in boundary_names - if (shape := cache.get_shape(name)) is not None - } - for node in graph.node: - if ( - _is_standard_node(node, local_functions, _RANK_AGNOSTIC_OPS) - and node.input - and all(name in live_ranks for name in node.input) - ): - rank = max(live_ranks[name] for name in node.input) - live_ranks.update(dict.fromkeys(node.output, rank)) - if rank < 4 and ( - len(node.input) > 1 or any(name in native_broadcasts for name in node.input) - ): - native_broadcasts.update(node.output) rank_sources = {} contrib_ops = _boundary_contrib_ops(model) + native_sources = {} + inputs_promotable = all( + cache.get_shape(value.name) is not None + for value in graph.input + if value.name not in cache.init_names + ) for value in graph.output: - shape = cache.get_shape(value.name) + shape: list | None = cache.get_shape(value.name) + symbolic_output = shape is None + if ( + shape is None + and defer_rank_conversion + and inputs_promotable + and value.type.tensor_type.HasField("shape") + ): + shape = [ + dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None + for dim in value.type.tensor_type.shape.dim + ] 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 symbolic_output: + if source is None or not _is_standard_node(source, local_functions, _RANK_FLOW_OPS): + continue + native_shape = _shape_before_boundary_promotion(cache, source.input[0]) + info = cache.vi_map.get(source.input[0]) + rank = ( + len(native_shape) + if native_shape is not None + else ( + len(info.type.tensor_type.shape.dim) + if info is not None and info.type.tensor_type.HasField("shape") + else None + ) + ) + target = cache.get_init(source.input[1]) if source.op_type == "Reshape" else None + if rank not in (2, 3) and not ( + target is not None and len(target.dims) == 1 and target.dims[0] < 4 + ): + continue if source is None and len(shape) == 1: continue if source is not None: @@ -419,8 +482,10 @@ def _apply_non4d_model_outputs( source, local_functions, {source.op_type}, contrib_ops=contrib_ops ): continue - if _is_standard_node(source, local_functions, {"Reshape", "Expand", "Tile"}): - if _has_native_symbolic_data(cache, source.input[0], local_functions): + if _is_standard_node(source, local_functions, _RANK_FLOW_OPS): + if _has_native_symbolic_data( + cache, source.input[0], local_functions, native_sources + ): continue target = cache.get_shape(source.input[1]) if source.op_type == "Reshape" else None if ( @@ -472,6 +537,8 @@ def _apply_non4d_model_outputs( if delta < 0: raise ValueError(f"Output Squeeze input rank exceeds 4 for {value.name!r}") axes = list(range(delta)) + [axis + delta if axis >= 0 else axis for axis in axes] + if defer_rank_conversion and not prepared[value.name]: + continue plans.append((value, shape, existing, axes)) if plans: _require_adapter_opset(model) @@ -479,7 +546,7 @@ def _apply_non4d_model_outputs( return model for value, shape, existing, axes in plans: axes_name = None - if len(shape) > 1 or shape == [1]: + if len(shape) > 1 or shape == [1] or any(not isinstance(dim, int) for dim in shape): axes_name = add_initializer( graph, cache.unique_name(value.name + "_squeeze_axes"), diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py new file mode 100644 index 00000000000..4a771d8b35e --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_rank_conversion.py @@ -0,0 +1,830 @@ +# 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. + +"""Ordered rank and axis passes sharing the preceding native shape snapshot. + +Do not infer shapes between stages: original axes precede upstream rank promotion. +""" + +import math +from functools import partial + +import numpy as np +from onnx import TensorProto, defs, helper, numpy_helper + +from ..utils.graph_utils import ( + GraphCache, + _attribute, + _shape_before_boundary_promotion, + add_initializer, + add_value_info, + pad4d, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _RANK_AGNOSTIC_OPS, _subgraph_captures +from .dla_indexing import _insert_support_nodes, _known_rank, _operator_nodes, _prepared_index_data +from .dla_model_boundaries import ( + _REDUCTION_OPS, + _has_native_symbolic_data, + _is_standard_node, + _native_conv1d, +) + +__all__ = [] + + +def _runtime_input_names(model): + """Exclude fixed initializer declarations required by legacy ONNX IR.""" + names = {value.name for value in model.graph.input} + if model.ir_version <= 3: + names.difference_update(tensor.name for tensor in model.graph.initializer) + return names + + +def _rank_preparation_proofs( + model, cache, prepared_data_ops, *, restoring_outputs=False, before_boundary=None +): + """Reuse the cached leading-layout proof for rank and axis stages.""" + # Standalone stages retain deferred inputs; registered calls inspect current layouts. + if before_boundary is None: + before_boundary = prepared_data_ops is None + if prepared_data_ops is None: + prepared_data_ops = ( + _RANK_AGNOSTIC_OPS + | _REDUCTION_OPS + | { + "Non4dInitializers", + "Reshape", + "Expand", + "Tile", + "Transpose", + "Slice", + "Flatten", + "LpNormalization", + "ArgMax", + "TopK", + "Concat", + "Split", + "QuantizeLinear", + "DequantizeLinear", + "Gather", + "GatherElements", + "MatMulToTransposeConvTranspose", + "RemoveIntermediarySqueezeAndUnsqueeze", + } + ) + return _prepared_index_data( + model, + cache, + prepared_data_ops, + [node.input[0] for node in model.graph.node if node.input] + + [value.name for value in model.graph.output], + before_boundary=before_boundary, + restoring_outputs=restoring_outputs, + ) + + +def _promoted_input_shape(cache, name, local_functions, native_sources, prepared): + """Distinguish native symbolic paths from inputs prepared for four-dimensional data.""" + shape = _shape_before_boundary_promotion(cache, name) + snapshot = cache._initializer_vi_map.get(name) + rank = len(shape) if shape is not None else _known_rank(cache, name) + if snapshot is not None and snapshot.type.tensor_type.HasField("shape"): + rank = len(snapshot.type.tensor_type.shape.dim) + current, seen = name, set() + while rank is None and current not in seen: + seen.add(current) + producer = cache.get_producer(current) + if producer is None or not _is_standard_node( + producer, + local_functions, + {"Reshape", "Slice", "QuantizeLinear", "DequantizeLinear"} | _REDUCTION_OPS, + ): + break + keepdims = _attribute(producer, "keepdims") + if producer.op_type in _REDUCTION_OPS and keepdims is not None and keepdims.i == 0: + break + if producer.op_type == "Reshape": + target_shape = cache.get_shape(producer.input[1]) + if target_shape is not None and len(target_shape) == 1: + rank = target_shape[0] + break + current = producer.input[0] + rank = _known_rank(cache, current) + native = ( + rank is not None + and rank < 4 + and ( + not prepared[name] + or _has_native_symbolic_data(cache, name, local_functions, native_sources) + ) + ) + return shape, rank, native + + +def _apply_non4d_slice_axis(model, *, prepared_data_ops=None): + """Remap fixed Slice axes before their data rank is promoted.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs( + model, + cache, + prepared_data_ops, + before_boundary=prepared_data_ops is None or "Non4dModelInputs" in prepared_data_ops, + ) + runtime_inputs = _runtime_input_names(model) + plans = [] + for node in _operator_nodes(model, {"Slice"}): + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native: + continue + if rank is None: + raise ValueError(f"{node.name!r}: Slice input rank must be known") + if rank not in (2, 3): + continue + if len(node.input) < 4 or not node.input[3]: + starts_shape = cache.get_shape(node.input[1]) + if starts_shape is None: + raise ValueError( + f"{node.name!r}: implicit Slice axes require a known starts length" + ) + values = np.arange(math.prod(starts_shape), dtype=np.int64) + 4 - rank + plans.append((node, node.input[1], values)) + continue + axes = cache.get_init(node.input[3]) + if axes is None or axes.name in runtime_inputs: + raise ValueError(f"{node.name!r}: Slice axes must be fixed") + values = numpy_helper.to_array(axes).reshape(-1) + if values.size and np.any(values >= 0): + plans.append((node, axes.name, np.where(values >= 0, values + 4 - rank, values))) + views = {} + for node, name, values in plans: + key = (name, tuple(values)) + if key not in views: + views[key] = add_initializer( + model.graph, + cache.unique_name(name + "_axes_only"), + np.asarray(values, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + if len(node.input) < 4: + node.input.append(views[key]) + else: + node.input[3] = views[key] + _add_legacy_initializer_inputs(model) + return model + + +def _apply_non4d_reshape(model): + """Prepare fixed Reshape targets for the deferred four-dimensional layout.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + runtime_inputs = _runtime_input_names(model) + reshapes = _operator_nodes(model, {"Reshape"}) + gathers = {id(node) for node in _operator_nodes(model, {"Gather"})} + squeezes = {id(node) for node in _operator_nodes(model, {"Squeeze"})} + replacements = {} + users = {} + captured = _subgraph_captures(model.graph) + for node in reshapes: + if any( + index == 1 and id(user) in gathers + for user, index in cache.get_consumers(node.output[0]) + ): + continue + tensor = cache.get_init(node.input[1]) + target_shape = cache.get_shape(node.input[1]) + native_input = _shape_before_boundary_promotion(cache, node.input[0]) + if target_shape == [4] and native_input is not None and len(native_input) == 4: + continue + rank = len(native_input) if native_input is not None else _known_rank(cache, node.input[0]) + if rank is None: + raise ValueError(f"{node.name!r}: Reshape input rank must be known") + if rank < 4 and _has_native_symbolic_data( + cache, node.input[0], local_functions, native_sources + ): + continue + if tensor is None or tensor.name in runtime_inputs: + raise ValueError(f"{node.name!r}: Reshape target must be fixed") + values = numpy_helper.to_array(tensor).reshape(-1).tolist() + if len(values) < 4: + native_rank = len(values) + values = pad4d(values) + output_shape = cache.get_shape(node.output[0]) + if output_shape is None: + aliases = [ + cache.get_shape(user.output[0]) + for user, index in cache.get_consumers(node.output[0]) + if index == 0 + and id(user) in squeezes + and user.output[0] in cache.graph_output_names + ] + aliases = [ + shape for shape in aliases if shape is not None and len(shape) == native_rank + ] + if aliases and all(shape == aliases[0] for shape in aliases): + output_shape = aliases[0] + if output_shape: + output_shape = pad4d(output_shape) + values = [ + output_shape[i] if value == 0 else value for i, value in enumerate(values) + ] + if 0 in values: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") + elif len(values) == 5: + if 0 in values: + output_shape = cache.get_shape(node.output[0]) + if output_shape is None or len(output_shape) != 5: + raise ValueError(f"{node.name!r}: cannot resolve 0 in 5D Reshape target") + values = [ + output_shape[i] if value == 0 else value for i, value in enumerate(values) + ] + if 1 not in values: + raise ValueError(f"{node.name!r}: 5D Reshape shape has no unary dimension") + values.remove(1) + elif len(values) == 4 and 0 in values: + output_shape = cache.get_shape(node.output[0]) + if output_shape is None or len(output_shape) != 4: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") + values = [output_shape[i] if value == 0 else value for i, value in enumerate(values)] + if 0 in values: + raise ValueError(f"{node.name!r}: cannot resolve 0 in Reshape target") + else: + continue + replacement = numpy_helper.from_array(np.asarray(values, dtype=np.int64), tensor.name) + if tensor.name in replacements and replacements[tensor.name] != replacement: + raise ValueError(f"Shared Reshape target {tensor.name!r} requires different shapes") + replacements[tensor.name] = replacement + users.setdefault(tensor.name, []).append(node) + for name, replacement in replacements.items(): + converted = {id(user) for user in users[name]} + private = name in cache.graph_output_names | cache.graph_input_names | captured or any( + index != 1 or id(user) not in converted for user, index in cache.get_consumers(name) + ) + if private: + target = add_initializer( + model.graph, + cache.unique_name(name + "_reshape_shape"), + numpy_helper.to_array(replacement), + dtype=TensorProto.INT64, + with_value_info=False, + ) + for user in users[name]: + user.input[1] = target + else: + tensor = cache.get_init(name) + assert tensor is not None + tensor.CopyFrom(replacement) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_non4d_initializers(model): + """Promote fixed weight layouts while preserving runtime input defaults.""" + expand, skip, front, avoid = set(), set(), set(), set() + supported = _operator_nodes(model, {node.op_type for node in model.graph.node}) + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): + supported += _operator_nodes( + model, {"QuantizeLinear", "DequantizeLinear", "GroupQueryAttention"}, ("com.microsoft",) + ) + selected = {id(node) for node in supported} + cache = GraphCache(model.graph) + avoid.update(name for node in supported if _native_conv1d(node, cache) for name in node.input) + protected = _subgraph_captures(model.graph) | {value.name for value in model.graph.output} + protected.update( + name for node in model.graph.node if id(node) not in selected for name in node.input + ) + protected.update(name for node in supported if node.op_type == "Gemm" for name in node.input) + used = ( + {name for node in supported for name in node.input} + if any(tensor.name in protected for tensor in model.graph.initializer) + else set() + ) + for node in supported: + if node.op_type in {"Div", "Sub", "Mul", "Add", "LSTM"}: + expand.update(node.input) + if node.op_type in {"MatMul", "Gemm"}: + skip.update(node.input) + if node.op_type in {"Where", "Gather", "Split"}: + front.update(node.input) + if node.op_type in {"Unsqueeze", "GroupQueryAttention"}: + avoid.update(node.input) + updates = [] + avoid.update(_runtime_input_names(model)) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + broadcasts = _RANK_AGNOSTIC_OPS - {"Clip"} + arithmetic = {"Add", "Sub", "Mul", "Div"} + private_users = {} + for tensor in model.graph.initializer: + consumers = cache.get_consumers(tensor.name) + if any( + user.op_type == "Shape" + or (port > 0 and user.op_type not in broadcasts | {"Concat", "MatMul", "LSTM"}) + for user, port in consumers + ): + readers = [(user, port) for user, port in consumers if user.op_type in arithmetic] + if readers: + private_users[tensor.name] = readers + native_initializers = { + name + for node in supported + if node.op_type in broadcasts + and len(node.input) > 1 + and any( + value not in cache.init_names + and _has_native_symbolic_data(cache, value, local_functions, native_sources) + for value in node.input + if value + ) + for name in node.input + if name in cache.init_names + } + for name in native_initializers: + tensor = cache.get_init(name) + assert tensor is not None + if not 0 < len(tensor.dims) < 4: + continue + for user, port in cache.get_consumers(name): + if ( + id(user) in selected + and (port == 0 or user.op_type in {"Concat", "MatMul"}) + and user.op_type not in {"Gemm", "Unsqueeze", "GroupQueryAttention", "Shape"} + and not _native_conv1d(user, cache) + and not (user.op_type in broadcasts and len(user.input) > 1) + and not ( + user.op_type == "Identity" + and all(output in protected for output in user.output) + ) + ): + raise ValueError( + f"Initializer {name!r} has incompatible native symbolic and rank-prepared usage" + ) + avoid.update(native_initializers) + expand_or_front = expand | front + for tensor in model.graph.initializer: + if tensor.name in avoid: + continue + rank = len(tensor.dims) + promote = ( + (rank == 1 and tensor.name in expand) + or (rank == 2 and tensor.name in expand_or_front) + or rank == 3 + ) + if tensor.name in protected: + if tensor.name in used and promote: + raise ValueError( + f"Initializer {tensor.name!r} has incompatible public, opaque or captured usage" + ) + continue + if promote: + shape = pad4d(list(tensor.dims)) + updates.append((tensor, shape)) + elif rank == 2 and tensor.name not in skip: + raise ValueError(f"Initializer {tensor.name!r} has unsupported rank 2 usage") + value_info = {value.name for value in model.graph.value_info} + native_inputs = { + node.input[0] + for node in model.graph.node + if node.input and node.op_type not in _RANK_AGNOSTIC_OPS + } + legacy_inputs = ( + {value.name: value for value in model.graph.input} if model.ir_version <= 3 else {} + ) + views = [] + for tensor, shape in updates: + if tensor.name in private_users: + axes = add_initializer( + model.graph, + cache.unique_name(tensor.name + "_4d_axes"), + np.arange(4 - len(tensor.dims), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + view = cache.unique_name(tensor.name + "_4d") + views.append( + helper.make_node( + "Unsqueeze", [tensor.name, axes], [view], name=cache.unique_name(view + "_view") + ) + ) + add_value_info(model.graph, view, tensor.data_type, shape) + for user, port in private_users[tensor.name]: + user.input[port] = view + continue + if tensor.name in native_inputs and tensor.name not in value_info: + # Later axis stages read native ranks from the same deferred metadata snapshot. + model.graph.value_info.append( + helper.make_tensor_value_info(tensor.name, tensor.data_type, list(tensor.dims)) + ) + del tensor.dims[:] + tensor.dims.extend(shape) + if tensor.name in legacy_inputs: + legacy_inputs[tensor.name].type.tensor_type.shape.CopyFrom( + helper.make_tensor_value_info( + tensor.name, tensor.data_type, shape + ).type.tensor_type.shape + ) + if views: + _insert_support_nodes(model.graph, [*views, *model.graph.node]) + _add_legacy_initializer_inputs(model) + return model + + +def _rank_one_path_widened(cache, name, rank_preserving, reshapes): + """Recognize earlier 4D initializer and Reshape promotion through rank-preserving paths.""" + pending, seen = [name], set() + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + tensor = cache.get_init(current) + if tensor is not None: + if len(tensor.dims) == 4: + return True + continue + producer = cache.get_producer(current) + if producer is not None and id(producer) in reshapes: + if cache.get_shape(producer.input[1]) == [4]: + return True + continue + if producer is not None and id(producer) in rank_preserving: + pending.extend(producer.input) + return False + + +def _pad_shape_inputs(model, op_type, suffix, *, prepared_data_ops=None): + """Preserve native rank-four layouts and plan fixed control padding otherwise.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + runtime_inputs = _runtime_input_names(model) + rank_preserving = {id(node) for node in _operator_nodes(model, _RANK_AGNOSTIC_OPS)} + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + plans = [] + for node in _operator_nodes(model, {op_type}): + tensor = cache.get_init(node.input[1]) + native = _shape_before_boundary_promotion(cache, node.input[0]) + rank = len(native) if native is not None else _known_rank(cache, node.input[0]) + if rank == 4 and cache.get_shape(node.input[1]) == [4]: + continue + if rank is None: + raise ValueError(f"{node.name!r}: {op_type} input rank must be known") + if op_type == "Tile" and rank == 1: + if not _rank_one_path_widened(cache, node.input[0], rank_preserving, reshapes): + continue + if op_type == "Tile" and rank < 4 and not prepared[node.input[0]]: + continue + if rank < 4 and _has_native_symbolic_data( + cache, node.input[0], local_functions, native_sources + ): + continue + if tensor is None or tensor.name in runtime_inputs: + raise ValueError(f"{node.name!r}: {op_type} target must be fixed") + values = numpy_helper.to_array(tensor).reshape(-1).tolist() + if len(values) < 4: + plans.append((node, tensor.name, pad4d(values))) + elif len(values) > 4: + raise ValueError(f"{node.name!r}: {op_type} target has rank {len(values)}") + views = {} + for node, name, values in plans: + if name not in views: + views[name] = add_initializer( + model.graph, + cache.unique_name(name + suffix), + np.asarray(values, dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.input[1] = views[name] + _add_legacy_initializer_inputs(model) + return model + + +_apply_non4d_expand = partial(_pad_shape_inputs, op_type="Expand", suffix="_expand_shape") +_apply_non4d_tile = partial(_pad_shape_inputs, op_type="Tile", suffix="_tile_shape") + + +def _collapsed_reshape_axes(cache, name, shape, reshapes): + """Identify singleton axes consistent with an actually collapsed native Reshape.""" + producer = cache.get_producer(name) + target = ( + cache.get_init_array(producer.input[1]) + if producer is not None and id(producer) in reshapes + else None + ) + if target is None or target.shape != (4,): + return [] + values = target.tolist() + return [ + axis + for axis, dim in enumerate(shape) + if dim == 1 + and all( + value in (-1, native) for value, native in zip(values, shape[:axis] + shape[axis + 1 :]) + ) + ] + + +def _apply_non4d_transpose(model, *, prepared_data_ops=None): + """Remap transpose permutations around added or removed unary axes.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + plans = [] + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + for node in _operator_nodes(model, {"Transpose"}): + perm = _attribute(node, "perm") + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native: + continue + if perm is None and rank is None: + raise ValueError(f"{node.name!r}: default Transpose requires a known input rank") + values = list(perm.ints) if perm is not None else list(reversed(range(rank))) + if len(values) in (2, 3): + delta = 4 - len(values) + values = list(range(delta)) + [v + delta for v in values] + elif len(values) == 5: + if shape is None or len(shape) != 5 or 1 not in shape: + raise ValueError(f"{node.name!r}: 5D Transpose requires a known unary dimension") + removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not removed_axes: + continue + if len(removed_axes) != 1: + raise ValueError(f"{node.name!r}: ambiguous collapsed 5D Transpose layout") + removed = removed_axes[0] + values = [v - int(v > removed) for v in values if v != removed] + else: + continue + plans.append((node, perm, values)) + for node, perm, values in plans: + if perm is None: + node.attribute.append(helper.make_attribute("perm", values)) + else: + del perm.ints[:] + perm.ints.extend(values) + return model + + +def _specialized_collapsed_axis(cache, node, shape, value, reshapes): + """Map a specialized operator axis only when its native singleton was actually collapsed.""" + removed = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not removed: + return value + if len(removed) != 1 or value % 5 in removed: + raise ValueError( + f"{node.name!r}: cannot map axis {value} through an ambiguous collapsed Reshape" + ) + native = value % 5 + return native - int(native > removed[0]) - (4 if value < 0 else 0) + + +def _apply_non4d_lpnorm(model, *, prepared_data_ops=None): + """Preserve normalization axes across leading-dimension promotion.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + plans = [] + for node in _operator_nodes(model, {"LpNormalization"}): + axis = _attribute(node, "axis") + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native: + continue + if rank is None or (rank == 5 and shape is None): + raise ValueError(f"{node.name!r}: LpNormalization input layout must be known") + value = axis.i if axis is not None else -1 + if shape and len(shape) == 5: + value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) + elif value >= 0 and rank in (2, 3): + value = -1 if value == rank - 1 else value + 4 - rank + plans.append((node, axis, value)) + for node, axis, value in plans: + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name += "_axis_updated" + return model + + +def _apply_flatten(model, *, prepared_data_ops=None): + """Replace Flatten with its four-dimensional equivalent shape.""" + graph = model.graph + cache = GraphCache(graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + plans = [] + for node in _operator_nodes(model, {"Flatten"}): + shape, _, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native and prepared_data_ops is not None: + continue + if any(reader.op_type == "Gemm" for reader, _ in cache.get_consumers(node.output[0])): + raise ValueError(f"{node.name!r}: Flatten output still feeds an unconverted Gemm") + axis = _attribute(node, "axis") + axis = axis.i if axis is not None else 1 + if shape is None or any(d <= 0 for d in shape): + raise ValueError(f"{node.name!r}: Flatten requires a known input shape") + if not -len(shape) <= axis <= len(shape): + raise ValueError(f"{node.name!r}: invalid Flatten axis {axis}") + if axis < 0: + axis += len(shape) + plans.append((node, math.prod(shape[:axis]))) + targets = {} + for node, batch in plans: + if batch not in targets: + hint = "flatten_reshape_shape" if batch == 1 else f"{node.name}_reshape_shape" + targets[batch] = add_initializer( + graph, + cache.unique_name(hint), + np.asarray([1, 1, batch, -1], dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.CopyFrom( + helper.make_node( + "Reshape", + [node.input[0], targets[batch]], + list(node.output), + name=cache.unique_name((node.name or node.output[0]) + "_reshape"), + ) + ) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_argmax(model, *, prepared_data_ops=None): + """Plan native ArgMax layouts before preparing private output restoration views.""" + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + plans = {} + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + for node in _operator_nodes(model, {"ArgMax"}): + shape, rank, native = _promoted_input_shape( + cache, node.input[0], local_functions, native_sources, prepared + ) + if native: + continue + keepdims = _attribute(node, "keepdims") + if rank == 4 and (keepdims is None or keepdims.i != 0): + continue + if not shape or any(dim <= 0 for dim in shape): + raise ValueError(f"{node.name!r}: ArgMax requires a static input shape") + axis = _attribute(node, "axis") + value = axis.i if axis is not None else 0 + if not -len(shape) <= value < len(shape): + raise ValueError(f"{node.name!r}: invalid ArgMax axis {value}") + output_shape = [dim for index, dim in enumerate(shape) if index != value % len(shape)] + if len(shape) == 5: + value = _specialized_collapsed_axis(cache, node, shape, value, reshapes) + elif value >= 0: + value += {2: 2, 3: 1}.get(len(shape), 0) + plans[id(node)] = (value, output_shape) + nodes = [] + for node in model.graph.node: + nodes.append(node) + if id(node) not in plans: + continue + value, output_shape = plans[id(node)] + axis = _attribute(node, "axis") + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + keepdims = _attribute(node, "keepdims") + restore_shape = keepdims is not None and keepdims.i == 0 + if keepdims is not None: + keepdims.i = 1 + original = node.output[0] + node.name = cache.unique_name((node.name or original) + "_axis_updated") + if restore_shape: + target = add_initializer( + model.graph, + cache.unique_name(node.name + "_reshape_shape"), + np.asarray(pad4d(output_shape), dtype=np.int64), + dtype=TensorProto.INT64, + with_value_info=False, + ) + node.output[0] = cache.unique_name(original + "_keepdims") + nodes.append( + helper.make_node( + "Reshape", + [node.output[0], target], + [original], + name=cache.unique_name(node.name + "_reshape"), + ) + ) + _insert_support_nodes(model.graph, nodes) + _add_legacy_initializer_inputs(model) + return model + + +def _apply_update_axis(model, *, prepared_data_ops=None): + """Remap supported operation axes from their native input ranks.""" + quantizers = {"QuantizeLinear", "DequantizeLinear"} + update = {"LogSoftmax", "Softmax", "TopK", "Concat", "Split"} | quantizers + skip = {"Gather", "GatherElements", "LpNormalization", "ArgMax", "Slice", "Flatten"} + skip |= {"ReduceMax", "ReduceMin", "ReduceSum", "ReduceMean"} + cache = GraphCache(model.graph) + local_functions = {(fn.domain, fn.name, fn.overload) for fn in model.functions} + native_sources = {} + prepared = _rank_preparation_proofs(model, cache, prepared_data_ops) + plans = [] + reshapes = {id(node) for node in _operator_nodes(model, {"Reshape"})} + selected = _operator_nodes(model, {node.op_type for node in model.graph.node}) + if any(op.domain == "com.microsoft" and op.version == 1 for op in model.opset_import): + selected += _operator_nodes(model, quantizers, ("com.microsoft",)) + versions = {op.domain: op.version for op in model.opset_import} + for node in selected: + axis = _attribute(node, "axis") + shape, rank, native = ( + _promoted_input_shape(cache, node.input[0], local_functions, native_sources, prepared) + if node.input + else (None, None, False) + ) + if native or "axis_updated" in node.name: + continue + if node.op_type in update and (rank is None or (rank == 5 and shape is None)): + raise ValueError(f"{node.name!r}: {node.op_type} input layout must be known") + if ( + axis is None + and node.op_type in update + and node.domain in {"", "ai.onnx", "com.microsoft"} + and ( + rank == 5 + or ( + rank in (2, 3) + and (node.op_type not in quantizers or cache.get_shape(node.input[1]) != []) + ) + ) + ): + if node.domain == "com.microsoft": + value = 1 + else: + schema = defs.get_schema( + node.op_type, versions.get(node.domain, versions.get("", 1)), node.domain + ) + default = schema.attributes.get("axis") + if default is not None and default.default_value.type: + value = default.default_value.i + else: + continue + elif axis is None or node.op_type in skip or (axis.i < 0 and node.op_type not in update): + continue + else: + value = axis.i + if node.op_type not in update: + if rank is not None and rank not in (2, 3, 5): + continue + raise ValueError(f"{node.op_type} is unsupported for DLA axis transformation") + original = value + if shape and rank == 5: + removed_axes = _collapsed_reshape_axes(cache, node.input[0], shape, reshapes) + if not -5 <= value < 5 or not removed_axes or value % 5 in removed_axes: + raise ValueError( + f"{node.name!r}: cannot map 5D axis {value} without an unambiguous collapsed Reshape" + ) + native_axis = value % 5 + value = native_axis - int(native_axis > removed_axes[0]) + if original < 0: + value -= 4 + elif value >= 0 and rank is not None: + value += {2: 2, 3: 1}.get(rank, 0) + if value != original: + plans.append((node, axis, value)) + for node, axis, value in plans: + if axis is None: + node.attribute.append(helper.make_attribute("axis", value)) + else: + axis.i = value + node.name += "_axis_updated" + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index e9c26b59172..b14fe04c876 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -53,6 +53,18 @@ _BoundaryPreflightError, _require_output_rank_conversion, ) +from .dla_transforms.dla_rank_conversion import ( + _apply_argmax, + _apply_flatten, + _apply_non4d_expand, + _apply_non4d_initializers, + _apply_non4d_lpnorm, + _apply_non4d_reshape, + _apply_non4d_slice_axis, + _apply_non4d_tile, + _apply_non4d_transpose, + _apply_update_axis, +) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes @@ -99,6 +111,7 @@ def _apply_constants_to_initializers(model, **options): ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), ("ReshapeReduceSum", _apply_reshape_reduce_sum), + ("Non4dSliceAxis", _apply_non4d_slice_axis), ("Gather", _apply_gather), ("DecomposeLstm", partial(_apply_decompose_lstm, defer_rank_conversion=True)), ( @@ -107,6 +120,8 @@ def _apply_constants_to_initializers(model, **options): ), ("Non4dModelInputs", _apply_non4d_model_inputs), ("Non4dModelOutputs", partial(_apply_non4d_model_outputs, defer_rank_conversion=True)), + ("Non4DReshape", _apply_non4d_reshape), + ("ArgMax", _apply_argmax), ("StandaloneReduceOperations", _apply_standalone_reduce_operations), ( "RemoveIntermediarySqueezeAndUnsqueeze", @@ -114,7 +129,14 @@ def _apply_constants_to_initializers(model, **options): ), ("GroupQueryAttention", _apply_group_query_attention), ("GatherElements", _apply_gather_elements), + ("Non4dInitializers", _apply_non4d_initializers), + ("Non4DExpand", _apply_non4d_expand), + ("Non4DTile", _apply_non4d_tile), + ("Flatten", _apply_flatten), + ("Non4DTranspose", _apply_non4d_transpose), + ("Non4DLpNorm", _apply_non4d_lpnorm), ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), + ("UpdateAxis", _apply_update_axis), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) @@ -135,9 +157,20 @@ def _boundary_stage_options(name, quant_types, verbose, prepared_input_ops, prep if name == "Non4dModelInputs": return {"prepared_input_ops": prepared_input_ops} if name == "Non4dModelOutputs": - return {"prepared_output_ops": prepared_output_ops} + return {"prepared_output_ops": prepared_output_ops, "prepared_data_ops": prepared_input_ops} if name == "RemoveIntermediarySqueezeAndUnsqueeze": return {"prepared_consumer_ops": prepared_input_ops} + if name in { + "Non4dSliceAxis", + "Non4DTranspose", + "Non4DLpNorm", + "ArgMax", + "Flatten", + "UpdateAxis", + "Non4DExpand", + "Non4DTile", + }: + return {"prepared_data_ops": prepared_input_ops} if name == "FixInstanceNormChannelMismatch": return {"prepared_consumer_ops": prepared_input_ops} return {} @@ -178,6 +211,7 @@ def _validate_boundary_contract( model, defer_rank_conversion=True, prepared_output_ops=prepared_output_ops, + prepared_data_ops=prepared_input_ops, validate_only=True, ) if "RemoveIntermediarySqueezeAndUnsqueeze" in prepared_input_ops: diff --git a/modelopt/onnx/graph_surgery/utils/graph_utils.py b/modelopt/onnx/graph_surgery/utils/graph_utils.py index d2a59549c1a..3a4a25a1138 100644 --- a/modelopt/onnx/graph_surgery/utils/graph_utils.py +++ b/modelopt/onnx/graph_surgery/utils/graph_utils.py @@ -744,6 +744,7 @@ class GraphCache: __slots__ = ( "_arr_cache", "_graph", + "_initializer_vi_map", "_names", "consumer_map", "graph_input_names", @@ -772,6 +773,9 @@ def rebuild(self) -> None: self._names.update([node.name, *node.input, *node.output]) self.init_map: dict[str, onnx.TensorProto] = {init.name: init for init in graph.initializer} self.init_names: set[str] = set(self.init_map) + self._initializer_vi_map = { + vi.name: vi for vi in graph.value_info if vi.name in self.init_names + } self.producer_map: dict[str, onnx.NodeProto] = {} for node in graph.node: for out in node.output: @@ -993,6 +997,14 @@ def _allocated_boundary_name(name, hint): def _shape_before_boundary_promotion(cache, name): """Recover native shapes at generated input and output boundary views.""" + snapshot = cache._initializer_vi_map.get(name) + if snapshot is not None and snapshot.type.tensor_type.HasField("shape"): + dimensions = snapshot.type.tensor_type.shape.dim + return ( + [dimension.dim_value for dimension in dimensions] + if all(dimension.HasField("dim_value") for dimension in dimensions) + else None + ) producer = cache.get_producer(name) if ( producer is not None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py new file mode 100644 index 00000000000..1f138505cd5 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_rank_conversion.py @@ -0,0 +1,2290 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Numerical checks for independently reviewable rank and boundary passes.""" + +import copy +import importlib + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_model_boundaries import ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, +) +from modelopt.onnx.graph_surgery.dla_transforms.dla_rank_conversion import ( + _apply_argmax, + _apply_flatten, + _apply_non4d_expand, + _apply_non4d_initializers, + _apply_non4d_lpnorm, + _apply_non4d_reshape, + _apply_non4d_slice_axis, + _apply_non4d_tile, + _apply_non4d_transpose, + _apply_update_axis, +) +from modelopt.onnx.graph_surgery.make_dla_compatible import ( + DLA_TRANSFORMS, + _transform_make_dla_compatible, +) + +dla_pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") + + +@pytest.mark.parametrize("existing_metadata", [False, True]) +def test_final_initializer_metadata_matches_materialized_layout(existing_metadata): + """Deferred native initializer ranks expire after axis conversion, before final inference.""" + weight = numpy_helper.from_array(np.arange(24, dtype=np.float32).reshape(2, 3, 4), "weight") + metadata = ( + [helper.make_tensor_value_info("weight", TensorProto.FLOAT, [2, 3, 4])] + if existing_metadata + else [] + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["weight"], ["y"], perm=[0, 2, 1])], + "initializer_metadata", + [], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4, 3])], + [weight], + value_info=metadata, + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=9, + ) + expected = _execute(model, {}) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + tensors = {tensor.name: tensor for tensor in transformed.graph.initializer} + for info in transformed.graph.value_info: + if info.name in tensors: + assert [dim.dim_value for dim in info.type.tensor_type.shape.dim] == list( + tensors[info.name].dims + ) + for actual, reference in zip(_execute(transformed, {}), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Slice", "Reshape", "Tile"]) +@pytest.mark.parametrize("ir_version", [3, 9]) +def test_initializer_promotion_preserves_shared_structural_controls(operation, ir_version): + """Arithmetic readers receive a private view while structural controls keep their native rank.""" + shape = [1, 1, 2, 4] if operation == "Tile" else [1, 1, 2, 3] + control = {"Slice": [1], "Reshape": [6], "Tile": [1, 1, 1, 2]}[operation] + initializers = [numpy_helper.from_array(np.asarray(control, np.int64), "control")] + inputs = ["x", "control"] + if operation == "Slice": + initializers.extend( + [ + numpy_helper.from_array(np.array([2], np.int64), "ends"), + numpy_helper.from_array(np.array([2], np.int64), "axes"), + ] + ) + inputs.extend(["ends", "axes"]) + result_shape = {"Slice": [1, 1, 1, 3], "Reshape": [6], "Tile": [1, 1, 2, 8]}[operation] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "control"], ["sum"]), + helper.make_node(operation, inputs, ["result"]), + ], + "shared_structural_control", + [helper.make_tensor_value_info("x", TensorProto.INT64, shape)], + [ + helper.make_tensor_value_info("sum", TensorProto.INT64, shape), + helper.make_tensor_value_info("result", TensorProto.INT64, result_shape), + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 21)], + ir_version=ir_version, + ) + if ir_version <= 3: + model.graph.input.extend( + helper.make_tensor_value_info(tensor.name, tensor.data_type, list(tensor.dims)) + for tensor in model.graph.initializer + ) + feeds = {"x": np.arange(np.prod(shape), dtype=np.int64).reshape(shape)} + expected = _execute(model, feeds) + tensor = model.graph.initializer[0] + before = tensor.SerializeToString() + transformed = _apply_non4d_initializers(model) + assert tensor is transformed.graph.initializer[0] + assert tensor.SerializeToString() == before + add = next(node for node in transformed.graph.node if node.op_type == "Add") + assert add.input[1] != "control" + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Transpose", "Reshape", "Tile", "Expand"]) +def test_rank_conversion_preserves_operator_values(operation): + """Execute native operators across explicitly deferred boundary and rank stages.""" + shape = [1, 3] if operation == "Expand" else [2, 3] + inputs, initializers, attributes = ["x"], [], {} + if operation == "Transpose": + attributes["perm"] = [1, 0] + output_shape = [3, 2] + transform = _apply_non4d_transpose + else: + values = {"Reshape": [3, 2], "Tile": [2, 1], "Expand": [2, 3]}[operation] + output_shape = {"Reshape": [3, 2], "Tile": [4, 3], "Expand": [2, 3]}[operation] + inputs.append("parameters") + initializers.append(numpy_helper.from_array(np.array(values, dtype=np.int64), "parameters")) + transform = { + "Reshape": _apply_non4d_reshape, + "Tile": _apply_non4d_tile, + "Expand": _apply_non4d_expand, + }[operation] + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"], name="operator", **attributes)], + "rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + transformed = copy.deepcopy(model) + for apply in (_apply_non4d_model_inputs, _apply_non4d_model_outputs, transform): + options = ( + { + "defer_rank_conversion": True, + "prepared_output_ops": {"Non4DReshape"} if operation == "Reshape" else set(), + } + if apply is _apply_non4d_model_outputs + else {} + ) + if apply is _apply_non4d_model_inputs: + options = {"prepared_input_ops": {operation}} + transformed = apply(transformed, **options) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + expected = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data}) + actual = ort.InferenceSession( + transformed.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, {"x": data}) + np.testing.assert_array_equal(actual[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["Flatten", "ArgMax"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_registered_rank_reducers_restore_native_public_outputs(operation, tail): + """Allow deferred output restoration when its actual later conversion stage is installed.""" + shape = [1, 48] if operation == "Flatten" else [1, 4, 4] + attributes = {"axis": 1} if operation == "Flatten" else {"axis": 1, "keepdims": 0} + nodes = [helper.make_node(operation, ["x"], ["middle" if tail else "y"], **attributes)] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "registered_rank_reducer", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 3, 4, 4])], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT if operation == "Flatten" else TensorProto.INT64, shape + ) + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(48, dtype=np.float32).reshape(1, 3, 4, 4)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_symbolic_native_transpose_keeps_its_permutation(explicit, batch): + """Keep rank-two permutation semantics when symbolic graph inputs remain native.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["x"], ["y"], **({"perm": [1, 0]} if explicit else {}))], + "native_transpose", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [3, "N"])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 3, dtype=np.float32).reshape(batch, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + "operation", ["Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Slice", "Transpose"] +) +@pytest.mark.parametrize("negative", [False, True]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_native_axis_chain_preserves_concrete_intermediate_rank(operation, negative, tail): + """Retain native axes after symbolic data acquires concrete Reshape metadata.""" + axis = -1 if negative else 1 + controls = [numpy_helper.from_array(np.asarray([2, 3], np.int64), "target")] + inputs = ["middle"] + attributes = {"axis": axis} + output = [2, 3] + dtype = TensorProto.FLOAT + if operation == "ArgMax": + attributes["keepdims"] = 0 + output, dtype = [2], TensorProto.INT64 + elif operation == "Transpose": + attributes = {} if negative else {"perm": [1, 0]} + output = [3, 2] + elif operation == "Slice": + attributes = {} + output = [2, 2] + for name, values in [("starts", [0]), ("ends", [2]), ("axes", [axis])]: + inputs.append(name) + controls.append(numpy_helper.from_array(np.asarray(values, np.int64), name)) + nodes = [ + helper.make_node("Reshape", ["x", "target"], ["middle"]), + helper.make_node(operation, inputs, ["result" if tail else "y"], **attributes), + ] + if tail: + nodes.append(helper.make_node("Identity", ["result"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_axis_chain", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", dtype, output)], + controls, + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(1, 7, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +def test_converted_reshape_target_preserves_skipped_native_owner(): + """Give the converted Reshape a private target when a native 4D owner is skipped.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["a", "target"], ["ay"]), + helper.make_node("Reshape", ["b", "target"], ["by"]), + ], + "shared_target", + [ + helper.make_tensor_value_info("a", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("b", TensorProto.FLOAT, [5, 6, 1, 1]), + ], + [ + helper.make_tensor_value_info("ay", TensorProto.FLOAT, [2, 1, 1, 3]), + helper.make_tensor_value_info("by", TensorProto.FLOAT, [5, 1, 1, 6]), + ], + [numpy_helper.from_array(np.array([0, 1, 1, -1], np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "a": np.arange(6, dtype=np.float32).reshape(2, 3), + "b": np.arange(30, dtype=np.float32).reshape(5, 6, 1, 1), + } + expected = ort.InferenceSession(model.SerializeToString()).run(None, feeds) + converted = _apply_non4d_reshape(copy.deepcopy(model)) + assert converted.graph.node[0].input[1] != converted.graph.node[1].input[1] + onnx.checker.check_model(converted, full_check=True) + for actual, reference in zip( + ort.InferenceSession(converted.SerializeToString()).run(None, feeds), expected + ): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("shape", [[4], [3, 4], [2, 3, 4]]) +def test_initializer_rank_conversion_preserves_overridable_defaults(shape): + """Preserve a runtime input's rank and default payload while preparing fixed weights.""" + default = np.ones(shape, dtype=np.float32) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Add", ["x", "weight"], ["y"])], + "runtime_default", + [ + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + for name in ("x", "weight") + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(default, "weight")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + transformed = _apply_non4d_initializers(copy.deepcopy(model)) + assert transformed.graph.initializer[0] == model.graph.initializer[0] + onnx.checker.check_model(transformed, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + original = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + converted = ort.InferenceSession( + transformed.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + for feeds in ({"x": data}, {"x": data, "weight": -default}): + np.testing.assert_array_equal(converted.run(None, feeds)[0], original.run(None, feeds)[0]) + + +def _execute(model, feeds): + """Fully check and execute the small native-rank contract fixtures.""" + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ).run(None, feeds) + + +@pytest.mark.parametrize("target", [[0, -1], [0, 0, -1], [0, -1, 0]]) +@pytest.mark.parametrize("collision", [False, True]) +def test_reshape_resolves_copied_dimensions_across_output_aliases(target, collision): + """Resolve copied dimensions even when output adapters allocate a suffixed private name.""" + shape = [2, 3, 4] + output_shape = [2, 12] if len(target) == 2 else shape + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "copied_shape", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(target, np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + if collision: + model.graph.input.append(helper.make_tensor_value_info("other", TensorProto.FLOAT, [1])) + model.graph.node.append(helper.make_node("Identity", ["other"], ["y_squeeze_input"])) + model.graph.output.append( + helper.make_tensor_value_info("y_squeeze_input", TensorProto.FLOAT, [1]) + ) + feeds["other"] = np.asarray([17], np.float32) + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("axis", [0, 1, 2, -1]) +@pytest.mark.parametrize("value_info", [False, True, "symbolic"]) +@pytest.mark.parametrize("ir_version", [3, 9]) +@pytest.mark.parametrize("entrypoint", ["stages", "pipeline"]) +def test_registered_axis_update_reads_native_initializer_rank( + axis, value_info, ir_version, entrypoint +): + """Keep Concat axes correct with fixed inputs that lack optional value metadata.""" + shape = [2, 3, 4] + output_shape = shape.copy() + output_shape[axis] *= 2 + values = np.arange(24, dtype=np.float32).reshape(shape) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Concat", ["first", "second"], ["y"], axis=axis)], + "initializer_axis", + [], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(values, "first"), numpy_helper.from_array(-values, "second")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = ir_version + if ir_version <= 3: + model.graph.input.extend( + helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + for name in ["first", "second"] + ) + if value_info: + model.graph.value_info.extend( + helper.make_tensor_value_info( + name, TensorProto.FLOAT, ["a", "b", "c"] if value_info == "symbolic" else shape + ) + for name in ["first", "second"] + ) + expected = _execute(model, {})[0] + if entrypoint == "pipeline": + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + else: + transformed = _apply_non4d_model_outputs(copy.deepcopy(model), defer_rank_conversion=True) + for name, transform in DLA_TRANSFORMS: + if name in {"Non4dInitializers", "UpdateAxis"}: + transformed = transform(transformed) + transformed.graph.ClearField("value_info") + np.testing.assert_array_equal(_execute(transformed, {})[0], expected) + + +def _collapsed_normalization_model( + shape, axis, operation="Softmax", opset=17, target=None, data_shape=None +): + """Build a native 5D normalization between a rank-expanding and rank-restoring Reshape.""" + collapsed = shape.copy() + collapsed.remove(1) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "native_shape"], ["native"]), + helper.make_node( + operation, + ["native"], + ["normalized"], + **({} if axis is None else {"axis": axis}), + ), + helper.make_node("Reshape", ["normalized", "output_shape"], ["y"]), + ], + "collapsed_normalization", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape or collapsed)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, collapsed)], + [ + numpy_helper.from_array(np.asarray(target or shape, np.int64), "native_shape"), + numpy_helper.from_array(np.asarray(collapsed, np.int64), "output_shape"), + ], + value_info=[ + helper.make_tensor_value_info(n, TensorProto.FLOAT, shape) + for n in ["native", "normalized"] + ], + ), + opset_imports=[helper.make_opsetid("", opset)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model + + +@pytest.mark.parametrize( + ("shape", "axis"), + [ + ([2, 3, 4, 5, 1], 1), + ([2, 3, 4, 5, 1], -4), + ([2, 3, 1, 4, 5], 3), + ([2, 3, 1, 4, 5], -2), + ([1, 2, 3, 4, 5], 2), + ([1, 2, 3, 4, 5], -3), + ], +) +def test_five_dimensional_axis_tracks_the_removed_unary_position(shape, axis): + """Execute the registered rank collapse and its downstream native axis mapping.""" + model = _collapsed_normalization_model(shape, axis) + input_shape = [d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim] + feeds = {"x": np.random.default_rng(20).standard_normal(input_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize( + ("shape", "axis"), [([1, 2, 3, 4, 5], 0), ([1, 2, 3, 4, 5], -5), ([2, 3, 4, 5, 6], 2)] +) +def test_invalid_five_dimensional_axis_refuses_before_earlier_axis_mutation(shape, axis): + """Validate every axis mapping before changing any previously supported operation.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Softmax", ["first"], ["first_out"], axis=1), + helper.make_node("Softmax", ["second"], ["second_out"], axis=axis), + ], + "axis_refusal", + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, shape), + ], + [ + helper.make_tensor_value_info("first_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("second_out", TensorProto.FLOAT, shape), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="5D axis"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("axis", [0, 1, 2]) +def test_axis_update_recovers_native_rank_when_boundary_name_is_allocated(axis): + """Preserve axis semantics when an unrelated public tensor occupies the adapter hint.""" + shape = [2, 3, 4] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Softmax", ["x"], ["y"], axis=axis), + helper.make_node("Identity", ["other"], ["x_unsqueeze_input"]), + ], + "allocated_input", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("other", TensorProto.FLOAT, [1]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + helper.make_tensor_value_info("x_unsqueeze_input", TensorProto.FLOAT, [1]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.random.default_rng(22).standard_normal(shape).astype(np.float32), + "other": np.asarray([17], np.float32), + } + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-6, atol=1e-7) + + +def _structural_control_model(operation, *, runtime=False, public=False): + """Create fixed shape controls with optional runtime or public owners.""" + values = [1, 1] if operation == "Tile" else [2, 3] + nodes = [helper.make_node(operation, ["x", "target"], ["y"])] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])] + initializers = [numpy_helper.from_array(np.asarray(values, np.int64), "target")] + if runtime: + inputs.append(helper.make_tensor_value_info("target", TensorProto.INT64, [2])) + nodes.insert(0, helper.make_node(operation, ["x", "earlier"], ["first"])) + initializers.append(numpy_helper.from_array(np.asarray(values, np.int64), "earlier")) + if public: + collision = "target_" + operation.lower() + "_shape" + initializers.append(numpy_helper.from_array(np.asarray([17], np.int64), collision)) + outputs.extend( + [ + helper.make_tensor_value_info("target", TensorProto.INT64, [2]), + helper.make_tensor_value_info(collision, TensorProto.INT64, [1]), + ] + ) + model = helper.make_model( + helper.make_graph(nodes, "structural_control", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + return model + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +def test_runtime_shape_defaults_refuse_before_any_structural_control_mutation(operation): + """Validate runtime shape controls before changing an earlier fixed-target node.""" + model = _structural_control_model(operation, runtime=True) + before = model.SerializeToString() + transform = { + "Reshape": _apply_non4d_reshape, + "Expand": _apply_non4d_expand, + "Tile": _apply_non4d_tile, + }[operation] + with pytest.raises(ValueError, match="fixed"): + transform(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile"]) +def test_private_shape_views_preserve_public_controls_and_existing_names(operation): + """Preserve public shape values when an allocated control hint is already occupied.""" + model = _structural_control_model(operation, public=True) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_argmax_keeps_untouched_constant_payload_in_place(): + """Preserve an unrelated Constant payload while preparing ArgMax.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Constant", + [], + ["constant"], + value=numpy_helper.from_array(np.arange(8, dtype=np.float32)), + ), + helper.make_node("ArgMax", ["x"], ["y"], axis=3, keepdims=1), + ], + "argmax_payload", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3])], + [ + helper.make_tensor_value_info("y", TensorProto.INT64, [1, 1, 2, 1]), + helper.make_tensor_value_info("constant", TensorProto.FLOAT, [8]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3)} + expected = _execute(model, feeds) + constant = model.graph.node[0] + transformed = _apply_argmax(model) + assert transformed.graph.node[0] is constant + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("collision", [False, True]) +@pytest.mark.parametrize("legacy_ir", [False, True]) +def test_unnamed_argmax_outputs_preserve_distinct_shapes_and_public_aliases(collision, legacy_ir): + """Restore two unnamed reductions while respecting an occupied output-adapter hint.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("ArgMax", ["x"], ["first"], axis=1, keepdims=0), + helper.make_node("ArgMax", ["x"], ["second"], axis=2, keepdims=0), + ], + "argmax_alias", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])], + [ + helper.make_tensor_value_info("first", TensorProto.INT64, [2, 4]), + helper.make_tensor_value_info("second", TensorProto.INT64, [2, 3]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + if collision: + model.graph.initializer.append( + numpy_helper.from_array(np.asarray([17], np.int64), "first_squeeze_input") + ) + model.graph.output.append( + helper.make_tensor_value_info("first_squeeze_input", TensorProto.INT64, [1]) + ) + if legacy_ir: + model.ir_version = 3 + model.graph.input.extend( + helper.make_tensor_value_info(t.name, t.data_type, list(t.dims)) + for t in model.graph.initializer + ) + feeds = {"x": np.random.default_rng(23).standard_normal((2, 3, 4)).astype(np.float32)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_late_unknown_argmax_shape_refuses_without_changing_an_earlier_node(): + """Validate all native ArgMax dimensions before allocating or mutating adapters.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("ArgMax", ["first"], ["first_out"], axis=1, keepdims=0), + helper.make_node("ArgMax", ["second"], ["second_out"], axis=1, keepdims=0), + ], + "argmax_refusal", + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info( + "second", TensorProto.FLOAT, [None, None, None, None] + ), + ], + [ + helper.make_tensor_value_info("first_out", TensorProto.INT64, [2, 4]), + helper.make_tensor_value_info("second_out", TensorProto.INT64, [None, None, None]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="static"): + _apply_argmax(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +def test_default_transpose_preserves_native_reversed_axes(shape): + """Execute the implicit permutation after leading singleton promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Transpose", ["x"], ["y"])], + "default_transpose", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, list(reversed(shape)))], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(np.prod(shape), dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +def test_unnamed_flatten_nodes_keep_distinct_batch_shapes_and_public_controls(): + """Allocate private targets for unrelated unnamed Flatten operations.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Flatten", ["x"], ["first"]), + helper.make_node("Flatten", ["z"], ["second"]), + ], + "flatten_names", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, [3, 2, 4]), + ], + [ + helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 12]), + helper.make_tensor_value_info("second", TensorProto.FLOAT, [3, 8]), + helper.make_tensor_value_info("_reshape_shape", TensorProto.INT64, [1]), + ], + [numpy_helper.from_array(np.asarray([17], np.int64), "_reshape_shape")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "z": np.arange(24, dtype=np.float32).reshape(3, 2, 4), + } + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Flatten", "Transpose"]) +def test_late_invalid_layout_refuses_before_earlier_node_mutation(operation): + """Plan every layout before mutating earlier valid operators.""" + shape = [None, 3, 4] if operation == "Flatten" else [2, 3, 4, 5, 6] + first_attrs = {} if operation == "Flatten" else {"perm": [1, 0]} + second_attrs = {} if operation == "Flatten" else {"perm": [4, 3, 2, 1, 0]} + model = helper.make_model( + helper.make_graph( + [ + helper.make_node(operation, ["x"], ["first"], **first_attrs), + helper.make_node(operation, ["z"], ["second"], **second_attrs), + ], + "layout_refusal", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("z", TensorProto.FLOAT, shape), + ], + [ + helper.make_tensor_value_info( + "first", TensorProto.FLOAT, [2, 3] if operation == "Flatten" else [3, 2] + ), + helper.make_tensor_value_info( + "second", + TensorProto.FLOAT, + [None, 12] if operation == "Flatten" else list(reversed(shape)), + ), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="known"): + {"Flatten": _apply_flatten, "Transpose": _apply_non4d_transpose}[operation](model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Softmax", "LogSoftmax"]) +def test_implicit_legacy_axis_tracks_removed_leading_unary_dimension(operation): + """Read the imported default axis after an actual registered Reshape collapse.""" + model = _collapsed_normalization_model([1, 2, 3, 4, 5], None, operation, 11) + feeds = {"x": np.random.default_rng(25).standard_normal((2, 3, 4, 5)).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +def test_implicit_axis_on_removed_dimension_refuses_without_mutation(): + """Reject an omitted default axis when its unary dimension would disappear.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Softmax", ["x"], ["y"])], + "implicit_refusal", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4, 5, 1])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4, 5, 1])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match="5D axis"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("owner", ["function", "capture"]) +@pytest.mark.parametrize("shared", [False, True]) +def test_opaque_initializer_owners_preserve_rank_or_refuse_shared_layout(owner, shared): + """Protect fixed native weights owned by an opaque function or nested scope.""" + shape = [2, 3, 4] + initializers = [numpy_helper.from_array(np.ones(shape, np.float32), "weight")] + functions = [] + opsets = [helper.make_opsetid("", 17)] + if owner == "function": + functions.append( + helper.make_function( + "opaque", + "Layout", + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + ) + ) + nodes = [helper.make_node("Layout", ["weight"], ["y"], domain="opaque")] + opsets.append(helper.make_opsetid("opaque", 1)) + else: + branch = helper.make_graph( + [helper.make_node("Identity", ["weight"], ["out"])], + "branch", + [], + [helper.make_tensor_value_info("out", TensorProto.FLOAT, shape)], + ) + initializers.append(numpy_helper.from_array(np.asarray(True), "condition")) + nodes = [ + helper.make_node("If", ["condition"], ["y"], then_branch=branch, else_branch=branch) + ] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, shape)] + if shared: + nodes.insert(0, helper.make_node("Identity", ["earlier"], ["early"])) + nodes.append(helper.make_node("Add", ["weight", "weight"], ["sum"])) + initializers.insert(0, numpy_helper.from_array(np.ones(shape, np.float32), "earlier")) + outputs.append(helper.make_tensor_value_info("sum", TensorProto.FLOAT, shape)) + model = helper.make_model( + helper.make_graph(nodes, "opaque_weight", [], outputs, initializers), + opset_imports=opsets, + functions=functions, + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + if shared: + with pytest.raises(ValueError, match="opaque or captured"): + _apply_non4d_initializers(model) + else: + _apply_non4d_initializers(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + "operation", ["Flatten", "Transpose", "LpNormalization", "Softmax", "Slice"] +) +def test_rank_passes_leave_model_local_operators_unchanged(operation): + """Skip opaque local functions whose names match standard layout operators.""" + function = helper.make_function( + "opaque", + operation, + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x"], ["y"], domain="opaque")], + "opaque_rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3, 4])], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("opaque", 1)], + functions=[function], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + for transform in [ + _apply_flatten, + _apply_non4d_transpose, + _apply_non4d_lpnorm, + _apply_update_axis, + _apply_non4d_slice_axis, + ]: + transform(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("runtime", [False, True]) +def test_shared_slice_axes_preserve_private_views_and_transactional_runtime_refusal(runtime): + """Reuse fixed axis views and reject later runtime axes before changing earlier nodes.""" + nodes = [ + helper.make_node("Slice", ["x", "starts", "ends", "axes"], ["first"]), + helper.make_node( + "Slice", ["x", "starts", "ends", "runtime_axes" if runtime else "axes"], ["second"] + ), + ] + initializers = [ + numpy_helper.from_array(np.asarray(v, np.int64), n) + for n, v in [("starts", [0]), ("ends", [2]), ("axes", [1]), ("axes_axes_only", [17])] + ] + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4])] + if runtime: + inputs.append(helper.make_tensor_value_info("runtime_axes", TensorProto.INT64, [1])) + initializers.append(numpy_helper.from_array(np.asarray([1], np.int64), "runtime_axes")) + model = helper.make_model( + helper.make_graph( + nodes, + "slice_controls", + inputs, + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [2, 2, 4]) + for n in ["first", "second"] + ] + + [helper.make_tensor_value_info("axes_axes_only", TensorProto.INT64, [1])], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + if runtime: + before = model.SerializeToString() + with pytest.raises(ValueError, match="fixed"): + _apply_non4d_slice_axis(model) + assert model.SerializeToString() == before + else: + feeds = {"x": np.arange(24, dtype=np.float32).reshape(2, 3, 4)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + slices = [n for n in transformed.graph.node if n.op_type == "Slice"] + assert slices[0].input[3] == slices[1].input[3] + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_implicit_split_axis_preserves_rank_three_values(): + """Map the standard default Split axis during leading singleton promotion.""" + shape = [2, 3, 4] + model = helper.make_model( + helper.make_graph( + [helper.make_node("Split", ["x"], ["first", "second"])], + "split_axis", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [1, 3, 4]) + for n in ["first", "second"] + ], + ), + opset_imports=[helper.make_opsetid("", 13)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(24, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(transformed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize( + ("rank", "empty", "runtime", "fixed_data"), + [ + (2, False, False, False), + (3, True, False, False), + (3, False, True, False), + (3, True, True, False), + (3, False, False, True), + ], +) +def test_implicit_slice_axes_preserve_registered_data_layouts(rank, empty, runtime, fixed_data): + """Derive implicit axes from the fixed starts length before promoting data.""" + shape = [2, 3] if rank == 2 else [2, 3, 4] + output_shape = [2, 2] if rank == 2 else [2, 2, 4] + names = ["x", "starts", "ends", ""] if empty else ["x", "starts", "ends"] + initializers = [numpy_helper.from_array(np.asarray([2, 3], np.int64), "ends")] + inputs = [] if fixed_data else [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + feeds = {} if fixed_data else {"x": data} + if fixed_data: + initializers.append(numpy_helper.from_array(data, "x")) + if runtime: + inputs.append(helper.make_tensor_value_info("starts", TensorProto.INT64, [2])) + feeds["starts"] = np.asarray([0, 1], np.int64) + else: + initializers.append(numpy_helper.from_array(np.asarray([0, 1], np.int64), "starts")) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Slice", names, ["y"])], + "implicit_slice", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) + + +def test_unknown_implicit_slice_length_refuses_before_earlier_mutation(): + """Reject an unknown runtime starts length before allocating an earlier fixed axis view.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "fixed", "ends"], ["first"]), + helper.make_node("Slice", ["x", "dynamic", "ends"], ["second"]), + ], + "unknown_starts", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]), + helper.make_tensor_value_info("dynamic", TensorProto.INT64, [None]), + ], + [ + helper.make_tensor_value_info(n, TensorProto.FLOAT, [None, None]) + for n in ["first", "second"] + ], + [ + numpy_helper.from_array(np.asarray([0, 0], np.int64), "fixed"), + numpy_helper.from_array(np.asarray([2, 3], np.int64), "ends"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + before = model.SerializeToString() + with pytest.raises(ValueError, match="known starts length"): + _apply_non4d_slice_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("producer", ["input", "Unsqueeze"]) +def test_native_five_dimensional_axes_refuse_without_a_collapsed_reshape(producer): + """Keep earlier axes unchanged when a later native input never lost a dimension.""" + shape = [1, 2, 3, 4, 5] + nodes = [helper.make_node("Softmax", ["first"], ["first_out"], axis=1)] + inputs = [helper.make_tensor_value_info("first", TensorProto.FLOAT, [2, 3, 4])] + initializers = [] + if producer == "input": + inputs.append(helper.make_tensor_value_info("native", TensorProto.FLOAT, shape)) + else: + inputs.append(helper.make_tensor_value_info("x", TensorProto.FLOAT, shape[1:])) + initializers.append(numpy_helper.from_array(np.asarray([0], np.int64), "axes")) + nodes.append(helper.make_node("Unsqueeze", ["x", "axes"], ["native"])) + nodes.append(helper.make_node("Softmax", ["native"], ["y"], axis=2)) + model = helper.make_model( + helper.make_graph( + nodes, + "native_five", + inputs, + [ + helper.make_tensor_value_info("first_out", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("y", TensorProto.FLOAT, shape), + ], + initializers, + value_info=[helper.make_tensor_value_info("native", TensorProto.FLOAT, shape)], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + with pytest.raises(ValueError, match="collapsed Reshape"): + _apply_update_axis(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize( + ("shape", "target", "data_shape", "axis"), + [ + ([1, 2, 1, 3, 4], [-1, 2, 1, 3, 4], [1, 2, 3, 4], 1), + ([1, 2, 4, 3, 5], [1, 2, 0, 3, 5], [2, 3, 4, 5], 2), + ], +) +def test_five_dimensional_axis_tracks_the_actually_removed_target_dimension( + shape, target, data_shape, axis +): + """Respect inferred and copied dimensions when locating the singleton actually removed.""" + model = _collapsed_normalization_model(shape, axis, target=target, data_shape=data_shape) + feeds = {"x": np.random.default_rng(29).standard_normal(data_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize("operation", ["ArgMax", "LpNormalization"]) +@pytest.mark.parametrize("axis", [2, -3]) +@pytest.mark.parametrize("collapse", [False, True]) +def test_specialized_axes_follow_only_actually_collapsed_reshapes(operation, axis, collapse): + """Remap specialized axes after a singleton collapse while keeping untouched native 5D axes.""" + shape = [1, 2, 3, 4, 5] + data_shape = shape[1:] if collapse else shape + reduced = [1, 2, 4, 5] + output_shape = reduced if operation == "ArgMax" else data_shape + inits = [] + nodes = [] + source = "x" + if collapse: + inits.append(numpy_helper.from_array(np.array(shape, np.int64), "native_shape")) + nodes.append(helper.make_node("Reshape", ["x", "native_shape"], ["native"])) + source = "native" + nodes.append( + helper.make_node( + operation, + [source], + ["result" if collapse else "y"], + axis=axis, + **({"keepdims": 0} if operation == "ArgMax" else {}), + ) + ) + if collapse: + inits.append(numpy_helper.from_array(np.array(output_shape, np.int64), "output_shape")) + nodes.append(helper.make_node("Reshape", ["result", "output_shape"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "specialized_axes", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, data_shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMax" else TensorProto.FLOAT, + output_shape, + ) + ], + inits, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.random.default_rng(73).standard_normal(data_shape).astype(np.float32)} + expected = _execute(model, feeds)[0] + converted = copy.deepcopy(model) + if collapse: + converted = _transform_make_dla_compatible(converted, verbose=False) + else: + converted = (_apply_argmax if operation == "ArgMax" else _apply_non4d_lpnorm)(converted) + np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +def test_public_initializer_outputs_keep_their_native_dimensions(shape): + """Preserve fixed public values that have no internal layout-changing consumer.""" + data = np.arange(np.prod(shape), dtype=np.float32).reshape(shape) + model = helper.make_model( + helper.make_graph( + [], + "public_weight", + [], + [helper.make_tensor_value_info("weight", TensorProto.FLOAT, shape)], + [numpy_helper.from_array(data, "weight")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + expected = _execute(model, {})[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, {})[0], expected) + + +@pytest.mark.parametrize("target", [[0, 1, 1, -1], [1, 0, 1, -1]]) +def test_four_dimensional_reshape_targets_resolve_native_copied_dimensions(target): + """Resolve copied positions before input promotion shifts native dimensions.""" + output_shape = [2, 1, 1, 3] if target[0] == 0 else [1, 3, 1, 2] + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "copied_four", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.asarray(target, np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +def test_runtime_four_dimensional_reshape_target_preserves_runtime_overrides(): + """Leave runtime targets native when neither input nor output rank needs promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Reshape", ["x", "target"], ["y"])], + "runtime_four", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 2, 3]), + helper.make_tensor_value_info("target", TensorProto.INT64, [4]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None, None, None, None])], + [numpy_helper.from_array(np.asarray([1, 1, 2, 3], np.int64), "target")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3), + "target": np.asarray([1, 2, 1, 3], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + ("producer", "copied"), + [("Reshape", False), ("Reshape", True), ("Squeeze", False), ("Unsqueeze", False)], +) +def test_native_index_views_survive_rank_preparation(producer, copied): + """Retain scalar and vector control ranks after upstream structural nodes are promoted.""" + source_shape = [] if producer == "Unsqueeze" else ([3, 1] if copied else [1, 3]) + count = 1 if producer == "Unsqueeze" else 3 + values = ( + np.asarray(2, np.int64) + if not source_shape + else np.arange(3, dtype=np.int64).reshape(source_shape) + ) + controls = np.asarray([0] if producer != "Reshape" or copied else [-1], np.int64) + nodes = [ + helper.make_node(producer, ["source", "control"], ["indices"]), + helper.make_node("Gather", ["x", "indices"], ["y"], axis=0), + ] + model = helper.make_model( + helper.make_graph( + nodes, + "native_index_views", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [5, 4]), + helper.make_tensor_value_info("source", TensorProto.INT64, source_shape), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [count, 4])], + [numpy_helper.from_array(controls, "control")], + value_info=[helper.make_tensor_value_info("indices", TensorProto.INT64, [count])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + onnx.checker.check_model(model, full_check=True) + feeds = {"x": np.arange(20, dtype=np.float32).reshape(5, 4), "source": values} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +@pytest.mark.parametrize("batch", [1, 2]) +@pytest.mark.parametrize("static_output", [False, True]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_structural_controls_keep_native_data_rank(operation, batch, static_output, tail): + """Retain symbolic native layouts, including fixed public outputs and unary readers.""" + target = { + "Tile": [0, 1] if static_output else [2, 1], + "Expand": [2, 3] if static_output else [1, 3], + "Reshape": [3, 1] if static_output else [3, -1], + }[operation] + if static_output and operation == "Reshape" and batch != 1: + target = [3, 2] + output_shape = ( + {"Tile": [0, 3], "Expand": [2, 3], "Reshape": [3, batch]}[operation] + if static_output + else ([3, "N"] if operation == "Reshape" else ["M", 3]) + ) + nodes = [helper.make_node(operation, ["x", "control"], ["middle" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Identity", ["middle"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "symbolic_control", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + [numpy_helper.from_array(np.array(target, np.int64), "control")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 3, dtype=np.float32).reshape(batch, 3)} + expected = _execute(model, feeds)[0] + original_control = model.graph.initializer[0].SerializeToString() + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + assert converted.graph.initializer[0].SerializeToString() == original_control + + +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +def test_symbolic_native_control_chain_preserves_concrete_intermediate_rank(operation): + """Keep native controls after a symbolic input acquires concrete intermediate metadata.""" + control = {"Tile": [1, 2], "Expand": [2, 3], "Reshape": [3, 2]}[operation] + output = {"Tile": [2, 6], "Expand": [2, 3], "Reshape": [3, 2]}[operation] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Reshape", ["x", "first_shape"], ["middle"]), + helper.make_node(operation, ["middle", "control"], ["y"]), + ], + "native_control_chain", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output)], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "first_shape"), + numpy_helper.from_array(np.array(control, np.int64), "control"), + ], + value_info=[helper.make_tensor_value_info("middle", TensorProto.FLOAT, [2, 3])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + for actual, original in zip(converted.graph.initializer, model.graph.initializer): + assert actual.SerializeToString() == original.SerializeToString() + + +@pytest.mark.parametrize("operation", ["Tile", "Reshape"]) +def test_runtime_slice_extents_retain_promoted_control_layouts(operation): + """Follow already-promoted data through unknown Slice extents before adapting controls.""" + values = [1, 2, 1] if operation == "Tile" else [4, -1] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "starts", "ends"], ["middle"]), + helper.make_node(operation, ["middle", "control"], ["y"]), + ], + "promoted_slice_control", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("starts", TensorProto.INT64, [2]), + ], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, ["A", "B", 4] if operation == "Tile" else ["A", "B"] + ) + ], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "ends"), + numpy_helper.from_array(np.array(values, np.int64), "control"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "starts": np.array([0, 1], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize( + "operation", ["Softmax", "LogSoftmax", "LpNormalization", "Transpose", "Slice"] +) +def test_runtime_slice_extents_retain_promoted_axes(operation): + """Map native axes using known rank when runtime Slice extents remain symbolic.""" + inputs, attributes = ["middle"], {"axis": 1} + if operation == "Transpose": + attributes = {"perm": [1, 2, 0]} + elif operation == "Slice": + inputs += ["starts", "ends", "axes"] + attributes = {} + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Slice", ["x", "starts", "ends"], ["middle"]), + helper.make_node(operation, inputs, ["y"], **attributes), + ], + "promoted_slice_axes", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("starts", TensorProto.INT64, [2]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["A", "B", "C"])], + [ + numpy_helper.from_array(np.array([2, 3], np.int64), "ends"), + numpy_helper.from_array(np.array([1, 2], np.int64), "axes"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 16, + "starts": np.array([0, 1], np.int64), + } + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(converted, feeds)[0], expected, rtol=1e-6, atol=1e-7) + + +@pytest.mark.parametrize("operation", ["Slice", "NonZero", "Range", "Reshape"]) +@pytest.mark.parametrize("tail", [False, True]) +def test_symbolic_outputs_restore_only_promoted_producer_paths(operation, tail): + """Reuse runtime-dependent native outputs while restoring a genuinely padded Reshape.""" + output = "middle" if tail else "y" + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, [6])] + initializers = [] + if operation in {"Slice", "Reshape"}: + inputs.append(helper.make_tensor_value_info("starts", TensorProto.INT64, [1])) + initializers.append(numpy_helper.from_array(np.array([6], np.int64), "ends")) + nodes = [ + helper.make_node( + "Slice", ["x", "starts", "ends"], ["sliced" if operation == "Reshape" else output] + ) + ] + if operation == "Reshape": + initializers.append(numpy_helper.from_array(np.array([-1], np.int64), "shape")) + nodes.append(helper.make_node("Reshape", ["sliced", "shape"], [output])) + cases = [ + {"x": np.arange(6, dtype=np.float32), "starts": np.array([start], np.int64)} + for start in [0, 2] + ] + elif operation == "NonZero": + nodes = [helper.make_node("NonZero", ["x"], [output])] + cases = [ + {"x": np.asarray(values, np.float32)} + for values in [[0, 0, 1, 0, 2, 0], [1, 2, 3, 4, 5, 0]] + ] + else: + inputs = [helper.make_tensor_value_info("stop", TensorProto.FLOAT, [])] + initializers = [ + numpy_helper.from_array(np.float32(value), name) + for name, value in [("start", 0), ("step", 1)] + ] + nodes = [helper.make_node("Range", ["start", "stop", "step"], [output])] + cases = [{"stop": np.asarray(stop, np.float32)} for stop in [3, 5]] + if tail: + nodes.append(helper.make_node("Identity", [output], ["y"])) + dtype = TensorProto.INT64 if operation == "NonZero" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + nodes, + "native_dynamic_output", + inputs, + [ + helper.make_tensor_value_info( + "y", dtype, [1, "N"] if operation == "NonZero" else ["N"] + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, reference in zip(cases, expected): + np.testing.assert_array_equal(_execute(changed, feeds)[0], reference) + + +@pytest.mark.parametrize(("initializer", "shared"), [(False, False), (True, False), (True, True)]) +@pytest.mark.parametrize("tail", [False, True]) +def test_rank_one_tile_data_and_repeats_remain_native(initializer, shared, tail): + """Execute rank-one Tile for graph inputs and fixed data without padding repeats.""" + data = np.array([1, 2, 3], np.float32) + nodes = [helper.make_node("Tile", ["x", "repeats"], ["tiled" if tail else "y"])] + if tail: + nodes.append(helper.make_node("Relu", ["tiled"], ["y"])) + inputs = [] if initializer else [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])] + outputs = [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])] + if shared: + nodes.append(helper.make_node("Add", ["other", "x"], ["added"])) + inputs.append(helper.make_tensor_value_info("other", TensorProto.FLOAT, [1, 1, 2, 3])) + outputs.append(helper.make_tensor_value_info("added", TensorProto.FLOAT, [1, 1, 2, 3])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_tile", + inputs, + outputs, + [numpy_helper.from_array(np.array([2], np.int64), "repeats")] + + ([numpy_helper.from_array(data, "x")] if initializer else []), + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {} if initializer else {"x": data} + if shared: + feeds["other"] = np.arange(6, dtype=np.float32).reshape(1, 1, 2, 3) + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Add", "Sub", "Mul", "Div"]) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_one_tile_tracks_widened_elementwise_initializer(operation, unary): + """Keep Tile repeats consistent with rank widened through an elementwise operand.""" + nodes = [helper.make_node(operation, ["x", "weight"], ["combined"])] + if unary: + nodes.append(helper.make_node("Relu", ["combined"], ["activated"])) + nodes.append(helper.make_node("Tile", ["activated" if unary else "combined", "repeats"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "widened_tile", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [6])], + [ + numpy_helper.from_array(np.array([1, 2, 3], np.float32), "weight"), + numpy_helper.from_array(np.array([2], np.int64), "repeats"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.array([3, 5, 7], np.float32)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize( + ("producer", "shape"), + [("Reshape", [2, 3]), ("Reshape", [1, 2, 3, 1]), ("ArgMax", [2, 3]), ("ReduceSum", [2, 3])], +) +@pytest.mark.parametrize("unary", [False, True]) +def test_rank_one_tile_tracks_registered_rank_restoring_reshape(producer, shape, unary): + """Execute vector Tile after earlier registered stages create a leading 4D view.""" + nodes = [] + initializers = [numpy_helper.from_array(np.array([2], np.int64), "repeats")] + if producer == "Reshape": + initializers.append(numpy_helper.from_array(np.array([-1], np.int64), "target")) + nodes.append(helper.make_node("Reshape", ["x", "target"], ["vector"])) + length = 6 + elif producer == "ArgMax": + nodes.append(helper.make_node("ArgMax", ["x"], ["vector"], axis=0, keepdims=0)) + length = 3 + else: + initializers.append(numpy_helper.from_array(np.array([0], np.int64), "axes")) + nodes.append(helper.make_node("ReduceSum", ["x", "axes"], ["vector"], keepdims=0)) + length = 3 + if unary: + nodes.append(helper.make_node("Identity", ["vector"], ["middle"])) + nodes.append(helper.make_node("Tile", ["middle" if unary else "vector", "repeats"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "restored_tile", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if producer == "ArgMax" else TensorProto.FLOAT, + [2 * length], + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["GatherElements", "ReduceSum"]) +@pytest.mark.parametrize("reshape", [False, True]) +def test_rank_preparation_preserves_unconverted_batched_matrix_consumers(operation, reshape): + """Prove runtime matrix rank from a padded view or boundary without requiring Conv eligibility.""" + shape = [2, 5, 4] + source_shape = [1, *shape] if reshape else shape + nodes = [] + initializers = [ + numpy_helper.from_array(np.tile(np.eye(4, dtype=np.float32), (2, 1, 1)), "weight") + ] + if reshape: + initializers.append(numpy_helper.from_array(np.asarray(shape, np.int64), "native_target")) + nodes.append(helper.make_node("Reshape", ["x", "native_target"], ["matrix"])) + nodes.append(helper.make_node("MatMul", ["matrix" if reshape else "x", "weight"], ["data"])) + if operation == "GatherElements": + initializers.append( + numpy_helper.from_array(np.arange(40, dtype=np.int64).reshape(shape) % 5, "control") + ) + nodes.append(helper.make_node(operation, ["data", "control"], ["selected"], axis=1)) + target = [1, *shape] + else: + initializers.append(numpy_helper.from_array(np.asarray([1], np.int64), "control")) + nodes.append(helper.make_node(operation, ["data", "control"], ["selected"], keepdims=0)) + target = [1, 1, 2, 4] + initializers.append(numpy_helper.from_array(np.asarray(target, np.int64), "target")) + nodes.append(helper.make_node("Reshape", ["selected", "target"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "matrix_rank_proof", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, target)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(40, dtype=np.float32).reshape(source_shape)} + expected = _execute(model, feeds) + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert any(node.op_type == "MatMul" for node in converted.graph.node) + for actual, reference in zip(_execute(converted, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("operation", ["Add", "Sub", "Mul", "Div", "Max"]) +@pytest.mark.parametrize("bias_shape", [[4], [3, 4], [1, 3, 4]]) +def test_symbolic_elementwise_initializers_keep_native_rank(operation, bias_shape): + """Reuse one converted model for multiple batches without widening native symbolic data.""" + nodes = [] + source = "x" + if len(bias_shape) == 2: + nodes.append(helper.make_node("Relu", [source], ["native"])) + source = "native" + operands = [source, "bias"] if len(bias_shape) != 3 else ["bias", source] + nodes.append(helper.make_node(operation, operands, ["y"])) + bias = np.arange(np.prod(bias_shape), dtype=np.float32).reshape(bias_shape) + 1 + model = helper.make_model( + helper.make_graph( + nodes, + "symbolic_bias", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4])], + [numpy_helper.from_array(bias, "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + payload = model.graph.initializer[0].SerializeToString() + cases = [ + {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4) + 1} for batch in [1, 2] + ] + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert ( + next(t for t in changed.graph.initializer if t.name == "bias").SerializeToString() + == payload + ) + for feeds, reference in zip(cases, expected): + np.testing.assert_array_equal(_execute(changed, feeds)[0], reference) + + +@pytest.mark.parametrize("public", [False, True]) +@pytest.mark.parametrize("source", ["constant", "default", "runtime"]) +def test_symbolic_bias_can_stay_native_for_a_shared_static_broadcast(public, source): + """Retain shared constant bytes and both native and already-4D broadcast outputs.""" + outputs = [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_y", TensorProto.FLOAT, [1, 2, 3, 4]), + ] + if public: + outputs.append(helper.make_tensor_value_info("bias", TensorProto.FLOAT, [1, 3, 4])) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "bias"], ["y"]), + helper.make_node("Add", ["static_x", "bias"], ["static_y"]), + ], + "shared_symbolic_bias", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 2, 3, 4]), + ], + outputs, + [numpy_helper.from_array(np.ones([1, 3, 4], np.float32), "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + if source != "constant": + model.graph.input.append( + helper.make_tensor_value_info("bias", TensorProto.FLOAT, [1, 3, 4]) + ) + # Make the boundary planner inspect the bias as the data operand. + model.graph.node[0].input[:] = ["bias", "x"] + if source == "default": + model.ir_version = 3 + else: + del model.graph.initializer[:] + cases = [ + { + "x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4), + "static_x": np.arange(24, dtype=np.float32).reshape(1, 2, 3, 4), + **({"bias": np.ones([1, 3, 4], np.float32)} if source == "runtime" else {}), + } + for batch in [1, 2] + ] + expected = [_execute(model, feeds) for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, references in zip(cases, expected): + for actual, reference in zip(_execute(changed, feeds), references): + np.testing.assert_array_equal(actual, reference) + + +def test_symbolic_bias_preserves_a_shared_quantization_control(): + """Keep a bias native when another consumer requires its original per-axis scale layout.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "scale"], ["y"]), + helper.make_node( + "QuantizeLinear", ["static_x", "scale", "zero"], ["quantized"], axis=3 + ), + ], + "shared_symbolic_scale", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 1, 3, 4]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("quantized", TensorProto.UINT8, [1, 1, 3, 4]), + ], + [ + numpy_helper.from_array(np.array([1, 2, 3, 4], np.float32), "scale"), + numpy_helper.from_array(np.zeros(4, np.uint8), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = { + "x": np.arange(24, dtype=np.float32).reshape(2, 3, 4), + "static_x": np.arange(12, dtype=np.float32).reshape(1, 1, 3, 4), + } + expected = _execute(model, feeds) + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for actual, reference in zip(_execute(changed, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +def test_symbolic_bias_refuses_conflicting_rank_preparation_before_mutation(): + """Refuse a shared native/forced-rank initializer instead of leaving one consumer invalid.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["x", "bias"], ["y"]), + helper.make_node("Concat", ["static_x", "bias"], ["joined"], axis=0), + ], + "mixed_symbolic_rank", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("static_x", TensorProto.FLOAT, [1, 3, 4]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4]), + helper.make_tensor_value_info("joined", TensorProto.FLOAT, [2, 3, 4]), + ], + [numpy_helper.from_array(np.ones([1, 3, 4], np.float32), "bias")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + _execute( + model, {"x": np.ones([2, 3, 4], np.float32), "static_x": np.ones([1, 3, 4], np.float32)} + ) + before = model.SerializeToString() + with pytest.raises(ValueError, match=r"native symbolic.*rank-prepared"): + _apply_non4d_initializers(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("operation", ["Tile", "Expand", "Reshape"]) +def test_unknown_structural_data_rank_refuses_before_mutation(operation): + """Refuse unresolved intermediate ranks before allocating any control view.""" + values = [2, 1] if operation == "Tile" else [2, 3] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Identity", ["x"], ["hidden"]), + helper.make_node(operation, ["hidden", "control"], ["y"]), + ], + "unknown_control_rank", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [ + helper.make_tensor_value_info( + "y", TensorProto.FLOAT, [4, 3] if operation == "Tile" else [2, 3] + ) + ], + [numpy_helper.from_array(np.array(values, np.int64), "control")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + _execute(model, {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}) + before = model.SerializeToString() + callback = { + "Tile": _apply_non4d_tile, + "Expand": _apply_non4d_expand, + "Reshape": _apply_non4d_reshape, + }[operation] + with pytest.raises(ValueError, match="input rank"): + callback(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("batch", [1, 2]) +def test_symbolic_input_output_remains_native(batch): + """Leave symbolic input ranks intact instead of restoring an unpromoted output.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Softmax", ["x"], ["y"], axis=-1)], + "symbolic_native", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, ["N", 3, 4])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(batch * 12, dtype=np.float32).reshape(batch, 3, 4)} + expected = _execute(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_allclose(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["ArgMin", "Hardmax", "LayerNormalization"]) +@pytest.mark.parametrize("symbolic", [False, True]) +def test_axis_update_preserves_native_four_dimensional_operators(operation, symbolic): + """Allow axis-bearing operators that already consume their final rank.""" + shape = [2, 3, 4, 5] + declared_shape = ["N", *shape[1:]] if symbolic else shape + inputs = ["x"] + initializers = [] + output_shape = [declared_shape[0], 1, 4, 5] if operation == "ArgMin" else declared_shape + if operation == "LayerNormalization": + inputs.append("scale") + initializers.append(numpy_helper.from_array(np.ones(shape[1:], np.float32), "scale")) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"], axis=1)], + "native_four_dimensional", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, declared_shape)], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMin" else TensorProto.FLOAT, + output_shape, + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + feeds = {"x": np.arange(120, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = _apply_update_axis(copy.deepcopy(model)) + assert transformed == model + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("keepdims", [None, 1]) +@pytest.mark.parametrize("batch", [1, 2]) +def test_registered_argmax_keeps_symbolic_four_dimensional_layout(keepdims, batch): + """Leave rank-four ArgMax native when its symbolic dimensions need no restoration.""" + attributes = {} if keepdims is None else {"keepdims": keepdims} + model = helper.make_model( + helper.make_graph( + [helper.make_node("ArgMax", ["x"], ["y"], axis=1, **attributes)], + "native_argmax", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, ["N", 3, 4, 5])], + [helper.make_tensor_value_info("y", TensorProto.INT64, ["N", 1, 4, 5])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(batch * 60, dtype=np.float32).reshape(batch, 3, 4, 5)} + expected = _execute(model, feeds)[0] + converted = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert converted.graph.node[0] == model.graph.node[0] + np.testing.assert_array_equal(_execute(converted, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["Reshape", "Expand", "Tile", "Slice", "Add"]) +def test_legacy_initializer_controls_remain_fixed(operation): + """Execute IR3 initializer declarations through boundary and rank conversion.""" + shape = [2, 3] + controls = { + "Reshape": [3, 2], + "Expand": [2, 3], + "Tile": [2, 1], + "Slice": [0], + "Add": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], + } + output_shape = {"Reshape": [3, 2], "Tile": [4, 3], "Slice": [1, 3]}.get(operation, shape) + inputs = ["x", "control"] + initializers = [ + numpy_helper.from_array( + np.asarray(controls[operation], np.float32 if operation == "Add" else np.int64), + "control", + ) + ] + if operation == "Slice": + for name, values in (("ends", [1]), ("axes", [0])): + initializers.append(numpy_helper.from_array(np.asarray(values, np.int64), name)) + inputs.append(name) + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, inputs, ["y"])], + "legacy_controls", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, shape)] + + [ + helper.make_tensor_value_info(t.name, t.data_type, list(t.dims)) + for t in initializers + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, output_shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 + feeds = {"x": np.arange(6, dtype=np.float32).reshape(shape)} + expected = _execute(model, feeds) + transformed = copy.deepcopy(model) + transforms = ( + _apply_non4d_model_inputs, + _apply_non4d_model_outputs, + { + "Reshape": _apply_non4d_reshape, + "Expand": _apply_non4d_expand, + "Tile": _apply_non4d_tile, + "Slice": _apply_non4d_slice_axis, + "Add": _apply_non4d_initializers, + }[operation], + ) + for apply in transforms: + transformed = apply( + transformed, + **( + { + "defer_rank_conversion": True, + "prepared_output_ops": {"Non4DReshape"} if operation == "Reshape" else set(), + } + if apply is _apply_non4d_model_outputs + else {"prepared_input_ops": {operation}} + if apply is _apply_non4d_model_inputs + else {} + ), + ) + if operation == "Add": + assert list(transformed.graph.initializer[0].dims) == [1, 1, 2, 3] + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected[0]) + + +@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) +@pytest.mark.parametrize("shape", [[2, 3], [2, 3, 4]]) +@pytest.mark.parametrize("per_channel", [False, True, "singleton"]) +@pytest.mark.parametrize("domain", ["", "com.microsoft"]) +def test_registered_implicit_quantizer_axis_preserves_channel_values( + operation, shape, per_channel, domain +): + """Map every vector scale's channel axis, including a single-channel vector.""" + shape = shape.copy() + if per_channel == "singleton": + shape[1] = 1 + scales = np.asarray( + [0.25, 0.5, 1.0] if per_channel is True else [0.5] if per_channel else 0.5, np.float32 + ) + zeros = np.asarray([0, 1, 2] if per_channel is True else [0] if per_channel else 0, np.uint8) + input_type = TensorProto.FLOAT if operation == "QuantizeLinear" else TensorProto.UINT8 + output_type = TensorProto.UINT8 if operation == "QuantizeLinear" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x", "scale", "zero"], ["y"], domain=domain)], + "implicit_quantizer", + [helper.make_tensor_value_info("x", input_type, shape)], + [helper.make_tensor_value_info("y", output_type, shape)], + [numpy_helper.from_array(scales, "scale"), numpy_helper.from_array(zeros, "zero")], + ), + opset_imports=[helper.make_opsetid("", 17)] + + ([helper.make_opsetid(domain, 1)] if domain else []), + ) + model.ir_version = 9 + data = ( + np.arange(np.prod(shape)) + .reshape(shape) + .astype(np.float32 if input_type == TensorProto.FLOAT else np.uint8) + ) + feeds = {"x": data} + expected = _execute(model, feeds)[0] + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + quantizer = next(node for node in transformed.graph.node if node.op_type == operation) + if per_channel: + assert next(attr.i for attr in quantizer.attribute if attr.name == "axis") == 5 - len(shape) + else: + assert all(attr.name != "axis" for attr in quantizer.attribute) + np.testing.assert_array_equal(_execute(transformed, feeds)[0], expected) + + +@pytest.mark.parametrize("operation", ["QuantizeLinear", "DequantizeLinear"]) +def test_legacy_quantizer_schema_keeps_implicit_scalar_axis(operation): + """Keep pre-axis quantizer schemas byte-identical during axis preparation.""" + input_type = TensorProto.FLOAT if operation == "QuantizeLinear" else TensorProto.UINT8 + output_type = TensorProto.UINT8 if operation == "QuantizeLinear" else TensorProto.FLOAT + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["x", "scale", "zero"], ["y"])], + "legacy_quantizer", + [helper.make_tensor_value_info("x", input_type, [2, 3])], + [helper.make_tensor_value_info("y", output_type, [2, 3])], + [ + numpy_helper.from_array(np.asarray(0.5, np.float32), "scale"), + numpy_helper.from_array(np.asarray(0, np.uint8), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 10)], + ) + model.ir_version = 9 + original = model.SerializeToString() + _apply_update_axis(model) + assert model.SerializeToString() == original + feeds = { + "x": np.arange(6) + .reshape(2, 3) + .astype(np.float32 if input_type == TensorProto.FLOAT else np.uint8) + } + _execute(model, feeds) + + +@pytest.mark.parametrize("operation", ["Expand", "Tile"]) +@pytest.mark.parametrize("source", ["input", "default", "shape", "symbolic"]) +def test_native_four_dimensional_runtime_controls_need_no_padding(operation, source): + """Execute native rank-four runtime layouts without freezing their controls.""" + shape = [1, 1, 2, 3] + first = [2, 1, 2, 3] if operation == "Expand" else [1, 1, 2, 1] + second = [3, 1, 2, 3] if operation == "Expand" else [2, 1, 1, 1] + declared = ["batch", 1, 2, 3] if source == "symbolic" else shape + inputs = [helper.make_tensor_value_info("x", TensorProto.FLOAT, declared)] + nodes = [] + initializers = [] + if source == "shape": + nodes.append(helper.make_node("Shape", ["x"], ["control"])) + else: + inputs.append(helper.make_tensor_value_info("control", TensorProto.INT64, [4])) + if source == "default": + initializers.append(numpy_helper.from_array(np.asarray(first, np.int64), "control")) + nodes.append(helper.make_node(operation, ["x", "control"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_runtime_controls", + inputs, + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [None] * 4)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 9 + data = np.arange(6, dtype=np.float32).reshape(shape) + feeds = ( + [{"x": data}] + if source in {"shape", "default"} + else [{"x": data, "control": np.asarray(first, np.int64)}] + ) + if source != "shape": + feeds.append({"x": data, "control": np.asarray(second, np.int64)}) + expected = [_execute(model, values)[0] for values in feeds] + if source == "shape": + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + else: + # ORT symbolic inference cannot resolve unconstrained runtime Expand controls. + transformed = dict(DLA_TRANSFORMS)["Non4D" + operation](copy.deepcopy(model)) + assert transformed.SerializeToString() == model.SerializeToString() + if source == "shape": + assert any(node.op_type == "Shape" for node in transformed.graph.node) + else: + assert any(value.name == "control" for value in transformed.graph.input) + if initializers: + assert ( + next(t for t in transformed.graph.initializer if t.name == "control") + == model.graph.initializer[0] + ) + for values, reference in zip(feeds, expected): + np.testing.assert_array_equal(_execute(transformed, values)[0], reference) + + +@pytest.mark.parametrize( + "operation", + ["Slice", "Transpose", "Softmax", "LogSoftmax", "LpNormalization", "ArgMax", "Expand", "Tile"], +) +@pytest.mark.parametrize("rank", [2, 3]) +@pytest.mark.parametrize("source", ["unsqueeze", "unary", "reshape"]) +def test_rank_one_adapters_keep_native_controls_unless_reshape_promotes(operation, rank, source): + """Distinguish native rank-one input adapters from an explicitly padded Reshape.""" + shape = [1] * (rank - 1) + [6] + initializers = [ + numpy_helper.from_array(np.array(list(range(rank - 1)), np.int64), "unsqueeze_axes") + ] + if source == "reshape": + initializers.append(numpy_helper.from_array(np.array(shape, np.int64), "target")) + nodes = [helper.make_node("Reshape", ["x", "target"], ["adapted"])] + else: + nodes = [helper.make_node("Unsqueeze", ["x", "unsqueeze_axes"], ["adapted"])] + data = "adapted" + if source == "unary": + nodes.append(helper.make_node("Relu", [data], ["unary"])) + data = "unary" + operands, attributes = [data], {"axis": rank - 1} + output_shape = shape.copy() + if operation == "Slice": + operands += ["starts", "ends", "slice_axes"] + attributes = {} + for name, values in [("starts", [1]), ("ends", [5]), ("slice_axes", [rank - 1])]: + initializers.append(numpy_helper.from_array(np.array(values, np.int64), name)) + output_shape[-1] = 4 + elif operation == "Transpose": + attributes = {"perm": list(reversed(range(rank)))} + output_shape.reverse() + elif operation == "ArgMax": + attributes["keepdims"] = 1 + output_shape[-1] = 1 + elif operation in {"Expand", "Tile"}: + output_shape[0] = 2 + values = output_shape if operation == "Expand" else [2] + [1] * (rank - 1) + initializers.append(numpy_helper.from_array(np.array(values, np.int64), "control")) + operands.append("control") + attributes = {} + nodes.append(helper.make_node(operation, operands, ["y"], **attributes)) + model = helper.make_model( + helper.make_graph( + nodes, + "rank_one_adapter_contract", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [6])], + [ + helper.make_tensor_value_info( + "y", + TensorProto.INT64 if operation == "ArgMax" else TensorProto.FLOAT, + output_shape, + ) + ], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + cases = [{"x": np.arange(1, 7, dtype=np.float32)}, {"x": np.arange(6, 0, -1, dtype=np.float32)}] + expected = [_execute(model, feeds)[0] for feeds in cases] + changed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + for feeds, reference in zip(cases, expected): + np.testing.assert_allclose(_execute(changed, feeds)[0], reference, rtol=1e-6, atol=1e-7) + + +def test_registered_slice_without_input_promotion_retains_native_axes(monkeypatch): + """Keep native Slice controls when the registered sequence omits input promotion.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Slice", ["x", "starts", "ends", "axes"], ["y"])], + "slice_without_input_promotion", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 2])], + [ + numpy_helper.from_array(np.array(values, np.int64), name) + for name, values in [("starts", [0]), ("ends", [2]), ("axes", [1])] + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(6, dtype=np.float32).reshape(2, 3)} + expected = _execute(model, feeds)[0] + monkeypatch.setattr( + dla_pipeline, + "DLA_TRANSFORMS", + tuple( + (name, callback) + for name, callback in dla_pipeline.DLA_TRANSFORMS + if name != "Non4dModelInputs" + ), + ) + changed = dla_pipeline._transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(changed, feeds)[0], expected) + + +@pytest.mark.parametrize("bad_first", [False, True]) +def test_flatten_refuses_native_gemm_before_mutating_other_branches(bad_first): + """Keep every branch unchanged when a native Gemm still needs a matrix input.""" + flatten = helper.make_node("Flatten", ["x"], ["flat"], name="gemm_flatten") + peer = helper.make_node("Flatten", ["x"], ["peer"], name="peer_flatten") + nodes = [flatten, peer] if bad_first else [peer, flatten] + nodes.append(helper.make_node("Gemm", ["flat", "weight"], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "native_gemm", + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3, 4]), + helper.make_tensor_value_info("weight", TensorProto.FLOAT, [12, 5]), + ], + [ + helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 5]), + helper.make_tensor_value_info("peer", TensorProto.FLOAT, [2, 12]), + ], + value_info=[helper.make_tensor_value_info("flat", TensorProto.FLOAT, [2, 12])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + original = model.SerializeToString() + with pytest.raises(ValueError, match="Gemm"): + _apply_flatten(model) + assert model.SerializeToString() == original + + +@pytest.mark.parametrize("native_shape", [[2, 12], [2, 3, 4]]) +@pytest.mark.parametrize("unary", [False, True]) +def test_internal_native_flatten_keeps_its_output_rank(native_shape, unary): + """Keep a native internal Flatten rank visible to a downstream Shape reader.""" + source_shape = [1] * (4 - len(native_shape)) + native_shape + nodes = [helper.make_node("Squeeze", ["x", "axes"], ["native"])] + data = "native" + if unary: + nodes.append(helper.make_node("Relu", [data], ["activated"])) + data = "activated" + nodes.extend( + [ + helper.make_node("Flatten", [data], ["flat"], axis=1), + helper.make_node("Shape", ["flat"], ["y"]), + ] + ) + model = helper.make_model( + helper.make_graph( + nodes, + "native_internal_flatten", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, source_shape)], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2])], + [numpy_helper.from_array(np.arange(4 - len(native_shape), dtype=np.int64), "axes")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(24, dtype=np.float32).reshape(source_shape)} + expected = _execute(model, feeds)[0] + result = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + np.testing.assert_array_equal(_execute(result, feeds)[0], expected) + assert any(node.op_type == "Flatten" for node in result.graph.node)