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

[Christian Sigg]

Author: Christian Sigg
Date: 2021-05-06T13:21:28+02:00
New Revision: a0d019fc89c57736e54a476aa4db63027a2dace2

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

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

Reviewed By: herhut

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
index e5278c7ad39c3..68417baf1459a 100644
--- a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp
@@ -78,8 +78,6 @@ struct GpuAsyncRegionPass::ThreadTokenCallback {
     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>();
 
@@ -96,7 +94,12 @@ struct GpuAsyncRegionPass::ThreadTokenCallback {
     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->getRegions(), newOp->getRegions()))
+      std::get<0>(pair).cloneInto(&std::get<1>(pair), mapping);
 
     // Replace the op with the async clone.
     auto results = newOp->getResults();