[mlir][linalg-transform] dyn_cast DestinationStyleOpInterface and early return #166299
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
hsjts0u
requested review from
ftynse,
nicolasvasilache and
rolfmorel
as code owners
Member
|
@llvm/pr-subscribers-mlir Author: Hsiang-Chieh Tsou (hsjts0u) ChangesUse module {
func.func @<!-- -->fuse(%arg0: tensor<4x4x4xf32>, %arg1: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {
%mul = "myop.a"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>
%add = "myop.b"(%mul, %mul) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>
return %add : tensor<4x4x4xf32>
}
transform.sequence failures(propagate) {
^bb0(%func: !transform.any_op):
%mul = transform.structured.match ops{["myop.a"]} in %func : (!transform.any_op) -> !transform.any_op
%add = transform.structured.match ops{["myop.b"]} in %func : (!transform.any_op) -> !transform.any_op
%loop, %tiled = transform.structured.tile_using_forall %add tile_sizes [1, 2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%mul_fused, %mul_containing = transform.structured.fuse_into_containing_op %mul into %tiled : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
}
}Full diff: https://github.com/llvm/llvm-project/pull/166299.diff 1 Files Affected:
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 3a433825fd31a..59629c422a034 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -997,8 +997,11 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
// Iterate over the outputs of the producer and over the loop bbArgs and
// check if any bbArg points to the same value as the producer output. In
// such case, make the producer output point to the bbArg directly.
- for (OpOperand &initOperandPtr :
- cast<DestinationStyleOpInterface>(clone).getDpsInitsMutable()) {
+ auto dpsInterface = dyn_cast<DestinationStyleOpInterface>(clone);
+ if (!dpsInterface)
+ return;
+
+ for (OpOperand &initOperandPtr : dpsInterface.getDpsInitsMutable()) {
Value producerOperand =
clone->getOperand(initOperandPtr.getOperandNumber());
for (BlockArgument containerIterArg :
|
Member
|
@llvm/pr-subscribers-mlir-linalg Author: Hsiang-Chieh Tsou (hsjts0u) ChangesUse module {
func.func @<!-- -->fuse(%arg0: tensor<4x4x4xf32>, %arg1: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {
%mul = "myop.a"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>
%add = "myop.b"(%mul, %mul) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>
return %add : tensor<4x4x4xf32>
}
transform.sequence failures(propagate) {
^bb0(%func: !transform.any_op):
%mul = transform.structured.match ops{["myop.a"]} in %func : (!transform.any_op) -> !transform.any_op
%add = transform.structured.match ops{["myop.b"]} in %func : (!transform.any_op) -> !transform.any_op
%loop, %tiled = transform.structured.tile_using_forall %add tile_sizes [1, 2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
%mul_fused, %mul_containing = transform.structured.fuse_into_containing_op %mul into %tiled : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
}
}Full diff: https://github.com/llvm/llvm-project/pull/166299.diff 1 Files Affected:
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 3a433825fd31a..59629c422a034 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -997,8 +997,11 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
// Iterate over the outputs of the producer and over the loop bbArgs and
// check if any bbArg points to the same value as the producer output. In
// such case, make the producer output point to the bbArg directly.
- for (OpOperand &initOperandPtr :
- cast<DestinationStyleOpInterface>(clone).getDpsInitsMutable()) {
+ auto dpsInterface = dyn_cast<DestinationStyleOpInterface>(clone);
+ if (!dpsInterface)
+ return;
+
+ for (OpOperand &initOperandPtr : dpsInterface.getDpsInitsMutable()) {
Value producerOperand =
clone->getOperand(initOperandPtr.getOperandNumber());
for (BlockArgument containerIterArg :
|
Contributor
Author
|
cc @pabloantoniom and @ftynse to review or tag relevant reviewers. Thanks! |
Contributor
Author
|
Bump on this |
Contributor
MaheshRavishankar
left a comment
MaheshRavishankar
left a comment
There was a problem hiding this comment.
Could you check the producerOp instead of the clone. That way we wont even clone if we dont need to .
ftynse
reviewed
Nov 5, 2025
Member
ftynse
left a comment
ftynse
left a comment
There was a problem hiding this comment.
For testing, you can add new ops into the test dialect and give them some trivial implementation of the interface, like do nothing or return failure. We only need to check there's no crash, not a particular behavior.
hsjts0u
force-pushed
the
user/htsou/dyn_cast_interface
branch
from
2c30f37 to
93b8734
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
hsjts0u
force-pushed
the
user/htsou/dyn_cast_interface
branch
2 times, most recently
from
9091c8f to
238d9eb
Compare
hsjts0u
force-pushed
the
user/htsou/dyn_cast_interface
branch
from
238d9eb to
72dd81d
Compare
ftynse
approved these changes
Nov 6, 2025
Contributor
Author
|
@ftynse Can you help merge this PR. Thank you! |
Collaborator
|
LLVM Buildbot has detected a new failure on builder Full details are available at: https://lab.llvm.org/buildbot/#/builders/169/builds/16808 Here is the relevant piece of the build log for the reference |
vinay-deshmukh
pushed a commit
to vinay-deshmukh/llvm-project
that referenced
this pull request
Nov 8, 2025
…ly return (llvm#166299) Use `dyn_cast` instead of `cast` and early return if op does not implement the `DestinationStyleOpInterface`. Before the change the following IR would cause a segfault when the transform interpreter is run, where `myop.a` and `myop.b` implement the `TilingInterface` and not the `DestinationStyleOpInterface`. Tried looking for ops in the upstream dialect that implement the `TilingInterface` and not the `DestinationStyleOpInterface` to add a test but could not find any. ```mlir module { func.func @fuse(%arg0: tensor<4x4x4xf32>, %arg1: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> { %mul = "myop.a"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32> %add = "myop.b"(%mul, %mul) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32> return %add : tensor<4x4x4xf32> } transform.sequence failures(propagate) { ^bb0(%func: !transform.any_op): %mul = transform.structured.match ops{["myop.a"]} in %func : (!transform.any_op) -> !transform.any_op %add = transform.structured.match ops{["myop.b"]} in %func : (!transform.any_op) -> !transform.any_op %loop, %tiled = transform.structured.tile_using_forall %add tile_sizes [1, 2, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op) %mul_fused, %mul_containing = transform.structured.fuse_into_containing_op %mul into %tiled : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op) } } ```
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.
Use
dyn_castinstead ofcastand early return if op does not implement theDestinationStyleOpInterface. Before the change the following IR would cause a segfault when the transform interpreter is run, wheremyop.aandmyop.bimplement theTilingInterfaceand not theDestinationStyleOpInterface. Tried looking for ops in the upstream dialect that implement theTilingInterfaceand not theDestinationStyleOpInterfaceto add a test but could not find any.