From 77f2c4fd44ee1bc19e907cdaa0b05fc3e2bda746 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:19 +0000 Subject: [PATCH 01/10] [DLA_transforms_14] [14/17] DLA boolean and comparison rewrites Signed-off-by: mgohil --- .../dla_transforms/dla_greater.py | 109 +++++ .../graph_surgery/dla_transforms/dla_not.py | 130 ++++++ .../graph_surgery/dla_transforms/dla_where.py | 173 ++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 6 + .../dla_transforms/test_dla_greater.py | 284 +++++++++++++ .../dla_transforms/test_dla_not.py | 276 +++++++++++++ .../dla_transforms/test_dla_where.py | 375 ++++++++++++++++++ 7 files changed, 1353 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_not.py create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_where.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py new file mode 100644 index 00000000000..0aa0c305b5b --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py @@ -0,0 +1,109 @@ +# 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. + +"""Cast INT32 / INT64 initializer inputs of ``Greater`` nodes to float32. + +``Greater`` on DLA requires float32 operands. When a comparison threshold is +stored as an integer initializer this transform replaces it with a float32 +copy so no runtime cast is needed. +""" + +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, + add_unique_initializers, + add_value_info, + insert_nodes_at_position, +) + +_INT_TYPES = frozenset({TensorProto.INT32, TensorProto.INT64}) + + +def _apply_greater(model: onnx.ModelProto) -> onnx.ModelProto: + """Ensure all inputs of every Greater node are float32. + + * Initializer inputs with INT32/INT64 dtype: replaced by a float32 initializer; + the original integer initializer is removed. + * Non-initializer (dynamic) inputs whose dtype is not float32: a + ``Cast(to=FLOAT)`` node is inserted before them. Shape inference must have + been run so that tensor dtypes are available in ``value_info``. + """ + graph = model.graph + cache = GraphCache(graph) + unique_init: set[str] = set() + cnt = 0 + converted = {} + + for node in cache.nodes_by_op("Greater"): + for idx, inp in enumerate(node.input): + if not inp: + continue + dtype = cache.get_dtype(inp) + if inp in converted: + node.input[idx] = converted[inp] + continue + if dtype == TensorProto.FLOAT: + continue + + if cache.is_init(inp): + if dtype not in _INT_TYPES: + continue + init_arr = cache.get_init_array(inp) + assert init_arr is not None, f"Initializer {inp!r} not found" + arr = init_arr.astype(np.float32) + new_name = cache.unique_name(f"{inp}_cast") + add_unique_initializers( + graph, unique_init, [numpy_helper.from_array(arr, name=new_name)] + ) + + node.input[idx] = new_name + converted[inp] = new_name + logger.debug( + "Greater %r input[%d]: initializer %r cast to float32.", node.name, idx, inp + ) + else: + # Dynamic tensor with non-float dtype — insert Cast before this node. + if dtype is None: + continue + cast_out = cache.unique_name(f"{inp}_to_float32") + cast_node = helper.make_node( + "Cast", + inputs=[inp], + outputs=[cast_out], + name=f"{inp}_cast_to_float32", + to=TensorProto.FLOAT, + ) + insert_nodes_at_position(graph, [cast_node], node) + node.input[idx] = cast_out + converted[inp] = cast_out + inp_shape = cache.get_shape(inp) + if inp_shape is not None: + add_value_info(graph, cast_out, TensorProto.FLOAT, inp_shape) + logger.debug( + "Greater %r input[%d]: inserted Cast for dynamic tensor %r.", + node.name, + idx, + inp, + ) + cnt += 1 + + logger.debug("greater: fixed %d input(s) to float32.", cnt) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py new file mode 100644 index 00000000000..c9a1ea6fcd2 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py @@ -0,0 +1,130 @@ +# 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. + +"""Replace ``Not`` nodes (preceded by a ``Cast``) with a float32-compatible ``Clip``/``Sub`` pattern. + +DLA does not support the ``Not`` op on boolean tensors. The common pattern is: + + X ──► Cast(to=BOOL) ──► Not ──► Y(BOOL) + +This is rewritten to: + + X ──► Cast(to=FLOAT) ──► Clip(0, 1) ──► Sub(1 - clipped) ──► Y(FLOAT) + +Only unshared Cast chains with provably binary source values are transformed. +``Not`` nodes with other producers (e.g. ``Equal``) or with graph-input operands are left in place. +""" + +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, + add_unique_initializers, + add_value_info, + batch_replace_nodes, +) + + +def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: + """Rewrite unshared casts of binary values; preserve other Not chains.""" + graph = model.graph + cache = GraphCache(graph) + node_idx = {id(n): i for i, n in enumerate(graph.node)} + unique_init: set[str] = set() + replacements: dict[int, tuple[list, list]] = {} + cnt = 0 + + for node in cache.nodes_by_op("Not"): + producer_node = cache.get_producer(node.input[0]) + if producer_node is None or producer_node.op_type != "Cast": + logger.debug("Not %r: input not from a Cast node — skipped.", node.name) + continue + if ( + len(cache.get_consumers(producer_node.output[0])) != 1 + or producer_node.output[0] in cache.graph_output_names + ): + continue + source = producer_node.input[0] + array = cache.get_init_array(source) + binary_source = cache.get_dtype(source) == TensorProto.BOOL or ( + array is not None and np.all((array == 0) | (array == 1)) + ) + boolean_users = {"And", "Or", "Xor", "Not"} + if not binary_source or any( + user.op_type in boolean_users or (user.op_type == "Where" and port == 0) + for user, port in cache.get_consumers(node.output[0]) + ): + continue + + # ── Redirect Cast: to=BOOL → to=FLOAT ──────────────────────────────── + for attr in producer_node.attribute: + if attr.name == "to": + attr.i = TensorProto.FLOAT + + cast_out_vi = cache.vi_map.get(producer_node.output[0]) + if cast_out_vi is not None: + cast_out_vi.type.tensor_type.elem_type = TensorProto.FLOAT + + # ── Unique initializer names scoped to this Not node ────────────────── + min_name = cache.unique_name(f"{node.name}_clip_min") + max_name = cache.unique_name(f"{node.name}_clip_max") + one_name = cache.unique_name(f"{node.name}_one") + + add_unique_initializers( + graph, + unique_init, + [ + numpy_helper.from_array(np.array(0.0, dtype=np.float32), name=min_name), + numpy_helper.from_array(np.array(1.0, dtype=np.float32), name=max_name), + numpy_helper.from_array(np.array(1.0, dtype=np.float32), name=one_name), + ], + ) + + clip_out = cache.unique_name(f"{node.name}_clipped") + clip_node = helper.make_node( + "Clip", + inputs=[producer_node.output[0], min_name, max_name], + outputs=[clip_out], + name=f"{node.name}_clip", + ) + clip_shape = cache.get_shape(producer_node.output[0]) + if clip_shape is not None: + add_value_info(graph, clip_out, TensorProto.FLOAT, clip_shape) + + # Update Not output dtype to FLOAT (graph output or intermediate value_info) + not_out_vi = cache.vi_map.get(node.output[0]) + if not_out_vi is not None: + not_out_vi.type.tensor_type.elem_type = TensorProto.FLOAT + + sub_node = helper.make_node( + "Sub", + inputs=[one_name, clip_out], + outputs=list(node.output), + name=f"{node.name}_sub", + ) + + replacements[node_idx[id(node)]] = ([node], [clip_node, sub_node]) + cnt += 1 + logger.debug("Not %r: replaced with Cast(FLOAT) → Clip → Sub.", node.name) + + batch_replace_nodes(graph, replacements) + + logger.debug("not: transformed %d Not node(s).", cnt) + return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py new file mode 100644 index 00000000000..9f4d8d64762 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.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. + +"""Replace ``Where`` nodes with a float-friendly ``Mul`` when one branch is an all-zero initializer. + +DLA does not support the ``Where`` op. When one of the two data branches is a static all-zero +tensor the operation simplifies to a scalar multiply: + + Where(cond, 0, y) → Mul(y, 1 - float(cond)) + Where(cond, x, 0) → Mul(x, float(cond)) + +The boolean condition is converted to float32 by: + +* **Static BOOL initializer** — replaced in-place by a pre-computed float32 initializer + (or its inverse, for the zero-true-branch case). The original bool initializer is removed + only if the ``Where`` node is its sole consumer. +* **Dynamic BOOL tensor** — a ``Cast(to=FLOAT)`` node is inserted before the ``Mul``. + For the zero-true-branch case a ``Sub(1, cast_out)`` is also inserted. +* **Non-BOOL condition** — used directly as the float mask (no cast needed). + +``Where`` nodes whose both data inputs are non-zero (or non-initializer) are left unchanged. +This C++-compatible arithmetic lowering assumes finite FP32 data values. +""" + +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, + add_unique_initializers, + add_value_info, + batch_replace_nodes, +) + + +def _zero_initializer_index(cache: GraphCache, names: list[str]) -> int | None: + """Return the index (in ``names``) of the first all-zero initializer, or ``None``.""" + for idx, name in enumerate(names): + if cache.is_init(name): + arr = cache.get_init_array(name) + if arr is not None and np.count_nonzero(arr) == 0: + return idx + return None + + +def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: + """Replace eligible ``Where`` nodes with ``Mul``-based equivalents. + + Only ``Where`` nodes where exactly one of the two data inputs is an all-zero + static initializer are transformed. All others are left untouched. + """ + graph = model.graph + cache = GraphCache(graph) + node_idx = {id(n): i for i, n in enumerate(graph.node)} + unique_init: set[str] = set() + replacements: dict[int, tuple[list, list]] = {} + cnt = 0 + + for node in cache.nodes_by_op("Where"): + condition = node.input[0] + x_name = node.input[1] # true branch + y_name = node.input[2] # false branch + + zero_idx = _zero_initializer_index(cache, [x_name, y_name]) + if zero_idx is None: + continue + + non_zero_name = y_name if zero_idx == 0 else x_name + if cache.get_dtype(non_zero_name) != TensorProto.FLOAT: + continue + shapes = [cache.get_shape(name) for name in [condition, x_name, y_name]] + if any(shape is None for shape in shapes): + continue + if np.broadcast_shapes(*shapes) != np.broadcast_shapes( + shapes[0], shapes[2 if zero_idx == 0 else 1] + ): + continue + condition_type = cache.get_dtype(condition) + nodes_to_add: list = [] + + if cache.is_init(condition): + if condition_type == TensorProto.BOOL: + cond_raw = cache.get_init_array(condition) + if cond_raw is None: + continue + cond_arr = cond_raw.astype(np.float32) + if zero_idx == 0: + # Where(cond, 0, y) → y * (1 - cond_float) + cond_arr = 1.0 - cond_arr + float_name = cache.unique_name(f"{condition}_float32") + add_unique_initializers( + graph, unique_init, [numpy_helper.from_array(cond_arr, name=float_name)] + ) + mul_inputs = [non_zero_name, float_name] + else: + # Non-bool static condition — use directly as the float mask. + mul_inputs = [non_zero_name, condition] + + else: + # ── Dynamic condition tensor ────────────────────────────────────── + if condition_type == TensorProto.BOOL: + float_name = cache.unique_name(f"{condition}_float32") + cast_node = helper.make_node( + "Cast", + inputs=[condition], + outputs=[float_name], + to=TensorProto.FLOAT, + name=f"{node.name}_cast_cond", + ) + nodes_to_add.append(cast_node) + cond_shape = cache.get_shape(condition) + if cond_shape is not None: + add_value_info(graph, float_name, TensorProto.FLOAT, cond_shape) + mask_name = float_name + else: + # Already a numeric type — use directly. + mask_name = condition + + if zero_idx == 0: + # Where(cond, 0, y) → y * (1 - mask) + one_name = cache.unique_name(f"{node.name}_one") + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(np.array(1.0, dtype=np.float32), name=one_name)], + ) + inv_name = cache.unique_name(f"{node.name}_inv_cond") + sub_node = helper.make_node( + "Sub", + inputs=[one_name, mask_name], + outputs=[inv_name], + name=f"{node.name}_inv", + ) + nodes_to_add.append(sub_node) + cond_shape = cache.get_shape(condition) + if cond_shape is not None: + add_value_info(graph, inv_name, TensorProto.FLOAT, cond_shape) + mul_inputs = [non_zero_name, inv_name] + else: + # Where(cond, x, 0) → x * mask + mul_inputs = [non_zero_name, mask_name] + + mul_node = helper.make_node( + "Mul", + inputs=mul_inputs, + outputs=list(node.output), + name=f"{node.name}_mul", + ) + nodes_to_add.append(mul_node) + replacements[node_idx[id(node)]] = ([node], nodes_to_add) + cnt += 1 + logger.debug("Where %r: replaced by Mul (zero_idx=%d).", node.name, zero_idx) + + batch_replace_nodes(graph, replacements) + + logger.debug("where: transformed %d Where node(s).", cnt) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index b14fe04c876..ffa9f757b7c 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -35,6 +35,7 @@ from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model from .dla_transforms.dla_decompose_lstm import _apply_decompose_lstm from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers +from .dla_transforms.dla_greater import _apply_greater from .dla_transforms.dla_indexing import ( _apply_gather, _apply_gather_elements, @@ -53,6 +54,7 @@ _BoundaryPreflightError, _require_output_rank_conversion, ) +from .dla_transforms.dla_not import _apply_not from .dla_transforms.dla_rank_conversion import ( _apply_argmax, _apply_flatten, @@ -73,6 +75,7 @@ _apply_reshape_reduce_sum, _apply_reshape_transpose_reshape, ) +from .dla_transforms.dla_where import _apply_where from .utils.graph_utils import ( DEFAULT_QDQ_REMOVE_QUANT_TYPES, GraphCache, @@ -137,6 +140,9 @@ def _apply_constants_to_initializers(model, **options): ("Non4DLpNorm", _apply_non4d_lpnorm), ("FixInstanceNormChannelMismatch", _apply_fix_instancenorm_channels), ("UpdateAxis", _apply_update_axis), + ("Where", _apply_where), + ("Not", _apply_not), + ("Greater", _apply_greater), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py new file mode 100644 index 00000000000..54ae0a9333d --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py @@ -0,0 +1,284 @@ +# 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_greater`. + +"""ORT CPU parity tests for ``dla_greater``. + +Transform rule +-------------- +For every ``Greater`` node, any initializer input whose dtype is INT32 or +INT64 is replaced by a new float32 initializer (name ``{original}_cast``). +Float32 / other-dtype initializers and dynamic (non-initializer) inputs are +left untouched. + +Coverage +-------- +INT32 initializer threshold → cast to float32, ORT parity +INT64 initializer threshold → cast to float32, ORT parity +float32 initializer → unchanged (no cast) +both inputs dynamic → graph unchanged, ORT parity +both inputs INT32 initializers → both cast to float32 +""" + +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_greater 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" / "greater_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_greater + + +# ── 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 _init_dtype(model: onnx.ModelProto, name: str) -> int | None: + for init in model.graph.initializer: + if init.name == name: + return init.data_type + return None + + +_RNG = np.random.default_rng(2024) + + +# ── 1. INT32 threshold initializer → cast to float32 ───────────────────────── + + +def test_greater_int32_init_converted(): + """Greater(X_int32, int32_threshold) → threshold becomes float32; ORT parity. + + Both inputs are INT32 so the base model is valid for ORT. + After the transform the threshold becomes float32; X's declared type is updated + to float32 in the test (reflecting that upstream transforms will produce float32). + """ + shape = [3, 4] + threshold = np.array(5, dtype=np.int32) + thresh_init = numpy_helper.from_array(threshold, name="thresh") + + greater = helper.make_node("Greater", ["X", "thresh"], ["Y"], name="greater") + graph = helper.make_graph( + [greater], + "greater_i32", + [_vi("X", TensorProto.INT32, shape)], + [_vi("Y", TensorProto.BOOL, shape)], + initializer=[thresh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = (_RNG.standard_normal(shape) * 10).astype(np.int32) + _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 _init_dtype(tr, "thresh_cast") == TensorProto.FLOAT, "threshold must become float32" + # original INT32 init is left for _remove_unused_initializers to clean up + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 2. INT64 threshold initializer → cast to float32 ───────────────────────── + + +def test_greater_int64_init_converted(): + """Greater(X_int64, int64_threshold) → threshold becomes float32; ORT parity.""" + shape = [2, 6] + threshold = np.array(-3, dtype=np.int64) + thresh_init = numpy_helper.from_array(threshold, name="thresh64") + + greater = helper.make_node("Greater", ["X", "thresh64"], ["Y"], name="greater") + graph = helper.make_graph( + [greater], + "greater_i64", + [_vi("X", TensorProto.INT64, shape)], + [_vi("Y", TensorProto.BOOL, shape)], + initializer=[thresh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x = (_RNG.standard_normal(shape) * 5).astype(np.int64) + _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 _init_dtype(tr, "thresh64_cast") == TensorProto.FLOAT + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 3. float32 initializer → unchanged ─────────────────────────────────────── + + +def test_greater_float32_init_unchanged(): + """Greater(X_float32, float32_threshold) — already float32, no change at all.""" + shape = [2, 4] + threshold = np.array(0.5, dtype=np.float32) + thresh_init = numpy_helper.from_array(threshold, name="thresh_f32") + + greater = helper.make_node("Greater", ["X", "thresh_f32"], ["Y"], name="greater") + graph = helper.make_graph( + [greater], + "greater_f32", + [_vi("X", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.BOOL, shape)], + initializer=[thresh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_node_count = len(model.graph.node) + original_init_count = len(model.graph.initializer) + + 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}) + + # Both inputs already float32 — no Cast inserted, no new initializer, node inputs unchanged + assert len(tr.graph.node) == original_node_count + assert len(tr.graph.initializer) == original_init_count + greater_node = next(n for n in tr.graph.node if n.op_type == "Greater") + assert greater_node.input[1] == "thresh_f32" + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 4. Both inputs dynamic → graph unchanged ───────────────────────────────── + + +def test_greater_dynamic_float32_inputs_noop(): + """Greater(X_float32, Z_float32) — both already float32, no Cast inserted.""" + shape = [3, 4] + greater = helper.make_node("Greater", ["X", "Z"], ["Y"], name="greater") + graph = helper.make_graph( + [greater], + "greater_dynamic", + [_vi("X", TensorProto.FLOAT, shape), _vi("Z", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.BOOL, 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) * 5 + z = _RNG.standard_normal(shape).astype(np.float32) * 5 + _prepare_model(model) + ref = _run_ort(model, {"X": x, "Z": z}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x, "Z": z}) + + # Both inputs are float32 — shape inference informs transform to skip them + assert len(tr.graph.node) == original_node_count, ( + "no Cast nodes should be inserted for float32 inputs" + ) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 5. Both inputs are INT32 initializers → both cast ──────────────────────── + + +def test_greater_both_inputs_int32_converted(): + """Greater(int32_A, int32_B) — both initializers cast to float32.""" + a_init = numpy_helper.from_array(np.array([1, 2, 3, 4], dtype=np.int32), name="A") + b_init = numpy_helper.from_array(np.array([4, 3, 2, 1], dtype=np.int32), name="B") + + greater = helper.make_node("Greater", ["A", "B"], ["Y"], name="greater") + graph = helper.make_graph( + [greater], + "greater_both_int", + [], + [_vi("Y", TensorProto.BOOL, [4])], + initializer=[a_init, b_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {}) + + assert _init_dtype(tr, "A_cast") == TensorProto.FLOAT, "A must be cast to float32" + assert _init_dtype(tr, "B_cast") == TensorProto.FLOAT, "B must be cast to float32" + np.testing.assert_array_equal(ref["Y"], got["Y"]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py new file mode 100644 index 00000000000..92d561b9d3d --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py @@ -0,0 +1,276 @@ +# 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_not`. + +"""ORT CPU parity tests for ``dla_not``. + +Transform rule +-------------- +``Cast(to=BOOL) → Not`` is rewritten to +``Cast(to=FLOAT) → Clip(0, 1) → Sub(1 - clipped)``. + +Only ``Not`` nodes whose input is produced by a ``Cast`` node are touched; +all other ``Not`` nodes (e.g. ``Equal → Not``, graph-input → Not) are left in place. + +Coverage +-------- +INT32 input → Cast(BOOL) → Not → transformed, ORT parity +Non-Cast producer (Equal → Not) → skipped, graph unchanged +Graph-input (BOOL) → Not → skipped, graph unchanged +Two Cast→Not chains → both transformed, unique initializer names, ORT parity +Graph with no Not nodes → noop, ORT parity +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper + +from modelopt.onnx.graph_surgery.dla_transforms import dla_not 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" / "not_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_not + + +# ── 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} + + +_RNG = np.random.default_rng(2024) + + +# ── 1. INT32 → Cast(BOOL) → Not → transformed ──────────────────────────────── + + +def test_not_cast_to_bool_transformed(): + """Cast(to=BOOL) → Not: Cast redirected to FLOAT, Not replaced by Clip+Sub; ORT parity. + + Input uses only 0 / 1 values so that Cast(INT32→FLOAT) and Cast(INT32→BOOL) + produce numerically consistent boolean representations. + """ + shape = [2, 3] + cast_node = helper.make_node("Cast", ["X"], ["X_bool"], to=TensorProto.BOOL, name="cast") + not_node = helper.make_node("Not", ["X_bool"], ["Y"], name="not") + graph = helper.make_graph( + [cast_node, not_node], + "not_cast_bool", + [_vi("X", TensorProto.BOOL, shape)], + [_vi("Y", TensorProto.BOOL, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + # Only 0/1 so Not(Cast(x, BOOL)) == 1 - cast(x, FLOAT) for these values. + x = np.array([[0, 1, 0], [1, 1, 0]], dtype=np.bool_) + _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}) + + # Transformed output is FLOAT; compare against bool cast to float. + np.testing.assert_array_equal(ref["Y"].astype(np.float32), got["Y"]) + + # Cast node must now produce FLOAT, Not node must be gone. + cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] + assert len(cast_nodes) == 1 + for attr in cast_nodes[0].attribute: + if attr.name == "to": + assert attr.i == TensorProto.FLOAT, "Cast must target FLOAT after transform" + assert not any(n.op_type == "Not" for n in tr.graph.node), "Not node must be removed" + + +# ── 2. Equal → Not: non-Cast producer → skipped ────────────────────────────── + + +def test_not_non_cast_producer_skipped(): + """Not whose input comes from Equal (not Cast) → graph unchanged.""" + shape = [3] + eq_node = helper.make_node("Equal", ["X", "Z"], ["X_bool"], name="eq") + not_node = helper.make_node("Not", ["X_bool"], ["Y"], name="not") + graph = helper.make_graph( + [eq_node, not_node], + "not_equal_producer", + [_vi("X", TensorProto.FLOAT, shape), _vi("Z", TensorProto.FLOAT, shape)], + [_vi("Y", TensorProto.BOOL, 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) + z = _RNG.standard_normal(shape).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X": x, "Z": z}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X": x, "Z": z}) + + assert len(tr.graph.node) == original_node_count, ( + "graph must be unchanged for non-Cast producer" + ) + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 3. BOOL graph input → Not: no producer → skipped ───────────────────────── + + +def test_not_graph_input_skipped(): + """Not whose operand is a BOOL graph input (no producing node) → graph unchanged.""" + shape = [2, 3] + not_node = helper.make_node("Not", ["X"], ["Y"], name="not") + graph = helper.make_graph( + [not_node], + "not_graph_input", + [_vi("X", TensorProto.BOOL, shape)], + [_vi("Y", TensorProto.BOOL, 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.integers(0, 2, shape).astype(bool) + _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, "graph must be unchanged for graph-input Not" + np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +# ── 4. Two Cast→Not chains → both transformed, unique initializer names ─────── + + +def test_not_multiple_not_nodes(): + """Two independent Cast(BOOL)→Not chains are both transformed with unique initializer names.""" + shape = [2] + cast_a = helper.make_node("Cast", ["A"], ["A_bool"], to=TensorProto.BOOL, name="cast_a") + not_a = helper.make_node("Not", ["A_bool"], ["Y_A"], name="not_a") + cast_b = helper.make_node("Cast", ["B"], ["B_bool"], to=TensorProto.BOOL, name="cast_b") + not_b = helper.make_node("Not", ["B_bool"], ["Y_B"], name="not_b") + graph = helper.make_graph( + [cast_a, not_a, cast_b, not_b], + "not_multi", + [_vi("A", TensorProto.BOOL, shape), _vi("B", TensorProto.BOOL, shape)], + [_vi("Y_A", TensorProto.BOOL, shape), _vi("Y_B", TensorProto.BOOL, shape)], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + a = np.array([0, 1], dtype=np.bool_) + b = np.array([1, 0], dtype=np.bool_) + _prepare_model(model) + ref = _run_ort(model, {"A": a, "B": b}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"A": a, "B": b}) + + np.testing.assert_array_equal(ref["Y_A"].astype(np.float32), got["Y_A"]) + np.testing.assert_array_equal(ref["Y_B"].astype(np.float32), got["Y_B"]) + + # Both Not nodes must be gone; no initializer name collisions. + assert not any(n.op_type == "Not" for n in tr.graph.node), "all Not nodes must be removed" + init_names = [init.name for init in tr.graph.initializer] + assert len(init_names) == len(set(init_names)), "initializer names must be unique" + + +# ── 5. No Not nodes → noop ──────────────────────────────────────────────────── + + +def test_no_not_nodes_noop(): + """Graph with no Not nodes: transform is a noop, ORT parity preserved.""" + shape = [2, 3] + relu_node = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu_node], + "no_not", + [_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, "noop: node count must not change" + np.testing.assert_array_equal(ref["Y"], got["Y"]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py new file mode 100644 index 00000000000..ae0788e5e57 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -0,0 +1,375 @@ +# 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_where`. + +"""ORT CPU parity tests for ``dla_where``. + +Transform rule +-------------- +``Where(cond, x, y)`` is replaced by a ``Mul``-based equivalent when one of +{x, y} is an all-zero static initializer: + + Where(cond, 0, y) → Mul(y, 1 - float(cond)) + Where(cond, x, 0) → Mul(x, float(cond)) + +Coverage +-------- +BOOL init condition, true-branch zeros → Mul(y, 1-cond_float), ORT parity +BOOL init condition, false-branch zeros → Mul(x, cond_float), ORT parity +Dynamic BOOL condition, true-branch zeros → Cast+Sub+Mul, ORT parity +Dynamic BOOL condition, false-branch zeros → Cast+Mul, ORT parity +Non-BOOL float condition, false-branch zeros → Mul directly, ORT parity +Neither branch all-zeros → graph unchanged +No Where nodes → noop, ORT parity +""" + +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_where 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" / "where_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + + +_apply = _mod._apply_where + + +# ── 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} + + +_RNG = np.random.default_rng(2024) +_SHAPE = [2, 3] + + +# ── 1. Static BOOL condition, true-branch (x) is zeros ─────────────────────── + + +def test_where_bool_init_cond_x_zeros(): + """Where(bool_init, zeros, y) → Mul(y, 1-cond_float); ORT parity. + + Where(cond, 0, y) selects y when condition is False, so the mask is (1 - cond). + """ + cond_arr = np.array([[True, False, True], [False, True, False]]) + zeros_arr = np.zeros(_SHAPE, dtype=np.float32) + y_arr = _RNG.standard_normal(_SHAPE).astype(np.float32) + + cond_init = numpy_helper.from_array(cond_arr, name="cond") + x_init = numpy_helper.from_array(zeros_arr, name="x_zeros") + + where = helper.make_node("Where", ["cond", "x_zeros", "Y_in"], ["Z"], name="where") + graph = helper.make_graph( + [where], + "where_bool_init_x0", + [_vi("Y_in", TensorProto.FLOAT, _SHAPE)], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[cond_init, x_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {"Y_in": y_arr}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"Y_in": y_arr}) + + assert not any(n.op_type == "Where" for n in tr.graph.node), "Where must be removed" + np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) + + +# ── 2. Static BOOL condition, false-branch (y) is zeros ────────────────────── + + +def test_where_bool_init_cond_y_zeros(): + """Where(bool_init, x, zeros) → Mul(x, cond_float); ORT parity. + + Where(cond, x, 0) selects x when condition is True, so the mask is cond itself. + """ + cond_arr = np.array([[True, False, True], [False, True, False]]) + zeros_arr = np.zeros(_SHAPE, dtype=np.float32) + x_arr = _RNG.standard_normal(_SHAPE).astype(np.float32) + + cond_init = numpy_helper.from_array(cond_arr, name="cond2") + y_init = numpy_helper.from_array(zeros_arr, name="y_zeros") + + where = helper.make_node("Where", ["cond2", "X_in", "y_zeros"], ["Z"], name="where") + graph = helper.make_graph( + [where], + "where_bool_init_y0", + [_vi("X_in", TensorProto.FLOAT, _SHAPE)], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[cond_init, y_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + _prepare_model(model) + ref = _run_ort(model, {"X_in": x_arr}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X_in": x_arr}) + + assert not any(n.op_type == "Where" for n in tr.graph.node), "Where must be removed" + np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) + + +# ── 3. Dynamic BOOL condition, true-branch (x) is zeros ────────────────────── + + +def test_where_dynamic_bool_cond_x_zeros(): + """Where(bool_graph_input, zeros_init, y) → Cast+Sub+Mul; ORT parity.""" + zeros_arr = np.zeros(_SHAPE, dtype=np.float32) + x_init = numpy_helper.from_array(zeros_arr, name="x_zeros_dyn") + + where = helper.make_node("Where", ["cond", "x_zeros_dyn", "Y_in"], ["Z"], name="where") + graph = helper.make_graph( + [where], + "where_dyn_bool_x0", + [ + _vi("cond", TensorProto.BOOL, _SHAPE), + _vi("Y_in", TensorProto.FLOAT, _SHAPE), + ], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[x_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + cond_val = _RNG.integers(0, 2, _SHAPE).astype(bool) + y_arr = _RNG.standard_normal(_SHAPE).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"cond": cond_val, "Y_in": y_arr}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"cond": cond_val, "Y_in": y_arr}) + + assert not any(n.op_type == "Where" for n in tr.graph.node), "Where must be removed" + # Cast + Sub + Mul should be present + op_types = {n.op_type for n in tr.graph.node} + assert "Cast" in op_types, "Cast node expected for dynamic BOOL condition" + assert "Sub" in op_types, "Sub node expected for x-is-zero inversion" + assert "Mul" in op_types + np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) + + +# ── 4. Dynamic BOOL condition, false-branch (y) is zeros ───────────────────── + + +def test_where_dynamic_bool_cond_y_zeros(): + """Where(bool_graph_input, x, zeros_init) → Cast+Mul; ORT parity.""" + zeros_arr = np.zeros(_SHAPE, dtype=np.float32) + y_init = numpy_helper.from_array(zeros_arr, name="y_zeros_dyn") + + where = helper.make_node("Where", ["cond", "X_in", "y_zeros_dyn"], ["Z"], name="where") + graph = helper.make_graph( + [where], + "where_dyn_bool_y0", + [ + _vi("cond", TensorProto.BOOL, _SHAPE), + _vi("X_in", TensorProto.FLOAT, _SHAPE), + ], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[y_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + cond_val = _RNG.integers(0, 2, _SHAPE).astype(bool) + x_arr = _RNG.standard_normal(_SHAPE).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"cond": cond_val, "X_in": x_arr}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"cond": cond_val, "X_in": x_arr}) + + assert not any(n.op_type == "Where" for n in tr.graph.node), "Where must be removed" + op_types = {n.op_type for n in tr.graph.node} + assert "Cast" in op_types, "Cast node expected for dynamic BOOL condition" + assert "Sub" not in op_types, "Sub must NOT be inserted for y-is-zero case" + assert "Mul" in op_types + np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) + + +# ── 5. Non-BOOL float condition (dynamic), false-branch zeros ───────────────── + + +def test_where_float_cond_y_zeros(): + """Where with a FLOAT dynamic condition and zeros false-branch → Mul directly (no Cast). + + The float condition (0.0 / 1.0 values) acts directly as the mask. + """ + zeros_arr = np.zeros(_SHAPE, dtype=np.float32) + y_init = numpy_helper.from_array(zeros_arr, name="y_zeros_f") + + # Sigmoid produces values in (0, 1); use a Cast(BOOL) then Cast(FLOAT) workaround + # is not needed here — instead, we make a simple graph where the condition is FLOAT. + # ORT does NOT accept a FLOAT condition for Where, so we model it differently: + # build a graph with Greater producing a BOOL mask, then use that as condition. + greater = helper.make_node("Greater", ["X_in", "thresh"], ["cond_bool"], name="greater") + cast = helper.make_node( + "Cast", ["cond_bool"], ["cond_float"], to=TensorProto.FLOAT, name="cast" + ) + where = helper.make_node("Where", ["cond_bool", "X_in", "y_zeros_f"], ["Z"], name="where") + + thresh_init = numpy_helper.from_array(np.array(0.0, dtype=np.float32), name="thresh") + graph = helper.make_graph( + [greater, cast, where], + "where_float_cond_y0", + [_vi("X_in", TensorProto.FLOAT, _SHAPE)], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[y_init, thresh_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + + x_arr = _RNG.standard_normal(_SHAPE).astype(np.float32) + _prepare_model(model) + ref = _run_ort(model, {"X_in": x_arr}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {"X_in": x_arr}) + + assert not any(n.op_type == "Where" for n in tr.graph.node), "Where must be removed" + np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) + + +# ── 6. Neither branch is all-zeros → graph unchanged ───────────────────────── + + +def test_where_no_zero_branch_skipped(): + """Where with no all-zero branch is left unchanged.""" + x_arr = np.ones(_SHAPE, dtype=np.float32) + y_arr = np.full(_SHAPE, 2.0, dtype=np.float32) + cond_arr = np.array([[True, False, True], [False, True, False]]) + + cond_init = numpy_helper.from_array(cond_arr, name="cond_nz") + x_init = numpy_helper.from_array(x_arr, name="x_nz") + y_init = numpy_helper.from_array(y_arr, name="y_nz") + + where = helper.make_node("Where", ["cond_nz", "x_nz", "y_nz"], ["Z"], name="where") + graph = helper.make_graph( + [where], + "where_no_zero", + [], + [_vi("Z", TensorProto.FLOAT, _SHAPE)], + initializer=[cond_init, x_init, y_init], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir(model) + original_node_count = len(model.graph.node) + + _prepare_model(model) + ref = _run_ort(model, {}) + + tr = copy.deepcopy(model) + _prepare_model(tr) + tr = _apply(tr) + _prepare_model(tr) + got = _run_ort(tr, {}) + + assert len(tr.graph.node) == original_node_count, "graph must be unchanged" + np.testing.assert_array_equal(ref["Z"], got["Z"]) + + +# ── 7. No Where nodes → noop ────────────────────────────────────────────────── + + +def test_no_where_nodes_noop(): + """Graph without Where nodes: transform is a noop, ORT parity preserved.""" + shape = [2, 3] + relu_node = helper.make_node("Relu", ["X"], ["Y"], name="relu") + graph = helper.make_graph( + [relu_node], + "no_where", + [_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, "noop: node count must not change" + np.testing.assert_array_equal(ref["Y"], got["Y"]) From 5239c23b8980f90cb53349a3398e864423879339 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:16:20 +0000 Subject: [PATCH 02/10] Preserve comparison ordering and boolean initializer ownership Signed-off-by: mgohil --- .../dla_transforms/dla_greater.py | 152 ++++++++------- .../graph_surgery/dla_transforms/dla_not.py | 52 ++++-- .../graph_surgery/dla_transforms/dla_where.py | 61 +++--- .../dla_transforms/test_dla_greater.py | 173 ++++++++++++++++-- .../dla_transforms/test_dla_not.py | 153 ++++++++++++++-- .../dla_transforms/test_dla_where.py | 128 +++++++++++-- 6 files changed, 544 insertions(+), 175 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py index 0aa0c305b5b..2572519a212 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py @@ -13,12 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Cast INT32 / INT64 initializer inputs of ``Greater`` nodes to float32. - -``Greater`` on DLA requires float32 operands. When a comparison threshold is -stored as an integer initializer this transform replaces it with a float32 -copy so no runtime cast is needed. -""" +"""Prepare FLOAT operands only where conversion preserves comparison ordering.""" from __future__ import annotations @@ -26,84 +21,87 @@ import onnx from onnx import TensorProto, helper, numpy_helper -from ...logging_config import logger -from ..utils.graph_utils import ( - GraphCache, - add_unique_initializers, - add_value_info, - insert_nodes_at_position, -) +from ..utils.graph_utils import GraphCache, add_unique_initializers, add_value_info +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_indexing import _insert_support_nodes, _operator_nodes +from .dla_rank_conversion import _runtime_input_names + +__all__ = [] _INT_TYPES = frozenset({TensorProto.INT32, TensorProto.INT64}) +_EXACT_TYPES = frozenset( + { + TensorProto.FLOAT, + TensorProto.FLOAT16, + TensorProto.INT8, + TensorProto.UINT8, + TensorProto.INT16, + TensorProto.UINT16, + } +) def _apply_greater(model: onnx.ModelProto) -> onnx.ModelProto: - """Ensure all inputs of every Greater node are float32. - - * Initializer inputs with INT32/INT64 dtype: replaced by a float32 initializer; - the original integer initializer is removed. - * Non-initializer (dynamic) inputs whose dtype is not float32: a - ``Cast(to=FLOAT)`` node is inserted before them. Shape inference must have - been run so that tensor dtypes are available in ``value_info``. - """ + """Convert comparisons whose ordering survives FLOAT operand conversion.""" graph = model.graph cache = GraphCache(graph) - unique_init: set[str] = set() - cnt = 0 - converted = {} - - for node in cache.nodes_by_op("Greater"): - for idx, inp in enumerate(node.input): - if not inp: - continue - dtype = cache.get_dtype(inp) - if inp in converted: - node.input[idx] = converted[inp] + selected = {id(node) for node in _operator_nodes(model, {"Greater"})} + runtime = _runtime_input_names(model) + unique_init = set(cache.init_names) + converted, safe_thresholds = {}, {} + nodes = [] + for node in list(graph.node): + if id(node) not in selected: + nodes.append(node) + continue + types = {cache.get_dtype(name) for name in node.input} + safe = types <= _EXACT_TYPES + if types <= _INT_TYPES: + for name in node.input: + if name not in safe_thresholds: + tensor = cache.get_init(name) if name not in runtime else None + array = numpy_helper.to_array(tensor) if tensor is not None else None + # Interior thresholds retain separation from either adjacent integer. + safe_thresholds[name] = array is not None and ( + array.size == 0 or (array.min() > -(2**24) and array.max() < 2**24) + ) + safe |= safe_thresholds[name] + if not safe: + nodes.append(node) + continue + for index, name in enumerate(node.input): + if cache.get_dtype(name) == TensorProto.FLOAT: continue - if dtype == TensorProto.FLOAT: - continue - - if cache.is_init(inp): - if dtype not in _INT_TYPES: - continue - init_arr = cache.get_init_array(inp) - assert init_arr is not None, f"Initializer {inp!r} not found" - arr = init_arr.astype(np.float32) - new_name = cache.unique_name(f"{inp}_cast") - add_unique_initializers( - graph, unique_init, [numpy_helper.from_array(arr, name=new_name)] - ) - - node.input[idx] = new_name - converted[inp] = new_name - logger.debug( - "Greater %r input[%d]: initializer %r cast to float32.", node.name, idx, inp + if name not in converted: + output = cache.unique_name( + name + "_cast" + if cache.is_init(name) and name not in runtime + else name + "_to_float32" ) - else: - # Dynamic tensor with non-float dtype — insert Cast before this node. - if dtype is None: - continue - cast_out = cache.unique_name(f"{inp}_to_float32") - cast_node = helper.make_node( - "Cast", - inputs=[inp], - outputs=[cast_out], - name=f"{inp}_cast_to_float32", - to=TensorProto.FLOAT, - ) - insert_nodes_at_position(graph, [cast_node], node) - node.input[idx] = cast_out - converted[inp] = cast_out - inp_shape = cache.get_shape(inp) - if inp_shape is not None: - add_value_info(graph, cast_out, TensorProto.FLOAT, inp_shape) - logger.debug( - "Greater %r input[%d]: inserted Cast for dynamic tensor %r.", - node.name, - idx, - inp, - ) - cnt += 1 - - logger.debug("greater: fixed %d input(s) to float32.", cnt) + if cache.is_init(name) and name not in runtime: + array = cache.get_init_array(name) + assert array is not None + add_unique_initializers( + graph, + unique_init, + [numpy_helper.from_array(array.astype(np.float32), output)], + ) + else: + nodes.append( + helper.make_node( + "Cast", + [name], + [output], + name=cache.unique_name(name + "_cast_to_float32"), + to=TensorProto.FLOAT, + ) + ) + shape = cache.get_shape(name) + if shape is not None: + add_value_info(graph, output, TensorProto.FLOAT, shape) + converted[name] = output + node.input[index] = converted[name] + nodes.append(node) + _insert_support_nodes(graph, nodes) + _add_legacy_initializer_inputs(model) return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py index c9a1ea6fcd2..d29923d8b03 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py @@ -40,36 +40,63 @@ add_value_info, batch_replace_nodes, ) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures +from .dla_indexing import _operator_nodes +from .dla_rank_conversion import _runtime_input_names + +__all__ = [] + + +def _binary_source(cache: GraphCache, name: str, runtime: set[str]) -> bool: + """Prove binary values without retaining a decoded fixed tensor.""" + if cache.get_dtype(name) == TensorProto.BOOL: + return True + tensor = cache.get_init(name) if name not in runtime else None + if tensor is None: + return False + array = numpy_helper.to_array(tensor) + return all( + np.all((chunk == 0) | (chunk == 1)) + for chunk in np.nditer( + array, flags=["external_loop", "buffered", "zerosize_ok"], buffersize=65536, order="C" + ) + ) def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: """Rewrite unshared casts of binary values; preserve other Not chains.""" graph = model.graph cache = GraphCache(graph) + if next((op.version for op in model.opset_import if op.domain in {"", "ai.onnx"}), 1) < 11: + return model node_idx = {id(n): i for i, n in enumerate(graph.node)} + captured = _subgraph_captures(graph) + casts = {id(node) for node in _operator_nodes(model, {"Cast"})} + runtime = _runtime_input_names(model) + binary = {} unique_init: set[str] = set() replacements: dict[int, tuple[list, list]] = {} cnt = 0 - for node in cache.nodes_by_op("Not"): + for node in _operator_nodes(model, {"Not"}): producer_node = cache.get_producer(node.input[0]) - if producer_node is None or producer_node.op_type != "Cast": + if producer_node is None or id(producer_node) not in casts: logger.debug("Not %r: input not from a Cast node — skipped.", node.name) continue if ( len(cache.get_consumers(producer_node.output[0])) != 1 - or producer_node.output[0] in cache.graph_output_names + or producer_node.output[0] in cache.graph_output_names | captured + or node.output[0] in captured ): continue source = producer_node.input[0] - array = cache.get_init_array(source) - binary_source = cache.get_dtype(source) == TensorProto.BOOL or ( - array is not None and np.all((array == 0) | (array == 1)) - ) - boolean_users = {"And", "Or", "Xor", "Not"} + if source not in binary: + binary[source] = _binary_source(cache, source, runtime) + binary_source = binary[source] if not binary_source or any( - user.op_type in boolean_users or (user.op_type == "Where" and port == 0) - for user, port in cache.get_consumers(node.output[0]) + user.op_type != "Cast" or id(user) not in casts + for user, _port in cache.get_consumers(node.output[0]) ): continue @@ -102,7 +129,7 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: "Clip", inputs=[producer_node.output[0], min_name, max_name], outputs=[clip_out], - name=f"{node.name}_clip", + name=cache.unique_name(f"{node.name}_clip"), ) clip_shape = cache.get_shape(producer_node.output[0]) if clip_shape is not None: @@ -117,7 +144,7 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: "Sub", inputs=[one_name, clip_out], outputs=list(node.output), - name=f"{node.name}_sub", + name=cache.unique_name(f"{node.name}_sub"), ) replacements[node_idx[id(node)]] = ([node], [clip_node, sub_node]) @@ -126,5 +153,6 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: batch_replace_nodes(graph, replacements) + _add_legacy_initializer_inputs(model) logger.debug("not: transformed %d Not node(s).", cnt) return model diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py index 9f4d8d64762..74021fca565 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -13,25 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Replace ``Where`` nodes with a float-friendly ``Mul`` when one branch is an all-zero initializer. +"""Lower zero-branch selections to arithmetic under the existing finite-FP32 assumption. -DLA does not support the ``Where`` op. When one of the two data branches is a static all-zero -tensor the operation simplifies to a scalar multiply: - - Where(cond, 0, y) → Mul(y, 1 - float(cond)) - Where(cond, x, 0) → Mul(x, float(cond)) - -The boolean condition is converted to float32 by: - -* **Static BOOL initializer** — replaced in-place by a pre-computed float32 initializer - (or its inverse, for the zero-true-branch case). The original bool initializer is removed - only if the ``Where`` node is its sole consumer. -* **Dynamic BOOL tensor** — a ``Cast(to=FLOAT)`` node is inserted before the ``Mul``. - For the zero-true-branch case a ``Sub(1, cast_out)`` is also inserted. -* **Non-BOOL condition** — used directly as the float mask (no cast needed). - -``Where`` nodes whose both data inputs are non-zero (or non-initializer) are left unchanged. -This C++-compatible arithmetic lowering assumes finite FP32 data values. +Fixed BOOL conditions get private FLOAT masks; runtime conditions use Cast nodes. +Overridable zero branches remain native, and initializer owners are preserved. """ from __future__ import annotations @@ -47,14 +32,22 @@ add_value_info, batch_replace_nodes, ) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_indexing import _operator_nodes +from .dla_rank_conversion import _runtime_input_names + +__all__ = [] -def _zero_initializer_index(cache: GraphCache, names: list[str]) -> int | None: +def _zero_initializer_index( + cache: GraphCache, names: list[str], runtime: set[str], zeros: dict[str, bool] +) -> int | None: """Return the index (in ``names``) of the first all-zero initializer, or ``None``.""" for idx, name in enumerate(names): - if cache.is_init(name): - arr = cache.get_init_array(name) - if arr is not None and np.count_nonzero(arr) == 0: + if cache.is_init(name) and name not in runtime: + if name not in zeros: + zeros[name] = np.count_nonzero(numpy_helper.to_array(cache.get_init(name))) == 0 + if zeros[name]: return idx return None @@ -71,13 +64,15 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: unique_init: set[str] = set() replacements: dict[int, tuple[list, list]] = {} cnt = 0 + runtime = _runtime_input_names(model) + zeros = {} - for node in cache.nodes_by_op("Where"): + for node in _operator_nodes(model, {"Where"}): condition = node.input[0] x_name = node.input[1] # true branch y_name = node.input[2] # false branch - zero_idx = _zero_initializer_index(cache, [x_name, y_name]) + zero_idx = _zero_initializer_index(cache, [x_name, y_name], runtime, zeros) if zero_idx is None: continue @@ -94,7 +89,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: condition_type = cache.get_dtype(condition) nodes_to_add: list = [] - if cache.is_init(condition): + if cache.is_init(condition) and condition not in runtime: if condition_type == TensorProto.BOOL: cond_raw = cache.get_init_array(condition) if cond_raw is None: @@ -102,10 +97,13 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: cond_arr = cond_raw.astype(np.float32) if zero_idx == 0: # Where(cond, 0, y) → y * (1 - cond_float) - cond_arr = 1.0 - cond_arr + np.subtract(1.0, cond_arr, out=cond_arr) float_name = cache.unique_name(f"{condition}_float32") - add_unique_initializers( - graph, unique_init, [numpy_helper.from_array(cond_arr, name=float_name)] + graph.initializer.add( + name=float_name, + data_type=TensorProto.FLOAT, + dims=cond_arr.shape, + raw_data=cond_arr.astype(" onnx.ModelProto: inputs=[condition], outputs=[float_name], to=TensorProto.FLOAT, - name=f"{node.name}_cast_cond", + name=cache.unique_name(f"{node.name}_cast_cond"), ) nodes_to_add.append(cast_node) cond_shape = cache.get_shape(condition) @@ -145,7 +143,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: "Sub", inputs=[one_name, mask_name], outputs=[inv_name], - name=f"{node.name}_inv", + name=cache.unique_name(f"{node.name}_inv"), ) nodes_to_add.append(sub_node) cond_shape = cache.get_shape(condition) @@ -160,7 +158,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: "Mul", inputs=mul_inputs, outputs=list(node.output), - name=f"{node.name}_mul", + name=cache.unique_name(f"{node.name}_mul"), ) nodes_to_add.append(mul_node) replacements[node_idx[id(node)]] = ([node], nodes_to_add) @@ -169,5 +167,6 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: batch_replace_nodes(graph, replacements) + _add_legacy_initializer_inputs(model) logger.debug("where: transformed %d Where node(s).", cnt) return model diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py index 54ae0a9333d..776aa6f3370 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py @@ -39,19 +39,17 @@ 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 onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_greater as _mod -pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "greater_test_models" _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -63,36 +61,42 @@ def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Declare a typed tensor shape.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir(model: onnx.ModelProto) -> None: + """Select an IR version supported by the CPU runtime.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _prepare_model(model: onnx.ModelProto) -> None: + """Populate metadata using optional runtime symbolic inference.""" 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 + # Optional ONNX Runtime tooling is absent in some CPU installations. + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + except ImportError: + return + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-untyped] - + """Fully check and execute all public outputs in one CPU session.""" _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} + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + sess = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + return dict(zip([o.name for o in model.graph.output], sess.run(None, feeds))) def _init_dtype(model: onnx.ModelProto, name: str) -> int | None: + """Read a named initializer element type.""" for init in model.graph.initializer: if init.name == name: return init.data_type @@ -282,3 +286,136 @@ def test_greater_both_inputs_int32_converted(): assert _init_dtype(tr, "A_cast") == TensorProto.FLOAT, "A must be cast to float32" assert _init_dtype(tr, "B_cast") == TensorProto.FLOAT, "B must be cast to float32" np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +@pytest.mark.parametrize("runtime_both", [False, True]) +def test_greater_preserves_large_integer_ordering(runtime_both): + """Keep adjacent large integers distinguishable instead of rounding both to FLOAT.""" + shape = [3] + values = np.asarray([2**24, 2**24 + 1, 2**24 + 2], np.int64) + threshold = np.full(shape, 2**24, np.int64) + inputs = [_vi("x", TensorProto.INT64, shape)] + initializers = [] + feeds = {"x": values} + if runtime_both: + inputs.append(_vi("threshold", TensorProto.INT64, shape)) + feeds["threshold"] = threshold + else: + initializers.append(numpy_helper.from_array(threshold, "threshold")) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Greater", ["x", "threshold"], ["y"])], + "integer_ordering", + inputs, + [_vi("y", TensorProto.BOOL, shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + expected = _run_ort(copy.deepcopy(model), feeds) + transformed = _apply(copy.deepcopy(model)) + np.testing.assert_array_equal(_run_ort(transformed, feeds)["y"], expected["y"]) + + +def test_greater_preserves_runtime_initializer_threshold(): + """Do not fold an overridable comparison threshold into its default value.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Greater", ["x", "threshold"], ["y"])], + "runtime_threshold", + [_vi("x", TensorProto.INT32, [3]), _vi("threshold", TensorProto.INT32, [])], + [_vi("y", TensorProto.BOOL, [3])], + [numpy_helper.from_array(np.asarray(0, np.int32), "threshold")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = {"x": np.asarray([1, 2, 3], np.int32), "threshold": np.asarray(2, np.int32)} + expected = _run_ort(copy.deepcopy(model), feeds) + np.testing.assert_array_equal(_run_ort(_apply(copy.deepcopy(model)), feeds)["y"], expected["y"]) + + +@pytest.mark.parametrize("domain", ["custom", ""]) +def test_greater_preserves_opaque_local_function(domain): + """Leave model-local comparisons and their native operand types untouched.""" + function = helper.make_function( + domain, + "Greater", + ["a", "b"], + ["c"], + [helper.make_node("Less", ["a", "b"], ["c"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Greater", ["x", "threshold"], ["y"], domain=domain)], + "opaque", + [_vi("x", TensorProto.INT32, [2])], + [_vi("y", TensorProto.BOOL, [2])], + [numpy_helper.from_array(np.asarray(1, np.int32), "threshold")], + ), + functions=[function], + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("custom", 1)], + ) + _clamp_ir(model) + original = model.SerializeToString() + feeds = {"x": np.asarray([0, 2], np.int32)} + expected = ReferenceEvaluator(model).run(None, feeds)[0] + _apply(model) + assert model.SerializeToString() == original + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_greater_allocates_distinct_names_and_legacy_controls(legacy): + """Execute repeated unnamed rewrites with occupied node names and legacy declarations.""" + inputs = [_vi("x", TensorProto.INT32, [2])] + initializers = [numpy_helper.from_array(np.asarray(1, np.int32), "threshold")] + nodes = [helper.make_node("Greater", ["x", "threshold"], [out]) for out in ("a", "b")] + nodes.append(helper.make_node("Identity", ["x"], ["spare"], name="x_cast_to_float32")) + outputs = [_vi(out, TensorProto.BOOL, [2]) for out in ("a", "b")] + outputs.append(_vi("spare", TensorProto.INT32, [2])) + feeds = {"x": np.asarray([0, 2], np.int32)} + if legacy: + inputs.extend(_vi(t.name, t.data_type, list(t.dims)) for t in initializers) + model = helper.make_model( + helper.make_graph(nodes, "names_and_legacy", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 if legacy else 9 + expected = _run_ort(copy.deepcopy(model), feeds) + transformed = _apply(copy.deepcopy(model)) + names = [node.name for node in transformed.graph.node if node.name] + assert len(names) == len(set(names)) + for name, actual in _run_ort(transformed, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) + + +@pytest.mark.parametrize("fixed", [False, True]) +def test_greater_half_precision_operands_convert_together(fixed): + """Preserve finite and nonfinite half comparisons with matching FLOAT operand types.""" + shape = [4] + inputs = [_vi("x", TensorProto.FLOAT16, shape)] + initializers = [] + threshold = np.asarray([0, 0, 0, 0], np.float16) + feeds = {"x": np.asarray([-1, 1, np.nan, np.inf], np.float16)} + if fixed: + initializers.append(numpy_helper.from_array(threshold, "threshold")) + else: + inputs.append(_vi("threshold", TensorProto.FLOAT16, shape)) + feeds["threshold"] = threshold + model = helper.make_model( + helper.make_graph( + [helper.make_node("Greater", ["x", "threshold"], ["y"])], + "half_operands", + inputs, + [_vi("y", TensorProto.BOOL, shape)], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + expected = _run_ort(copy.deepcopy(model), feeds) + np.testing.assert_array_equal(_run_ort(_apply(copy.deepcopy(model)), feeds)["y"], expected["y"]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py index 92d561b9d3d..e0eecf1d739 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py @@ -40,19 +40,17 @@ from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx import pytest -from onnx import TensorProto, helper +from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_not as _mod -pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "not_test_models" _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -64,33 +62,38 @@ def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Declare a typed tensor shape.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir(model: onnx.ModelProto) -> None: + """Select an IR version supported by the CPU runtime.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _prepare_model(model: onnx.ModelProto) -> None: + """Populate metadata using optional runtime symbolic inference.""" 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 + # Optional ONNX Runtime tooling is absent in some CPU installations. + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + except ImportError: + return + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-untyped] - + """Fully check and execute all public outputs in one CPU session.""" _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} + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + sess = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + return dict(zip([o.name for o in model.graph.output], sess.run(None, feeds))) _RNG = np.random.default_rng(2024) @@ -274,3 +277,115 @@ def test_no_not_nodes_noop(): assert len(tr.graph.node) == original_node_count, "noop: node count must not change" np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +@pytest.mark.parametrize("owner", ["capture", "identity", "override"]) +def test_not_preserves_boolean_owners_and_runtime_sources(owner): + """Keep captured Casts, boolean aliases and overridable numeric sources native.""" + initializers = [] + source_type = TensorProto.INT32 if owner == "override" else TensorProto.BOOL + inputs = [_vi("x", source_type, [2])] + nodes = [ + helper.make_node("Cast", ["x"], ["casted"], to=TensorProto.BOOL), + helper.make_node("Not", ["casted"], ["y"]), + ] + outputs = [_vi("y", TensorProto.BOOL, [2])] + feeds = { + "x": np.asarray([-2, 3], np.int32) if owner == "override" else np.asarray([False, True]) + } + if owner == "override": + initializers.append(numpy_helper.from_array(np.asarray([0, 1], np.int32), "x")) + if owner == "identity": + nodes.append(helper.make_node("Identity", ["y"], ["alias"])) + outputs = [_vi("alias", TensorProto.BOOL, [2])] + if owner == "capture": + inputs.append(_vi("choose", TensorProto.BOOL, [])) + feeds["choose"] = np.asarray(True) + branch = helper.make_graph( + [helper.make_node("Identity", ["casted"], ["captured"])], + "capture", + [], + [_vi("captured", TensorProto.BOOL, [2])], + ) + nodes.append( + helper.make_node( + "If", ["choose"], ["capture_out"], then_branch=branch, else_branch=branch + ) + ) + outputs.append(_vi("capture_out", TensorProto.BOOL, [2])) + model = helper.make_model( + helper.make_graph(nodes, "not_ownership", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + expected = _run_ort(copy.deepcopy(model), feeds) + transformed = _apply(copy.deepcopy(model)) + actual = _run_ort(transformed, feeds) + for name in expected: + np.testing.assert_array_equal(actual[name], expected[name]) + + +@pytest.mark.parametrize("domain", ["custom", ""]) +def test_not_preserves_opaque_local_function(domain): + """Do not reinterpret a local Not function as the standard boolean operator.""" + function = helper.make_function( + domain, + "Not", + ["a"], + ["b"], + [helper.make_node("Identity", ["a"], ["b"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["x"], ["casted"], to=TensorProto.BOOL), + helper.make_node("Not", ["casted"], ["y"], domain=domain), + ], + "opaque", + [_vi("x", TensorProto.BOOL, [2])], + [_vi("y", TensorProto.BOOL, [2])], + ), + functions=[function], + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("custom", 1)], + ) + _clamp_ir(model) + original = model.SerializeToString() + feeds = {"x": np.asarray([False, True])} + expected = ReferenceEvaluator(model).run(None, feeds)[0] + _apply(model) + assert model.SerializeToString() == original + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_not_allocates_distinct_names_and_legacy_controls(legacy): + """Execute repeated unnamed rewrites with occupied node names and legacy declarations.""" + inputs = [_vi("x", TensorProto.BOOL, [2])] + initializers = [] + nodes = [] + for index, output in enumerate(("a", "b")): + nodes.extend( + [ + helper.make_node("Cast", ["x"], [f"cast_{index}"], to=TensorProto.BOOL), + helper.make_node("Not", [f"cast_{index}"], [output]), + ] + ) + nodes.append(helper.make_node("Identity", ["x"], ["spare"], name="_clip")) + outputs = [_vi(out, TensorProto.BOOL, [2]) for out in ("a", "b", "spare")] + feeds = {"x": np.asarray([False, True])} + if legacy: + inputs.extend(_vi(t.name, t.data_type, list(t.dims)) for t in initializers) + model = helper.make_model( + helper.make_graph(nodes, "names_and_legacy", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 if legacy else 9 + expected = _run_ort(copy.deepcopy(model), feeds) + transformed = _apply(copy.deepcopy(model)) + names = [node.name for node in transformed.graph.node if node.name] + assert len(names) == len(set(names)) + for name, actual in _run_ort(transformed, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py index ae0788e5e57..0cd99562b40 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -42,19 +42,17 @@ 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 onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_where as _mod -pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "where_test_models" _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -66,33 +64,38 @@ def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Declare a typed tensor shape.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir(model: onnx.ModelProto) -> None: + """Select an IR version supported by the CPU runtime.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _prepare_model(model: onnx.ModelProto) -> None: + """Populate metadata using optional runtime symbolic inference.""" 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 + # Optional ONNX Runtime tooling is absent in some CPU installations. + from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference + except ImportError: + return + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + if inferred is not None: + del model.graph.value_info[:] + model.graph.value_info.extend(inferred.graph.value_info) def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-untyped] - + """Fully check and execute all public outputs in one CPU session.""" _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} + onnx.checker.check_model(model, full_check=True) + options = ort.SessionOptions() + options.intra_op_num_threads = 1 + sess = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + return dict(zip([o.name for o in model.graph.output], sess.run(None, feeds))) _RNG = np.random.default_rng(2024) @@ -373,3 +376,92 @@ def test_no_where_nodes_noop(): assert len(tr.graph.node) == original_node_count, "noop: node count must not change" np.testing.assert_array_equal(ref["Y"], got["Y"]) + + +@pytest.mark.parametrize("overridden", ["condition", "zeros"]) +@pytest.mark.parametrize("zero_index", [0, 1]) +def test_where_respects_overridable_initializer_inputs(overridden, zero_index): + """Use runtime condition defaults and decline nonconstant zero branches.""" + inputs = [_vi("data", TensorProto.FLOAT, [3]), _vi("condition", TensorProto.BOOL, [3])] + initializers = [ + numpy_helper.from_array(np.asarray([False, False, False]), "condition"), + numpy_helper.from_array(np.zeros(3, np.float32), "zeros"), + ] + if overridden == "zeros": + inputs.append(_vi("zeros", TensorProto.FLOAT, [3])) + branches = ["zeros", "data"] if zero_index == 0 else ["data", "zeros"] + model = helper.make_model( + helper.make_graph( + [helper.make_node("Where", ["condition", *branches], ["y"])], + "runtime_where", + inputs, + [_vi("y", TensorProto.FLOAT, [3])], + initializers, + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = { + "data": np.asarray([2, 3, 4], np.float32), + "condition": np.asarray([True, False, True]), + } + if overridden == "zeros": + feeds["zeros"] = np.asarray([9, 8, 7], np.float32) + expected = _run_ort(copy.deepcopy(model), feeds) + np.testing.assert_array_equal(_run_ort(_apply(copy.deepcopy(model)), feeds)["y"], expected["y"]) + + +@pytest.mark.parametrize("domain", ["custom", ""]) +def test_where_preserves_opaque_local_function(domain): + """Leave a local Where function and its fixed owners intact.""" + function = helper.make_function( + domain, + "Where", + ["condition", "a", "b"], + ["c"], + [helper.make_node("Identity", ["b"], ["c"])], + [helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Where", ["condition", "zeros", "data"], ["y"], domain=domain)], + "opaque", + [_vi("condition", TensorProto.BOOL, [2]), _vi("data", TensorProto.FLOAT, [2])], + [_vi("y", TensorProto.FLOAT, [2])], + [numpy_helper.from_array(np.zeros(2, np.float32), "zeros")], + ), + functions=[function], + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("custom", 1)], + ) + _clamp_ir(model) + original = model.SerializeToString() + feeds = {"condition": np.asarray([False, True]), "data": np.asarray([2, 3], np.float32)} + expected = ReferenceEvaluator(model).run(None, feeds)[0] + _apply(model) + assert model.SerializeToString() == original + onnx.checker.check_model(model, full_check=True) + np.testing.assert_array_equal(ReferenceEvaluator(model).run(None, feeds)[0], expected) + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_where_allocates_distinct_names_and_legacy_controls(legacy): + """Execute repeated unnamed rewrites with occupied node names and legacy declarations.""" + inputs = [_vi("x", TensorProto.FLOAT, [2]), _vi("condition", TensorProto.BOOL, [2])] + initializers = [numpy_helper.from_array(np.zeros(2, np.float32), "zeros")] + nodes = [helper.make_node("Where", ["condition", "x", "zeros"], [out]) for out in ("a", "b")] + nodes.append(helper.make_node("Identity", ["x"], ["spare"], name="_mul")) + outputs = [_vi(out, TensorProto.FLOAT, [2]) for out in ("a", "b", "spare")] + feeds = {"x": np.asarray([2, 3], np.float32), "condition": np.asarray([False, True])} + if legacy: + inputs.extend(_vi(t.name, t.data_type, list(t.dims)) for t in initializers) + model = helper.make_model( + helper.make_graph(nodes, "names_and_legacy", inputs, outputs, initializers), + opset_imports=[helper.make_opsetid("", 17)], + ) + model.ir_version = 3 if legacy else 9 + expected = _run_ort(copy.deepcopy(model), feeds) + transformed = _apply(copy.deepcopy(model)) + names = [node.name for node in transformed.graph.node if node.name] + assert len(names) == len(set(names)) + for name, actual in _run_ort(transformed, feeds).items(): + np.testing.assert_array_equal(actual, expected[name]) From bce52d59440d589d1a838d28eda170400b841b55 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 03:21:08 +0000 Subject: [PATCH 03/10] Finish comparisons after preceding FLOAT data preparation Signed-off-by: mgohil --- .../dla_transforms/dla_greater.py | 4 +++- .../dla_transforms/test_dla_greater.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py index 2572519a212..51d5fd11dc6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_greater.py @@ -40,6 +40,8 @@ } ) +_SUPPORTED_TYPES = _EXACT_TYPES | _INT_TYPES + def _apply_greater(model: onnx.ModelProto) -> onnx.ModelProto: """Convert comparisons whose ordering survives FLOAT operand conversion.""" @@ -56,7 +58,7 @@ def _apply_greater(model: onnx.ModelProto) -> onnx.ModelProto: continue types = {cache.get_dtype(name) for name in node.input} safe = types <= _EXACT_TYPES - if types <= _INT_TYPES: + if not safe and types <= _SUPPORTED_TYPES: for name in node.input: if name not in safe_thresholds: tensor = cache.get_init(name) if name not in runtime else None diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py index 776aa6f3370..7398afd2a6a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_greater.py @@ -419,3 +419,27 @@ def test_greater_half_precision_operands_convert_together(fixed): _clamp_ir(model) expected = _run_ort(copy.deepcopy(model), feeds) np.testing.assert_array_equal(_run_ort(_apply(copy.deepcopy(model)), feeds)["y"], expected["y"]) + + +def test_greater_finishes_partially_prepared_integer_operands(): + """Finish a preceding FLOAT data adapter while retaining a small fixed integer threshold.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node("Greater", ["x", "threshold"], ["y"])], + "prepared_data_adapter", + [_vi("x", TensorProto.INT64, [4])], + [_vi("y", TensorProto.BOOL, [4])], + [numpy_helper.from_array(np.asarray(1, np.int64), "threshold")], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + feeds = {"x": np.asarray([-(2**63), 0, 2, 2**63 - 1], np.int64)} + expected = _run_ort(copy.deepcopy(model), feeds) + prepared = copy.deepcopy(model) + prepared.graph.node[0].input[0] = "prepared_float" + prepared.graph.node.insert( + 0, helper.make_node("Cast", ["x"], ["prepared_float"], to=TensorProto.FLOAT) + ) + prepared.graph.value_info.append(_vi("prepared_float", TensorProto.FLOAT, [4])) + np.testing.assert_array_equal(_run_ort(_apply(prepared), feeds)["y"], expected["y"]) From 330ecccf3dc52849100f964ca189b0ac089564b0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 05:00:20 +0000 Subject: [PATCH 04/10] Clarify the Greater-produced BOOL Where fixture Signed-off-by: mgohil --- .../dla_transforms/test_dla_where.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py index 0cd99562b40..8202837aefe 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -34,7 +34,7 @@ BOOL init condition, false-branch zeros → Mul(x, cond_float), ORT parity Dynamic BOOL condition, true-branch zeros → Cast+Sub+Mul, ORT parity Dynamic BOOL condition, false-branch zeros → Cast+Mul, ORT parity -Non-BOOL float condition, false-branch zeros → Mul directly, ORT parity +Greater-produced BOOL condition, false-branch zeros: Cast+Mul, ORT parity Neither branch all-zeros → graph unchanged No Where nodes → noop, ORT parity """ @@ -263,31 +263,21 @@ def test_where_dynamic_bool_cond_y_zeros(): np.testing.assert_allclose(ref["Z"], got["Z"], rtol=1e-5, atol=1e-6) -# ── 5. Non-BOOL float condition (dynamic), false-branch zeros ───────────────── +# 5. Greater-produced BOOL condition, false-branch zeros -def test_where_float_cond_y_zeros(): - """Where with a FLOAT dynamic condition and zeros false-branch → Mul directly (no Cast). - - The float condition (0.0 / 1.0 values) acts directly as the mask. - """ +def test_where_greater_bool_cond_y_zeros(): + """Lower a Greater-produced BOOL condition with a zero false branch.""" zeros_arr = np.zeros(_SHAPE, dtype=np.float32) y_init = numpy_helper.from_array(zeros_arr, name="y_zeros_f") - # Sigmoid produces values in (0, 1); use a Cast(BOOL) then Cast(FLOAT) workaround - # is not needed here — instead, we make a simple graph where the condition is FLOAT. - # ORT does NOT accept a FLOAT condition for Where, so we model it differently: - # build a graph with Greater producing a BOOL mask, then use that as condition. greater = helper.make_node("Greater", ["X_in", "thresh"], ["cond_bool"], name="greater") - cast = helper.make_node( - "Cast", ["cond_bool"], ["cond_float"], to=TensorProto.FLOAT, name="cast" - ) where = helper.make_node("Where", ["cond_bool", "X_in", "y_zeros_f"], ["Z"], name="where") thresh_init = numpy_helper.from_array(np.array(0.0, dtype=np.float32), name="thresh") graph = helper.make_graph( - [greater, cast, where], - "where_float_cond_y0", + [greater, where], + "where_greater_bool_cond_y0", [_vi("X_in", TensorProto.FLOAT, _SHAPE)], [_vi("Z", TensorProto.FLOAT, _SHAPE)], initializer=[y_init, thresh_init], From 33a85985369a3a15e58b7b48e1c72b3399df88db Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 17:15:00 +0000 Subject: [PATCH 05/10] Preserve boolean representation for Not string consumers Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_not.py | 5 ++- .../dla_transforms/test_dla_not.py | 41 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py index d29923d8b03..f2416aa1bef 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py @@ -39,6 +39,7 @@ add_unique_initializers, add_value_info, batch_replace_nodes, + get_node_attr_i, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_graph_cleanup import _subgraph_captures @@ -95,7 +96,9 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: binary[source] = _binary_source(cache, source, runtime) binary_source = binary[source] if not binary_source or any( - user.op_type != "Cast" or id(user) not in casts + user.op_type != "Cast" + or id(user) not in casts + or get_node_attr_i(user, "to", 0) == TensorProto.STRING for user, _port in cache.get_consumers(node.output[0]) ): continue diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py index e0eecf1d739..f1a0db95546 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py @@ -279,6 +279,47 @@ def test_no_not_nodes_noop(): np.testing.assert_array_equal(ref["Y"], got["Y"]) +@pytest.mark.parametrize( + "target", [TensorProto.STRING, TensorProto.FLOAT, TensorProto.INT32, TensorProto.BOOL] +) +@pytest.mark.parametrize("runtime", ["reference", "ort"]) +def test_not_preserves_downstream_cast_representation(target, runtime): + """Keep string representations stable and continue lowering eligible numeric Cast consumers.""" + model = helper.make_model( + helper.make_graph( + [ + helper.make_node("Cast", ["x"], ["casted"], to=TensorProto.BOOL), + helper.make_node("Not", ["casted"], ["inverted"]), + helper.make_node("Cast", ["inverted"], ["y"], to=target), + ], + "cast_representation", + [_vi("x", TensorProto.BOOL, [2])], + [_vi("y", target, [2])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ) + _clamp_ir(model) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.asarray([False, True])} + expected = ( + ReferenceEvaluator(model).run(None, feeds)[0] + if runtime == "reference" + else _run_ort(model, feeds)["y"] + ) + before = model.SerializeToString() + changed = _apply(copy.deepcopy(model)) + actual = ( + ReferenceEvaluator(changed).run(None, feeds)[0] + if runtime == "reference" + else _run_ort(changed, feeds)["y"] + ) + np.testing.assert_array_equal(actual, expected) + if target == TensorProto.STRING: + assert changed.SerializeToString() == before + else: + assert not any(node.op_type == "Not" for node in changed.graph.node) + + @pytest.mark.parametrize("owner", ["capture", "identity", "override"]) def test_not_preserves_boolean_owners_and_runtime_sources(owner): """Keep captured Casts, boolean aliases and overridable numeric sources native.""" From d29988350999430ea7d85872614d0bc9e2ce9bdc Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 20:53:57 +0000 Subject: [PATCH 06/10] Preserve public boolean Not output contracts Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_not.py | 3 +- .../dla_transforms/test_dla_not.py | 55 +++++++++++++++---- 2 files changed, 46 insertions(+), 12 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py index f2416aa1bef..9edc711cbd6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py @@ -88,7 +88,7 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: if ( len(cache.get_consumers(producer_node.output[0])) != 1 or producer_node.output[0] in cache.graph_output_names | captured - or node.output[0] in captured + or node.output[0] in cache.graph_output_names | captured ): continue source = producer_node.input[0] @@ -138,7 +138,6 @@ def _apply_not(model: onnx.ModelProto) -> onnx.ModelProto: if clip_shape is not None: add_value_info(graph, clip_out, TensorProto.FLOAT, clip_shape) - # Update Not output dtype to FLOAT (graph output or intermediate value_info) not_out_vi = cache.vi_map.get(node.output[0]) if not_out_vi is not None: not_out_vi.type.tensor_type.elem_type = TensorProto.FLOAT diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py index f1a0db95546..d2ed3ad0f38 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py @@ -110,12 +110,12 @@ def test_not_cast_to_bool_transformed(): """ shape = [2, 3] cast_node = helper.make_node("Cast", ["X"], ["X_bool"], to=TensorProto.BOOL, name="cast") - not_node = helper.make_node("Not", ["X_bool"], ["Y"], name="not") + not_node = helper.make_node("Not", ["X_bool"], ["Y_bool"], name="not") graph = helper.make_graph( - [cast_node, not_node], + [cast_node, not_node, helper.make_node("Cast", ["Y_bool"], ["Y"], to=TensorProto.FLOAT)], "not_cast_bool", [_vi("X", TensorProto.BOOL, shape)], - [_vi("Y", TensorProto.BOOL, shape)], + [_vi("Y", TensorProto.FLOAT, shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) @@ -136,7 +136,7 @@ def test_not_cast_to_bool_transformed(): # Cast node must now produce FLOAT, Not node must be gone. cast_nodes = [n for n in tr.graph.node if n.op_type == "Cast"] - assert len(cast_nodes) == 1 + assert len(cast_nodes) == 2 for attr in cast_nodes[0].attribute: if attr.name == "to": assert attr.i == TensorProto.FLOAT, "Cast must target FLOAT after transform" @@ -216,14 +216,16 @@ def test_not_multiple_not_nodes(): """Two independent Cast(BOOL)→Not chains are both transformed with unique initializer names.""" shape = [2] cast_a = helper.make_node("Cast", ["A"], ["A_bool"], to=TensorProto.BOOL, name="cast_a") - not_a = helper.make_node("Not", ["A_bool"], ["Y_A"], name="not_a") + not_a = helper.make_node("Not", ["A_bool"], ["Y_A_bool"], name="not_a") + output_a = helper.make_node("Cast", ["Y_A_bool"], ["Y_A"], to=TensorProto.FLOAT) cast_b = helper.make_node("Cast", ["B"], ["B_bool"], to=TensorProto.BOOL, name="cast_b") - not_b = helper.make_node("Not", ["B_bool"], ["Y_B"], name="not_b") + not_b = helper.make_node("Not", ["B_bool"], ["Y_B_bool"], name="not_b") + output_b = helper.make_node("Cast", ["Y_B_bool"], ["Y_B"], to=TensorProto.FLOAT) graph = helper.make_graph( - [cast_a, not_a, cast_b, not_b], + [cast_a, not_a, output_a, cast_b, not_b, output_b], "not_multi", [_vi("A", TensorProto.BOOL, shape), _vi("B", TensorProto.BOOL, shape)], - [_vi("Y_A", TensorProto.BOOL, shape), _vi("Y_B", TensorProto.BOOL, shape)], + [_vi("Y_A", TensorProto.FLOAT, shape), _vi("Y_B", TensorProto.FLOAT, shape)], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir(model) @@ -411,11 +413,16 @@ def test_not_allocates_distinct_names_and_legacy_controls(legacy): nodes.extend( [ helper.make_node("Cast", ["x"], [f"cast_{index}"], to=TensorProto.BOOL), - helper.make_node("Not", [f"cast_{index}"], [output]), + helper.make_node("Not", [f"cast_{index}"], [output + "_bool"]), + helper.make_node("Cast", [output + "_bool"], [output], to=TensorProto.FLOAT), ] ) nodes.append(helper.make_node("Identity", ["x"], ["spare"], name="_clip")) - outputs = [_vi(out, TensorProto.BOOL, [2]) for out in ("a", "b", "spare")] + outputs = [ + _vi("a", TensorProto.FLOAT, [2]), + _vi("b", TensorProto.FLOAT, [2]), + _vi("spare", TensorProto.BOOL, [2]), + ] feeds = {"x": np.asarray([False, True])} if legacy: inputs.extend(_vi(t.name, t.data_type, list(t.dims)) for t in initializers) @@ -430,3 +437,31 @@ def test_not_allocates_distinct_names_and_legacy_controls(legacy): assert len(names) == len(set(names)) for name, actual in _run_ort(transformed, feeds).items(): np.testing.assert_array_equal(actual, expected[name]) + + +@pytest.mark.parametrize("reader", [None, TensorProto.FLOAT, TensorProto.INT32]) +def test_public_not_result_retains_boolean_type_and_values(reader): + """Keep declared BOOL outputs native even when numeric Cast readers are eligible.""" + nodes = [ + helper.make_node("Cast", ["x"], ["boolean"], to=TensorProto.BOOL), + helper.make_node("Not", ["boolean"], ["public"]), + ] + outputs = [_vi("public", TensorProto.BOOL, [2])] + if reader is not None: + nodes.append(helper.make_node("Cast", ["public"], ["numeric"], to=reader)) + outputs.append(_vi("numeric", reader, [2])) + model = helper.make_model( + helper.make_graph(nodes, "public_not", [_vi("x", TensorProto.BOOL, [2])], outputs), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + model = onnx.shape_inference.infer_shapes(model) + feeds = {"x": np.asarray([False, True])} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + changed = _apply(copy.deepcopy(model)) + assert changed.SerializeToString() == before + actual = _run_ort(changed, feeds) + assert actual["public"].dtype == np.bool_ + for name in expected: + np.testing.assert_array_equal(actual[name], expected[name]) From 987a478566df90808bf3c30edc8b11fa48bf55e0 Mon Sep 17 00:00:00 2001 From: mgohil Date: Wed, 28 Oct 2026 03:16:44 +0000 Subject: [PATCH 07/10] Document signed-zero limits of arithmetic Where lowering Signed-off-by: mgohil --- modelopt/onnx/graph_surgery/dla_transforms/dla_where.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py index 74021fca565..651513c448e 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -17,6 +17,7 @@ Fixed BOOL conditions get private FLOAT masks; runtime conditions use Cast nodes. Overridable zero branches remain native, and initializer owners are preserved. +Masked multiplication does not preserve the sign of selected zero values. """ from __future__ import annotations From 2debf2b9ca4a8718940ae986395a8c6d8f72b52d Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:03:36 +0000 Subject: [PATCH 08/10] Use materialized promoted shapes for Where eligibility Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_where.py | 29 +++++++++++-- .../dla_transforms/test_dla_where.py | 41 +++++++++++++++++++ 2 files changed, 67 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py index 651513c448e..28aa26e2d82 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -29,9 +29,11 @@ from ...logging_config import logger from ..utils.graph_utils import ( GraphCache, + _attribute, add_unique_initializers, add_value_info, batch_replace_nodes, + pad4d, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs from .dla_indexing import _operator_nodes @@ -68,6 +70,27 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: runtime = _runtime_input_names(model) zeros = {} + transposes = {id(node) for node in _operator_nodes(model, {"Transpose"})} + + def current_shape(name): + """Read materialized layouts while rank stages defer intermediate value-info updates.""" + tensor = cache.get_init(name) + if tensor is not None and name not in runtime: + return list(tensor.dims) + producer = cache.get_producer(name) + if producer is not None and id(producer) in transposes: + perm = _attribute(producer, "perm") + source = cache.get_shape(producer.input[0]) + if ( + perm is not None + and sorted(perm.ints) == [0, 1, 2, 3] + and source is not None + and len(source) <= 4 + ): + source = pad4d(source) + return [source[index] for index in perm.ints] + return cache.get_shape(name) + for node in _operator_nodes(model, {"Where"}): condition = node.input[0] x_name = node.input[1] # true branch @@ -80,7 +103,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: non_zero_name = y_name if zero_idx == 0 else x_name if cache.get_dtype(non_zero_name) != TensorProto.FLOAT: continue - shapes = [cache.get_shape(name) for name in [condition, x_name, y_name]] + shapes = [current_shape(name) for name in [condition, x_name, y_name]] if any(shape is None for shape in shapes): continue if np.broadcast_shapes(*shapes) != np.broadcast_shapes( @@ -123,7 +146,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: name=cache.unique_name(f"{node.name}_cast_cond"), ) nodes_to_add.append(cast_node) - cond_shape = cache.get_shape(condition) + cond_shape = current_shape(condition) if cond_shape is not None: add_value_info(graph, float_name, TensorProto.FLOAT, cond_shape) mask_name = float_name @@ -147,7 +170,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: name=cache.unique_name(f"{node.name}_inv"), ) nodes_to_add.append(sub_node) - cond_shape = cache.get_shape(condition) + cond_shape = current_shape(condition) if cond_shape is not None: add_value_info(graph, inv_name, TensorProto.FLOAT, cond_shape) mul_inputs = [non_zero_name, inv_name] diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py index 8202837aefe..9cfffd82ea5 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -50,6 +50,7 @@ from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms import dla_where as _mod +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") @@ -455,3 +456,43 @@ def test_where_allocates_distinct_names_and_legacy_controls(legacy): assert len(names) == len(set(names)) for name, actual in _run_ort(transformed, feeds).items(): np.testing.assert_array_equal(actual, expected[name]) + + +@pytest.mark.parametrize("source_kind", ["initializer", "Transpose"]) +@pytest.mark.parametrize("zero_branch", [0, 1]) +def test_registered_where_uses_materialized_promoted_shapes(source_kind, zero_branch): + """Lower selections after rank preparation even when intermediate metadata is deferred.""" + nodes, inputs, feeds = [], [], {} + values = np.arange(1, 7, dtype=np.float32).reshape(3, 2) + initializers = [ + numpy_helper.from_array(np.array(True), "condition"), + numpy_helper.from_array(np.zeros((3, 2), np.float32), "zero"), + ] + if source_kind == "initializer": + initializers.append(numpy_helper.from_array(values, "data")) + else: + inputs.append(_vi("x", TensorProto.FLOAT, [2, 3])) + feeds["x"] = values.T.copy() + nodes.append(helper.make_node("Transpose", ["x"], ["data"], perm=[1, 0])) + branches = ["data", "data"] + branches[zero_branch] = "zero" + nodes.append(helper.make_node("Where", ["condition", *branches], ["y"])) + model = helper.make_model( + helper.make_graph( + nodes, + "promoted_where", + inputs, + [_vi("y", TensorProto.FLOAT, [3, 2])], + initializers, + value_info=[_vi("data", TensorProto.FLOAT, [3, 2])], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + expected = _run_ort(model, feeds) + result = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + assert not any(node.op_type == "Where" for node in result.graph.node) + onnx.checker.check_model(result, full_check=True) + actual = _run_ort(result, feeds) + for name in expected: + np.testing.assert_array_equal(actual[name], expected[name]) From 91f942ef24c294a255ea8e5fc88de561e329851e Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:10:11 +0000 Subject: [PATCH 09/10] Track proven Where layouts through deferred unary paths Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_where.py | 45 ++++++++++--------- .../dla_transforms/test_dla_where.py | 16 ++++--- 2 files changed, 33 insertions(+), 28 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py index 28aa26e2d82..8c10b85fa16 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -29,14 +29,13 @@ from ...logging_config import logger from ..utils.graph_utils import ( GraphCache, - _attribute, add_unique_initializers, add_value_info, batch_replace_nodes, pad4d, ) from .dla_constants_to_initializers import _add_legacy_initializer_inputs -from .dla_indexing import _operator_nodes +from .dla_indexing import _operator_nodes, _prepared_index_data from .dla_rank_conversion import _runtime_input_names __all__ = [] @@ -70,28 +69,30 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: runtime = _runtime_input_names(model) zeros = {} - transposes = {id(node) for node in _operator_nodes(model, {"Transpose"})} + selections = _operator_nodes(model, {"Where"}) + if selections: + for tensor in graph.initializer: + if tensor.name in cache.vi_map and tensor.name not in runtime: + cache.vi_map[tensor.name] = helper.make_tensor_value_info( + tensor.name, tensor.data_type, list(tensor.dims) + ) + prepared = ( + _prepared_index_data( + model, + cache, + (), + [name for node in selections for name in node.input], + ) + if selections + else {} + ) def current_shape(name): - """Read materialized layouts while rank stages defer intermediate value-info updates.""" - tensor = cache.get_init(name) - if tensor is not None and name not in runtime: - return list(tensor.dims) - producer = cache.get_producer(name) - if producer is not None and id(producer) in transposes: - perm = _attribute(producer, "perm") - source = cache.get_shape(producer.input[0]) - if ( - perm is not None - and sorted(perm.ints) == [0, 1, 2, 3] - and source is not None - and len(source) <= 4 - ): - source = pad4d(source) - return [source[index] for index in perm.ints] - return cache.get_shape(name) - - for node in _operator_nodes(model, {"Where"}): + """Use proven leading layouts without changing metadata needed by other passes.""" + shape = cache.get_shape(name) + return pad4d(shape) if shape is not None and prepared[name] else shape + + for node in selections: condition = node.input[0] x_name = node.input[1] # true branch y_name = node.input[2] # false branch diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py index 9cfffd82ea5..1660ee5cf7c 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -460,13 +460,15 @@ def test_where_allocates_distinct_names_and_legacy_controls(legacy): @pytest.mark.parametrize("source_kind", ["initializer", "Transpose"]) @pytest.mark.parametrize("zero_branch", [0, 1]) -def test_registered_where_uses_materialized_promoted_shapes(source_kind, zero_branch): +@pytest.mark.parametrize("unary", [False, True]) +def test_registered_where_uses_materialized_promoted_shapes(source_kind, zero_branch, unary): """Lower selections after rank preparation even when intermediate metadata is deferred.""" nodes, inputs, feeds = [], [], {} - values = np.arange(1, 7, dtype=np.float32).reshape(3, 2) + shape = [1, 3, 2] if source_kind == "initializer" and unary else [3, 2] + values = np.arange(1, 7, dtype=np.float32).reshape(shape) initializers = [ numpy_helper.from_array(np.array(True), "condition"), - numpy_helper.from_array(np.zeros((3, 2), np.float32), "zero"), + numpy_helper.from_array(np.zeros(shape, np.float32), "zero"), ] if source_kind == "initializer": initializers.append(numpy_helper.from_array(values, "data")) @@ -474,7 +476,9 @@ def test_registered_where_uses_materialized_promoted_shapes(source_kind, zero_br inputs.append(_vi("x", TensorProto.FLOAT, [2, 3])) feeds["x"] = values.T.copy() nodes.append(helper.make_node("Transpose", ["x"], ["data"], perm=[1, 0])) - branches = ["data", "data"] + if unary: + nodes.append(helper.make_node("Relu", ["data"], ["activated"])) + branches = ["activated" if unary else "data"] * 2 branches[zero_branch] = "zero" nodes.append(helper.make_node("Where", ["condition", *branches], ["y"])) model = helper.make_model( @@ -482,9 +486,9 @@ def test_registered_where_uses_materialized_promoted_shapes(source_kind, zero_br nodes, "promoted_where", inputs, - [_vi("y", TensorProto.FLOAT, [3, 2])], + [_vi("y", TensorProto.FLOAT, shape)], initializers, - value_info=[_vi("data", TensorProto.FLOAT, [3, 2])], + value_info=[_vi("data", TensorProto.FLOAT, shape)], ), opset_imports=[helper.make_opsetid("", 17)], ir_version=9, From ec5f81eb2bde6a1ccd8136a8383dde7f1b848f2c Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 15:40:45 +0000 Subject: [PATCH 10/10] Reuse fixed Boolean Where masks for matching polarity Signed-off-by: mgohil --- .../graph_surgery/dla_transforms/dla_where.py | 33 +++++++----- .../dla_transforms/test_dla_where.py | 54 +++++++++++++++++++ 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py index 8c10b85fa16..5fdd3a91907 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -68,6 +68,7 @@ def _apply_where(model: onnx.ModelProto) -> onnx.ModelProto: cnt = 0 runtime = _runtime_input_names(model) zeros = {} + static_masks = {} selections = _operator_nodes(model, {"Where"}) if selections: @@ -116,20 +117,24 @@ def current_shape(name): if cache.is_init(condition) and condition not in runtime: if condition_type == TensorProto.BOOL: - cond_raw = cache.get_init_array(condition) - if cond_raw is None: - continue - cond_arr = cond_raw.astype(np.float32) - if zero_idx == 0: - # Where(cond, 0, y) → y * (1 - cond_float) - np.subtract(1.0, cond_arr, out=cond_arr) - float_name = cache.unique_name(f"{condition}_float32") - graph.initializer.add( - name=float_name, - data_type=TensorProto.FLOAT, - dims=cond_arr.shape, - raw_data=cond_arr.astype("