core: new inference system for constraints #3456

alexarice · 2024-11-15T14:44:08Z

I've gone through a few different refactors this week and have eventually settled on this slightly smaller refactor for the constraint system. The key addition is that each constraint now returns a dictionary of "resolvers" (name can maybe be improved) which resolve a variable from the given attribute. This allows inference to be decoupled entirely from verification, and inference should no longer be quite as fragile/maybe slightly quicker.

In this version, the resolvers are not used for verification, as I had originally planned, and instead the previous verification framework is left as is. I believe this change can enable the following changes in the future:

Allow constraints on properties to resolve variables in a clean way (see core: introduce TypedAttributeConstraint #3318)
Allow the introduction of "length/int constraints" for managing the length of variadics (see misc: raise error if variadic results types aren't referenced in assembly format #3416)
Potentially facilitate extensions to assembly format (i.e. directly supplying a variable)
Resolvers could be used to automate construction of initialiser methods for operation.

PS. as of time of writing it appears this saves 4 lines of code

codecov · 2024-11-15T14:49:50Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.31%. Comparing base (59e59ad) to head (5ad1495).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3456      +/-   ##
==========================================
+ Coverage   90.29%   90.31%   +0.02%     
==========================================
  Files         462      462              
  Lines       57937    58039     +102     
  Branches     5566     5562       -4     
==========================================
+ Hits        52313    52417     +104     
+ Misses       4189     4188       -1     
+ Partials     1435     1434       -1

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

xdsl/irdl/constraints.py

tests/irdl/test_declarative_assembly_format.py

xdsl/irdl/constraints.py

superlopuh

Overall looks good, although I'm still not sure what the motivation for this is. Is it better error messages? Or is there more downstream? Also, my first instinct would be to somehow be a bit more lazy about the resolver situation, would it make sense to make them into lazy model objects, that let the user iterate over the keys, and define each function only once, instead of creating wrappers of wrappers. Not 100% sure if it would work, but might be worth a try in this PR?

alexarice · 2024-11-15T17:20:46Z

I tried to put some motivation in the PR description and there is a bit more discussion in #3318 about needing to verify properties during parsing with the old system (I have rebased that PR onto this PR).

superlopuh · 2024-11-15T17:22:30Z

Thank you, happy to approve once the CI passes

xdsl/irdl/constraints.py

alexarice · 2024-11-15T17:29:55Z

Also, my first instinct would be to somehow be a bit more lazy about the resolver situation, would it make sense to make them into lazy model objects, that let the user iterate over the keys, and define each function only once, instead of creating wrappers of wrappers. Not 100% sure if it would work, but might be worth a try in this PR?

I'm also not sure how happy I am about stacks of wrappers. Another option which could be worth trying is to have a resolver be a list of "accessors" where each accessor is a piece of data (such as the attribute name, or operand index). I haven't tried putting this into action so I'm not sure how messy it would get.

Could you expand on what you mean by "lazy model objects"?

superlopuh · 2024-11-15T17:46:21Z

Something like

class AbstractResolver(abc.ABC):
    @abstractmethod
    def iter_variables() -> Iterator[str]: ...

     @abstractmethod
    def resolve(attr: Attribute) -> Attribute | Sequence[Attribute]:
        ...

@dataclass
class WhateverResolver(AbstractResolver):

    a: AbstractResolver
    b: AbstractResolver

    def iter_variables() -> Iterator[str]:
        yield from a.iter_variables()
        yield from b.iter_variables()

    def resolve(attribute: Attribute) -> Attribute | Sequence[Attribute]:
        return a.resolve(attribute) if creative_condition else b.resolve(attribute)

alexarice · 2024-11-15T19:49:15Z

I like that idea, I'll give it a go next week

@superlopuh

This implements something halfway between what I originally had in #3456 and what @superlopuh suggested in #3456 (comment) . I think this a strict improvement over the original which uses callables and lambda functions and will merge it into that PR if people agree

alexarice · 2024-11-18T16:34:59Z

Have merged #3464 here, so this is ready for review again (Think there are no outstanding open comments)

superlopuh · 2024-11-18T19:32:01Z

xdsl/irdl/declarative_assembly_format_parser.py

+        resolvers = dict[str, Resolver[ParsingState]]()
        for i, (_, operand_def) in enumerate(self.op_def.operands):
            if self.seen_operand_types[i]:
-                seen_variables |= operand_def.constr.get_resolved_variables()
+                resolvers.update(
+                    {
+                        v: self._OperandResultResolver(i, True, r)
+                        for v, r in operand_def.constr.get_resolvers().items()
+                    }
+                )
        for i, (_, result_def) in enumerate(self.op_def.results):
            if self.seen_result_types[i]:
-                seen_variables |= result_def.constr.get_resolved_variables()
-        return seen_variables
+                resolvers.update(
+                    {
+                        v: self._OperandResultResolver(i, False, r)
+                        for v, r in result_def.constr.get_resolvers().items()
+                    }
+                )
+        return resolvers


What happens if there's an overlap in vs? Worth a comment if that's not possible? My instinct here would be to create two dictionaries, probably with two dictionary comprehensions, and then take their union. (Maybe assert that the intersection of keys is empty)

If there's an overlap in vs then the second one encountered will be used. This would be happen if a variable can be inferred from two different places in the assembly format. Which resolver is used shouldn't make any difference for any operation that successfully verifies.

how do you feel about a resolver that asserts at runtime that the two resolved results are the same?

I imagine this would be easy to do with the new class-based system. I'm not sure how I feel about putting what is effectively verification as a check in parsing.

superlopuh · 2024-11-18T21:04:56Z

xdsl/dialects/bufferization.py

+        if isa(attr, TensorType[Attribute]):
+            memref_type = MemRefType(attr.element_type, attr.shape)
+            return self.memref_var_constr.verify(memref_type, constraint_context)
+        if isa(attr, UnrankedTensorType[Attribute]):
+            memref_type = UnrankedMemrefType.from_type(attr.element_type)
+
+        raise VerifyException(f"Expected TensorType or UnrankedTensorType, got {attr}")


the logic here seems to be changed

It is, though it appears I forgot to amend the doc comment. In my opinion the previous constraint was an abuse of the particular workings of the previous system. The one I have written here is a bit more well-behaved with respect to inference, and can be used to achieve the same effect as the old one. I'll update the doc comment.

Would it make sense to merge this change as a separate PR, while we discuss the details of this one?

Sure, I'll get that ready first thing tomorrow

xdsl/irdl/constraints.py

alexarice · 2024-11-19T18:57:21Z

I believe all comments have been resolved.

I've renamed Resolvers to VarExtractors and given them a (admittedly small) docstring.
I have also introduced a MergeExtractor which is used to record that two extractors have been provided for the same variable. This performs a runtime check to determine that each extractor would return the same thing. I have added 3 tests for each case where this is used, using slightly contrived operations.
I have rebased on top of the change to the constraint in bufferization.

alexarice · 2024-11-19T18:58:02Z

Forgot to change the can_infer method argument name, will do that now

tests/dialects/test_bufferization.py

xdsl/irdl/declarative_assembly_format_parser.py

xdsl/irdl/constraints.py

alexarice · 2024-11-19T19:43:57Z

@superlopuh do you feel the extractor merging is worth the added complexity? My examples in the tests feel a bit contrived

superlopuh · 2024-11-19T19:59:04Z

I feel like it's OK, I'm just imagining the headache that would emerge from trying to debug this, and it all seems worth it.

Co-authored-by: Sasha Lopoukhine <[email protected]>

alexarice · 2024-11-20T08:05:26Z

@math-fehr @superlopuh are you both happy for this to be merged? Do you want to merge other conflicting PRs first?

superlopuh · 2024-11-20T08:23:46Z

I just merged the small conflicting PR, hopefully will be easy to resolve. I think the best order now is for you to merge this and for me to resolve my draft PRs.

alexarice added the core xDSL core (ir, textual format, ...) label Nov 15, 2024

alexarice requested review from superlopuh, compor, AntonLydike, math-fehr, PapyChacal and kimxworrall November 15, 2024 14:44

alexarice self-assigned this Nov 15, 2024

alexarice mentioned this pull request Nov 15, 2024

core: introduce TypedAttributeConstraint #3318

Merged

superlopuh reviewed Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

tests/irdl/test_declarative_assembly_format.py Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 15, 2024

View reviewed changes

alexarice force-pushed the alexarice/constraints-v1.5 branch from 1fab125 to d1001b6 Compare November 15, 2024 17:14

alexarice commented Nov 15, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

alexarice mentioned this pull request Nov 18, 2024

Use classes in place of callables for resolvers #3464

Merged

superlopuh reviewed Nov 18, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh mentioned this pull request Nov 19, 2024

core: (irdl) ParamAttrConstraint can infer recursively #3477

Merged

alexarice added 2 commits November 19, 2024 15:47

core: new inference system for constraints

bdc8420

Fix bufferization

aeadd84

alexarice force-pushed the alexarice/constraints-v1.5 branch from a002ad5 to aeadd84 Compare November 19, 2024 15:59

alexarice added 2 commits November 19, 2024 17:20

Rename some classes

f50c014

Add merge-extractors and tests

c437da0

rename argument of can_infer

c66995b

superlopuh reviewed Nov 19, 2024

View reviewed changes

tests/dialects/test_bufferization.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 19, 2024

View reviewed changes

xdsl/irdl/declarative_assembly_format_parser.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 19, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 19, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 19, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh reviewed Nov 19, 2024

View reviewed changes

xdsl/irdl/constraints.py Outdated Show resolved Hide resolved

superlopuh approved these changes Nov 19, 2024

View reviewed changes

alexarice and others added 6 commits November 19, 2024 21:04

Update xdsl/irdl/constraints.py

f17a9f7

Co-authored-by: Sasha Lopoukhine <[email protected]>

Update tests/dialects/test_bufferization.py

fb6a520

Co-authored-by: Sasha Lopoukhine <[email protected]>

Update xdsl/irdl/constraints.py

4703d6b

Co-authored-by: Sasha Lopoukhine <[email protected]>

Update xdsl/irdl/constraints.py

cc87241

Co-authored-by: Sasha Lopoukhine <[email protected]>

Update xdsl/irdl/constraints.py

a1d38d7

Co-authored-by: Sasha Lopoukhine <[email protected]>

More renaming

badc71b

alexarice added 2 commits November 20, 2024 08:54

Merge branch 'main' into alexarice/constraints-v1.5

f57346e

Replace dict() by {}

5ad1495

alexarice merged commit d334411 into main Nov 20, 2024
15 checks passed

alexarice deleted the alexarice/constraints-v1.5 branch November 20, 2024 11:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

core: new inference system for constraints #3456

core: new inference system for constraints #3456

alexarice commented Nov 15, 2024 •

edited

Loading

codecov bot commented Nov 15, 2024 •

edited

Loading

superlopuh left a comment

alexarice commented Nov 15, 2024

superlopuh commented Nov 15, 2024

alexarice commented Nov 15, 2024

superlopuh commented Nov 15, 2024 •

edited

Loading

alexarice commented Nov 15, 2024

alexarice commented Nov 18, 2024

superlopuh Nov 18, 2024

alexarice Nov 18, 2024

superlopuh Nov 18, 2024

alexarice Nov 18, 2024

superlopuh Nov 18, 2024

alexarice Nov 18, 2024

superlopuh Nov 18, 2024

alexarice Nov 18, 2024

alexarice commented Nov 19, 2024

alexarice commented Nov 19, 2024

alexarice commented Nov 19, 2024

superlopuh commented Nov 19, 2024

alexarice commented Nov 20, 2024

superlopuh commented Nov 20, 2024

core: new inference system for constraints #3456

core: new inference system for constraints #3456

Conversation

alexarice commented Nov 15, 2024 • edited Loading

codecov bot commented Nov 15, 2024 • edited Loading

Codecov Report

superlopuh left a comment

Choose a reason for hiding this comment

alexarice commented Nov 15, 2024

superlopuh commented Nov 15, 2024

alexarice commented Nov 15, 2024

superlopuh commented Nov 15, 2024 • edited Loading

alexarice commented Nov 15, 2024

alexarice commented Nov 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alexarice commented Nov 19, 2024

alexarice commented Nov 19, 2024

alexarice commented Nov 19, 2024

superlopuh commented Nov 19, 2024

alexarice commented Nov 20, 2024

superlopuh commented Nov 20, 2024

alexarice commented Nov 15, 2024 •

edited

Loading

codecov bot commented Nov 15, 2024 •

edited

Loading

superlopuh commented Nov 15, 2024 •

edited

Loading