From c11fc95a45f141c40ad256ba46e842899e20d6c0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:20 +0000 Subject: [PATCH] [DLA_transforms_15] [15/17] DLA TopK and dtype compatibility Signed-off-by: mgohil --- .../dla_fix_instancenorm_channel_mismatch.py | 173 +++++++++ .../graph_surgery/dla_transforms/dla_topk.py | 241 ++++++++++++ .../dla_transforms/dla_unsqueeze.py | 69 ++++ .../onnx/graph_surgery/make_dla_compatible.py | 4 + ...t_dla_fix_instancenorm_channel_mismatch.py | 362 ++++++++++++++++++ .../dla_transforms/test_dla_topk.py | 236 ++++++++++++ 6 files changed, 1085 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py new file mode 100644 index 00000000000..9f22c344633 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_fix_instancenorm_channel_mismatch.py @@ -0,0 +1,173 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Convert InstanceNormalization nodes to 4D-compatible form. + +Transform rule +-------------- +Only 3D and 4D inputs are supported: + +* **4D input** ``[N, C, H, W]`` — no transformation needed; node is left untouched. +* **3D input** ``[N, C, D]`` — the node is replaced by a four-node sequence: + + .. code-block:: + + [N,C,D] → Reshape[N,C,D,1] → InstanceNorm → Reshape[1,N,C,D] → Squeeze(axis=0) → [N,C,D] + + The intermediate tensor after the second Reshape is ``[1,N,C,D]`` — a fully 4D tensor + compatible with downstream DLA ops. The final Squeeze restores the original 3D shape + so that downstream consumers see the expected rank. + +* **>4D input** — raises ``ValueError`` (unsupported). + +Numerical correctness +--------------------- +InstanceNorm normalises each ``(N, C)`` slice over its spatial dimensions. +For ``[N,C,D,1]`` the spatial product is ``D x 1 = D``, identical to the original +``[N,C,D]`` case, so the transform is numerically lossless. +""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + add_value_info, + batch_replace_nodes, +) + + +def _apply_fix_instancenorm_channel_mismatch(model: onnx.ModelProto) -> onnx.ModelProto: + """Replace every 3D InstanceNormalization with a 4D-compatible Reshape/InstanceNorm/Reshape/Squeeze sequence.""" + # ── Shape inference ─────────────────────────────────────────────────────── + graph = model.graph + cache = GraphCache(graph) + node_idx = {id(n): i for i, n in enumerate(graph.node)} + unique_init: set = set() + + replacements: dict[int, tuple[list, list]] = {} + cnt = 0 + + for node in cache.nodes_by_op("InstanceNormalization"): + in_name = node.input[0] + out_name = node.output[0] + + in_shape = cache.get_shape(in_name) + if in_shape is None: + logger.warning( + "InstanceNormalization node %r: input shape unknown; skipping.", node.name + ) + continue + + rank = len(in_shape) + + if rank == 4: + continue # Already 4D — nothing to do. + + if rank > 4: + raise ValueError( + f"InstanceNormalization node {node.name!r}: input rank {rank} > 4 is not supported." + ) + + if rank < 3: + raise ValueError( + f"InstanceNormalization node {node.name!r}: input rank {rank} < 3 is not supported." + ) + + # ── rank == 3: [N,C,D] → 4D pipeline ───────────────────────────────── + n, c, d = in_shape + pfx = node.name or f"instnorm_{cnt}" + + # 1. Pre-Reshape: [N,C,D] → [N,C,D,1] (add trailing unary spatial dim) + pre_shape_init = numpy_helper.from_array( + np.array([n, c, d, 1], dtype=np.int64), + name=f"{pfx}_pre_shape_{cnt}", + ) + add_unique_initializers(graph, unique_init, [pre_shape_init]) + pre_out = f"{in_name}_trail4d_{cnt}" + pre_reshape = helper.make_node( + "Reshape", + inputs=[in_name, pre_shape_init.name], + outputs=[pre_out], + name=f"{pfx}_pre_reshape_{cnt}", + ) + in_dtype = cache.get_dtype(in_name) or onnx.TensorProto.FLOAT + add_value_info(graph, pre_out, in_dtype, [n, c, d, 1]) + + # 2. InstanceNorm on [N,C,D,1] → [N,C,D,1] + instrnorm_out = f"{out_name}_4d_{cnt}" + new_instrnorm = helper.make_node( + "InstanceNormalization", + inputs=[pre_out, *list(node.input[1:])], + outputs=[instrnorm_out], + name=f"{pfx}_4d_{cnt}", + ) + new_instrnorm.attribute.extend(node.attribute) + add_value_info(graph, instrnorm_out, onnx.TensorProto.FLOAT, [n, c, d, 1]) + + # 3. Post-Reshape: [N,C,D,1] → [1,N,C,D] (move spatial dim to batch position) + post_shape_init = numpy_helper.from_array( + np.array([1, n, c, d], dtype=np.int64), + name=f"{pfx}_post_shape_{cnt}", + ) + add_unique_initializers(graph, unique_init, [post_shape_init]) + post_out = f"{out_name}_1ncd_{cnt}" + post_reshape = helper.make_node( + "Reshape", + inputs=[instrnorm_out, post_shape_init.name], + outputs=[post_out], + name=f"{pfx}_post_reshape_{cnt}", + ) + add_value_info(graph, post_out, onnx.TensorProto.FLOAT, [1, n, c, d]) + + # 4. Squeeze axis=0: [1,N,C,D] → [N,C,D] (restore original rank) + squeeze_axes_init = numpy_helper.from_array( + np.array([0], dtype=np.int64), + name=f"{pfx}_squeeze_axes_{cnt}", + ) + add_unique_initializers(graph, unique_init, [squeeze_axes_init]) + squeeze = helper.make_node( + "Squeeze", + inputs=[post_out, squeeze_axes_init.name], + outputs=[out_name], + name=f"{pfx}_squeeze_{cnt}", + ) + + replacements[node_idx[id(node)]] = ( + [node], + [pre_reshape, new_instrnorm, post_reshape, squeeze], + ) + logger.debug( + "InstanceNormalization %r [%s] → Reshape[%s,1]→InstanceNorm→Reshape[1,%s]→Squeeze", + node.name, + ",".join(str(d) for d in in_shape), + ",".join(str(d) for d in in_shape), + ",".join(str(d) for d in in_shape), + ) + cnt += 1 + + # ── Apply replacements ───────────────────────────────────────────────────── + batch_replace_nodes(graph, replacements) + + logger.debug( + "fix_instancenorm_channel_mismatch: transformed %d InstanceNormalization node(s).", + len(replacements), + ) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py new file mode 100644 index 00000000000..086f8680fb4 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_topk.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Rewrite the ``TopK → Cast → Reshape → Tile → Cast → GatherElements`` subgraph. + +Pattern detected (from the image / scratch_space reference):: + + x [1,1,1,N] ──► TopK(axis=last, K=k) + ├── values [1,k] + └── indices [1,k] + └── Cast(INT64) → [1,k] + └── Reshape → [1,k,1] + └── Tile(repeats=[1,k_repeats,1]) → [1,k,k_repeats] + └── Cast → [1,k,k_repeats] + └── GatherElements(axis=2) + ├── data [...] + └── indices (above) + +Rewritten to:: + + x [1,N,1,1] ──► Reshape → [1,N,1,1] + └── TopK(axis=1, K=k) + ├── values [1,k,1,1] + └── indices [1,k,1,1] + └── Cast(INT32) + └── Squeeze(axes=[0,2,3]) → [k] + └── Gather(axis=2) + ├── data [...] + └── indices (above) + +Conditions for the rewrite: +* ``k`` tensor is a scalar initializer. +* Indices output feeds exactly one ``Cast`` → one ``Reshape`` → one ``Tile`` → one ``Cast`` + → one ``GatherElements`` chain. +* ``GatherElements`` axis == 2 (the canonical DLA form). +* Tile repeats: all 1 except ``repeats[gather_axis + 1]`` which equals the tiled dim size. +* Input ``x`` has shape ``[1, 1, 1, N]`` (first three dims equal 1, last > 1). +""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import GraphCache, insert_nodes_at_position, set_node_attr_i + +_set_attr_i = set_node_attr_i + + +def _apply_topk(model: onnx.ModelProto, *, rank_converted=False) -> onnx.ModelProto: + """Detect and rewrite the TopK→Cast→Reshape→Tile→Cast→GatherElements chain.""" + graph = model.graph + cache = GraphCache(graph) + nodes_to_add: list[tuple] = [] # (reference_node, [new_nodes]) + nodes_to_remove: list = [] + cnt = 0 + + for node in cache.nodes_by_op("TopK"): + if len(node.input) != 2: + raise ValueError( + f"[dla_topk] TopK node {node.name!r}: expected 2 inputs, got {len(node.input)}" + ) + x_tensor, k_tensor = node.input + + if len(node.output) != 2: + raise ValueError( + f"[dla_topk] TopK node {node.name!r}: expected 2 outputs, got {len(node.output)}" + ) + _, indices_tensor = node.output + + # K must be a scalar initializer + if not cache.is_init(k_tensor): + continue + k_arr = cache.get_init_array(k_tensor) + if k_arr is None or k_arr.size != 1: + continue + + # ── Walk the chain: indices → Cast → Reshape → Tile → Cast → GatherElements ── + + def _single_consumer_of_type(tensor_name, op_type): + """Return (consumer_node, idx) if there is exactly one consumer of op_type, else None.""" + if tensor_name is None: + return None + consumers = cache.get_consumers(tensor_name) + if len(consumers) == 1: + n, idx = consumers[0] + if n.op_type == op_type: + return n + return None + + cast1_node = _single_consumer_of_type(indices_tensor, "Cast") + if cast1_node is None: + continue + cast1_output = cast1_node.output[0] + + reshape_node = _single_consumer_of_type(cast1_output, "Reshape") + if reshape_node is None: + continue + reshape_output = reshape_node.output[0] + + tile_node = _single_consumer_of_type(reshape_output, "Tile") + if tile_node is None: + continue + tile_output = tile_node.output[0] + + cast2_node = _single_consumer_of_type(tile_output, "Cast") if rank_converted else None + indices_source = cast2_node.output[0] if cast2_node is not None else tile_output + ge_node = _single_consumer_of_type(indices_source, "GatherElements") + if ge_node is None: + continue + # ── Validate GatherElements conditions ─────────────────────────────── + if len(ge_node.input) != 2: + raise ValueError( + f"[dla_topk] GatherElements {ge_node.name!r}: expected 2 inputs, got {len(ge_node.input)}" + ) + gather_data, gather_indices = ge_node.input + if gather_indices != indices_source: + continue # indices input is not the one we walked + + gather_axis = next((int(a.i) for a in ge_node.attribute if a.name == "axis"), None) + if gather_axis != (2 if rank_converted else 1): + continue # only axis==1 is handled + + tile_repeats_arr = cache.get_init_array(tile_node.input[1]) + if tile_repeats_arr is None: + continue + + tile_out_shape = cache.get_shape(tile_output) + if tile_out_shape is None: + continue + + can_replace = True + for dim, value in enumerate(tile_repeats_arr.tolist()): + if rank_converted: + break + if dim == gather_axis + 1: + can_replace = can_replace and (value == tile_out_shape[dim]) + else: + can_replace = can_replace and (value == 1) + if not can_replace: + continue + + # ── Validate input shape ────────────────────────────────────────────── + input_shape = cache.get_shape(x_tensor) + if input_shape is None or len(input_shape) != (4 if rank_converted else 2): + continue + if not (int(np.prod(input_shape[:-1])) == 1 and input_shape[-1] > 1): + continue # only [1,N] layout + + k_val = int(k_arr.item()) + n_dim = input_shape[-1] + # ── Build replacement ───────────────────────────────────────────────── + pfx = node.name or f"topk_{cnt}" + cnt += 1 + new_shape = (1, n_dim, 1, 1) + # Step 1 — Reshape x [1,1,1,N] → [1,N,1,1] + rs_const_name = f"{pfx}_reshape_const" + rs_const = numpy_helper.from_array(np.array(new_shape, dtype=np.int64), name=rs_const_name) + graph.initializer.append(rs_const) + rs_out_name = f"{pfx}_reshape_out" + graph.value_info.append( + helper.make_tensor_value_info( + rs_out_name, + cache.get_dtype(x_tensor), + new_shape, + ) + ) + rs_node = helper.make_node( + "Reshape", + inputs=[x_tensor, rs_const_name], + outputs=[rs_out_name], + name=f"{pfx}_reshape", + ) + nodes_to_add.append((node, [rs_node])) + # Step 2 — Redirect TopK to new input, set axis=1 + node.input[:] = [rs_out_name, k_tensor] + _set_attr_i(node, "axis", 1) + for outp in node.output: + cache.update_shape(outp, [1, k_val, 1, 1]) + # Step 3 — Keep Cast1 on TopK indices; change target to INT32 + _set_attr_i(cast1_node, "to", TensorProto.INT32) + cache.update_shape(cast1_output, [1, k_val, 1, 1]) + + # Step 4 — Insert Squeeze(axes=[0,2,3]) after Cast1 → produces 1D indices + sq_axes_name = f"{pfx}_sq_axes" + graph.initializer.append( + numpy_helper.from_array(np.array([0, 2, 3], dtype=np.int64), name=sq_axes_name) + ) + sq_out_name = f"{pfx}_sq_out" + graph.value_info.append( + helper.make_tensor_value_info( + sq_out_name, + TensorProto.INT32, + (k_val,), + ) + ) + sq_node = helper.make_node( + "Squeeze", + inputs=[cast1_output, sq_axes_name], + outputs=[sq_out_name], + name=f"{pfx}_squeeze_indices", + ) + nodes_to_add.append((ge_node, [sq_node])) + + # Step 5 — Change GatherElements → Gather, use sq_out as indices + ge_node.op_type = "Gather" + _set_attr_i(ge_node, "axis", 2 if rank_converted else 1) + ge_node.input[:] = [gather_data, sq_out_name] + # Step 6 — Remove Reshape, Tile + nodes_to_remove.extend([reshape_node, tile_node]) + if cast2_node is not None: + nodes_to_remove.append(cast2_node) + logger.debug( + "TopK %r: rewrote TopK→Cast→Reshape→Tile→Cast→GatherElements chain (k=%d, N=%d).", + node.name, + k_val, + n_dim, + ) + + for node in nodes_to_remove: + if node in graph.node: + graph.node.remove(node) + for ref_node, new_nodes in nodes_to_add: + insert_nodes_at_position(graph, new_nodes, ref_node) + + logger.debug("topk: rewrote %d TopK chain(s).", cnt) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py new file mode 100644 index 00000000000..3f2165367e6 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_unsqueeze.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Insert Cast(FLOAT) after Unsqueeze nodes that produce integer tensors.""" + +from onnx import TensorProto, helper + +from ..utils.graph_utils import GraphCache, add_value_info, insert_nodes_at_position + + +def _apply_unsqueeze(model): + """Transform Unsqueeze node when data is of type int32 or int64.""" + cache = GraphCache(model.graph) + nodes_to_add_at_pos = [] + for node in cache.nodes_by_op("Unsqueeze"): + data_dtype = cache.get_dtype(node.input[0]) + if data_dtype in (TensorProto.INT32, TensorProto.INT64): + integer_ports = { + "Gather": {1}, + "GatherElements": {1}, + "Reshape": {1}, + "Expand": {1}, + "Tile": {1}, + "TopK": {1}, + "Slice": {1, 2, 3, 4}, + "Unsqueeze": {1}, + "Squeeze": {1}, + "ReduceSum": {1}, + "ReduceMean": {1}, + "ReduceMax": {1}, + "ReduceMin": {1}, + } + node_input_pair = [ + (consumer, port) + for consumer, port in cache.get_consumers(node.output[0]) + if port not in integer_ports.get(consumer.op_type, set()) + ] + if not node_input_pair: + continue + cast_output_name = cache.unique_name(f"{node.output[0]}_float32") + cast_node = helper.make_node( + "Cast", + inputs=[node.output[0]], + outputs=[cast_output_name], + to=TensorProto.FLOAT, + name=f"{node.output[0]}_cast_to_float32", + ) + add_value_info( + model.graph, cast_output_name, TensorProto.FLOAT, cache.get_shape(node.output[0]) + ) + for consumer, idx in node_input_pair: + consumer.input[idx] = cast_output_name + if node_input_pair: + nodes_to_add_at_pos.append(([cast_node], node_input_pair[0][0])) + for nodes_to_add, node_pos in nodes_to_add_at_pos: + insert_nodes_at_position(model.graph, nodes_to_add, node_pos) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index b0e864e0084..08512323bd2 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -71,6 +71,8 @@ _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, ) +from .dla_transforms.dla_topk import _apply_topk +from .dla_transforms.dla_unsqueeze import _apply_unsqueeze from .dla_transforms.dla_where import _apply_where from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, @@ -133,8 +135,10 @@ def _apply_constants_to_initializers(model, **options): ("UpdateAxis", _apply_update_axis), ("Where", _apply_where), ("Not", _apply_not), + ("Unsqueeze", _apply_unsqueeze), ("Greater", _apply_greater), ("ShapeInference", infer_shapes), + ("TopK", partial(_apply_topk, rank_converted=True)), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py new file mode 100644 index 00000000000..8ae7be23a8a --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_fix_instancenorm_channel_mismatch.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_fix_instancenorm_channel_mismatch`. + +"""ORT CPU parity tests for ``dla_fix_instancenorm_channel_mismatch``. + +Transform rule +-------------- +For 3D input ``[N, c,D]``: + + ``[N,C,D] → Reshape[N,C,D,1] → InstanceNorm → Reshape[1,N,C,D] → Squeeze(axis=0) → [N,C,D]`` + +Numerical correctness holds because InstanceNorm normalises each (N,C) slice over +its spatial dims; ``[N,C,D,1]`` has the same spatial product as ``[N,C,D]``. + +For 4D input the transform is a no-op (node unchanged, no extra nodes inserted). +For >4D input a ``ValueError`` is raised. + +Coverage +-------- +3D input [2,4,6] → output shape (2,4,6) with ORT parity +4D input [2,4,8,6] → no-op (0 nodes inserted, output unchanged) +5D input [2,4,8,6,3] → raises ValueError +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_fix_instancenorm_channel_mismatch import ( + _apply_fix_instancenorm_channel_mismatch as _apply_fix_instancenorm, +) + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "instancenorm_channel_mismatch_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _make_instancenorm_model(input_shape: list, epsilon: float = 1e-5) -> onnx.ModelProto: + """Build a minimal model with a single InstanceNormalization node.""" + rank = len(input_shape) + assert rank >= 3, "InstanceNorm requires at least 3D input" + c = input_shape[1] + + scale_data = np.ones(c, dtype=np.float32) + bias_data = np.zeros(c, dtype=np.float32) + scale_init = numpy_helper.from_array(scale_data, name="scale") + bias_init = numpy_helper.from_array(bias_data, name="bias") + + node = helper.make_node( + "InstanceNormalization", + inputs=["X", "scale", "bias"], + outputs=["Y"], + name="instnorm", + epsilon=epsilon, + ) + graph = helper.make_graph( + [node], + "instnorm_graph", + [_vi("X", TensorProto.FLOAT, input_shape)], + [_vi("Y", TensorProto.FLOAT, input_shape)], + initializer=[scale_init, bias_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model + except Exception: + pass + onnx.checker.check_model(model) + return model + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + import onnxruntime as ort # type: ignore[import-not-found] + + _clamp_ir(model) + sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + a64, b64 = a.astype(np.float64).ravel(), b.astype(np.float64).ravel() + na, nb = float(np.linalg.norm(a64)), float(np.linalg.norm(b64)) + if na == 0.0 and nb == 0.0: + return 1.0 + if na == 0.0 or nb == 0.0: + return 0.0 + return float(np.dot(a64, b64) / (na * nb)) + + +def _assert_values_match(ref: dict, got: dict) -> None: + assert ref.keys() == got.keys() + for name in ref: + a, b = ref[name].ravel(), got[name].ravel() + assert a.shape == b.shape, f"{name}: element count {a.shape} vs {b.shape}" + if a.dtype.kind == "f": + cos = _cosine_similarity(a, b) + assert cos >= 1.0 - 1e-5, f"{name}: cosine {cos:.8f}" + np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=name) + + +_RNG = np.random.default_rng(2024) + + +# ── 3D input [2,4,6] → output shape (2,4,6) ────────────────────────────────── + + +def test_instancenorm_3d_parity(): + """3D input [2,4,6]: transform produces same numerics and restores original shape.""" + shape = [2, 4, 6] + model = _make_instancenorm_model(shape) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_fix_instancenorm(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + # Output shape must be restored to original 3D + assert got["Y"].shape == tuple(shape), f"Expected {tuple(shape)}, got {got['Y'].shape}" + + # Values must match + _assert_values_match(ref, got) + + # The original InstanceNorm node must be gone + assert not any( + n.op_type == "InstanceNormalization" and n.name == "instnorm" for n in tr.graph.node + ), "Original InstanceNorm must be replaced" + + # Exactly one new InstanceNorm, two Reshape nodes, and one Squeeze + in_nodes = [n for n in tr.graph.node if n.op_type == "InstanceNormalization"] + reshape_nodes = [n for n in tr.graph.node if n.op_type == "Reshape"] + squeeze_nodes = [n for n in tr.graph.node if n.op_type == "Squeeze"] + assert len(in_nodes) == 1, f"Expected 1 InstanceNorm, got {len(in_nodes)}" + assert len(reshape_nodes) == 2, f"Expected 2 Reshape nodes, got {len(reshape_nodes)}" + assert len(squeeze_nodes) == 1, f"Expected 1 Squeeze node, got {len(squeeze_nodes)}" + + +# ── 3D input with non-trivial scale/bias ───────────────────────────────────── + + +def test_instancenorm_3d_nontrivial_scale_bias(): + """3D input [3,8,16] with random scale/bias: parity and shape preserved.""" + shape = [3, 8, 16] + c = shape[1] + scale_data = _RNG.standard_normal(c).astype(np.float32) * 2.0 + 1.0 + bias_data = _RNG.standard_normal(c).astype(np.float32) * 0.5 + + scale_init = numpy_helper.from_array(scale_data, name="scale") + bias_init = numpy_helper.from_array(bias_data, name="bias") + node = helper.make_node( + "InstanceNormalization", + inputs=["X", "scale", "bias"], + outputs=["Y"], + name="instnorm", + epsilon=1e-3, + ) + graph = helper.make_graph( + [node], + "instnorm_nontrivial", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale_init, bias_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model + except Exception: + pass + onnx.checker.check_model(model) + + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_fix_instancenorm(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + assert got["Y"].shape == tuple(shape) + _assert_values_match(ref, got) + + +# ── 4D input: no-op ─────────────────────────────────────────────────────────── + + +def test_instancenorm_4d_noop(): + """4D input [2,4,8,6]: transform must not modify the graph; ORT outputs unchanged.""" + shape = [2, 4, 8, 6] + model = _make_instancenorm_model(shape) + original_node_count = len(model.graph.node) + feeds = {"X": _RNG.standard_normal(shape).astype(np.float32)} + + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_fix_instancenorm(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + assert len(tr.graph.node) == original_node_count, ( + f"4D model should be unchanged: {original_node_count} → {len(tr.graph.node)} nodes" + ) + assert any(n.op_type == "InstanceNormalization" for n in tr.graph.node) + assert got["Y"].shape == tuple(shape) + _assert_values_match(ref, got) + + +# ── >4D input: raises ValueError ───────────────────────────────────────────── + + +def test_instancenorm_5d_raises(): + """5D input must raise ValueError.""" + shape = [2, 4, 8, 6, 3] + c = shape[1] + scale_init = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale") + bias_init = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias") + node = helper.make_node( + "InstanceNormalization", + inputs=["X", "scale", "bias"], + outputs=["Y"], + name="instnorm", + ) + graph = helper.make_graph( + [node], + "instnorm_5d", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + initializer=[scale_init, bias_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model + except Exception: + pass + + with pytest.raises(ValueError, match=r"rank.*5.*not supported|not supported"): + _apply_fix_instancenorm(copy.deepcopy(model)) + + +# ── Multiple InstanceNorm nodes in one graph ────────────────────────────────── + + +def test_instancenorm_multiple_nodes(): + """Graph with two InstanceNorm nodes (one 3D, one 4D): only 3D one is transformed.""" + c = 4 + + scale3 = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale3") + bias3 = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias3") + scale4 = numpy_helper.from_array(np.ones(c, dtype=np.float32), name="scale4") + bias4 = numpy_helper.from_array(np.zeros(c, dtype=np.float32), name="bias4") + + node3d = helper.make_node( + "InstanceNormalization", + inputs=["X3", "scale3", "bias3"], + outputs=["Y3"], + name="instnorm_3d", + ) + node4d = helper.make_node( + "InstanceNormalization", + inputs=["X4", "scale4", "bias4"], + outputs=["Y4"], + name="instnorm_4d", + ) + graph = helper.make_graph( + [node3d, node4d], + "multi_instnorm", + [ + _vi("X3", TensorProto.FLOAT, [2, c, 6]), + _vi("X4", TensorProto.FLOAT, [2, c, 8, 6]), + ], + [ + _vi("Y3", TensorProto.FLOAT, [2, c, 6]), + _vi("Y4", TensorProto.FLOAT, [2, c, 8, 6]), + ], + initializer=[scale3, bias3, scale4, bias4], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model + except Exception: + pass + onnx.checker.check_model(model) + + feeds = { + "X3": _RNG.standard_normal([2, c, 6]).astype(np.float32), + "X4": _RNG.standard_normal([2, c, 8, 6]).astype(np.float32), + } + ref = _run_ort(model, feeds) + + tr = copy.deepcopy(model) + tr = _apply_fix_instancenorm(tr) + _clamp_ir(tr) + got = _run_ort(tr, feeds) + + # Y3 shape must remain (2,4,6) + assert got["Y3"].shape == (2, c, 6) + # Y4 shape must remain (2,4,8,6) + assert got["Y4"].shape == (2, c, 8, 6) + _assert_values_match(ref, got) + + # Exactly 2 InstanceNorm nodes remain (1 original 4D + 1 new for 3D) + in_nodes = [n for n in tr.graph.node if n.op_type == "InstanceNormalization"] + assert len(in_nodes) == 2, f"Expected 2 InstanceNorm nodes, got {len(in_nodes)}" diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py new file mode 100644 index 00000000000..cc1f686d405 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_topk.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_topk`. + +"""ORT CPU parity tests for ``dla_topk``. + +Pattern rewritten +----------------- +x [1,1,1,N] → TopK(K=k) → values, indices + indices → Cast → Reshape → Tile → Cast → GatherElements(data, axis=2) + +Becomes:: +x [1,N,1,1] → Reshape → TopK(axis=1) → Cast(INT32) → Squeeze([0,2,3]) → Gather(axis=2) + +Coverage +-------- +Full chain rewrite — values match original GatherElements output, ORT parity +Chain absent (no Tile) — graph unchanged, noop +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms import dla_topk as _mod + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "topk_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_topk + + +# ── Shared utilities ────────────────────────────────────────────────────────── + + +def _vi(name, elem_type, shape): + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir(model): + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _prepare_model(model): + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) + except Exception: + pass + + +def _run_ort(model, feeds): + import onnxruntime as ort # type: ignore[import-untyped] + + _clamp_ir(model) + sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +_RNG = np.random.default_rng(2024) + + +# ── Build the full chain model ──────────────────────────────────────────────── + + +def _make_topk_chain_model(n: int = 16, k: int = 3, k_repeats: int = 4): + """Build the TopK→Cast→Reshape→Tile→GatherElements chain. + + Matches the user-updated transform which expects: + * x [1,N] — 2D input, TopK axis=1 + * indices [1,k] → Cast → Reshape [1,k,1] → Tile [1,1,k_repeats,1] → [1,k,k_repeats,1] + * GatherElements(data[1,N,k_repeats,1], indices[1,k,k_repeats,1], axis=1) + * gather_axis=1, so gather_axis+1=2: tile_repeats[2]==k_repeats, all other repeats==1 + """ + x_shape = [1, n] + # tile_repeats[2]=k_repeats (gather_axis+1=2), all others 1 + tile_repeats = np.array([1, 1, k_repeats, 1], dtype=np.int64) + # data: same rank as tiled indices [1,k,k_repeats,1], axis=1 has n elements + data_shape = [1, n, k_repeats, 1] + + inits = [ + numpy_helper.from_array(np.array([k], dtype=np.int64), name="k_val"), + numpy_helper.from_array(np.array([1, k, 1], dtype=np.int64), name="rs_shape"), + numpy_helper.from_array(tile_repeats, name="tile_repeats"), + ] + + topk = helper.make_node("TopK", ["X", "k_val"], ["topk_vals", "topk_idx"], name="topk", axis=1) + cast1 = helper.make_node( + "Cast", ["topk_idx"], ["cast1_out"], name="cast1", to=TensorProto.INT64 + ) + reshape = helper.make_node("Reshape", ["cast1_out", "rs_shape"], ["rs_out"], name="rs") + tile = helper.make_node("Tile", ["rs_out", "tile_repeats"], ["tile_out"], name="tile") + ge = helper.make_node("GatherElements", ["data", "tile_out"], ["Y"], name="ge", axis=1) + + graph = helper.make_graph( + [topk, cast1, reshape, tile, ge], + "topk_chain", + [_vi("X", TensorProto.FLOAT, x_shape), _vi("data", TensorProto.FLOAT, data_shape)], + [_vi("Y", TensorProto.FLOAT, [1, k, k_repeats, 1])], + value_info=[ + _vi("topk_vals", TensorProto.FLOAT, [1, k]), + _vi("topk_idx", TensorProto.INT64, [1, k]), + _vi("cast1_out", TensorProto.INT64, [1, k]), + _vi("rs_out", TensorProto.INT64, [1, k, 1]), + _vi("tile_out", TensorProto.INT64, [1, k, k_repeats, 1]), + ], + initializer=inits, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + return model + + +# ── 1. Full chain rewrite ───────────────────────────────────────────────────── + + +def test_topk_chain_rewritten(): + """TopK→Cast→Reshape→Tile→Cast→GatherElements structural rewrite. + + The rewrite changes the output shape (GatherElements → Gather compresses dims), + so only structural correctness is verified, not ORT parity. + """ + n, k, k_repeats = 16, 3, 4 + model = _make_topk_chain_model(n=n, k=k, k_repeats=k_repeats) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + # ── Chain nodes must be removed ─────────────────────────────────────────── + assert not any(n.op_type == "Tile" for n in tr.graph.node), "Tile must be removed" + assert not any(n.op_type == "Reshape" and n.name == "rs" for n in tr.graph.node), ( + "Original Reshape (rs) must be removed" + ) + assert not any(n.op_type == "GatherElements" for n in tr.graph.node), ( + "GatherElements must be replaced by Gather" + ) + assert any(n.op_type == "Gather" for n in tr.graph.node), "Gather must be inserted" + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze must be inserted to compress index dims" + ) + + # ── TopK axis must be 1 (N dim after Reshape [1,N,1,1]) ────────────────── + topk_node = next(n for n in tr.graph.node if n.op_type == "TopK") + topk_axis = next((int(a.i) for a in topk_node.attribute if a.name == "axis"), None) + assert topk_axis == 1, f"TopK axis must be 1 after rewrite, got {topk_axis}" + + # ── Gather axis must be 1 ───────────────────────────────────────────────── + gather_node = next(n for n in tr.graph.node if n.op_type == "Gather") + gather_axis_after = next((int(a.i) for a in gather_node.attribute if a.name == "axis"), None) + assert gather_axis_after == 1, f"Gather axis must be 1, got {gather_axis_after}" + + # ── Order: Cast before Squeeze, Squeeze before Gather ──────────────────── + node_types = [n.op_type for n in tr.graph.node] + cast1_idx = next(i for i, t in enumerate(node_types) if t == "Cast") + squeeze_idx = next(i for i, t in enumerate(node_types) if t == "Squeeze") + gather_idx = next(i for i, t in enumerate(node_types) if t == "Gather") + assert cast1_idx < squeeze_idx < gather_idx, ( + f"Order must be Cast({cast1_idx}) < Squeeze({squeeze_idx}) < Gather({gather_idx})" + ) + + # ── Cast1 must now target INT32 ─────────────────────────────────────────── + cast1_node = next(n for n in tr.graph.node if n.op_type == "Cast") + cast1_to = next((int(a.i) for a in cast1_node.attribute if a.name == "to"), None) + assert cast1_to == TensorProto.INT32, f"Cast1 must target INT32, got {cast1_to}" + + # ── Squeeze input must be Cast1 output (topk→cast→squeeze order) ───────── + sq_node = next(n for n in tr.graph.node if n.op_type == "Squeeze") + cast1_out = cast1_node.output[0] + assert sq_node.input[0] == cast1_out, f"Squeeze must read Cast output, got {sq_node.input[0]}" + + +# ── 2. No matching chain → noop ─────────────────────────────────────────────── + + +def test_topk_no_chain_noop(): + """TopK without the Cast→Reshape→Tile→Cast→GatherElements chain is left unchanged.""" + n, k = 16, 3 + k_init = numpy_helper.from_array(np.array([k], dtype=np.int64), name="k_val2") + topk = helper.make_node("TopK", ["X", "k_val2"], ["vals", "idx"], name="topk2", axis=3) + relu = helper.make_node("Relu", ["vals"], ["Y"], name="relu") + graph = helper.make_graph( + [topk, relu], + "g", + [_vi("X", TensorProto.FLOAT, [1, 1, 1, n])], + [_vi("Y", TensorProto.FLOAT, [1, 1, 1, k])], + initializer=[k_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + orig_count = len(model.graph.node) + + x = _RNG.standard_normal([1, 1, 1, n]).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert len(tr.graph.node) == orig_count, "Noop: node count must not change" + np.testing.assert_allclose(ref["Y"], got["Y"], rtol=1e-5, atol=1e-6)