From b52a6ae95780ce2405a2eeb102303ca7afbbb138 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 10 Nov 2026 15:32:27 -0800 Subject: [PATCH] Preserve overridable initializer defaults in rewrite rules Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- onnxscript/rewriter/_ir_utils.py | 10 +- onnxscript/rewriter/_matcher.py | 3 +- .../rewriter/_overridable_initializer_test.py | 393 ++++++++++++++++++ onnxscript/rewriter/_rewrite_rule.py | 2 +- .../rewriter/ort_fusions/cos_sin_cache.py | 8 +- .../ort_fusions/fused_matmul_rule_sets.py | 7 +- .../instance_to_group_normalization.py | 9 +- .../rewriter/rules/common/_basic_rules.py | 30 +- .../rules/common/_broadcast_to_matmul.py | 7 +- .../rewriter/rules/common/_collapse_slices.py | 8 +- .../rules/common/_fuse_pad_into_conv.py | 19 +- .../common/_materialize_reshape_shape.py | 3 + .../rewriter/rules/common/_min_max_to_clip.py | 3 +- .../rules/common/_min_max_to_clip_test.py | 12 +- .../rules/common/_redundant_scatter_nd.py | 4 +- .../rules/common/_remove_optional_bias.py | 3 +- 16 files changed, 462 insertions(+), 59 deletions(-) create mode 100644 onnxscript/rewriter/_overridable_initializer_test.py diff --git a/onnxscript/rewriter/_ir_utils.py b/onnxscript/rewriter/_ir_utils.py index 85c796a701..a095ebf4c2 100644 --- a/onnxscript/rewriter/_ir_utils.py +++ b/onnxscript/rewriter/_ir_utils.py @@ -52,6 +52,9 @@ def visit(node: ir.Node, depth): def get_const_value(value: ir.Value) -> ir.TensorProtocol | None: + """Return a runtime constant, excluding overridable graph-input defaults.""" + if value.is_graph_input(): + return None node = value.producer() if node is not None: optimizer.basic_constant_propagation([node]) @@ -61,11 +64,8 @@ def get_const_value(value: ir.Value) -> ir.TensorProtocol | None: def get_numpy_value(val: ir.Value | None) -> np.ndarray | None: """Convenience wrapper to get (optional) numpy value from an optional IR Value. - This is intended for use in optimizations/rewriting. Note that this does not - yet handle the distinction between inputs with default values (values that are - both graph inputs and graph initializers), which should not be treated as a - constant, and true constant values. The caller should make the distinction, as - a value does not contain enough information to determine this. (TODO) + This is intended for use in optimizations/rewriting. Graph inputs with + initializer defaults are not constants, since callers can override them. """ if val is None: return None diff --git a/onnxscript/rewriter/_matcher.py b/onnxscript/rewriter/_matcher.py index 36f1ba14f4..35828fa180 100644 --- a/onnxscript/rewriter/_matcher.py +++ b/onnxscript/rewriter/_matcher.py @@ -15,6 +15,7 @@ import onnx_ir as ir import onnxscript.rewriter._basics as _basics +import onnxscript.rewriter._ir_utils as _ir_utils import onnxscript.rewriter._pattern_ir as _pattern_ir @@ -81,7 +82,7 @@ def _match_constant(self, pattern_constant: _pattern_ir.Constant, value: ir.Valu if subgraph replacement happens. But subsequent DCE will remove the constant node if it is not used elsewhere. """ - constant_value = value.const_value + constant_value = _ir_utils.get_const_value(value) if constant_value is None: return self.fail( f"Value {value.name} is not a constant, expecting {pattern_constant.value}.", diff --git a/onnxscript/rewriter/_overridable_initializer_test.py b/onnxscript/rewriter/_overridable_initializer_test.py new file mode 100644 index 0000000000..b3b081b1af --- /dev/null +++ b/onnxscript/rewriter/_overridable_initializer_test.py @@ -0,0 +1,393 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +from __future__ import annotations + +import unittest + +import numpy as np +import onnx +import onnx_ir as ir +import onnxruntime as ort +import parameterized + +from onnxscript import optimizer, rewriter +from onnxscript.rewriter import _ir_utils +from onnxscript.rewriter.ort_fusions import ( + cos_sin_cache, + fused_matmul_rule_sets, + instance_to_group_normalization, +) +from onnxscript.rewriter.rules.common import ( + _basic_rules, + _broadcast_to_matmul, + _collapse_slices, + _fuse_pad_into_conv, + _materialize_reshape_shape, + _min_max_to_clip, + _redundant_scatter_nd, + _remove_optional_bias, +) + + +def _session(model: onnx.ModelProto) -> ort.InferenceSession: + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + options.log_severity_level = 3 + return ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) + + +class OverridableInitializerTest(unittest.TestCase): + @parameterized.parameterized.expand( + [ + ("Add", 0.0), + ("Sub", 0.0), + ("Mul", 1.0), + ("Div", 1.0), + ] + ) + def test_public_optimizer_preserves_overrides(self, op_type, default): + for representation in ("overridable", "initializer", "Constant"): + for reverse in (False, True) if op_type in ("Add", "Mul") else (False,): + with self.subTest(representation=representation, reverse=reverse): + tensor = onnx.numpy_helper.from_array(np.array(default, np.float32), "c") + inputs = [onnx.helper.make_tensor_value_info("x", 1, [2])] + nodes = [] + if representation == "overridable": + inputs.append(onnx.helper.make_tensor_value_info("c", 1, [])) + if representation == "Constant": + nodes.append( + onnx.helper.make_node("Constant", [], ["c"], value=tensor) + ) + nodes.append( + onnx.helper.make_node( + op_type, ["c", "x"] if reverse else ["x", "c"], ["y"] + ) + ) + model = onnx.helper.make_model( + onnx.helper.make_graph( + nodes, + "arithmetic", + inputs, + [onnx.helper.make_tensor_value_info("y", 1, [2])], + [] if representation == "Constant" else [tensor], + ), + ir_version=9, + opset_imports=[onnx.helper.make_opsetid("", 18)], + ) + original = model.SerializeToString() + optimized = optimizer.optimize(model) + self.assertEqual(model.SerializeToString(), original) + onnx.checker.check_model(optimized) + baseline, actual = _session(model), _session(optimized) + feeds = {"x": np.array([2.0, 4.0], np.float32)} + np.testing.assert_array_equal( + actual.run(None, feeds)[0], baseline.run(None, feeds)[0] + ) + if representation == "overridable": + self.assertEqual([v.name for v in optimized.graph.input], ["x", "c"]) + self.assertEqual(optimized.graph.initializer[0], tensor) + self.assertIn(op_type, [n.op_type for n in optimized.graph.node]) + feeds["c"] = np.array(3.0, np.float32) + np.testing.assert_array_equal( + actual.run(None, feeds)[0], baseline.run(None, feeds)[0] + ) + self.assertFalse( + np.array_equal( + actual.run(None, feeds)[0], + actual.run(None, {"x": feeds["x"]})[0], + ) + ) + else: + self.assertNotIn(op_type, [n.op_type for n in optimized.graph.node]) + + @parameterized.parameterized.expand( + [ + ( + "expand", + "float[1] x", + "int64[1] c = {1}", + "y = Expand(x, c)", + _basic_rules.ExpandIdentity.rule(), + {"x": np.array([2.0], np.float32)}, + [3], + ), + ( + "reshape", + "float[2,3] x", + "int64[2] c = {2,3}, int64[2] s = {3,2}", + "t = Reshape(x, s)\n y = Reshape(t, c)", + _basic_rules.ReshapeReshape.rule(), + {"x": np.arange(6, dtype=np.float32).reshape(2, 3)}, + [3, 2], + ), + ( + "slice", + "float[4] x", + "int64[1] c = {0}, int64[1] e = {4}, int64[1] a = {0}, int64[1] s = {1}", + "y = Slice(x, c, e, a, s)", + _collapse_slices.collapse_slice_rule, + {"x": np.arange(4, dtype=np.float32)}, + [1], + ), + ( + "min", + "float[4] x", + "float c = {1}, float b = {6}", + "t = Min(x, c)\n y = Min(t, b)", + _min_max_to_clip.min_min_rule, + {"x": np.array([-3.0, 0.0, 4.0, 8.0], np.float32)}, + 5.0, + ), + ( + "max", + "float[4] x", + "float c = {1}, float b = {6}", + "t = Max(x, c)\n y = Max(t, b)", + _min_max_to_clip.max_max_rule, + {"x": np.array([-3.0, 0.0, 4.0, 8.0], np.float32)}, + 10.0, + ), + ( + "clip", + "float[4] x", + "float c = {1}, float b = {6}", + "t = Max(x, c)\n y = Min(t, b)", + _min_max_to_clip.max_min_rule, + {"x": np.array([-3.0, 0.0, 4.0, 8.0], np.float32)}, + 5.0, + ), + ( + "bias", + "float[2,2] x", + "float[2] c = {0,0}, float[2,2] w = {1,0,0,1}", + "y = Gemm(x, w, c)", + _remove_optional_bias.remove_optional_bias_from_gemm_rule, + {"x": np.arange(4, dtype=np.float32).reshape(2, 2)}, + [2.0, 3.0], + ), + ( + "scatter", + "float[2] x, float[2] updates", + "int64[2,1] c = {0,1}", + "y = ScatterND(x, c, updates)", + _redundant_scatter_nd.no_op_static_scatter_nd_rule, + { + "x": np.array([1.0, 2.0], np.float32), + "updates": np.array([4.0, 5.0], np.float32), + }, + [[1], [0]], + ), + ( + "pad", + "float[1,1,2] x", + "float c = {0}, int64[6] pads = {0,0,1,0,0,1}, float[1,1,1] w = {1}", + "t = Pad(x, pads, c)\n y = Conv(t, w)", + _fuse_pad_into_conv.FuseConvPad.rule(), + {"x": np.array([[[1.0, 2.0]]], np.float32)}, + 3.0, + ), + ( + "matmul_shape", + "float[1,2,2] x, float[1,2,2] z", + "int64[3] c = {1,2,2}, int64[3] s = {1,2,2}", + "a = Reshape(x, s)\n b = Reshape(z, s)\n t = MatMul(a, b)\n y = Reshape(t, c)", + _broadcast_to_matmul.two_reshapes_matmul_reshape_rule, + { + "x": np.arange(4, dtype=np.float32).reshape(1, 2, 2), + "z": np.arange(4, dtype=np.float32).reshape(1, 2, 2), + }, + [2, 1, 2], + ), + ( + "ort_divisor", + "float[2,2] x, float[2,2] z", + "float c = {2}", + "t = MatMul(x, z)\n y = Div(t, c)", + fused_matmul_rule_sets.FusedMatMulDiv1.rule(), + { + "x": np.arange(4, dtype=np.float32).reshape(2, 2), + "z": np.arange(4, dtype=np.float32).reshape(2, 2), + }, + 4.0, + ), + ] + ) + def test_rule_preserves_default_and_override( + self, _, inputs, initializers, body, rule, feeds, override + ): + model = ir.from_onnx_text( + f""" + + graph ({inputs}) => (float[] y) <{initializers}> {{ {body} }} + """ + ) + # Leave result dimensions unconstrained so runtime overrides may change them. + model.graph.outputs[0].shape = None + if model.graph[-1].op_type == "ScatterND": + model.graph.outputs[0].shape = ir.Shape([2]) + baseline = _session(ir.to_proto(model)) + default = model.graph.initializers["c"] + default_tensor = default.const_value + assert default_tensor is not None + for overridable in (True, False): + with self.subTest(overridable=overridable): + candidate = model.clone() + if overridable: + candidate.graph.inputs.append(candidate.graph.initializers["c"]) + count = rewriter.RewriteRuleSet([rule]).apply_to_model(candidate) + self.assertEqual(count, 0 if overridable else 1) + if overridable: + candidate = rewriter.rewrite(candidate) + actual = _session(ir.to_proto(candidate)) + np.testing.assert_allclose( + actual.run(None, feeds)[0], baseline.run(None, feeds)[0] + ) + if overridable: + self.assertEqual([v.name for v in candidate.graph.inputs], [*feeds, "c"]) + candidate_tensor = candidate.graph.initializers["c"].const_value + assert candidate_tensor is not None + np.testing.assert_array_equal( + candidate_tensor.numpy(), default_tensor.numpy() + ) + source = model.clone() + source.graph.inputs.append(source.graph.initializers["c"]) + overridden_feeds = { + **feeds, + "c": np.array(override, dtype=default_tensor.dtype.numpy()), + } + expected = _session(ir.to_proto(source)).run(None, overridden_feeds)[0] + np.testing.assert_allclose(actual.run(None, overridden_feeds)[0], expected) + self.assertFalse(np.array_equal(expected, baseline.run(None, feeds)[0])) + + def test_materialize_preserves_overridable_shape(self): + model = ir.from_onnx_text(""" + + graph (float[2,3] x, int64[2] shape) => (float[2,3] y) + + { y = Reshape(x, shape) } + """) + self.assertEqual(_materialize_reshape_shape.rules.apply_to_model(model), 0) + + @parameterized.parameterized.expand( + [(name,) for name in ("begin0", "end0", "axes0", "begin1", "end1", "axes1")] + ) + def test_slices_split_rejects_each_overridable_parameter(self, name): + model = ir.from_onnx_text(""" + + graph (float[4] x) => (float[?] y0, float[?] y1) + + { + y0 = Slice(x, begin0, end0, axes0) + y1 = Slice(x, begin1, end1, axes1) + } + """) + for overridable in (True, False): + with self.subTest(overridable=overridable): + candidate = model.clone() + if overridable: + candidate.graph.inputs.append(candidate.graph.initializers[name]) + result = _basic_rules.SlicesSplit().check( + None, candidate.graph.inputs[0], **candidate.graph.initializers + ) + self.assertEqual(bool(result), not overridable) + + def test_constant_helpers_preserve_defaults(self): + model = ir.from_onnx_text(""" + + graph (float c) => (float y) + { y = Identity(c) } + """) + value = model.graph.inputs[0] + default = value.const_value + self.assertIsNone(_ir_utils.get_const_value(value)) + self.assertIsNone(_ir_utils.get_numpy_value(value)) + self.assertIsNone(_ir_utils.get_singleton_value(value)) + self.assertFalse(_ir_utils.is_singleton_value(value, 1)) + self.assertIs(value.const_value, default) + model.graph.inputs.remove(value) + self.assertIs(_ir_utils.get_const_value(value), default) + + def test_cos_sin_cache_rejects_overridable_frequencies(self): + model = ir.from_onnx_text(""" + + graph (int64[1,2] positions) => (int64[1,2] y) + + { y = Identity(positions) } + """) + for const_freqs, name in ((False, "inv_freq"), (True, "freqs")): + with self.subTest(const_freqs=const_freqs): + rule = cos_sin_cache.CosSinCacheFusion("test", const_freqs=const_freqs) + args = ( + None, + model.graph.initializers["inv_freq"], + model.graph.inputs[0], + model.graph.initializers["freqs"], + model.graph.initializers["axes"], + ) + self.assertTrue(rule.check(*args)) + value = model.graph.initializers[name] + model.graph.inputs.append(value) + self.assertFalse(rule.check(*args)) + model.graph.inputs.remove(value) + + @parameterized.parameterized.expand( + [ + (name,) + for name in ( + "adjusted_input_shape", + "original_input_shape", + "weight_for_norm", + "bias_for_norm", + ) + ] + ) + def test_instance_norm_rejects_overridable_parameters(self, name): + model = ir.from_onnx_text(""" + + graph (float[1,2,2,2] x) => (float[1,2,2,2] y) + + { y = Identity(x) } + """) + check = instance_to_group_normalization.check_if_simulated_instance_norm_is_used + self.assertTrue(check(None, model.graph.inputs[0], **model.graph.initializers)) + model.graph.inputs.append(model.graph.initializers[name]) + self.assertFalse(check(None, model.graph.inputs[0], **model.graph.initializers)) + + def test_nested_graph_capture_is_not_a_constant(self): + model = ir.from_onnx_text(""" + + graph (float[2] x, float c, bool condition) => (float[2] y) + + { + y = If< + then_branch: graph = then_graph () => (float[2] t) { t = Add(x, c) }, + else_branch: graph = else_graph () => (float[2] e) { e = Sub(x, c) } + >(condition) + } + """) + rewriter.RewriteRuleSet(rewriter._DEFAULT_REWRITE_RULES).apply_to_model(model) + node = model.graph[0] + self.assertEqual(node.attributes["then_branch"].as_graph()[0].op_type, "Add") + self.assertEqual(node.attributes["else_branch"].as_graph()[0].op_type, "Sub") + actual = _session(ir.to_proto(model)) + for condition in (True, False): + for c in (None, np.array(3.0, np.float32)): + feeds = { + "x": np.array([2.0, 4.0], np.float32), + "condition": np.array(condition), + } + if c is not None: + feeds["c"] = c + expected = feeds["x"] + (0 if c is None else c * (1 if condition else -1)) + np.testing.assert_array_equal(actual.run(None, feeds)[0], expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/onnxscript/rewriter/_rewrite_rule.py b/onnxscript/rewriter/_rewrite_rule.py index abafb4019f..423439bc63 100644 --- a/onnxscript/rewriter/_rewrite_rule.py +++ b/onnxscript/rewriter/_rewrite_rule.py @@ -572,7 +572,7 @@ def copy_value(value: ir.Value | None) -> ir.Value | None: if value is None: return None if value not in value_map: - const_value = value.const_value + const_value = _ir_utils.get_const_value(value) if const_value is not None: # create a Constant node to represent the value value_attr = ir.AttrTensor("value", const_value) diff --git a/onnxscript/rewriter/ort_fusions/cos_sin_cache.py b/onnxscript/rewriter/ort_fusions/cos_sin_cache.py index 8e6ec1d9da..1d3da44601 100644 --- a/onnxscript/rewriter/ort_fusions/cos_sin_cache.py +++ b/onnxscript/rewriter/ort_fusions/cos_sin_cache.py @@ -72,9 +72,9 @@ def _compute_dynamic_freqs(self, op, inv_freq, position_ids, dtype): if self._max_pos_id is not None: # Use max_pos_id from the model metadata max_pos_id = self._max_pos_id - elif position_ids.const_value is not None: + elif (position_ids_value := _ir_utils.get_numpy_value(position_ids)) is not None: # Calculate max_pos_id from the position_ids tensor - max_pos_id = int(np.max(position_ids.const_value.numpy())) + max_pos_id = int(np.max(position_ids_value)) else: # Dynamically compute max_pos_id from position_ids using ONNX ops inv_freq = op.Reshape(inv_freq, op.Constant(value_ints=[1, -1])) @@ -165,7 +165,7 @@ def check( check_result = pattern.MatchResult() # TODO(rama): handle redundant reshape/expand if self._const_freqs: - if (freqs.const_value is None) or not _ir_utils.has_rank(freqs, 3): + if (_ir_utils.get_const_value(freqs) is None) or not _ir_utils.has_rank(freqs, 3): return check_result.fail("freqs is not a constant or not 3D.", freqs) else: return check_result @@ -184,7 +184,7 @@ def check( if not _ir_utils.has_rank(expanded_inv_freq, 3): return check_result.fail("expanded_inv_freq is not 3D.", expanded_inv_freq) # TODO: check expanded_inv_freq shape - if inv_freq.const_value is None: # TODO: should this be inv_freq_shape? + if _ir_utils.get_const_value(inv_freq) is None: return check_result.fail("inv_freq is not a constant.", inv_freq) if inv_freq_shape[0] != 1 or inv_freq_shape[2] != 1: return check_result.fail("inv_freq is not of shape [1, ., 1].", inv_freq) diff --git a/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py b/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py index c3de7004ff..30ee29c846 100644 --- a/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py +++ b/onnxscript/rewriter/ort_fusions/fused_matmul_rule_sets.py @@ -31,9 +31,8 @@ def pattern(self, op, x, y, cst): def check(self, context, x, y, cst) -> orp.MatchResult: check_result = orp.MatchResult() - if cst.const_value is None: + if (value := _ir_utils.get_numpy_value(cst)) is None: return check_result.fail("Divisor is not a constant value.") - value = cst.const_value.numpy() if value.size > 1: return check_result.fail("Divisor is not a scalar value.") return check_result @@ -52,9 +51,9 @@ def pattern(self, op, x, y, cst): def check(self, context, x, y, cst, **_) -> orp.MatchResult: check_result = orp.MatchResult() - if cst.const_value is None: + if (value := _ir_utils.get_numpy_value(cst)) is None: return check_result.fail("Divisor is not a constant value.") - if cst.const_value.numpy().size > 1: + if value.size > 1: return check_result.fail("Divisor is not a scalar value.") return check_result diff --git a/onnxscript/rewriter/ort_fusions/instance_to_group_normalization.py b/onnxscript/rewriter/ort_fusions/instance_to_group_normalization.py index 8ea43e4b84..9a89c550f9 100644 --- a/onnxscript/rewriter/ort_fusions/instance_to_group_normalization.py +++ b/onnxscript/rewriter/ort_fusions/instance_to_group_normalization.py @@ -7,6 +7,7 @@ import numpy as np import onnx +from onnxscript.rewriter import _ir_utils from onnxscript.rewriter._rewrite_rule import RewriteRule, RewriteRuleSet logger = logging.getLogger(__name__) @@ -40,12 +41,12 @@ def check_if_simulated_instance_norm_is_used( Returns: bool: True if the simulated instance normalization is used, False otherwise. """ - weight_for_norm_const_value = weight_for_norm.const_value + weight_for_norm_const_value = _ir_utils.get_const_value(weight_for_norm) if weight_for_norm_const_value is None: return False weight_for_norm = weight_for_norm_const_value.numpy() - bias_for_norm_const_value = bias_for_norm.const_value + bias_for_norm_const_value = _ir_utils.get_const_value(bias_for_norm) if bias_for_norm_const_value is None: return False bias_for_norm = bias_for_norm_const_value.numpy() @@ -72,7 +73,7 @@ def check_if_simulated_instance_norm_is_used( if not all(dim == 1 for dim in bias_full_shape[1:]): return False - adjusted_input_shape_const_value = adjusted_input_shape.const_value + adjusted_input_shape_const_value = _ir_utils.get_const_value(adjusted_input_shape) g = weight_for_norm.shape[0] if ( @@ -82,7 +83,7 @@ def check_if_simulated_instance_norm_is_used( return False # NOTE: Restrict the rule to only support constant shape - original_input_shape_const_value = original_input_shape.const_value + original_input_shape_const_value = _ir_utils.get_const_value(original_input_shape) if ( original_input_shape_const_value is None or original_input_shape_const_value.numpy().tolist() != input_x.shape diff --git a/onnxscript/rewriter/rules/common/_basic_rules.py b/onnxscript/rewriter/rules/common/_basic_rules.py index 66531daa2a..6c7f336b1c 100644 --- a/onnxscript/rewriter/rules/common/_basic_rules.py +++ b/onnxscript/rewriter/rules/common/_basic_rules.py @@ -102,15 +102,15 @@ def rewrite(self, op, x: ir.Value, shape: ir.Value): def check(self, context, x, shape) -> MatchResult: check_result = MatchResult() - if shape.const_value is None: + if (shape_value := ir_utils.get_numpy_value(shape)) is None: # Shape is not a constant and cannot be guessed. return check_result.fail("Shape is not a constant and cannot be guessed.") if (x_shape := x.shape) is None: # We don't know the shape of the input return check_result.fail("Input shape is not known.") - if x_shape.dims != tuple(shape.const_value.numpy().tolist()): + if x_shape.dims != tuple(shape_value.tolist()): return check_result.fail( - f"Input shape {x_shape.dims} does not match the shape {shape.const_value.numpy().tolist()}." + f"Input shape {x_shape.dims} does not match the shape {shape_value.tolist()}." ) return check_result @@ -171,12 +171,12 @@ def pattern(self, op, x, begin0, end0, axes0, begin1, end1, axes1): def check(self, context, x, begin0, end0, axes0, begin1, end1, axes1) -> MatchResult: check_result = MatchResult() if ( - axes0.const_value is None - or axes1.const_value is None - or axes0.const_value.numpy().tolist() != axes1.const_value.numpy().tolist() + (axes0_value := ir_utils.get_numpy_value(axes0)) is None + or (axes1_value := ir_utils.get_numpy_value(axes1)) is None + or axes0_value.tolist() != axes1_value.tolist() ): return check_result.fail("Axes are not equal or not constant.") - axes = axes0.const_value.numpy().tolist() + axes = axes0_value.tolist() if len(axes) != 1: return check_result.fail("Axes has more than one dimension.") if x.shape is not None: @@ -186,18 +186,18 @@ def check(self, context, x, begin0, end0, axes0, begin1, end1, axes1) -> MatchRe if axes[0] != -1 and axes[0] != rank - 1: return check_result.fail("Axes is not -1 or last dimension.") if ( - begin0.const_value is None - or end0.const_value is None - or begin1.const_value is None - or end1.const_value is None + (begin0_value := ir_utils.get_numpy_value(begin0)) is None + or (end0_value := ir_utils.get_numpy_value(end0)) is None + or (begin1_value := ir_utils.get_numpy_value(begin1)) is None + or (end1_value := ir_utils.get_numpy_value(end1)) is None ): return check_result.fail("Begin or end are not constant values.") - if begin0.const_value.numpy().tolist() != [0]: + if begin0_value.tolist() != [0]: return check_result.fail("First begin value is not 0.") e0, b1, e1 = ( - end0.const_value.numpy().tolist(), - begin1.const_value.numpy().tolist(), - end1.const_value.numpy().tolist(), + end0_value.tolist(), + begin1_value.tolist(), + end1_value.tolist(), ) if e0[0] != b1[0]: return check_result.fail("End0 is not equal to Begin1.") diff --git a/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py b/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py index fb62db216e..5aa4361cf4 100644 --- a/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py +++ b/onnxscript/rewriter/rules/common/_broadcast_to_matmul.py @@ -8,6 +8,7 @@ import onnx_ir as ir +from onnxscript.rewriter import _ir_utils from onnxscript.rewriter._basics import MatchFailureError from onnxscript.rewriter._rewrite_rule import RewriteRule, RewriteRuleSet @@ -235,7 +236,7 @@ def check_if_not_need_reshape( """ input_a_shape = input_a.shape input_b_shape = input_b.shape - shape_c_tensor = shape_c.const_value + shape_c_tensor = _ir_utils.get_const_value(shape_c) if shape_c_tensor is None: logger.info("The value 'shape_c' is not statically known.") return False @@ -271,7 +272,7 @@ def check_if_not_need_reshape( matched_reshape_b = matched_node.inputs[1].producer() if shape_b is not None else None allowzero_a = _get_allowzero(matched_reshape_a) if shape_a is not None: - shape_a_tensor = shape_a.const_value + shape_a_tensor = _ir_utils.get_const_value(shape_a) if shape_a_tensor is None or len(shape_a_tensor.shape) != 1: logger.info("The value 'shape_a' is not statically known or 1-D.") return False @@ -285,7 +286,7 @@ def check_if_not_need_reshape( allowzero_b = _get_allowzero(matched_reshape_b) if shape_b is not None: - shape_b_tensor = shape_b.const_value + shape_b_tensor = _ir_utils.get_const_value(shape_b) if shape_b_tensor is None or len(shape_b_tensor.shape) != 1: logger.info("The value 'shape_b' is not statically known or 1-D.") return False diff --git a/onnxscript/rewriter/rules/common/_collapse_slices.py b/onnxscript/rewriter/rules/common/_collapse_slices.py index 7f45997de3..baacdb24b6 100644 --- a/onnxscript/rewriter/rules/common/_collapse_slices.py +++ b/onnxscript/rewriter/rules/common/_collapse_slices.py @@ -25,10 +25,10 @@ def _check_if_redundant_slice( """If the starts is 0, and the ends is equal to or grater than the shape of the specified axis, then the slice is redundant.""" del context # Reserved for future extensions - starts_const = starts.const_value - ends_const = ends.const_value - axes_const = axes.const_value - steps_const = steps.const_value + starts_const = _ir_utils.get_const_value(starts) + ends_const = _ir_utils.get_const_value(ends) + axes_const = _ir_utils.get_const_value(axes) + steps_const = _ir_utils.get_const_value(steps) if starts_const is None or ends_const is None or axes_const is None or steps_const is None: logger.info("The value 'start', 'end', 'axis', 'step' is not statically known.") diff --git a/onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py b/onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py index 3ba37fb5b1..9234a8ed59 100644 --- a/onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py +++ b/onnxscript/rewriter/rules/common/_fuse_pad_into_conv.py @@ -14,6 +14,7 @@ import numpy as np import onnx_ir as ir +from onnxscript.rewriter import _ir_utils from onnxscript.rewriter import pattern as orp @@ -128,24 +129,26 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc ) # Pad constraints: inputs - if (pads := pad_node.inputs[1]).const_value is None: + pads = pad_node.inputs[1] + assert pads is not None + if (pads_value := _ir_utils.get_numpy_value(pads)) is None: return check_result.fail(f"{pads.name} is not a constant/initializer.") if len(pad_node.inputs) > 2 and (constant_value := pad_node.inputs[2]) is not None: - if constant_value.const_value is None: + if (fill_value := _ir_utils.get_numpy_value(constant_value)) is None: return check_result.fail( f"{constant_value.name} is not a constant/initializer." ) - elif constant_value.const_value.numpy().item() != 0: + elif fill_value.item() != 0: return check_result.fail(f"{constant_value.name} must be equal to 0.") if len(pad_node.inputs) > 3 and (axes := pad_node.inputs[3]) is not None: - if axes.const_value is None: + if (axes_value := _ir_utils.get_numpy_value(axes)) is None: return check_result.fail(f"{axes.name} is not a constant/initializer.") - axes_list = [x if x >= 0 else x_rank + x for x in axes.const_value.numpy()] + axes_list = [x if x >= 0 else x_rank + x for x in axes_value] else: axes_list = list(range(x_rank)) # Pad constraints: values - self._pads_list = fill_pads_with_axes(pads.const_value.numpy(), axes_list, x_rank) + self._pads_list = fill_pads_with_axes(pads_value.tolist(), axes_list, x_rank) if np.any(self._pads_list[:2] + self._pads_list[x_rank : x_rank + 2]): self._pads_list = None return check_result.fail(f"{pads.name} must be zero in non-spatial dimensions.") @@ -200,9 +203,9 @@ def check(self, context, x: ir.Value, pad: ir.Value, conv: ir.Value) -> orp.Matc # only agree when x_zero_point is absent or zero. conv_node = conv.producer() if len(conv_node.inputs) > 2 and (x_zero_point := conv_node.inputs[2]) is not None: - if x_zero_point.const_value is None: + if (zero_point_value := _ir_utils.get_numpy_value(x_zero_point)) is None: return check_result.fail(f"{x_zero_point.name} is not a constant/initializer.") - if np.any(x_zero_point.const_value.numpy() != 0): + if np.any(zero_point_value != 0): return check_result.fail(f"{x_zero_point.name} must be equal to 0.") return check_result diff --git a/onnxscript/rewriter/rules/common/_materialize_reshape_shape.py b/onnxscript/rewriter/rules/common/_materialize_reshape_shape.py index 32379a1897..8309b5d138 100644 --- a/onnxscript/rewriter/rules/common/_materialize_reshape_shape.py +++ b/onnxscript/rewriter/rules/common/_materialize_reshape_shape.py @@ -30,6 +30,9 @@ def pattern(self, op, data, shape): def check(self, context, data: ir.Value, shape: ir.Value) -> MatchResult: check_result = MatchResult() + if shape.is_graph_input() and shape.const_value is not None: + return check_result.fail("Shape input has an overridable initializer default.") + # Shape input must not already be a constant if ir_utils.get_numpy_value(shape) is not None: return check_result.fail("Shape input is already a constant.") diff --git a/onnxscript/rewriter/rules/common/_min_max_to_clip.py b/onnxscript/rewriter/rules/common/_min_max_to_clip.py index c0aff9e8b9..70f3a9fc0c 100644 --- a/onnxscript/rewriter/rules/common/_min_max_to_clip.py +++ b/onnxscript/rewriter/rules/common/_min_max_to_clip.py @@ -30,6 +30,7 @@ import numpy as np import onnx_ir as ir +from onnxscript.rewriter import _ir_utils from onnxscript.rewriter._basics import MatchResult from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet @@ -95,7 +96,7 @@ def check(self, context, out1, out2, **_): # Ensure all inputs except the first are constants for input_ in first_node.inputs[1:] + second_node.inputs[1:]: - if ir.convenience.get_const_tensor(input_) is None: + if _ir_utils.get_const_value(input_) is None: return check_result.fail(f"{input_.name} is not a constant.") # If scalars are required (Clip fusion), enforce scalar-ness diff --git a/onnxscript/rewriter/rules/common/_min_max_to_clip_test.py b/onnxscript/rewriter/rules/common/_min_max_to_clip_test.py index dd09078a9e..e2b8f65d48 100644 --- a/onnxscript/rewriter/rules/common/_min_max_to_clip_test.py +++ b/onnxscript/rewriter/rules/common/_min_max_to_clip_test.py @@ -177,7 +177,7 @@ def test_failure_fuse_successive_min_or_max_non_constant(self, _, op_type, rewri ("max_graph_input", "Max"), ] ) - def test_successful_fuse_successive_min_or_max_graph_inputs_as_constants(self, _, op_type): + def test_preserve_successive_min_or_max_with_overridable_defaults(self, _, op_type): base_model = ir.from_onnx_text(f""" < ir_version: 10, opset_import: ["" : 20] > test_model (float[N, 32, 14, 17] X, float[1] cst1, float[1] cst2) => (float[N, ?, ?, ?] Y) @@ -187,7 +187,7 @@ def test_successful_fuse_successive_min_or_max_graph_inputs_as_constants(self, _ Y = {op_type}(x1, cst2) }} """) - self.run_test(base_model, expected_op_types=[op_type]) + self.run_test(base_model, expected_op_types=[op_type, op_type]) class TestMinMaxToClip(_TestMinMaxToClipBase): @@ -216,7 +216,7 @@ def test_successful_min_max_to_clip_constants(self): """) self.run_test(base_model, expected_op_types=["Constant", "Clip"]) - def test_successful_min_max_to_clip_graph_inputs_as_constants(self): + def test_preserve_min_max_with_overridable_defaults(self): base_model = ir.from_onnx_text(""" < ir_version: 10, opset_import: ["" : 20] > test_model (float[N, 32, 14, 17] X, float[1] min, float[1] max) => (float [N, ?, ?, ?] Y) @@ -226,7 +226,7 @@ def test_successful_min_max_to_clip_graph_inputs_as_constants(self): Y = Max(x1, max) } """) - self.run_test(base_model, expected_op_types=["Clip"]) + self.run_test(base_model, expected_op_types=["Min", "Max"]) def test_failure_min_max_to_clip_invalid_bounds(self): """Min node should have the max value and Max node should have the min value.""" @@ -293,7 +293,7 @@ def test_successful_max_min_to_clip_constants(self): """) self.run_test(base_model, expected_op_types=["Constant", "Clip"]) - def test_successful_max_min_to_clip_graph_inputs_as_constants(self): + def test_preserve_max_min_with_overridable_defaults(self): base_model = ir.from_onnx_text(""" < ir_version: 10, opset_import: ["" : 20] > test_model (float[N, 32, 14, 17] X, float[1] min, float[1] max) => (float [N, ?, ?, ?] Y) @@ -303,7 +303,7 @@ def test_successful_max_min_to_clip_graph_inputs_as_constants(self): Y = Min(x1, min) } """) - self.run_test(base_model, expected_op_types=["Clip"]) + self.run_test(base_model, expected_op_types=["Max", "Min"]) def test_successful_max_min_to_clip_check_bounds(self): base_model = ir.from_onnx_text(""" diff --git a/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py b/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py index 6880cb4dff..566aa7c26c 100644 --- a/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py +++ b/onnxscript/rewriter/rules/common/_redundant_scatter_nd.py @@ -112,13 +112,13 @@ def check(self, context, data, indices, updates, **_): ) # 3. the indices is referring to the whole data, which is from 0 to data.shape[0] - if indices.const_value is None: + if (indices_value := _ir_utils.get_numpy_value(indices)) is None: return result.fail("The value 'indices' is not statically known.", indices) first_dim = data.shape[0] if not isinstance(first_dim, int): return result.fail("The first dimension of 'data' is not statically known.", data) expected_indices = [[i] for i in range(first_dim)] - actual_indices = indices.const_value.numpy().tolist() + actual_indices = indices_value.tolist() if actual_indices != expected_indices: return result.fail("The 'indices' is not referring to the whole data.", indices) diff --git a/onnxscript/rewriter/rules/common/_remove_optional_bias.py b/onnxscript/rewriter/rules/common/_remove_optional_bias.py index 0ba2e3426f..e0309856ed 100644 --- a/onnxscript/rewriter/rules/common/_remove_optional_bias.py +++ b/onnxscript/rewriter/rules/common/_remove_optional_bias.py @@ -9,6 +9,7 @@ import numpy as np import onnx_ir as ir +from onnxscript.rewriter import _ir_utils from onnxscript.rewriter._basics import MatchResult from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase, RewriteRuleSet @@ -37,7 +38,7 @@ def check(self, context, b: ir.Value, **_) -> MatchResult: check_result = MatchResult() # Check if bias is a constant/initializer - bias_tensor = ir.convenience.get_const_tensor(b) + bias_tensor = _ir_utils.get_const_value(b) if bias_tensor is None: return check_result.fail("Bias is not a constant/initializer.")