[MLIR][Transform][Python] transform.foreach wrapper and .owner OpViews #171544
+107
−6
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Friendlier wrapper and makes it so that OpResult.owner returns the relevant OpView instead of Operation (Like OpResultList etc).
rolfmorel
requested review from
ftynse,
makslevental and
stellaraccident
as code owners
Member
|
@llvm/pr-subscribers-mlir Author: Rolf Morel (rolfmorel) ChangesFriendlier wrapper and makes it so that OpResult.owner returns the relevant OpView instead of Operation (Like OpResultList etc). Full diff: https://github.com/llvm/llvm-project/pull/171544.diff 4 Files Affected:
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 2e0c2b895216f..eed12fe4380bb 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1519,12 +1519,12 @@ class PyOpResult : public PyConcreteValue<PyOpResult> {
static void bindDerived(ClassTy &c) {
c.def_prop_ro(
"owner",
- [](PyOpResult &self) -> nb::typed<nb::object, PyOperation> {
+ [](PyOpResult &self) -> nb::typed<nb::object, PyOpView> {
assert(mlirOperationEqual(self.getParentOperation()->get(),
mlirOpResultGetOwner(self.get())) &&
"expected the owner of the value in Python to match that in "
"the IR");
- return self.getParentOperation().getObject();
+ return self.getParentOperation()->createOpView();
},
"Returns the operation that produces this result.");
c.def_prop_ro(
diff --git a/mlir/python/mlir/dialects/memref.py b/mlir/python/mlir/dialects/memref.py
index bc9a3a52728ad..91185b37a5b5f 100644
--- a/mlir/python/mlir/dialects/memref.py
+++ b/mlir/python/mlir/dialects/memref.py
@@ -14,8 +14,10 @@
def _is_constant_int_like(i):
return (
isinstance(i, Value)
- and isinstance(i.owner, Operation)
- and isinstance(i.owner.opview, ConstantOp)
+ and (
+ (isinstance(i.owner, Operation) and isinstance(i.owner.opview, ConstantOp))
+ or isinstance(i.owner, ConstantOp)
+ )
and _is_integer_like_type(i.type)
)
diff --git a/mlir/python/mlir/dialects/transform/__init__.py b/mlir/python/mlir/dialects/transform/__init__.py
index b3dd79c7dbd79..fbe4078782997 100644
--- a/mlir/python/mlir/dialects/transform/__init__.py
+++ b/mlir/python/mlir/dialects/transform/__init__.py
@@ -310,6 +310,8 @@ def __init__(
sym_visibility=sym_visibility,
arg_attrs=arg_attrs,
res_attrs=res_attrs,
+ loc=loc,
+ ip=ip,
)
self.regions[0].blocks.append(*input_types)
@@ -468,6 +470,54 @@ def apply_registered_pass(
).result
+@_ods_cext.register_operation(_Dialect, replace=True)
+class ForeachOp(ForeachOp):
+ def __init__(
+ self,
+ results: Sequence[Type],
+ targets: Sequence[Union[Operation, Value, OpView]],
+ *,
+ with_zip_shortest: Optional[bool] = False,
+ loc=None,
+ ip=None,
+ ):
+ targets = [_get_op_result_or_value(target) for target in targets]
+ super().__init__(
+ results_=results,
+ targets=targets,
+ with_zip_shortest=with_zip_shortest,
+ loc=loc,
+ ip=ip,
+ )
+ self.regions[0].blocks.append(*[target.type for target in targets])
+
+ @property
+ def body(self) -> Block:
+ return self.regions[0].blocks[0]
+
+ @property
+ def bodyTargets(self) -> BlockArgumentList:
+ return self.regions[0].blocks[0].arguments
+
+
+def foreach(
+ results: Sequence[Type],
+ targets: Sequence[Union[Operation, Value, OpView]],
+ *,
+ with_zip_shortest: Optional[bool] = False,
+ loc=None,
+ ip=None,
+) -> Union[OpResult, OpResultList, ForeachOp]:
+ results = ForeachOp(
+ results=results,
+ targets=targets,
+ with_zip_shortest=with_zip_shortest,
+ loc=loc,
+ ip=ip,
+ ).results
+ return results if len(results) > 1 else (results[0] if len(results) == 1 else op)
+
+
AnyOpTypeT = NewType("AnyOpType", AnyOpType)
diff --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py
index f58442d04fc66..dfcc890b83ffc 100644
--- a/mlir/test/python/dialects/transform.py
+++ b/mlir/test/python/dialects/transform.py
@@ -401,3 +401,55 @@ def testApplyRegisteredPassOp(module: Module):
options={"exclude": (symbol_a, symbol_b)},
)
transform.YieldOp()
+
+
+# CHECK-LABEL: TEST: testForeachOp
+@run
+def testForeachOp(module: Module):
+ # CHECK: transform.sequence
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [transform.AnyOpType.get()],
+ transform.AnyOpType.get(),
+ )
+ with InsertionPoint(sequence.body):
+ # CHECK: {{.*}} = foreach %{{.*}} : !transform.any_op -> !transform.any_op
+ foreach1 = transform.ForeachOp(
+ (transform.AnyOpType.get(),), (sequence.bodyTarget,)
+ )
+ with InsertionPoint(foreach1.body):
+ # CHECK: transform.yield {{.*}} : !transform.any_op
+ transform.yield_(foreach1.bodyTargets)
+
+ a_val = transform.get_operand(
+ transform.AnyValueType.get(), foreach1.result, [0]
+ )
+ a_param = transform.param_constant(
+ transform.AnyParamType.get(), StringAttr.get("a_param")
+ )
+
+ # CHECK: {{.*}} = foreach %{{.*}}, %{{.*}}, %{{.*}} : !transform.any_op, !transform.any_value, !transform.any_param -> !transform.any_value, !transform.any_param
+ foreach2 = transform.foreach(
+ (transform.AnyValueType.get(), transform.AnyParamType.get()),
+ (sequence.bodyTarget, a_val, a_param),
+ )
+ with InsertionPoint(foreach2.owner.body):
+ # CHECK: transform.yield {{.*}} : !transform.any_value, !transform.any_param
+ transform.yield_(foreach2.owner.bodyTargets[1:3])
+
+ another_param = transform.param_constant(
+ transform.AnyParamType.get(), StringAttr.get("another_param")
+ )
+ params = transform.merge_handles([a_param, another_param])
+
+ # CHECK: {{.*}} = foreach %{{.*}}, %{{.*}}, %{{.*}} with_zip_shortest : !transform.any_op, !transform.any_param, !transform.any_param -> !transform.any_op
+ foreach3 = transform.foreach(
+ (transform.AnyOpType.get(),),
+ (foreach1.result, foreach2[1], params),
+ with_zip_shortest=True,
+ )
+ with InsertionPoint(foreach3.owner.body):
+ # CHECK: transform.yield {{.*}} : !transform.any_op
+ transform.yield_((foreach3.owner.bodyTargets[0],))
+
+ transform.yield_((foreach3,))
|
rolfmorel
changed the title
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/116/builds/22238 Here is the relevant piece of the build log for the reference |
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/177/builds/25891 Here is the relevant piece of the build log for the reference |
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/138/builds/23132 Here is the relevant piece of the build log for the reference |
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/129/builds/34985 Here is the relevant piece of the build log for the reference |
Collaborator
|
Failures seem legit, reverting in #172225 for now to unbreak CI, feel free to reland with a fix! |
llvm-sync bot
pushed a commit
to arm/arm-toolchain
that referenced
this pull request
Dec 14, 2025
…r and .owner OpViews" (#172225) Reverts llvm/llvm-project#171544 ; bots are broken.
rolfmorel
added a commit
that referenced
this pull request
Dec 14, 2025
#172228) Friendlier wrapper for transform.foreach. To facilitate that friendliness, makes it so that OpResult.owner returns the relevant OpView instead of Operation. For good measure, also changes Value.owner to return OpView instead of Operation, thereby ensuring consistency. That is, makes it is so that all op-returning .owner accessors return OpView (and thereby give access to all goodies available on registered OpViews.) Reland of #171544 due to fixup for integration test.
anonymouspc
pushed a commit
to anonymouspc/llvm
that referenced
this pull request
Dec 15, 2025
llvm#171544) Friendlier wrapper for `transform.foreach`. To facilitate that friendliness, makes it so that `OpResult.owner` returns the relevant `OpView` instead of `Operation`. For good measure, also changes `Value.owner` to return `OpView` instead of `Operation`, thereby ensuring consistency. That is, makes it is so that all op-returning `.owner` accessors return `OpView` (and thereby give access to all goodies available on registered `OpView`s.)
anonymouspc
pushed a commit
to anonymouspc/llvm
that referenced
this pull request
Dec 15, 2025
…r OpViews" (llvm#172225) Reverts llvm#171544 ; bots are broken.
anonymouspc
pushed a commit
to anonymouspc/llvm
that referenced
this pull request
Dec 15, 2025
llvm#172228) Friendlier wrapper for transform.foreach. To facilitate that friendliness, makes it so that OpResult.owner returns the relevant OpView instead of Operation. For good measure, also changes Value.owner to return OpView instead of Operation, thereby ensuring consistency. That is, makes it is so that all op-returning .owner accessors return OpView (and thereby give access to all goodies available on registered OpViews.) Reland of llvm#171544 due to fixup for integration test.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Friendlier wrapper for
transform.foreach.To facilitate that friendliness, makes it so that
OpResult.ownerreturns the relevantOpViewinstead ofOperation. For good measure, also changesValue.ownerto returnOpViewinstead ofOperation, thereby ensuring consistency. That is, makes it is so that all op-returning.owneraccessors returnOpView(and thereby give access to all goodies available on registeredOpViews.)