support region op.

Author: linuxlonelyeagle

mlir/lib/Transforms/HoistPureOps.cpp

@@ -56,10 +56,34 @@ static Value getDomaincedValue(DominanceInfo &dominanceInfo, Value a, Value b) {
   }
 }
 
+static bool isOpContainBlock(Operation *op, Block *block) {
+  Operation *parentOp = block->getParentOp();
+  while (parentOp && parentOp != op) {
+    parentOp = parentOp->getParentOp();
+  }
+  return parentOp == op ? true : false;
+}
+
 /// Find the hoisting position for the pure op.
 static Value getDestPos(Operation *op) {
   DominanceInfo dominanceInfo(op);
   SmallVector<Value> operands(op->getOperands());
+  if (op->getNumRegions()) {
+    op->walk([&](Operation *operation) {
+      for (auto operand : operation->getOperands()) {
+        Operation *defineOp = operand.getDefiningOp();
+        if (!defineOp) {
+          BlockArgument argument = cast<BlockArgument>(operand);
+          if (!isOpContainBlock(op, argument.getOwner()))
+            operands.push_back(operand);
+          continue;
+        }
+        if (!isOpContainBlock(op, defineOp->getBlock())) {
+          operands.push_back(operand);
+        }
+      }
+    });
+  }
   if (operands.empty())
     return {};
   Value ret = operands[0];
@@ -71,13 +95,18 @@ static Value getDestPos(Operation *op) {
 
 /// Hoist single pure op.
 static void hoistPureOp(RewriterBase &rewriter, Operation *op) {
+  LDBG() << "hoistPureOp: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
   Value pos = getDestPos(op);
   if (!pos)
     return;
 
   if (Operation *defineOp = pos.getDefiningOp()) {
+    if (op == defineOp)
+      return;
+
     LDBG() << "move " << OpWithFlags(op, OpPrintingFlags().skipRegions())
-           << " after " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+           << " after "
+           << OpWithFlags(defineOp, OpPrintingFlags().skipRegions());
     rewriter.moveOpAfter(op, defineOp);
     return;
   }

mlir/test/Transforms/loop-invariant-code-motion.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s  -split-input-file -loop-invariant-code-motion | FileCheck %s
+// RUN: mlir-opt %s  -split-input-file -hoist-pure-ops | FileCheck %s
 
 func.func @nested_loops_both_having_invariant_code() {
   %m = memref.alloc() : memref<10xf32>