[mlir] Add support for ops with regions in 'gpu-async-region' rewriter.

[csigg]

Remove result mapping.

[herhut]

On what op is this used as there is no test here?

[csigg]

In D101757#2741553, @herhut wrote:

On what op is this used as there is no test here?

This is only used downstream.
I think it's trivial enough to punt on a test here, which would require an extra test op.

mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp

 
 private:
   // Replaces asyncOp with a clone that returns a token.
   LogicalResult rewriteAsyncOp(gpu::AsyncOpInterface asyncOp) {
     auto *op = asyncOp.getOperation();
     if (asyncOp.getAsyncToken())
       // TODO: Support ops that are already async.
       return op->emitOpError("is already async");
-    if (op->getNumRegions() > 0)
-      return op->emitOpError("regions are not supported");
 
     auto tokenType = builder.getType<gpu::AsyncTokenType>();
 
     // If there is no current token, insert a `gpu.wait async` without
     // dependencies to create one.
     if (!currentToken)
       currentToken = createWaitOp(op->getLoc(), tokenType, {});
     asyncOp.addAsyncDependency(currentToken);
 
     // Clone the op to return a token in addition to the other results.
     SmallVector<Type, 1> resultTypes;
     resultTypes.reserve(1 + op->getNumResults());
     copy(op->getResultTypes(), std::back_inserter(resultTypes));
     resultTypes.push_back(tokenType);
     auto *newOp = Operation::create(op->getLoc(), op->getName(), resultTypes,
                                     op->getOperands(), op->getAttrDictionary(),
-                                    op->getSuccessors());
+                                    op->getSuccessors(), op->getNumRegions());
+
+    // Clone regions into new op.
+    BlockAndValueMapping mapping;
+    for (auto pair : llvm::zip_first(op->getResults(), newOp->getResults()))
+      mapping.map(std::get<0>(pair), std::get<1>(pair));
+    for (auto pair : llvm::zip_first(op->getRegions(), newOp->getRegions()))
+      std::get<0>(pair).cloneInto(&std::get<1>(pair), mapping);
 
     // Replace the op with the async clone.
     auto results = newOp->getResults();
     currentToken = results.back();
     builder.insert(newOp);
     op->replaceAllUsesWith(results.drop_back());
     op->erase();