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..51d5fd11dc6 --- /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. + +"""Prepare FLOAT operands only where conversion preserves comparison ordering.""" + +from __future__ import annotations + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +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, + } +) + +_SUPPORTED_TYPES = _EXACT_TYPES | _INT_TYPES + + +def _apply_greater(model: onnx.ModelProto) -> onnx.ModelProto: + """Convert comparisons whose ordering survives FLOAT operand conversion.""" + graph = model.graph + cache = GraphCache(graph) + 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 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 + 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 name not in converted: + output = cache.unique_name( + name + "_cast" + if cache.is_init(name) and name not in runtime + else name + "_to_float32" + ) + 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 new file mode 100644 index 00000000000..9edc711cbd6 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_not.py @@ -0,0 +1,160 @@ +# 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, + get_node_attr_i, +) +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 _operator_nodes(model, {"Not"}): + producer_node = cache.get_producer(node.input[0]) + 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 | captured + or node.output[0] in cache.graph_output_names | captured + ): + continue + source = producer_node.input[0] + 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 != "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 + + # ── 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=cache.unique_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) + + 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=cache.unique_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) + + _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 new file mode 100644 index 00000000000..5fdd3a91907 --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_where.py @@ -0,0 +1,202 @@ +# 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. + +"""Lower zero-branch selections to arithmetic under the existing finite-FP32 assumption. + +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 + +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, + pad4d, +) +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_indexing import _operator_nodes, _prepared_index_data +from .dla_rank_conversion import _runtime_input_names + +__all__ = [] + + +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) 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 + + +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 + runtime = _runtime_input_names(model) + zeros = {} + static_masks = {} + + 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): + """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 + + zero_idx = _zero_initializer_index(cache, [x_name, y_name], runtime, zeros) + 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 = [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( + 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) and condition not in runtime: + if condition_type == TensorProto.BOOL: + mask_key = (condition, zero_idx == 0) + float_name = static_masks.get(mask_key) + if float_name is None: + 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(" 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: + # 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: + """Fully check and execute all public outputs in one CPU session.""" + _clamp_ir(model) + 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 + 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"]) + + +@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"]) + + +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"]) 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..d2ed3ad0f38 --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_not.py @@ -0,0 +1,467 @@ +# 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 + +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_not as _mod + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_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: + """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: + # 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: + """Fully check and execute all public outputs in one CPU session.""" + _clamp_ir(model) + 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) + + +# ── 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_bool"], name="not") + graph = helper.make_graph( + [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.FLOAT, 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) == 2 + 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_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_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, output_a, cast_b, not_b, output_b], + "not_multi", + [_vi("A", TensorProto.BOOL, shape), _vi("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) + + 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"]) + + +@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.""" + 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 + "_bool"]), + helper.make_node("Cast", [output + "_bool"], [output], to=TensorProto.FLOAT), + ] + ) + nodes.append(helper.make_node("Identity", ["x"], ["spare"], name="_clip")) + 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) + 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("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]) 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..fbbb29ac90a --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_where.py @@ -0,0 +1,556 @@ +# 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 +Greater-produced BOOL condition, false-branch zeros: Cast+Mul, ORT parity +Neither branch all-zeros → graph unchanged +No Where nodes → noop, ORT parity +""" + +from __future__ import annotations + +import copy + +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 +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible + +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_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: + """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: + # 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: + """Fully check and execute all public outputs in one CPU session.""" + _clamp_ir(model) + 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) +_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. Greater-produced BOOL condition, false-branch zeros + + +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") + + greater = helper.make_node("Greater", ["X_in", "thresh"], ["cond_bool"], name="greater") + 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, where], + "where_greater_bool_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"]) + + +@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]) + + +@pytest.mark.parametrize("source_kind", ["initializer", "Transpose"]) +@pytest.mark.parametrize("zero_branch", [0, 1]) +@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 = [], [], {} + 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(shape, 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])) + 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( + helper.make_graph( + nodes, + "promoted_where", + inputs, + [_vi("y", TensorProto.FLOAT, shape)], + initializers, + value_info=[_vi("data", TensorProto.FLOAT, shape)], + ), + 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]) + + +def test_shared_fixed_bool_masks_are_reused_by_polarity(): + """Share each direct or inverted FLOAT mask without changing BOOL ownership or values.""" + shape = [1, 1, 2, 3] + condition = np.array([True, False, True, False, True, False]).reshape(shape) + branches = [0, 1, 0, 1] + model = helper.make_model( + helper.make_graph( + [ + helper.make_node( + "Where", + ["condition", "zero" if side == 0 else "x", "x" if side == 0 else "zero"], + [f"y{i}"], + ) + for i, side in enumerate(branches) + ], + "shared_bool_masks", + [_vi("x", TensorProto.FLOAT, shape)], + [_vi(f"y{i}", TensorProto.FLOAT, shape) for i in range(len(branches))], + initializer=[ + numpy_helper.from_array(condition, "condition"), + numpy_helper.from_array(np.float32(0), "zero"), + ], + ), + opset_imports=[helper.make_opsetid("", 17)], + ir_version=9, + ) + feeds = {"x": np.arange(1, 7, dtype=np.float32).reshape(shape)} + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds) + runtime_expected = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + original = model.graph.initializer[0] + payload = original.SerializeToString() + _apply(model) + onnx.checker.check_model(model, full_check=True) + masks = [node.input[1] for node in model.graph.node if node.op_type == "Mul"] + assert len(masks) == 4 + assert masks[0] == masks[2] and masks[1] == masks[3] and masks[0] != masks[1] + assert ( + sum(t.data_type == TensorProto.FLOAT and t.name != "zero" for t in model.graph.initializer) + == 2 + ) + assert model.graph.initializer[0] is original + assert original.SerializeToString() == payload + for actual, reference in zip(ReferenceEvaluator(model).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + runtime_actual = ort.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ).run(None, feeds) + for actual, reference in zip(runtime_actual, runtime_expected): + np.testing.assert_array_equal(actual, reference)