[Mlir-commits] [mlir] 7365a11 - [mlir][linalg][transform] PackGreedilyOp/TileToForallOp: Store payload ops in SmallVector

[Matthias Springer]

Author: Matthias Springer
Date: 2023-04-20T17:42:48+09:00
New Revision: 7365a1127bfe2f6bdd2c04931eab04f4ca1df015

URL: https://github.com/llvm/llvm-project/commit/7365a1127bfe2f6bdd2c04931eab04f4ca1df015
DIFF: https://github.com/llvm/llvm-project/commit/7365a1127bfe2f6bdd2c04931eab04f4ca1df015.diff

LOG: [mlir][linalg][transform] PackGreedilyOp/TileToForallOp: Store payload ops in SmallVector

The TrackingListener removes ops from the internal transform dialect state when they are erased. At the same time, the `apply` interface method is iterating over all payload ops. Elements may not be removed from the state while iterating over the state.

Differential Revision: https://reviews.llvm.org/D148771

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 6907cf6ba115b..dbadabd66cba0 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -1439,7 +1439,11 @@ packMatmulGreedily(RewriterBase &rewriter, LinalgOp linalgOp,
 DiagnosedSilenceableFailure
 PackGreedilyOp::apply(transform::TransformResults &transformResults,
                       transform::TransformState &state) {
-  ArrayRef<Operation *> targetOps = state.getPayloadOps(getTarget());
+  ArrayRef<Operation *> targetOpsView = state.getPayloadOps(getTarget());
+  // Store payload ops into a separate SmallVector because the TrackingListener
+  // removes erased ops from the transform state.
+  SmallVector<Operation *> targetOps(targetOpsView.begin(),
+                                     targetOpsView.end());
 
   SmallVector<Operation *> results;
   TrackingListener listener(state, *this);
@@ -2773,7 +2777,10 @@ transform::TileToForallOp::apply(transform::TransformResults &transformResults,
   TrackingListener listener(state, *this);
   IRRewriter rewriter(getContext(), &listener);
   auto transformOp = cast<TransformOpInterface>(getOperation());
-  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  ArrayRef<Operation *> targetsView = state.getPayloadOps(getTarget());
+  // Store payload ops into a separate SmallVector because the TrackingListener
+  // removes erased ops from the transform state.
+  SmallVector<Operation *> targets(targetsView.begin(), targetsView.end());
 
   // Result payload ops.
   SmallVector<Operation *> tileOps;