From 2903df13874105d8d1e7a1191a9f9e8a4bb81817 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:04 +0000 Subject: [PATCH 1/6] [DLA_transforms_4] [4/17] Private ordered DLA scaffold and node ordering Signed-off-by: mgohil --- .../dla_transforms/dla_cast_to_fp32.py | 89 ++++ .../dla_transforms/dla_reorder_nodes.py | 155 ++++++ .../onnx/graph_surgery/make_dla_compatible.py | 91 ++++ .../dla_transforms/test_dla_cast_to_fp32.py | 481 ++++++++++++++++++ .../dla_transforms/test_dla_reorder_nodes.py | 101 ++++ 5 files changed, 917 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py create mode 100644 modelopt/onnx/graph_surgery/make_dla_compatible.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py new file mode 100644 index 00000000000..6e7a6be2e51 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py @@ -0,0 +1,89 @@ +# 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. + +"""Redirect all Cast nodes to produce float32 output. + +Skip conditions +--------------- +A Cast node is left untouched when its output is consumed by an op that +requires a specific non-float32 dtype: + +* **Gather / GatherElements** ``input[1]`` (indices) — must stay INT32. +* **GroupQueryAttention** inputs — require FP16. + +These Cast nodes are inserted by earlier transforms +(``dla-convert-ops-to-4d`` for Gather/GatherElements indices, or +externally for GQA) and must not be overridden. +""" + +from __future__ import annotations + +import onnx +from onnx import TensorProto + +from ...logging_config import logger + + +def _apply_cast_to_fp32(model: onnx.ModelProto) -> onnx.ModelProto: + """Change every eligible Cast node's ``to`` attribute to FLOAT (float32). + + A Cast node is skipped when its output feeds: + - ``input[1]`` (indices) of a Gather or GatherElements node, or + - any input of a GroupQueryAttention node. + """ + graph = model.graph + + # ── Build set of protected Cast outputs ────────────────────────────────── + # Tensor names produced by Cast nodes that must keep their current target + # dtype: INT32 for Gather/GatherElements indices, FP16 for GQA inputs. + protected: set[str] = set() + for node in graph.node: + if node.op_type in ("Gather", "GatherElements"): + # input[1] is the indices tensor — must stay INT32 + if len(node.input) > 1 and node.input[1]: + protected.add(node.input[1]) + elif node.op_type in {"Reshape", "Expand", "Tile", "TopK"}: + protected.update(node.input[1:2]) + elif node.op_type == "Slice": + protected.update(node.input[1:]) + elif node.op_type in {"Not", "And", "Or", "Xor"}: + protected.update(node.input) + elif node.op_type == "Where": + protected.update(node.input[:1]) + elif node.op_type == "GroupQueryAttention": + for inp in node.input: + if inp: + protected.add(inp) + + # ── Redirect eligible Cast nodes to FLOAT ──────────────────────────────── + cnt = 0 + for node in graph.node: + if node.op_type != "Cast": + continue + if not node.output or node.output[0] in protected: + continue + for attr in node.attribute: + if attr.name == "to" and attr.i != TensorProto.FLOAT: + attr.i = TensorProto.FLOAT + # Update value_info dtype for the Cast output tensor + for vi in graph.value_info: + if vi.name == node.output[0]: + vi.type.tensor_type.elem_type = TensorProto.FLOAT + break + cnt += 1 + break + + logger.debug("cast_to_fp32: redirected %d Cast node(s) to float32.", cnt) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py new file mode 100644 index 00000000000..ea9eabb3819 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py @@ -0,0 +1,155 @@ +# 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. + +"""Final priority topological node reorder for DLA ONNX models.""" + +from __future__ import annotations + +import heapq +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import onnx + +from ..utils.graph_utils import _node_inputs + +PRIORITIZED_OPS = { + "Constant", + "ConstantOfShape", + "QuantizeLinear", + "DequantizeLinear", +} + + +def _node_label(node: onnx.NodeProto, index: int) -> str: + return node.name or f"" + + +def _build_edges(nodes: list[onnx.NodeProto]) -> tuple[list[set[int]], list[list[int]]]: + output_to_node: dict[str, int] = {} + for index, node in enumerate(nodes): + for output in node.output: + if output: + output_to_node[output] = index + + dependencies = [set() for _ in nodes] + consumers: list[list[int]] = [[] for _ in nodes] + seen_edges: set[tuple[int, int]] = set() + for consumer_index, node in enumerate(nodes): + for input_name in _node_inputs(node): + producer_index = output_to_node.get(input_name) + if producer_index is None: + continue + + edge = (producer_index, consumer_index) + if edge in seen_edges: + continue + + seen_edges.add(edge) + dependencies[consumer_index].add(producer_index) + consumers[producer_index].append(consumer_index) + + return dependencies, consumers + + +def _unresolved_dependency_error(nodes: list[onnx.NodeProto], indegrees: list[int]) -> ValueError: + remaining = [ + _node_label(nodes[index], index) for index, indegree in enumerate(indegrees) if indegree > 0 + ] + return ValueError( + "Graph contains a cycle or unresolved dependency involving: " + ", ".join(remaining[:10]) + ) + + +def _compute_depths( + nodes: list[onnx.NodeProto], + dependencies: list[set[int]], + consumers: list[list[int]], +) -> list[int]: + depths = [0] * len(nodes) + indegrees = [len(node_dependencies) for node_dependencies in dependencies] + ready = [index for index, indegree in enumerate(indegrees) if indegree == 0] + + processed = 0 + read_index = 0 + while read_index < len(ready): + index = ready[read_index] + read_index += 1 + processed += 1 + + for consumer_index in consumers[index]: + depths[consumer_index] = max(depths[consumer_index], depths[index] + 1) + indegrees[consumer_index] -= 1 + if indegrees[consumer_index] == 0: + ready.append(consumer_index) + + if processed != len(nodes): + raise _unresolved_dependency_error(nodes, indegrees) + + return depths + + +def priority_topological_order( + graph: onnx.GraphProto, +) -> tuple[list[int], list[int], list[int]]: + """Return node indices in DLA priority topological order. + + The order always respects data dependencies. Among ready nodes, generic ops + keep the lowest priority, while Constant/QDQ-style support nodes are delayed + toward their first consumer depth so they appear nearer to their use. + """ + nodes = list(graph.node) + dependencies, consumers = _build_edges(nodes) + depths = _compute_depths(nodes, dependencies, consumers) + sink_priority = max(depths, default=0) + 1 + + def node_priority(index: int) -> int: + if nodes[index].op_type not in PRIORITIZED_OPS: + return 0 + if consumers[index]: + return depths[consumers[index][0]] + return sink_priority + + indegrees = [len(node_dependencies) for node_dependencies in dependencies] + ready: list[tuple[int, int]] = [] + for index, indegree in enumerate(indegrees): + if indegree == 0: + heapq.heappush(ready, (node_priority(index), index)) + + order: list[int] = [] + while ready: + _, index = heapq.heappop(ready) + order.append(index) + + for consumer_index in consumers[index]: + indegrees[consumer_index] -= 1 + if indegrees[consumer_index] == 0: + heapq.heappush(ready, (node_priority(consumer_index), consumer_index)) + + if len(order) != len(nodes): + raise _unresolved_dependency_error(nodes, indegrees) + + priorities = [node_priority(index) for index in range(len(nodes))] + return order, depths, priorities + + +def _apply_reorder_nodes(model: onnx.ModelProto) -> onnx.ModelProto: + """Reorder ``graph.node`` in-place using DLA priority topological order.""" + order, _, _ = priority_topological_order(model.graph) + nodes = list(model.graph.node) + sorted_nodes = [nodes[index] for index in order] + del model.graph.node[:] + model.graph.node.extend(sorted_nodes) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py new file mode 100644 index 00000000000..d091ea53b32 --- /dev/null +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DLA graph transforms in the order recorded by the C++ reference pipeline. + +DLA_TRANSFORMS is the single ordered list of supported passes. Individual reviews +can add passes alongside their focused tests. Shape inference is explicit: +inserting it between rank-conversion passes changes original shape metadata used +to adjust axes. +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +import onnx + +from ..logging_config import logger +from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32 +from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers +from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes +from .utils.graph_utils import ( + DEFAULT_QDQ_REMOVE_QUANT_TYPES, + infer_shapes, + parse_qdq_quantized_dtype_list, +) + +if TYPE_CHECKING: + from collections.abc import Iterable +__all__ = [] +DLA_TRANSFORMS = ( + ("Cast", _apply_cast_to_fp32), + ("ShapeInference", infer_shapes), + ("ShapeInference", infer_shapes), + ("RemoveUnusedInitializers", _remove_unused_initializers), +) + + +def _transform_make_dla_compatible( + model: onnx.ModelProto, + *, + verbose: bool = True, + qdq_quantized_dtypes: Iterable[int | str] | None = None, + remove_all_qdq: bool = False, +) -> onnx.ModelProto: + """Apply the logged DLA sequence in memory, then infer shapes and validate.""" + started = time.perf_counter() + if remove_all_qdq and qdq_quantized_dtypes is not None: + raise ValueError("Specify remove_all_qdq or qdq_quantized_dtypes, not both") + quant_types = ( + DEFAULT_QDQ_REMOVE_QUANT_TYPES + if qdq_quantized_dtypes is None + else parse_qdq_quantized_dtype_list(qdq_quantized_dtypes) + ) + if remove_all_qdq: + quant_types = frozenset(onnx.TensorProto.DataType.values()) + model = infer_shapes(model) + for name, transform in DLA_TRANSFORMS: + if verbose: + logger.info("[DLA pipeline] Applying %s", name) + options = {} + if name == "RemoveQDQ": + options = { + "keep_clip_after_inputs": False, + "max_chained_qdq_pairs": 5, + "quant_types_to_remove": quant_types, + } + elif name == "ConstantsToInitializers": + options = {"verbose": verbose} + result = transform(model, **options) + if result is not None: + model = result + model = infer_shapes(model) + model = _apply_reorder_nodes(model) + onnx.checker.check_model(model) + if verbose: + logger.info("[DLA pipeline] Done (%.1fs).", time.perf_counter() - started) + return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py new file mode 100644 index 00000000000..ee78a77bd6c --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py @@ -0,0 +1,481 @@ +# 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_cast_to_fp32`. + +"""ORT CPU parity tests for ``dla_cast_to_fp32``. + +Transform rule +-------------- +Every Cast node whose output is NOT consumed by a dtype-sensitive op has its +``to`` attribute changed to FLOAT (float32). + +Skip conditions (Cast node left untouched): + - Output feeds ``input[1]`` (indices) of a Gather node → must stay INT32 + - Output feeds ``input[1]`` (indices) of a GatherElements node → must stay INT32 + - Output feeds any input of a GroupQueryAttention node → must stay FP16 + +Coverage +-------- +Cast INT64→FLOAT redirected to FLOAT, ORT parity +Cast INT32→FLOAT redirected to FLOAT, ORT parity +Cast already targeting FLOAT → unchanged (cnt == 0) +Cast → Gather indices: skipped (output keeps INT32/INT64) +Cast → GatherElements indices: skipped +Cast → GroupQueryAttention input: skipped +Multiple Casts in one graph: eligible ones converted, protected ones skipped +No Cast nodes: graph unchanged, ORT output unchanged +""" + +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 import dla_cast_to_fp32 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" / "cast_to_fp32_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_cast_to_fp32 + + +# ── 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 _prepare_model(model: onnx.ModelProto) -> None: + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + inferred = SymbolicShapeInference.infer_shapes(model) + 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: onnx.ModelProto, feeds: dict) -> dict: + 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} + + +def _cast_to_attr(model: onnx.ModelProto, cast_name: str) -> int | None: + """Return the ``to`` attribute value of a Cast node by name.""" + for node in model.graph.node: + if node.op_type == "Cast" and node.name == cast_name: + for attr in node.attribute: + if attr.name == "to": + return int(attr.i) + return None + + +_RNG = np.random.default_rng(2024) + + +# ── 1. Cast INT64 → redirected to FLOAT, ORT parity ───────────────────────── + + +def test_cast_int64_redirected_to_float(): + """Cast(to=INT64) on a float input is redirected to FLOAT; ORT output matches.""" + shape = [2, 4] + cast_node = helper.make_node( + "Cast", + inputs=["X"], + outputs=["Y"], + name="cast_i64", + to=TensorProto.INT64, + ) + graph = helper.make_graph( + [cast_node], + "cast_int64", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + # After redirect: Cast(to=FLOAT) is identity for float input + ref = x.copy() + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert _cast_to_attr(tr, "cast_i64") == TensorProto.FLOAT + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + + +# ── 2. Cast INT32 → redirected to FLOAT, ORT parity ───────────────────────── + + +def test_cast_int32_redirected_to_float(): + """Cast(to=INT32) redirected to FLOAT.""" + shape = [3, 8] + cast_node = helper.make_node( + "Cast", + inputs=["X"], + outputs=["Y"], + name="cast_i32", + to=TensorProto.INT32, + ) + graph = helper.make_graph( + [cast_node], + "cast_int32", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).astype(np.float32) + ref = x.copy() + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x}) + + assert _cast_to_attr(tr, "cast_i32") == TensorProto.FLOAT + np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) + + +# ── 3. Cast already FLOAT → unchanged (cnt == 0) ───────────────────────────── + + +def test_cast_already_float_unchanged(): + """Cast(to=FLOAT) is already float32 — attribute must not be touched.""" + shape = [2, 4] + cast_node = helper.make_node( + "Cast", + inputs=["X"], + outputs=["Y"], + name="cast_fp32", + to=TensorProto.FLOAT, + ) + graph = helper.make_graph( + [cast_node], + "cast_fp32_graph", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = _RNG.standard_normal(shape).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 _cast_to_attr(tr, "cast_fp32") == TensorProto.FLOAT + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 4. Cast → Gather indices: skipped ──────────────────────────────────────── + + +def test_cast_to_gather_indices_skipped(): + """Cast feeding Gather input[1] (indices) must NOT be redirected to FLOAT.""" + data_shape = [4, 8] + # Cast INT64 → INT32 for the indices, then feed into Gather + cast_node = helper.make_node( + "Cast", + inputs=["idx_i64"], + outputs=["idx_i32"], + name="cast_idx", + to=TensorProto.INT32, + ) + gather_node = helper.make_node( + "Gather", + inputs=["X", "idx_i32"], + outputs=["Y"], + name="gather", + axis=0, + ) + graph = helper.make_graph( + [cast_node, gather_node], + "cast_gather", + [ + _vi("X", TensorProto.FLOAT, data_shape), + _vi("idx_i64", TensorProto.INT64, []), + ], + [_vi("Y", TensorProto.FLOAT, [data_shape[1]])], + value_info=[_vi("idx_i32", TensorProto.INT32, [])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + # Cast must still target INT32 — it was NOT redirected + assert _cast_to_attr(tr, "cast_idx") == TensorProto.INT32, ( + "Cast feeding Gather indices must keep INT32" + ) + + +# ── 5. Cast → GatherElements indices: skipped ──────────────────────────────── + + +def test_cast_to_gatherelements_indices_skipped(): + """Cast feeding GatherElements input[1] (indices) must NOT be redirected.""" + shape = [3, 4] + cast_node = helper.make_node( + "Cast", + inputs=["idx_i64"], + outputs=["idx_i32"], + name="cast_geidx", + to=TensorProto.INT32, + ) + ge_node = helper.make_node( + "GatherElements", + inputs=["X", "idx_i32"], + outputs=["Y"], + name="gather_elements", + axis=0, + ) + graph = helper.make_graph( + [cast_node, ge_node], + "cast_gatherelements", + [ + _vi("X", TensorProto.FLOAT, shape), + _vi("idx_i64", TensorProto.INT64, shape), + ], + [_vi("Y", TensorProto.FLOAT, shape)], + value_info=[_vi("idx_i32", TensorProto.INT32, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _cast_to_attr(tr, "cast_geidx") == TensorProto.INT32, ( + "Cast feeding GatherElements indices must keep INT32" + ) + + +# ── 6. Cast → GQA input: skipped ───────────────────────────────────────────── + + +def test_cast_to_gqa_input_skipped(): + """Cast feeding any GroupQueryAttention input must NOT be redirected.""" + shape = [2, 4, 8] + # Cast FLOAT → FLOAT16 for GQA input + cast_node = helper.make_node( + "Cast", + inputs=["X"], + outputs=["X_fp16"], + name="cast_gqa", + to=TensorProto.FLOAT16, + ) + # Minimal GQA node stub — ORT symbolic inference requires at least 3 inputs + gqa_node = helper.make_node( + "GroupQueryAttention", + inputs=["X_fp16", "", ""], + outputs=["Y"], + name="gqa", + domain="com.microsoft", + num_heads=2, + kv_num_heads=2, + ) + graph = helper.make_graph( + [cast_node, gqa_node], + "cast_gqa_graph", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT16, shape)], + value_info=[_vi("X_fp16", TensorProto.FLOAT16, shape)], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", _OPSET_VERSION), + helper.make_opsetid("com.microsoft", 1), + ], + ) + _clamp_ir(model) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + + assert _cast_to_attr(tr, "cast_gqa") == TensorProto.FLOAT16, ( + "Cast feeding GQA input must keep FLOAT16" + ) + + +# ── 7. Mixed graph: eligible Cast converted, protected Cast skipped ─────────── + + +def test_mixed_graph_eligible_converted_protected_skipped(): + """Graph with one eligible and one protected Cast; ORT parity on transformed model. + + Structure + --------- + Cast(eligible, INT64) → Reshape(same shape) → X_reshaped [graph output 1] + Cast(protected, INT32) → Gather(X, axis=0) → Y_gather [graph output 2] + + Reshape accepts any dtype so no type-op compatibility issue. + X_i64 is intentionally NOT added to value_info: after the transform its + dtype changes INT64→FLOAT and a stale INT64 declaration would cause ORT + to reject the transformed model. Omitting it lets ORT infer the type. + + ORT is run only on the transformed model; numpy reference is used for + parity (the base model would output INT64 for X_reshaped, the transformed + model outputs FLOAT — they are numerically the same for integer-valued x). + """ + data_shape = [4, 8] + shape_init = numpy_helper.from_array(np.array(data_shape, dtype=np.int64), name="reshape_shape") + + # ── Path 1: eligible cast → Reshape (dtype-agnostic) ───────────────────── + cast_eligible = helper.make_node( + "Cast", + inputs=["X"], + outputs=["X_i64"], + name="cast_eligible", + to=TensorProto.INT64, + ) + reshape_node = helper.make_node( + "Reshape", + inputs=["X_i64", "reshape_shape"], + outputs=["X_reshaped"], + name="reshape_node", + ) + + # ── Path 2: protected cast → Gather ────────────────────────────────────── + cast_protected = helper.make_node( + "Cast", + inputs=["idx_i64"], + outputs=["idx_i32"], + name="cast_protected", + to=TensorProto.INT32, + ) + gather_node = helper.make_node( + "Gather", + inputs=["X", "idx_i32"], + outputs=["Y_gather"], + name="gather", + axis=0, + ) + + graph = helper.make_graph( + [cast_eligible, reshape_node, cast_protected, gather_node], + "mixed_graph", + [ + _vi("X", TensorProto.FLOAT, data_shape), + _vi("idx_i64", TensorProto.INT64, []), + ], + [ + _vi("X_reshaped", TensorProto.FLOAT, data_shape), + _vi("Y_gather", TensorProto.FLOAT, [data_shape[1]]), + ], + initializer=[shape_init], + # X_i64 intentionally omitted from value_info — its dtype changes after + # the transform (INT64 → FLOAT) and a stale entry would cause ORT type errors. + value_info=[_vi("idx_i32", TensorProto.INT32, [])], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + # Use integer-valued input so INT64 cast is lossless and numpy ref is exact. + x = np.arange(np.prod(data_shape), dtype=np.float32).reshape(data_shape) + + # Numpy reference for the TRANSFORMED model: + # cast_eligible → FLOAT (identity), Reshape (no-op) → x + # Gather row 1 → x[1, :] + reshaped_ref = x.reshape(data_shape) + y_gather_ref = x[1, :] + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x, "idx_i64": np.array(1, dtype=np.int64)}) + + # ── Attribute checks ────────────────────────────────────────────────────── + assert _cast_to_attr(tr, "cast_eligible") == TensorProto.FLOAT, "eligible must become FLOAT" + assert _cast_to_attr(tr, "cast_protected") == TensorProto.INT32, "protected must stay INT32" + + # ── ORT parity (transformed model vs numpy) ─────────────────────────────── + np.testing.assert_allclose(got["X_reshaped"], reshaped_ref, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(got["Y_gather"], y_gather_ref, rtol=1e-5, atol=1e-5) + + +# ── 8. No Cast nodes: graph and ORT output unchanged ───────────────────────── + + +def test_no_cast_nodes_noop(): + """Graph with no Cast nodes: graph structure and ORT output unchanged.""" + shape = [2, 4] + relu = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu], + "no_cast", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_node_count = len(model.graph.node) + + x = _RNG.standard_normal(shape).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) == original_node_count + np.testing.assert_array_equal(ref["Y"], got["Y"]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py new file mode 100644 index 00000000000..522ee0ec4e4 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py @@ -0,0 +1,101 @@ +# 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. + +"""Tests for the final DLA priority topological node reorder.""" + +from __future__ import annotations + +import numpy as np +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_reorder_nodes import ( + _apply_reorder_nodes, + priority_topological_order, +) + + +def _vi(name: str, dtype: int = TensorProto.FLOAT, shape: list[int] | None = None): + return helper.make_tensor_value_info(name, dtype, shape or [1]) + + +def _node_names(model) -> list[str]: + return [node.name for node in model.graph.node] + + +def test_priority_reorder_preserves_dependencies_and_delays_support_nodes(): + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["relu", "bias"], ["Y"], name="add"), + helper.make_node( + "Constant", + [], + ["bias"], + name="const_bias", + value=numpy_helper.from_array( + np.array([1.0], dtype=np.float32), name="bias_value" + ), + ), + helper.make_node("QuantizeLinear", ["X", "scale", "zero"], ["xq"], name="q"), + helper.make_node("DequantizeLinear", ["xq", "scale", "zero"], ["xdq"], name="dq"), + helper.make_node("Relu", ["xdq"], ["relu"], name="relu"), + ], + "g", + [_vi("X")], + [_vi("Y")], + initializer=[ + numpy_helper.from_array(np.array([0.5], dtype=np.float32), "scale"), + numpy_helper.from_array(np.array([0], dtype=np.uint8), "zero"), + ], + ) + ) + + _apply_reorder_nodes(model) + + assert _node_names(model) == ["q", "dq", "relu", "const_bias", "add"] + + +def test_duplicate_producer_consumer_edges_are_counted_once(): + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Add", ["hidden", "hidden"], ["Y"], name="add"), + helper.make_node("Identity", ["X"], ["hidden"], name="identity"), + ], + "g", + [_vi("X")], + [_vi("Y")], + ) + ) + + _apply_reorder_nodes(model) + + assert _node_names(model) == ["identity", "add"] + + +def test_unresolved_dependency_error_uses_fallback_label_for_unnamed_nodes(): + graph = helper.make_graph( + [ + helper.make_node("Identity", ["b"], ["a"]), + helper.make_node("Identity", ["a"], ["b"], name="named_cycle"), + ], + "g", + [], + [_vi("a")], + ) + + with pytest.raises(ValueError, match=r""): + priority_topological_order(graph) From 224d78b6a85b417e434cbb4d57837204752e5875 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 14:52:24 +0000 Subject: [PATCH 2/6] Preserve ONNX type constraints in DLA Cast conversion Signed-off-by: mgohil --- .../dla_transforms/dla_cast_to_fp32.py | 174 ++++++++++---- .../dla_transforms/dla_reorder_nodes.py | 5 + .../onnx/graph_surgery/make_dla_compatible.py | 5 +- .../dla_transforms/test_dla_cast_to_fp32.py | 226 +++++++++++++++--- .../test_dla_pipeline_scaffold.py | 55 +++++ .../dla_transforms/test_dla_reorder_nodes.py | 5 + 6 files changed, 396 insertions(+), 74 deletions(-) create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py index 6e7a6be2e51..fc2acd7b226 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py @@ -13,19 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Redirect all Cast nodes to produce float32 output. +"""Redirect eligible standard Cast targets to FP32 using ONNX type constraints. -Skip conditions ---------------- -A Cast node is left untouched when its output is consumed by an op that -requires a specific non-float32 dtype: - -* **Gather / GatherElements** ``input[1]`` (indices) — must stay INT32. -* **GroupQueryAttention** inputs — require FP16. - -These Cast nodes are inserted by earlier transforms -(``dla-convert-ops-to-4d`` for Gather/GatherElements indices, or -externally for GQA) and must not be overridden. +Shared type variables connect operands and outputs across operations. Fixed +non-FP32 operands, integer or boolean controls, opaque operations and lexical +captures prevent conversion of their connected type groups. Compatible Casts +can convert together; their connected tensor metadata is updated as well. """ from __future__ import annotations @@ -34,56 +27,147 @@ from onnx import TensorProto from ...logging_config import logger +from ..utils.graph_utils import _captured_names + + +def _fp32_cast_outputs(model: onnx.ModelProto) -> set[str]: + """Find tensor type groups that can become FP32 without changing fixed operands.""" + graph = model.graph + candidates = { + node.output[0] + for node in graph.node + if node.op_type == "Cast" + and node.domain in {"", "ai.onnx"} + and not node.overload + and node.output + and any(attr.name == "to" and attr.i != TensorProto.FLOAT for attr in node.attribute) + } + if not candidates: + return set() + parents: dict[str, str] = {} + + def find(name: str) -> str: + """Find and compress a tensor's shared-type component.""" + parents.setdefault(name, name) + while parents[name] != name: + parents[name] = parents[parents[name]] + name = parents[name] + return name + + blocked: set[str] = set() + types = { + value.name: value.type.tensor_type.elem_type + for value in (*graph.input, *graph.value_info, *graph.output) + if value.type.HasField("tensor_type") + } + types.update((value.name, value.data_type) for value in graph.initializer) + fixed = {value.name for value in graph.input} | {value.name for value in graph.initializer} + fixed.update(value.values.name for value in graph.sparse_initializer) + types.update((value.values.name, value.values.data_type) for value in graph.sparse_initializer) + opsets = {opset.domain: opset.version for opset in model.opset_import} + opsets[""] = opsets.get("", opsets.get("ai.onnx", 1)) + schemas = {} + for node in graph.node: + for attribute in node.attribute: + subgraphs = ( + [attribute.g] if attribute.type == onnx.AttributeProto.GRAPH else attribute.graphs + ) + for subgraph in subgraphs: + blocked.update(_captured_names(subgraph)) + domain = "" if node.domain == "ai.onnx" else node.domain + key = (domain, node.op_type) + if key not in schemas: + try: + schemas[key] = onnx.defs.get_schema(node.op_type, opsets.get(domain, 1), domain) + except onnx.defs.SchemaError: + schemas[key] = None + schema = schemas[key] + if ( + schema is None + or node.overload + or any( + parameter.option == onnx.defs.OpSchema.FormalParameterOption.Variadic + and not parameter.is_homogeneous + for parameter in (*schema.inputs, *schema.outputs) + ) + ): + blocked.update(name for name in (*node.input, *node.output) if name) + continue + allowed = { + constraint.type_param_str: constraint.allowed_type_strs + for constraint in schema.type_constraints + } + groups: dict[str, list[str]] = {} + input_types: set[str] = set() + for values, parameters, is_output in ( + (node.input, schema.inputs, False), + (node.output, schema.outputs, True), + ): + for index, name in enumerate(values): + if not name: + continue + parameter = parameters[min(index, len(parameters) - 1)] + type_name = parameter.type_str + groups.setdefault(type_name, []).append(name) + if "tensor(float)" not in allowed.get(type_name, [type_name]): + blocked.add(name) + if not is_output: + input_types.add(type_name) + elif type_name not in input_types: + fixed.add(name) + if type_name == "tensor(float)": + types[name] = TensorProto.FLOAT + elif type_name in allowed and node.op_type != "Cast": + # Independent output variables may still default to an input's dtype. + blocked.update(name for name in node.input if name) + for names in groups.values(): + for name in names[1:]: + parents[find(name)] = find(names[0]) + if node.op_type == "Cast" and domain == "": + for attribute in node.attribute: + if attribute.name == "to": + types[node.output[0]] = attribute.i + elif node.op_type == "Constant" and domain == "": + for attribute in node.attribute: + if attribute.name == "value": + types[node.output[0]] = attribute.t.data_type + # Runtime attention inputs have implementation-specific precision requirements. + if node.op_type == "GroupQueryAttention": + blocked.update(name for name in node.input if name) + blocked.update(name for name in fixed - candidates if types.get(name) != TensorProto.FLOAT) + blocked_roots = {find(name) for name in blocked} + promoted_roots = {find(name) for name in candidates if find(name) not in blocked_roots} + return {name for name in parents if find(name) in promoted_roots} def _apply_cast_to_fp32(model: onnx.ModelProto) -> onnx.ModelProto: """Change every eligible Cast node's ``to`` attribute to FLOAT (float32). - A Cast node is skipped when its output feeds: - - ``input[1]`` (indices) of a Gather or GatherElements node, or - - any input of a GroupQueryAttention node. + Preserve connected type requirements, custom operations and subgraph captures. + Update intermediate and graph-output descriptors for converted type groups. """ graph = model.graph - - # ── Build set of protected Cast outputs ────────────────────────────────── - # Tensor names produced by Cast nodes that must keep their current target - # dtype: INT32 for Gather/GatherElements indices, FP16 for GQA inputs. - protected: set[str] = set() - for node in graph.node: - if node.op_type in ("Gather", "GatherElements"): - # input[1] is the indices tensor — must stay INT32 - if len(node.input) > 1 and node.input[1]: - protected.add(node.input[1]) - elif node.op_type in {"Reshape", "Expand", "Tile", "TopK"}: - protected.update(node.input[1:2]) - elif node.op_type == "Slice": - protected.update(node.input[1:]) - elif node.op_type in {"Not", "And", "Or", "Xor"}: - protected.update(node.input) - elif node.op_type == "Where": - protected.update(node.input[:1]) - elif node.op_type == "GroupQueryAttention": - for inp in node.input: - if inp: - protected.add(inp) - - # ── Redirect eligible Cast nodes to FLOAT ──────────────────────────────── + promoted = _fp32_cast_outputs(model) + value_info_by_name = {} + for value_info in graph.value_info: + value_info_by_name.setdefault(value_info.name, value_info) + outputs_by_name = {output.name: output for output in graph.output} cnt = 0 for node in graph.node: if node.op_type != "Cast": continue - if not node.output or node.output[0] in protected: + if not node.output or node.output[0] not in promoted: continue for attr in node.attribute: if attr.name == "to" and attr.i != TensorProto.FLOAT: attr.i = TensorProto.FLOAT - # Update value_info dtype for the Cast output tensor - for vi in graph.value_info: - if vi.name == node.output[0]: - vi.type.tensor_type.elem_type = TensorProto.FLOAT - break cnt += 1 break + for name in promoted: + for value_info in (value_info_by_name.get(name), outputs_by_name.get(name)): + if value_info is not None: + value_info.type.tensor_type.elem_type = TensorProto.FLOAT + logger.debug("cast_to_fp32: redirected %d Cast node(s) to float32.", cnt) return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py index ea9eabb3819..aebfd5a862e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py @@ -34,10 +34,12 @@ def _node_label(node: onnx.NodeProto, index: int) -> str: + """Identify a node in an error even when its optional name is absent.""" return node.name or f"" def _build_edges(nodes: list[onnx.NodeProto]) -> tuple[list[set[int]], list[list[int]]]: + """Collect unique dependencies, including values captured by nested subgraphs.""" output_to_node: dict[str, int] = {} for index, node in enumerate(nodes): for output in node.output: @@ -65,6 +67,7 @@ def _build_edges(nodes: list[onnx.NodeProto]) -> tuple[list[set[int]], list[list def _unresolved_dependency_error(nodes: list[onnx.NodeProto], indegrees: list[int]) -> ValueError: + """Describe nodes left blocked after a topological traversal.""" remaining = [ _node_label(nodes[index], index) for index, indegree in enumerate(indegrees) if indegree > 0 ] @@ -78,6 +81,7 @@ def _compute_depths( dependencies: list[set[int]], consumers: list[list[int]], ) -> list[int]: + """Compute each node's longest dependency depth and reject cycles.""" depths = [0] * len(nodes) indegrees = [len(node_dependencies) for node_dependencies in dependencies] ready = [index for index, indegree in enumerate(indegrees) if indegree == 0] @@ -116,6 +120,7 @@ def priority_topological_order( sink_priority = max(depths, default=0) + 1 def node_priority(index: int) -> int: + """Delay support nodes toward their first consumer's dependency depth.""" if nodes[index].op_type not in PRIORITIZED_OPS: return 0 if consumers[index]: diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index d091ea53b32..8431e2040e6 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -44,7 +44,6 @@ DLA_TRANSFORMS = ( ("Cast", _apply_cast_to_fp32), ("ShapeInference", infer_shapes), - ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) @@ -60,6 +59,10 @@ def _transform_make_dla_compatible( started = time.perf_counter() if remove_all_qdq and qdq_quantized_dtypes is not None: raise ValueError("Specify remove_all_qdq or qdq_quantized_dtypes, not both") + if (remove_all_qdq or qdq_quantized_dtypes is not None) and not any( + name == "RemoveQDQ" for name, _ in DLA_TRANSFORMS + ): + raise NotImplementedError("QDQ removal is not available in this pipeline") quant_types = ( DEFAULT_QDQ_REMOVE_QUANT_TYPES if qdq_quantized_dtypes is None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py index ee78a77bd6c..e2975f12f06 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py @@ -13,11 +13,6 @@ # 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_cast_to_fp32`. - """ORT CPU parity tests for ``dla_cast_to_fp32``. Transform rule @@ -45,7 +40,6 @@ from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx @@ -54,10 +48,7 @@ from modelopt.onnx.graph_surgery.dla_transforms import dla_cast_to_fp32 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" / "cast_to_fp32_test_models" +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -69,15 +60,19 @@ def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Build a typed tensor descriptor for a test graph.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir(model: onnx.ModelProto) -> None: + """Keep the test model compatible with the supported ORT versions.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _prepare_model(model: onnx.ModelProto) -> None: + """Populate intermediate shapes when the optional ORT tool is available.""" try: + # The tool submodule is optional; model execution only requires ORT itself. from onnxruntime.tools.symbolic_shape_infer import ( SymbolicShapeInference, # type: ignore[import-untyped] ) @@ -91,8 +86,7 @@ def _prepare_model(model: onnx.ModelProto) -> None: def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-untyped] - + """Execute a small test graph on the CPU and return named outputs.""" _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} @@ -128,10 +122,12 @@ def test_cast_int64_redirected_to_float(): [cast_node], "cast_int64", [_vi("X", TensorProto.FLOAT, shape)], - [_vi("Y", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.INT64, shape)], + value_info=[_vi("Y", TensorProto.INT64, shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) x = _RNG.standard_normal(shape).astype(np.float32) # After redirect: Cast(to=FLOAT) is identity for float input @@ -140,7 +136,7 @@ def test_cast_int64_redirected_to_float(): tr = copy.deepcopy(model) _prepare_model(tr) tr = _apply(tr) - _prepare_model(tr) + onnx.checker.check_model(tr, full_check=True) got = _run_ort(tr, {"X": x}) assert _cast_to_attr(tr, "cast_i64") == TensorProto.FLOAT @@ -164,10 +160,11 @@ def test_cast_int32_redirected_to_float(): [cast_node], "cast_int32", [_vi("X", TensorProto.FLOAT, shape)], - [_vi("Y", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.INT32, shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) x = _RNG.standard_normal(shape).astype(np.float32) ref = x.copy() @@ -175,7 +172,7 @@ def test_cast_int32_redirected_to_float(): tr = copy.deepcopy(model) _prepare_model(tr) tr = _apply(tr) - _prepare_model(tr) + onnx.checker.check_model(tr, full_check=True) got = _run_ort(tr, {"X": x}) assert _cast_to_attr(tr, "cast_i32") == TensorProto.FLOAT @@ -265,48 +262,221 @@ def test_cast_to_gather_indices_skipped(): # ── 5. Cast → GatherElements indices: skipped ──────────────────────────────── -def test_cast_to_gatherelements_indices_skipped(): - """Cast feeding GatherElements input[1] (indices) must NOT be redirected.""" +@pytest.mark.parametrize( + ("op_type", "index_dtype"), + [("GatherElements", TensorProto.INT32), ("GatherND", TensorProto.INT64)], +) +def test_cast_to_gather_indices_remain_executable(op_type, index_dtype): + """Index casts keep their required dtype and produce the same gathered values.""" shape = [3, 4] + index_shape = shape if op_type == "GatherElements" else [2, 1] cast_node = helper.make_node( "Cast", inputs=["idx_i64"], outputs=["idx_i32"], name="cast_geidx", - to=TensorProto.INT32, + to=index_dtype, ) ge_node = helper.make_node( - "GatherElements", + op_type, inputs=["X", "idx_i32"], outputs=["Y"], name="gather_elements", - axis=0, + **({"axis": 0} if op_type == "GatherElements" else {}), ) graph = helper.make_graph( [cast_node, ge_node], "cast_gatherelements", [ _vi("X", TensorProto.FLOAT, shape), - _vi("idx_i64", TensorProto.INT64, shape), + _vi("idx_i64", TensorProto.INT64, index_shape), ], - [_vi("Y", TensorProto.FLOAT, shape)], - value_info=[_vi("idx_i32", TensorProto.INT32, shape)], + [_vi("Y", TensorProto.FLOAT, shape if op_type == "GatherElements" else [2, 4])], + value_info=[_vi("idx_i32", index_dtype, index_shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + feeds = { + "X": np.arange(12, dtype=np.float32).reshape(shape), + "idx_i64": np.zeros(index_shape, dtype=np.int64), + } + ref = _run_ort(model, feeds) tr = copy.deepcopy(model) _prepare_model(tr) tr = _apply(tr) - - assert _cast_to_attr(tr, "cast_geidx") == TensorProto.INT32, ( - "Cast feeding GatherElements indices must keep INT32" - ) + onnx.checker.check_model(tr, full_check=True) + assert _cast_to_attr(tr, "cast_geidx") == index_dtype + np.testing.assert_array_equal(_run_ort(tr, feeds)["Y"], ref["Y"]) # ── 6. Cast → GQA input: skipped ───────────────────────────────────────────── +@pytest.mark.parametrize( + "op_type", ["Reshape", "Expand", "Tile", "TopK", "Slice", "Not", "And", "Or", "Xor", "Where"] +) +def test_protected_control_casts_remain_executable(op_type): + """Shape controls and boolean predicates retain their types and runtime values.""" + boolean = op_type in {"Not", "And", "Or", "Xor", "Where"} + dtype = TensorProto.BOOL if boolean else TensorProto.INT64 + control = [1, 0] if boolean else [2, 2] + initializers = [numpy_helper.from_array(np.arange(4, dtype=np.float32).reshape(2, 2), "data")] + inputs = ["data", "control"] + outputs = [_vi("Y", TensorProto.FLOAT, [2, 2])] + if boolean: + inputs = ["control"] + outputs = [_vi("Y", TensorProto.BOOL, [2])] + if op_type in {"And", "Or", "Xor"}: + inputs.append("other") + initializers = [numpy_helper.from_array(np.array([False, True]), "other")] + elif op_type == "Where": + inputs.extend(["data", "other"]) + initializers.append(numpy_helper.from_array(np.ones((2, 2), dtype=np.float32), "other")) + outputs = [_vi("Y", TensorProto.FLOAT, [2, 2])] + elif op_type == "Tile": + control = [1, 1] + elif op_type == "TopK": + control = [1] + outputs = [_vi("Y", TensorProto.FLOAT, [2, 1]), _vi("indices", TensorProto.INT64, [2, 1])] + elif op_type == "Slice": + control = [0] + inputs.append("ends") + initializers.append(numpy_helper.from_array(np.array([1], dtype=np.int64), "ends")) + outputs = [_vi("Y", TensorProto.FLOAT, [1, 2])] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["X"], ["control"], name="control_cast", to=dtype), + helper.make_node(op_type, inputs, [value.name for value in outputs]), + ], + "protected_controls", + [_vi("X", TensorProto.FLOAT, [len(control)])], + outputs, + initializer=initializers, + value_info=[_vi("control", dtype, [len(control)])], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.array(control, dtype=np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert _cast_to_attr(transformed, "control_cast") == dtype + for name, value in _run_ort(transformed, feeds).items(): + np.testing.assert_array_equal(value, expected[name]) + + +@pytest.mark.parametrize("peer", ["initializer", "input", "cast", "identity", "reshape"]) +def test_cast_conversion_respects_shared_operand_types(peer): + """Coupled operands stay compatible, including through intermediate operations.""" + nodes = [helper.make_node("Cast", ["X"], ["cast"], name="cast", to=TensorProto.INT64)] + inputs = [_vi("X", TensorProto.FLOAT, [2])] + initializers = [] + feeds = {"X": np.array([1, 3], dtype=np.float32)} + left = "cast" + if peer == "cast": + inputs.append(_vi("peer", TensorProto.FLOAT, [2])) + feeds["peer"] = np.array([2, 4], dtype=np.float32) + nodes.append(helper.make_node("Cast", ["peer"], ["right"], to=TensorProto.INT64)) + elif peer == "input": + inputs.append(_vi("right", TensorProto.INT64, [2])) + feeds["right"] = np.array([2, 4], dtype=np.int64) + else: + initializers.append(numpy_helper.from_array(np.array([2, 4], dtype=np.int64), "right")) + if peer == "identity": + nodes.append(helper.make_node("Identity", ["cast"], ["left"])) + left = "left" + elif peer == "reshape": + initializers.append(numpy_helper.from_array(np.array([2], dtype=np.int64), "shape")) + nodes.append(helper.make_node("Reshape", ["cast", "shape"], ["left"])) + left = "left" + nodes.append(helper.make_node("Add", [left, "right"], ["Y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "coupled_types", + inputs, + [_vi("Y", TensorProto.INT64, [2])], + initializer=initializers, + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds)["Y"] + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + target_dtype = TensorProto.FLOAT if peer == "cast" else TensorProto.INT64 + assert _cast_to_attr(transformed, "cast") == target_dtype + assert transformed.graph.output[0].type.tensor_type.elem_type == target_dtype + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) + + +def test_cast_preserves_implicitly_derived_output_dtype(): + """A generator's default output dtype can depend on a separately typed input.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["X"], ["cast"], name="cast", to=TensorProto.INT64), + helper.make_node("EyeLike", ["cast"], ["Y"]), + ], + "implicit_dtype", + [_vi("X", TensorProto.FLOAT, [2, 2])], + [_vi("Y", TensorProto.INT64, [2, 2])], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + model = onnx.shape_inference.infer_shapes(model) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.ones((2, 2), dtype=np.float32)} + expected = _run_ort(model, feeds)["Y"] + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert _cast_to_attr(transformed, "cast") == TensorProto.INT64 + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) + + +def test_custom_cast_function_keeps_its_target_type(): + """A custom Cast function is opaque even when it forwards a standard Cast attribute.""" + body = helper.make_node("Cast", ["X"], ["Y"]) + body.attribute.append(helper.make_attribute_ref("to", onnx.AttributeProto.INT)) + body.attribute[0].ref_attr_name = "to" + function = helper.make_function( + "local", + "Cast", + ["X"], + ["Y"], + [body], + [helper.make_opsetid("", _OPSET_VERSION)], + attributes=["to"], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Cast", ["X"], ["Y"], domain="local", to=TensorProto.INT64)], + "custom_cast", + [_vi("X", TensorProto.FLOAT, [2])], + [_vi("Y", TensorProto.INT64, [2])], + ), + functions=[function], + opset_imports=[helper.make_opsetid("", _OPSET_VERSION), helper.make_opsetid("local", 1)], + ) + _clamp_ir(model) + onnx.checker.check_model(model, full_check=True) + before = model.SerializeToString() + feeds = {"X": np.array([1.5, -2.5], dtype=np.float32)} + expected = _run_ort(model, feeds) + _apply(model) + assert model.SerializeToString() == before + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(_run_ort(model, feeds)["Y"], expected["Y"]) + + def test_cast_to_gqa_input_skipped(): """Cast feeding any GroupQueryAttention input must NOT be redirected.""" shape = [2, 4, 8] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py new file mode 100644 index 00000000000..b660666bfd6 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py @@ -0,0 +1,55 @@ +# 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. + +"""Validation for unavailable stages in the private DLA pipeline scaffold.""" + +import importlib + +import pytest +from onnx import TensorProto, helper + +pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") + + +@pytest.mark.parametrize( + "options", + [ + {"remove_all_qdq": True}, + {"qdq_quantized_dtypes": ["int16"]}, + {"remove_all_qdq": True, "qdq_quantized_dtypes": ["int16"]}, + ], +) +def test_pipeline_rejects_unavailable_qdq_options(monkeypatch, options): + """Unavailable QDQ options and conflicting options fail without changing the model.""" + monkeypatch.setattr( + pipeline, + "DLA_TRANSFORMS", + tuple(stage for stage in pipeline.DLA_TRANSFORMS if stage[0] != "RemoveQDQ"), + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Identity", ["X"], ["Y"])], + "g", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [1])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + before = model.SerializeToString() + error = ValueError if len(options) == 2 else NotImplementedError + with pytest.raises(error): + pipeline._transform_make_dla_compatible(model, verbose=False, **options) + assert model.SerializeToString() == before diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py index 522ee0ec4e4..68b9b3e0bd8 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_reorder_nodes.py @@ -28,14 +28,17 @@ def _vi(name: str, dtype: int = TensorProto.FLOAT, shape: list[int] | None = None): + """Build a tensor descriptor for a small ordering fixture.""" return helper.make_tensor_value_info(name, dtype, shape or [1]) def _node_names(model) -> list[str]: + """Read node names in their current serialized order.""" return [node.name for node in model.graph.node] def test_priority_reorder_preserves_dependencies_and_delays_support_nodes(): + """Support nodes stay close to their consumers without breaking dependencies.""" model = helper.make_model( helper.make_graph( [ @@ -69,6 +72,7 @@ def test_priority_reorder_preserves_dependencies_and_delays_support_nodes(): def test_duplicate_producer_consumer_edges_are_counted_once(): + """Using the same producer twice does not leave a consumer blocked.""" model = helper.make_model( helper.make_graph( [ @@ -87,6 +91,7 @@ def test_duplicate_producer_consumer_edges_are_counted_once(): def test_unresolved_dependency_error_uses_fallback_label_for_unnamed_nodes(): + """Cycle diagnostics identify unnamed nodes as well as named nodes.""" graph = helper.make_graph( [ helper.make_node("Identity", ["b"], ["a"]), From af9664a715c6c676e794543fbdeb689baa192a24 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 15:18:40 +0000 Subject: [PATCH 3/6] Exercise the DLA pipeline scaffold end to end Signed-off-by: mgohil --- .../test_dla_pipeline_scaffold.py | 49 +++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py index b660666bfd6..632a04456fd 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py @@ -17,10 +17,13 @@ import importlib +import numpy as np +import onnx import pytest from onnx import TensorProto, helper pipeline = importlib.import_module("modelopt.onnx.graph_surgery.make_dla_compatible") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for pipeline execution") @pytest.mark.parametrize( @@ -53,3 +56,49 @@ def test_pipeline_rejects_unavailable_qdq_options(monkeypatch, options): with pytest.raises(error): pipeline._transform_make_dla_compatible(model, verbose=False, **options) assert model.SerializeToString() == before + + +def test_pipeline_applies_cast_reorders_nodes_and_returns_model(): + """The real pipeline dispatches stages, orders support nodes and returns executable outputs.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Constant", + [], + ["K"], + name="constant", + value=helper.make_tensor("k", TensorProto.FLOAT, [2], [1.0, 2.0]), + ), + helper.make_node("Cast", ["X"], ["C"], name="cast", to=TensorProto.INT32), + helper.make_node("Identity", ["C"], ["Y"], name="consumer"), + helper.make_node("Identity", ["K"], ["Z"], name="constant_consumer"), + ], + "pipeline", + [helper.make_tensor_value_info("X", TensorProto.FLOAT, [2])], + [ + helper.make_tensor_value_info("Y", TensorProto.INT32, [2]), + helper.make_tensor_value_info("Z", TensorProto.FLOAT, [2]), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + onnx.checker.check_model(model, full_check=True) + result = pipeline._transform_make_dla_compatible(model, verbose=False) + names = [node.name for node in result.graph.node] + assert names.index("cast") < names.index("consumer") < names.index("constant_consumer") + if "constant" in names: + # Later slices turn Constants into initializers before the final reorder. + assert names.index("consumer") < names.index("constant") < names.index("constant_consumer") + cast = next(node for node in result.graph.node if node.name == "cast") + assert ( + next(attribute for attribute in cast.attribute if attribute.name == "to").i + == TensorProto.FLOAT + ) + onnx.checker.check_model(result, full_check=True) + x = np.array([1.5, -2.5], dtype=np.float32) + session = ort.InferenceSession(result.SerializeToString(), providers=["CPUExecutionProvider"]) + outputs = dict(zip([value.name for value in result.graph.output], session.run(None, {"X": x}))) + np.testing.assert_array_equal(outputs["Y"], x) + np.testing.assert_array_equal(outputs["Z"], np.array([1.0, 2.0], dtype=np.float32)) From 6c06aa4f52697175be6e9f67db960ac94dc8ce65 Mon Sep 17 00:00:00 2001 From: mgohil Date: Fri, 23 Oct 2026 15:36:05 +0000 Subject: [PATCH 4/6] Preserve integer-mode Mod operands in DLA Cast conversion Signed-off-by: mgohil --- .../dla_transforms/dla_cast_to_fp32.py | 11 +++++++ .../dla_transforms/test_dla_cast_to_fp32.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py index fc2acd7b226..962b6db6b7b 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py @@ -93,6 +93,17 @@ def find(name: str) -> str: ): blocked.update(name for name in (*node.input, *node.output) if name) continue + if ( + domain == "" + and node.op_type == "Mod" + and schema.since_version < 28 + and not any( + attribute.name == "fmod" and attribute.i == 1 for attribute in node.attribute + ) + ): + # Older Mod schemas require integer inputs in the default attribute mode. + blocked.update(name for name in (*node.input, *node.output) if name) + continue allowed = { constraint.type_param_str: constraint.allowed_type_strs for constraint in schema.type_constraints diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py index e2975f12f06..e2c58afee75 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py @@ -417,6 +417,38 @@ def test_cast_conversion_respects_shared_operand_types(peer): np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) +@pytest.mark.parametrize("fmod", [None, 0, 1]) +def test_cast_conversion_respects_mod_attribute_mode(fmod): + """Integer Mod modes keep integer casts while floating fmod can promote together.""" + source_dtype = TensorProto.DOUBLE if fmod == 1 else TensorProto.INT64 + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["X"], ["left"], name="left_cast", to=source_dtype), + helper.make_node("Cast", ["B"], ["right"], name="right_cast", to=source_dtype), + helper.make_node( + "Mod", ["left", "right"], ["Y"], **({} if fmod is None else {"fmod": fmod}) + ), + ], + "mod_mode", + [_vi("X", TensorProto.FLOAT, [2]), _vi("B", TensorProto.FLOAT, [2])], + [_vi("Y", source_dtype, [2])], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ir_version=_MAX_IR_VERSION_FOR_ORT, + ) + model = onnx.shape_inference.infer_shapes(model) + onnx.checker.check_model(model, full_check=True) + feeds = {"X": np.array([-5, 5], dtype=np.float32), "B": np.array([3, -3], dtype=np.float32)} + expected = _run_ort(model, feeds)["Y"] + transformed = _apply(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["Y"], expected) + target_dtype = TensorProto.FLOAT if fmod == 1 else TensorProto.INT64 + assert _cast_to_attr(transformed, "left_cast") == target_dtype + assert _cast_to_attr(transformed, "right_cast") == target_dtype + + def test_cast_preserves_implicitly_derived_output_dtype(): """A generator's default output dtype can depend on a separately typed input.""" model = helper.make_model( From 2b8baa996f48e8509e720a17dd09b35add783eff Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 06:44:37 +0000 Subject: [PATCH 5/6] Consolidate Cast tests and name pipeline tests after their module Signed-off-by: mgohil --- .../dla_transforms/test_dla_cast_to_fp32.py | 132 +++--------------- ...caffold.py => test_make_dla_compatible.py} | 0 2 files changed, 18 insertions(+), 114 deletions(-) rename tests/unit/onnx/graph_surgery/dla_transforms/{test_dla_pipeline_scaffold.py => test_make_dla_compatible.py} (100%) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py index e2c58afee75..99a1b7418aa 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py @@ -13,29 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""ORT CPU parity tests for ``dla_cast_to_fp32``. - -Transform rule --------------- -Every Cast node whose output is NOT consumed by a dtype-sensitive op has its -``to`` attribute changed to FLOAT (float32). - -Skip conditions (Cast node left untouched): - - Output feeds ``input[1]`` (indices) of a Gather node → must stay INT32 - - Output feeds ``input[1]`` (indices) of a GatherElements node → must stay INT32 - - Output feeds any input of a GroupQueryAttention node → must stay FP16 - -Coverage --------- -Cast INT64→FLOAT redirected to FLOAT, ORT parity -Cast INT32→FLOAT redirected to FLOAT, ORT parity -Cast already targeting FLOAT → unchanged (cnt == 0) -Cast → Gather indices: skipped (output keeps INT32/INT64) -Cast → GatherElements indices: skipped -Cast → GroupQueryAttention input: skipped -Multiple Casts in one graph: eligible ones converted, protected ones skipped -No Cast nodes: graph unchanged, ORT output unchanged -""" +"""Numerical and structural checks for dtype-safe Cast preparation.""" from __future__ import annotations @@ -108,22 +86,24 @@ def _cast_to_attr(model: onnx.ModelProto, cast_name: str) -> int | None: # ── 1. Cast INT64 → redirected to FLOAT, ORT parity ───────────────────────── -def test_cast_int64_redirected_to_float(): +@pytest.mark.parametrize( + ("dtype", "shape"), [(TensorProto.INT64, [2, 4]), (TensorProto.INT32, [3, 8])] +) +def test_integer_cast_redirected_to_float(dtype, shape): """Cast(to=INT64) on a float input is redirected to FLOAT; ORT output matches.""" - shape = [2, 4] cast_node = helper.make_node( "Cast", inputs=["X"], outputs=["Y"], - name="cast_i64", - to=TensorProto.INT64, + name="cast_integer", + to=dtype, ) graph = helper.make_graph( [cast_node], - "cast_int64", + "cast_integer", [_vi("X", TensorProto.FLOAT, shape)], - [_vi("Y", TensorProto.INT64, shape)], - value_info=[_vi("Y", TensorProto.INT64, shape)], + [_vi("Y", dtype, shape)], + value_info=[_vi("Y", dtype, shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) @@ -139,43 +119,7 @@ def test_cast_int64_redirected_to_float(): onnx.checker.check_model(tr, full_check=True) got = _run_ort(tr, {"X": x}) - assert _cast_to_attr(tr, "cast_i64") == TensorProto.FLOAT - np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) - - -# ── 2. Cast INT32 → redirected to FLOAT, ORT parity ───────────────────────── - - -def test_cast_int32_redirected_to_float(): - """Cast(to=INT32) redirected to FLOAT.""" - shape = [3, 8] - cast_node = helper.make_node( - "Cast", - inputs=["X"], - outputs=["Y"], - name="cast_i32", - to=TensorProto.INT32, - ) - graph = helper.make_graph( - [cast_node], - "cast_int32", - [_vi("X", TensorProto.FLOAT, shape)], - [_vi("Y", TensorProto.INT32, shape)], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) - _clamp_ir(model) - onnx.checker.check_model(model, full_check=True) - - x = _RNG.standard_normal(shape).astype(np.float32) - ref = x.copy() - - tr = copy.deepcopy(model) - _prepare_model(tr) - tr = _apply(tr) - onnx.checker.check_model(tr, full_check=True) - got = _run_ort(tr, {"X": x}) - - assert _cast_to_attr(tr, "cast_i32") == TensorProto.FLOAT + assert _cast_to_attr(tr, "cast_integer") == TensorProto.FLOAT np.testing.assert_allclose(got["Y"], ref, rtol=1e-5, atol=1e-5) @@ -215,61 +159,21 @@ def test_cast_already_float_unchanged(): np.testing.assert_array_equal(ref["Y"], got["Y"]) -# ── 4. Cast → Gather indices: skipped ──────────────────────────────────────── - - -def test_cast_to_gather_indices_skipped(): - """Cast feeding Gather input[1] (indices) must NOT be redirected to FLOAT.""" - data_shape = [4, 8] - # Cast INT64 → INT32 for the indices, then feed into Gather - cast_node = helper.make_node( - "Cast", - inputs=["idx_i64"], - outputs=["idx_i32"], - name="cast_idx", - to=TensorProto.INT32, - ) - gather_node = helper.make_node( - "Gather", - inputs=["X", "idx_i32"], - outputs=["Y"], - name="gather", - axis=0, - ) - graph = helper.make_graph( - [cast_node, gather_node], - "cast_gather", - [ - _vi("X", TensorProto.FLOAT, data_shape), - _vi("idx_i64", TensorProto.INT64, []), - ], - [_vi("Y", TensorProto.FLOAT, [data_shape[1]])], - value_info=[_vi("idx_i32", TensorProto.INT32, [])], - ) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) - _clamp_ir(model) - - tr = copy.deepcopy(model) - _prepare_model(tr) - tr = _apply(tr) - - # Cast must still target INT32 — it was NOT redirected - assert _cast_to_attr(tr, "cast_idx") == TensorProto.INT32, ( - "Cast feeding Gather indices must keep INT32" - ) - - # ── 5. Cast → GatherElements indices: skipped ──────────────────────────────── @pytest.mark.parametrize( ("op_type", "index_dtype"), - [("GatherElements", TensorProto.INT32), ("GatherND", TensorProto.INT64)], + [ + ("Gather", TensorProto.INT32), + ("GatherElements", TensorProto.INT32), + ("GatherND", TensorProto.INT64), + ], ) def test_cast_to_gather_indices_remain_executable(op_type, index_dtype): """Index casts keep their required dtype and produce the same gathered values.""" shape = [3, 4] - index_shape = shape if op_type == "GatherElements" else [2, 1] + index_shape = shape if op_type == "GatherElements" else [2] if op_type == "Gather" else [2, 1] cast_node = helper.make_node( "Cast", inputs=["idx_i64"], @@ -282,7 +186,7 @@ def test_cast_to_gather_indices_remain_executable(op_type, index_dtype): inputs=["X", "idx_i32"], outputs=["Y"], name="gather_elements", - **({"axis": 0} if op_type == "GatherElements" else {}), + **({"axis": 0} if op_type in {"Gather", "GatherElements"} else {}), ) graph = helper.make_graph( [cast_node, ge_node], diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_make_dla_compatible.py similarity index 100% rename from tests/unit/onnx/graph_surgery/dla_transforms/test_dla_pipeline_scaffold.py rename to tests/unit/onnx/graph_surgery/dla_transforms/test_make_dla_compatible.py From 2f348cf32e8cfc712060cd73dfa09e4dfba3fd48 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 10:40:27 +0000 Subject: [PATCH 6/6] Preserve SequenceInsert element types during Cast promotion Signed-off-by: mgohil --- .../dla_transforms/dla_cast_to_fp32.py | 2 ++ .../dla_transforms/test_dla_cast_to_fp32.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py index 962b6db6b7b..2a773614420 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py @@ -85,6 +85,8 @@ def find(name: str) -> str: if ( schema is None or node.overload + # SequenceInsert couples sequence and tensor types across separate schema variables. + or (domain == "" and node.op_type == "SequenceInsert") or any( parameter.option == onnx.defs.OpSchema.FormalParameterOption.Variadic and not parameter.is_homogeneous diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py index 99a1b7418aa..fbe61e881e3 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_cast_to_fp32.py @@ -559,6 +559,38 @@ def test_mixed_graph_eligible_converted_protected_skipped(): # ── 8. No Cast nodes: graph and ORT output unchanged ───────────────────────── +@pytest.mark.parametrize("dtype", [TensorProto.INT32, TensorProto.FLOAT16]) +def test_sequence_insert_preserves_inserted_tensor_type(dtype): + """Keep the tensor type coupled to its sequence element type.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["x"], ["inserted"], to=dtype, name="cast"), + helper.make_node("SequenceConstruct", ["seed"], ["sequence"]), + helper.make_node("SequenceInsert", ["sequence", "inserted"], ["y"]), + ], + "sequence_insert", + [_vi("x", TensorProto.FLOAT, [2]), _vi("seed", dtype, [2])], + [helper.make_tensor_sequence_value_info("y", dtype, [2])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = { + "x": np.array([2, 3], np.float32), + "seed": np.array([0, 1], helper.tensor_dtype_to_np_dtype(dtype)), + } + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds)["y"] + before = model.SerializeToString() + _apply(model) + assert model.SerializeToString() == before + onnx.checker.check_model(model, full_check=True) + for actual, reference in zip(_run_ort(model, feeds)["y"], expected): + np.testing.assert_array_equal(actual, reference) + assert actual.dtype == reference.dtype + + def test_no_cast_nodes_noop(): """Graph with no Cast nodes: graph structure and ORT output unchanged.""" shape = [2, 4]