Preserve unresolved function attributes in rewrite rules

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()
+```

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,

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

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.

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