[Mlir-commits] [mlir] [mlir][linalg-transform] dyn_cast DestinationStyleOpInterface and early return (PR #166299)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Hsiang-Chieh Tsou (hsjts0u)

<details>
<summary>Changes</summary>

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)
}
}
```

---
Full diff: https://github.com/llvm/llvm-project/pull/166299.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+5-2) 


``````````diff
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 :

``````````

</details>


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