[Mlir-commits] [mlir] [MLIR][transform][python] add sugared python abstractions for transform dialect (PR #75073)

[Maksim Levental]

Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>,
Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>,
Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>,
Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>,
Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>,
Martin =?utf-8?q?Lücke?= <martin.luecke at ed.ac.uk>
Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/75073 at github.com>


================
@@ -0,0 +1,148 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from __future__ import annotations
+from typing import Callable, Optional, Sequence
+
+from .... import ir
+from ....dialects import transform
+from ....dialects.transform import structured
+
+
+class Handle(ir.Value):
+    """
+    Base class for wrappers around different types of transform handle with
+    methods to chain further transforms.
+
+    The fields `children` and `parent` are used to capture the relation of
+    handles statically in order to enable further analysis. The payload
+    operation of a child handle is nested into a region of the payload operation
+    of the corresponding parent handle.
+    """
+
+    def __init__(
+        self,
+        v: ir.Value,
+        *,
+        parent: Optional[Handle] = None,
+        children: Optional[Sequence[Handle]] = None,
+    ):
+        super().__init__(v)
+        self.parent = parent
+        self.children = children if children is not None else []
+
+
+ at ir.register_value_caster(transform.AnyOpType.get_static_typeid())
+ at ir.register_value_caster(transform.OperationType.get_static_typeid())
+class OpHandle(Handle):
+    """
+    Wrapper around a transform operation handle with methods to chain further
+    transforms.
+    """
+
+    def __init__(
+        self,
+        v: ir.Value,
+        *,
+        parent: Optional[Handle] = None,
+        children: Optional[Sequence[Handle]] = None,
+    ):
+        super().__init__(v, parent=parent, children=children)
+
+    def match_ops(
+        self,
+        ops: str
+        | ir.OpView
+        | structured.MatchInterfaceEnum
+        | Sequence[str | ir.OpView],
+    ) -> OpHandle:
+        """
+        Emits a `transform.structured.MatchOp`.
+        Returns a handle to payload ops that match the given names, types, or
+        interface. If only a single type is given, the value wrapped by the
+        resulting handle is populated with the respective type.
+        """
+        # Handle interface.
+        if isinstance(ops, structured.MatchInterfaceEnum) or (
+            isinstance(ops, str) and ops in structured.MatchInterfaceEnum.__members__
+        ):
+            if isinstance(ops, str):
+                ops = structured.MatchInterfaceEnum[ops]
----------------
makslevental wrote:

Hmm makes me think we should change the emitted enums to handle strings as well like this (not in this PR obv).

https://github.com/llvm/llvm-project/pull/75073