From a65e473fc9fd1d2cbb72fb3b6aac8456800cd96d Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 10 Nov 2026 12:35:19 -0800 Subject: [PATCH] Reject unresolved attributes in value-dependent rewrite rules Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/tutorial/rewriter/conditional_rewrite.md | 18 +++ onnxscript/rewriter/_pattern_ir.py | 2 + onnxscript/rewriter/pattern_test.py | 117 ++++++++++++++++++ .../rewriter/rules/common/_basic_rules.py | 8 ++ .../rules/common/_basic_rules_test.py | 79 ++++++++++++ 5 files changed, 224 insertions(+) diff --git a/docs/tutorial/rewriter/conditional_rewrite.md b/docs/tutorial/rewriter/conditional_rewrite.md index 379788e657..01cf4fe25d 100644 --- a/docs/tutorial/rewriter/conditional_rewrite.md +++ b/docs/tutorial/rewriter/conditional_rewrite.md @@ -101,3 +101,21 @@ def advanced_condition_check(context, x, y, **_): ``` This context information enables condition functions to make decisions based on the broader graph structure, the specific nodes involved in the match, and relationships between matched patterns and the rest of the model. + +## Attribute references in functions + +Attributes inside an ONNX function can reference a function parameter rather than +hold a concrete value. Literal pattern constraints such as `perm=[1, 0]` do not +match these references. Attribute variables can still match references and pass +them unchanged to a replacement operator. + +When a condition needs an attribute's concrete value, check `attr.is_ref()` before +calling accessors such as `as_int()` or `as_ints()`. A reference is not an omitted +attribute: do not replace it with the operator's default. + +```python +def check_cast(context, x, to): + if to.is_ref(): + return False + return x.dtype == to.as_int() +``` diff --git a/onnxscript/rewriter/_pattern_ir.py b/onnxscript/rewriter/_pattern_ir.py index 73de2eb817..bef990e1a0 100644 --- a/onnxscript/rewriter/_pattern_ir.py +++ b/onnxscript/rewriter/_pattern_ir.py @@ -127,6 +127,8 @@ def __init__(self, value: SupportedAttrTypes): self._value = value def matches(self, attr: ir.Attr) -> bool: + if attr.is_ref(): + return False if attr.type in { ir.AttributeType.INTS, ir.AttributeType.FLOATS, diff --git a/onnxscript/rewriter/pattern_test.py b/onnxscript/rewriter/pattern_test.py index 81cb58ff13..4e0453bec0 100644 --- a/onnxscript/rewriter/pattern_test.py +++ b/onnxscript/rewriter/pattern_test.py @@ -227,6 +227,123 @@ def test_concat_in_function(self): class RewriteRuleTest(unittest.TestCase): + def test_literal_attribute_constraints_reject_references(self): + def make_target(expected, multi_output): + def target(op, x): + constrained = op.Custom(x, attr=expected, _domain="local") + return (op.Identity(x), constrained) if multi_output else constrained + + return target + + for expected, attribute_type in ( + (1, ir.AttributeType.INT), + (1.0, ir.AttributeType.FLOAT), + ("value", ir.AttributeType.STRING), + ([0, 1], ir.AttributeType.INTS), + ([0.0, 1.0], ir.AttributeType.FLOATS), + (["a", "b"], ir.AttributeType.STRINGS), + ): + for multi_output in (False, True): + with self.subTest(expected=expected, multi_output=multi_output): + target = make_target(expected, multi_output) + x = ir.Value(name="x") + identity = ir.Node("", "Identity", [x]) + constrained = ir.Node( + "local", + "Custom", + [x], + [ir.RefAttr("attr", "parameter", attribute_type)], + ) + graph = ir.Graph( + [x], + [identity.outputs[0], constrained.outputs[0]], + nodes=[identity, constrained], + opset_imports={"": 18, "local": 1}, + ) + function = ir.Function( + "local", + "function", + graph=graph, + attributes=[ir.Attr("parameter", attribute_type, None)], + ) + model = ir.Model( + ir.Graph([], [], nodes=[], opset_imports={"": 18}, name="main"), + ir_version=10, + functions=[function], + ) + root = identity if multi_output else constrained + self.assertFalse(pattern.Pattern(target).match(model, function, root)) + + def test_attribute_variable_preserves_reference(self): + def target(op, x, perm): + return op.Transpose(op.Identity(x), perm=perm) + + def replacement(op, x, perm): + return op.Transpose(x, perm=perm) + + x = ir.Value(name="x") + identity = ir.Node("", "Identity", [x]) + transpose = ir.Node( + "", + "Transpose", + identity.outputs, + [ir.RefAttr("perm", "permutation", ir.AttributeType.INTS)], + ) + function = ir.Function( + "local", + "function", + graph=ir.Graph( + [x], transpose.outputs, nodes=[identity, transpose], opset_imports={"": 18} + ), + attributes=[ir.Attr("permutation", ir.AttributeType.INTS, None)], + ) + model = ir.Model( + ir.Graph([], [], nodes=[], opset_imports={"": 18}, name="main"), + ir_version=10, + functions=[function], + ) + self.assertEqual(pattern.RewriteRule(target, replacement).apply_to_model(model), 1) + perm = function.graph[-1].attributes["perm"] + self.assertTrue(perm.is_ref()) + self.assertEqual(perm.ref_attr_name, "permutation") + onnx.checker.check_model(ir.to_proto(model)) + + def test_multi_output_literal_attribute_search_skips_references(self): + def target(op, x): + return op.Identity(x), op.Transpose(x, perm=[1, 0]) + + x = ir.Value(name="x") + root = ir.Node("", "Identity", [x]) + reference = ir.Node( + "", "Transpose", [x], [ir.RefAttr("perm", "permutation", ir.AttributeType.INTS)] + ) + wrong = ir.Node("", "Transpose", [x], [ir.AttrInt64s("perm", [0, 1])]) + correct = ir.Node("", "Transpose", [x], [ir.AttrInt64s("perm", [1, 0])]) + function = ir.Function( + "local", + "function", + graph=ir.Graph( + [x], + [root.outputs[0], correct.outputs[0]], + nodes=[root, reference, wrong, correct], + opset_imports={"": 18}, + ), + attributes=[ir.Attr("permutation", ir.AttributeType.INTS, None)], + ) + model = ir.Model( + ir.Graph([], [], nodes=[], opset_imports={"": 18}, name="main"), + ir_version=10, + functions=[function], + ) + matcher = pattern.Pattern(target) + match = matcher.match(model, function, root) + self.assertIsNotNone(match) + self.assertEqual(list(match.outputs), list(function.outputs)) + correct.attributes["perm"] = ir.AttrInt64s("perm", [0, 1]) + self.assertFalse(matcher.match(model, function, root)) + correct.attributes["perm"] = ir.AttrInt64s("perm", [1, 0]) + self.assertIsNotNone(matcher.match(model, function, root)) + def test_commute(self): def add_0(op, x): return x + 0 diff --git a/onnxscript/rewriter/rules/common/_basic_rules.py b/onnxscript/rewriter/rules/common/_basic_rules.py index 66531daa2a..c184071bdb 100644 --- a/onnxscript/rewriter/rules/common/_basic_rules.py +++ b/onnxscript/rewriter/rules/common/_basic_rules.py @@ -53,6 +53,8 @@ def rewrite(self, op, x: ir.Value, to: ir.Attr): def check(self, context, x, to) -> MatchResult: check_result = MatchResult() + if to.is_ref(): + return check_result.fail("Cast target type is a reference attribute.") if x.dtype != to.as_int(): return check_result.fail("Input and output types are not the same") return check_result @@ -130,6 +132,9 @@ def rewrite(self, op, x: ir.Value, shape_ignored: ir.Value, shape: ir.Value): def check(self, context, x, shape_ignored, shape) -> MatchResult: check_result = MatchResult() + allowzero = context.root.attributes.get("allowzero") + if allowzero is not None and allowzero.is_ref(): + return check_result.fail("Reshape allowzero is a reference attribute.") # Shape must be a constant. if (np_shape := ir_utils.get_numpy_value(shape)) is None: @@ -315,6 +320,9 @@ def rewrite(self, op, x: ir.Value): def check(self, context, x: ir.Value) -> MatchResult: check_result = MatchResult() + axis_attr = context.root.attributes.get("axis") + if axis_attr is not None and axis_attr.is_ref(): + return check_result.fail("Flatten axis is a reference attribute.") self._new_shape = np.array([-1, -1], "int64") # Convert axis in a positive value if possible. diff --git a/onnxscript/rewriter/rules/common/_basic_rules_test.py b/onnxscript/rewriter/rules/common/_basic_rules_test.py index 595e1b18bc..fdaa65a252 100644 --- a/onnxscript/rewriter/rules/common/_basic_rules_test.py +++ b/onnxscript/rewriter/rules/common/_basic_rules_test.py @@ -38,6 +38,65 @@ def clone_model(model: ir.Model) -> ir.Model: class BasicRulesTest(unittest.TestCase): + @parameterized.parameterized.expand( + [ + ("cast_unknown_type", "Cast", "to", None, _basic_rules.no_op_cast_rule), + ("cast_known_type", "Cast", "to", ir.DataType.FLOAT, _basic_rules.no_op_cast_rule), + ( + "flatten", + "Flatten", + "axis", + ir.DataType.FLOAT, + _basic_rules.flatten_to_reshape_rule, + ), + ( + "reshape", + "Reshape", + "allowzero", + ir.DataType.FLOAT, + _basic_rules.reshape_reshape_rule, + ), + ] + ) + def test_reference_attributes_in_function(self, _, op_type, attribute_name, dtype, rule): + x = ir.Value( + name="x", + type=ir.TensorType(dtype) if dtype is not None else None, + shape=ir.Shape([2, 3]), + ) + ref_attr = ir.RefAttr(attribute_name, "parameter", ir.AttributeType.INT) + nodes = [] + if op_type == "Reshape": + shape = ir.Node( + "", + "Constant", + [], + [ir.AttrTensor("value", ir.tensor(np.array([0, 3], dtype=np.int64)))], + ) + first = ir.Node("", "Reshape", [x, shape.outputs[0]]) + root = ir.Node("", "Reshape", [first.outputs[0], shape.outputs[0]], [ref_attr]) + nodes.extend([shape, first]) + else: + root = ir.Node("", op_type, [x], [ref_attr]) + nodes.append(root) + function = ir.Function( + "local", + "function", + graph=ir.Graph([x], root.outputs, nodes=nodes, opset_imports={"": 18}), + attributes=[ir.Attr("parameter", ir.AttributeType.INT, None)], + ) + model = ir.Model( + ir.Graph([], [], nodes=[], opset_imports={"": 18}, name="main"), + ir_version=10, + functions=[function], + ) + onnxscript.optimizer.basic_constant_propagation(function) + before = ir.to_proto(model).SerializeToString() + self.assertIsNone(rule.try_rewrite(model, function, root)) + self.assertEqual(rule.apply_to_model(model), 0) + self.assertEqual(ir.to_proto(model).SerializeToString(), before) + onnx.checker.check_model(ir.to_proto(model)) + def _get_random_inputs(self, model: onnx.ModelProto) -> dict[str, Any]: feeds: dict[str, Any] = {} for i in model.graph.input: @@ -478,6 +537,26 @@ def model3(X: ot.FLOAT[1, 1]): class ReshapeReshapeTest(unittest.TestCase): + def test_failed_match_does_not_contaminate_later_rewrites(self): + rule = _basic_rules.ReshapeReshape.rule(remove_nodes=False) + models = [ + self.create_model( + (2, 3, 4), + np.array([3, 8], dtype=np.int64), + np.array(shape, dtype=np.int64), + ) + for shape in ([6, 4], [0, -1], [4, 6]) + ] + feeds = {"X": np.arange(24, dtype=np.float32).reshape(2, 3, 4)} + for model, expected_count in zip(models, (1, 0, 1)): + original = clone_model(model) + self.assertEqual(rule.apply_to_model(model), expected_count) + testing.assert_numerically_equal( + original, model, feeds, atol=0, rtol=0, use_reference=True + ) + self.assertEqual(models[0].graph[-1].inputs[1].const_value.numpy().tolist(), [6, 4]) + self.assertEqual(models[-1].graph[-1].inputs[1].const_value.numpy().tolist(), [4, 6]) + @staticmethod def create_model( input_shape, shape1, shape2, allowzero1=0, allowzero2=0, infer_shape=False