From a0d019fc89c57736e54a476aa4db63027a2dace2 Mon Sep 17 00:00:00 2001 From: Christian Sigg Date: Wed, 5 May 2021 20:06:37 +0200 Subject: [PATCH] [mlir] Add support for ops with regions in 'gpu-async-region' rewriter. Reviewed By: herhut Differential Revision: https://reviews.llvm.org/D101757 --- mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp index e5278c7..68417ba 100644 --- a/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp +++ b/mlir/lib/Dialect/GPU/Transforms/AsyncRegionRewriter.cpp @@ -78,8 +78,6 @@ private: 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(); @@ -96,7 +94,12 @@ private: 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(); -- 2.7.4