[transform_ext] Move populate_pattern op to dialects dir by rolfmorel · Pull Request #81 · llvm/lighthouse

rolfmorel · 2026-03-16T23:27:07Z

Introduces op
transform_ext.populate_pattern PAT_NAME TARGET_OP_KIND PRIORITY where patterns can be registered on PopulatePatternOp via its name_to_pattern_rewrite class member.

Fixes #80.

Introduces op `transform_ext.populate_pattern TARGET_OP_KIND PAT_NAME PRIORITY` where patterns can be registered on `PopulatePatternOp` via its `name_to_pattern_rewrite` class member. Fixes #80.

rolfmorel · 2026-03-16T23:37:00Z

@charithaintc, this fixes #80. We should be able to merge it when @fschlimb and @adam-smnk are up in the morning.

(You can check out users/rolfmorel/fix-pattern-op in the meantime if you'd like to be unblocked today.)

fschlimb · 2026-03-17T09:52:24Z

lighthouse/dialects/__init__.py

+def register_and_load():
+    from . import transform_ext
+
+    transform_ext.register_and_load()


What's the benefit of this wrapper/indirection?

Because there's already a PR which extends multiple (sub)dialects: https://github.com/llvm/lighthouse/pull/63/changes#diff-c806f5426817d17a93adc42142d8d36986989d6f11083732266d2a127bcf8d9c

The indirection is so we have to only call one function once we have obtained the ir.Context().

fschlimb · 2026-03-17T09:52:52Z

lighthouse/dialects/transform_ext.py

+def register_and_load(context=None):
+    TransformExtensionDialect.load()


As in the tuning PR, this is so to have a unified interface to accomplish the respective registering and loading across the different (sub)dialect extensions: https://github.com/llvm/lighthouse/pull/63/changes#diff-c806f5426817d17a93adc42142d8d36986989d6f11083732266d2a127bcf8d9c

adam-smnk · 2026-03-17T09:53:20Z

Does this also act as a general registry i.e., allow to register all custom patterns in one module and make them accessible for direct use?
Or is it just updated mechanism only replacing pattern_rewrite_schedule?

fschlimb · 2026-03-17T09:56:01Z

lighthouse/dialects/transform_ext.py

+def populate_pattern(
+    pattern_name: str, op_kind: str, priority: int
+) -> PopulatePatternOp:
+    """Camelcase constructor for PopulatePatternOp."""


Suggested change

"""Camelcase constructor for PopulatePatternOp."""

"""snake_case constructor for PopulatePatternOp."""

fschlimb · 2026-03-17T10:12:10Z

lighthouse/dialects/transform_ext.py

+    # A mapping from pattern names to their corresponding rewrite functions.
+    # This should be populated by the users of this operation. In effect serves
+    # as a registry for rewrite patterns that can be applied by this operation.
+    name_to_rewrite_pattern = {}


I find this rather weird. Filling a specific singleton (that I don't control) and referring to my entry later feels odd. Suggest to not use a pattern with logically uses a global/singleton. And I don't see what problem it solves.

I feel that this could/should be auto-populated with what's available in lighthouse to skip manual registration.

I don't see what problem it solves.

The previous approach conflates an op definition with an instance of that op. That is, if you want to use multiple op instances in the IR, it would attempt to overwrite the previous op definition registration.

In general, on-demand usage of the factory pattern is not suited to the (current) op registration mechanism. The registration method for Python-defined ops assumes that all ops of a dialect are to be registered at the same time, before any instances of them are created. Invoking the factory function during IR construction violates this (when we have multiple ops in a dialect).

This PR's approach is set up to "have our cake and eat it". That is, we have the upfront definition of the relevant op(s) and can adjust the patterns that can be invoked during IR construction. That is, by modifying PopulatePatternOp.name_to_rewrite_pattern.

If we want to go for dedicated transform_ext.apply_patterns.PAT1 and transform_ext.apply_patterns.PAT2, etc. ops via a factory method, let me know. We would need to adjust the code to try to keep people from invoking the factory method after other ops from the same dialect have already been registered.

I don't know. It feels like an odd workaround. In my mind, what we should provide is something like

my_pattern = register_pattern("func.func", lambda op: ...) apply_my_pattern(my_pattern)

Can we avoid accessing a dict and instead have a registration function outside the class? Also asking the user to care about the name/key of the patterns doesn't seem needed.

rolfmorel · 2026-03-17T11:22:55Z

Does this also act as a general registry i.e., allow to register all custom patterns in one module and make them accessible for direct use?

PopulatePatternOp.name_to_rewrite_pattern is a registry, though this PR does not attempt a systematic approach for getting rewrite patterns registered with it. Presumably later on we would have a global registry in a dedicated place somewhere, once we start to collect more rewrite patterns in lighthouse. We could have a nice mechanism to have the rewrite patterns registered with it (e.g. a decorator).

Note that in the MLIR C++ codebase, RewritePatterns are not invokable by name. If anything, the transform.apply_patterns.PAT1 ops that live upstream are achieving just that: binding a name to a pattern. If we want, we could keep to that: a dedicated op for each pattern, which is also how the name for the pattern gets bound (though, like in upstream, the pattern would only be accessible by name from inside a schedule).

adam-smnk

I like the registry direction.
I think this provides greater flexibility and reduces duplication of explicit wrapper ops.
While it diverges from upstream (C++ driven) pattern registration, I can see potential benefits that make it worth exploring alternatives.

PR provides only initial skeleton but looks like a good starting point for further iterations.

charithaintc · 2026-03-17T15:49:46Z

@charithaintc, this fixes #80. We should be able to merge it when @fschlimb and @adam-smnk are up in the morning.

(You can check out users/rolfmorel/fix-pattern-op in the meantime if you'd like to be unblocked today.)

thanks @rolfmorel

rolfmorel requested a review from adam-smnk March 16, 2026 23:27

[transform_ext] Move populate_pattern op to dialects dir

302a972

Introduces op `transform_ext.populate_pattern TARGET_OP_KIND PAT_NAME PRIORITY` where patterns can be registered on `PopulatePatternOp` via its `name_to_pattern_rewrite` class member. Fixes #80.

rolfmorel force-pushed the users/rolfmorel/fix-pattern-op branch from a6b71bb to 302a972 Compare March 16, 2026 23:28

rolfmorel requested a review from fschlimb March 16, 2026 23:29

rolfmorel added 2 commits March 16, 2026 16:31

Register op as well for MPI example

3fdefb9

Change arg order

b18dad1

fschlimb reviewed Mar 17, 2026

View reviewed changes

adam-smnk approved these changes Mar 17, 2026

View reviewed changes

		def register_and_load(context=None):
		TransformExtensionDialect.load()

	"""Camelcase constructor for PopulatePatternOp."""
	"""snake_case constructor for PopulatePatternOp."""

Conversation

rolfmorel commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rolfmorel commented Mar 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adam-smnk commented Mar 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rolfmorel commented Mar 17, 2026

Uh oh!

adam-smnk left a comment

Choose a reason for hiding this comment

Uh oh!

charithaintc commented Mar 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

rolfmorel commented Mar 16, 2026 •

edited

Loading

rolfmorel commented Mar 16, 2026 •

edited

Loading