Fix rewrite rules to handle attribute references (RefAttr) properly

[Copilot]

This PR fixes an issue where rewrite rules would fail when encountering attribute references (RefAttr) inside functions, addressing cases where attribute values are not available as concrete values but reference function parameters.

Problem

Rewrite rules previously failed in two scenarios involving RefAttr:

1. Function extraction with as_function=True: When a rewrite rule tried to extract a matched subgraph into a function and encountered nodes with RefAttr, it would raise NotImplementedError: RefAttr not supported in the _copy_for_function utility.

2. Pattern matching with attribute access: Replacement patterns that tried to access .value on attributes would fail when the attribute was a RefAttr instead of a concrete attribute.

Solution

1. Fixed RefAttr handling in function extraction

Changed copy_attr_value function in _rewrite_rule.py to properly handle RefAttr:

def copy_attr_value(attr: ir.Attr) -> ir.Attr:
    if attr.is_ref():
        # RefAttr objects are immutable and can be shared as-is.
        # The referenced attribute parameter will be handled separately
        # when the function signature is created.
        return attr
    # ... rest of function

Instead of raising NotImplementedError, RefAttr objects are now returned as-is since they are immutable and can be safely shared in the new function context.

2. Fixed RefAttr handling in pattern matching

Updated generic_pattern_test.py to check for RefAttr before accessing .value:

if perm0 is not None:
    if perm0.is_ref():
        # Cannot optimize when attribute is a reference, as we don't have the concrete value
        return None
    perm0 = perm0.value

This allows replacement patterns to gracefully handle cases where optimization cannot proceed due to attribute references.

Testing

Added comprehensive test test_rewrite_rule_with_ref_attr that verifies:

• Rewrite rules with as_function=True work correctly when encountering RefAttr
• RefAttr is properly preserved in extracted functions
• The rewrite process completes without errors

Example

Before this fix:

# This would fail with NotImplementedError
rule = pattern.RewriteRule(pattern_func, replacement, as_function=True)
# Model with RefAttr would cause rule.apply_to_model(model) to fail

After this fix:

# This now works correctly
rule = pattern.RewriteRule(pattern_func, replacement, as_function=True)  
# Model with RefAttr is handled properly, RefAttr preserved in function
count = rule.apply_to_model(model)  # ✅ Success

Fixes #1442.

---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.