[Mlir-commits] [mlir] 5048933 - [mlir] Added the dialect inliner to the SCF dialect

[Feng Liu]

Author: Feng Liu
Date: 2020-06-23T10:49:05-07:00
New Revision: 5048933c47157a55349e623c289c9ed77d58d628

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

LOG: [mlir] Added the dialect inliner to the SCF dialect

 Currently no restrictions are added to the destination regions.

Subscribers: mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, aartbik, stephenneuendorffer, Joonsoo, grosul1, Kayjukh, jurahul, msifontes

Tags: #mlir

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/SCF.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/SCF.cpp b/mlir/lib/Dialect/SCF/SCF.cpp
index f980cdb96a6a..559a5f1673f6 100644
--- a/mlir/lib/Dialect/SCF/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/SCF.cpp
@@ -19,10 +19,45 @@
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Support/MathExtras.h"
+#include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace mlir::scf;
 
+//===----------------------------------------------------------------------===//
+// SCFDialect Dialect Interfaces
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct SCFInlinerInterface : public DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  // We don't have any special restrictions on what can be inlined into
+  // destination regions (e.g. while/conditional bodies). Always allow it.
+  bool isLegalToInline(Region *dest, Region *src,
+                       BlockAndValueMapping &valueMapping) const final {
+    return true;
+  }
+  // Operations in scf dialect are always legal to inline since they are
+  // pure.
+  bool isLegalToInline(Operation *, Region *,
+                       BlockAndValueMapping &) const final {
+    return true;
+  }
+  // Handle the given inlined terminator by replacing it with a new operation
+  // as necessary. Required when the region has only one block.
+  void handleTerminator(Operation *op,
+                        ArrayRef<Value> valuesToRepl) const final {
+    auto retValOp = dyn_cast<YieldOp>(op);
+    if (!retValOp)
+      return;
+
+    for (auto retValue : llvm::zip(valuesToRepl, retValOp.getOperands())) {
+      std::get<0>(retValue).replaceAllUsesWith(std::get<1>(retValue));
+    }
+  }
+};
+} // end anonymous namespace
+
 //===----------------------------------------------------------------------===//
 // SCFDialect
 //===----------------------------------------------------------------------===//
@@ -33,6 +68,7 @@ SCFDialect::SCFDialect(MLIRContext *context)
 #define GET_OP_LIST
 #include "mlir/Dialect/SCF/SCFOps.cpp.inc"
       >();
+  addInterfaces<SCFInlinerInterface>();
 }
 
 /// Default callback for IfOp builders. Inserts a yield without arguments.