[dialects][transform] Reorganize transform dialect extensions#86
[dialects][transform] Reorganize transform dialect extensions#86
Conversation
Refactors transform extension dialects into separate submodules and places their definitions into separate files. The refactor aligns lighthouse dialect structure closer to upstream bindings mirroring its nesting. That also makes the op discovery easier and allows for cleaner addition of new extensions. When possible only the snake case op wrappers are exposed to promote their use which simplifies reuse across code and IR forms.
| class TransformExtensionDialect(ext.Dialect, name="transform_ext"): | ||
| @classmethod | ||
| def load(cls, *args, **kwargs): | ||
| # Registers the dialect and its op classes and loads the dialect and ops into the context. | ||
| super().load(*args, **kwargs) | ||
|
|
||
| # Attach interfaces to just registered/loaded operations. | ||
| for op_cls in cls.operations: | ||
| if hasattr(op_cls, "attach_interface_impls"): | ||
| op_cls.attach_interface_impls() |
There was a problem hiding this comment.
Because this already seems to be boilerplate, maybe we can have our own DialectBase class somewhere.
As the point of this PR is to move code around, we could do it in a follow up.
There was a problem hiding this comment.
The code in lighthouse/dialects/smt_ext.py is only related to the SMT dialect itself. That is, those helpers are/could be useful anywhere you create SMT ops. At the moment they are used just inside sections of (schedule) code which generate regions which contain solely smt-dialect ops.
Could you please move/keep these in a smt_ext dialect, that's not treated as a namespace of transform_ext?
| sg_m, sg_n = layer_params["sg_m"], layer_params["sg_n"] | ||
|
|
||
| @td_smt_ext.constrain_params(wg_m, wg_n, sg_m, sg_n) | ||
| @smt_ext.constrain_params(wg_m, wg_n, sg_m, sg_n) |
There was a problem hiding this comment.
| @smt_ext.constrain_params(wg_m, wg_n, sg_m, sg_n) | |
| @td_ext_smt.constrain_params(wg_m, wg_n, sg_m, sg_n) |
I am chatting with @PragmaTwice to allow transform_ext.smt.constrain_params. We could also wait for that to merge upstream and change the above to @transform_ext.smt.constrain_params(...).
3 participants
Refactors transform extension dialects into separate submodules and places their definitions into separate files.
The refactor aligns lighthouse dialect structure closer to upstream bindings mirroring its nesting. That also makes the op discovery easier and allows for cleaner addition of new extensions.
When possible only the snake case op wrappers are exposed to promote their use which simplifies reuse across code and IR forms.