[Mlir-commits] [mlir] [MLIR][SCF] Add canonicalization pattern to fold away iter args of scf.forall (PR #90189)

[Matthias Springer]

================
@@ -1415,6 +1416,31 @@ InParallelOp ForallOp::getTerminator() {
   return cast<InParallelOp>(getBody()->getTerminator());
 }
 
+FailureOr<SubsetInsertionOpInterface>
+ForallOp::getStoreOpUser(BlockArgument bbArg) {
+  Value::user_range users = bbArg.getUsers();
+  bool foundUser = false;
+  SubsetInsertionOpInterface storeOp = nullptr;
+  for (Operation *userOp : users) {
+    if (auto parallelInsertSliceOp =
+            dyn_cast<tensor::ParallelInsertSliceOp>(userOp);
+        parallelInsertSliceOp && isa<InParallelOp>(userOp->getParentOp())) {
+      // Return failure in case we find more than one user of the block argument
+      // within scf.forall.in_parallel.
+      if (foundUser) {
+        return failure();
+      }
+      storeOp = cast<SubsetInsertionOpInterface>(userOp);
----------------
matthias-springer wrote:

This should not be necessary.

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