Fix loop restructure

Author: wsmoses

lib/polygeist/Passes/CanonicalizeFor.cpp

@@ -592,6 +592,13 @@ struct MoveWhileToFor : public OpRewritePattern<WhileOp> {
     if (!step)
       return failure();
 
+    // Cannot transform for if step is not loop-invariant
+    if (auto op = step.getDefiningOp()) {
+      if (loop->isAncestor(op)) {
+        return failure();
+      }
+    }
+
     bool negativeStep = false;
     if (auto cop = step.getDefiningOp<ConstantIntOp>()) {
       if (cop.value() < 0) {
@@ -905,6 +912,7 @@ struct MoveWhileDown2 : public OpRewritePattern<WhileOp> {
         return failure();
 
       SmallVector<std::pair<BlockArgument, Value>, 2> m;
+
       // The return results of the while which are used
       SmallVector<Value, 2> prevResults;
       // The corresponding value in the before which
@@ -958,7 +966,17 @@ struct MoveWhileDown2 : public OpRewritePattern<WhileOp> {
               }
             return failure();
           }
-          m.emplace_back(std::get<2>(pair), thenYielded);
+          // If the value yielded from then then is defined in the while before
+          // but not being moved down with the if, don't change anything.
+          if (!ifOp.getThenRegion().isAncestor(thenYielded.getParentRegion()) &&
+              op.getBefore().isAncestor(thenYielded.getParentRegion())) {
+            prevResults.push_back(std::get<0>(pair));
+            condArgs.push_back(thenYielded);
+          } else {
+            // Otherwise, mark the corresponding after argument to be replaced
+            // with the value yielded in the if statement.
+            m.emplace_back(std::get<2>(pair), thenYielded);
+          }
         } else {
           assert(prevResults.size() == condArgs.size());
           prevResults.push_back(std::get<0>(pair));
@@ -974,29 +992,35 @@ struct MoveWhileDown2 : public OpRewritePattern<WhileOp> {
       rewriter.updateRootInPlace(afterYield, [&] {
         afterYield.getResultsMutable().assign(yieldArgs);
       });
-
-      llvm::SetVector<Value> sv;
-      findValuesUsedBelow(ifOp, sv);
-
       Block *afterB = &op.getAfter().front();
 
-      for (auto v : sv) {
-        condArgs.push_back(v);
-        auto arg = afterB->addArgument(v.getType());
-        for (OpOperand &use : llvm::make_early_inc_range(v.getUses())) {
-          if (ifOp->isAncestor(use.getOwner()) || use.getOwner() == afterYield)
-            rewriter.updateRootInPlace(use.getOwner(), [&]() { use.set(arg); });
+      {
+        llvm::SetVector<Value> sv;
+        findValuesUsedBelow(ifOp, sv);
+
+        for (auto v : sv) {
+          condArgs.push_back(v);
+          auto arg = afterB->addArgument(v.getType());
+          for (OpOperand &use : llvm::make_early_inc_range(v.getUses())) {
+            if (ifOp->isAncestor(use.getOwner()) ||
+                use.getOwner() == afterYield)
+              rewriter.updateRootInPlace(use.getOwner(),
+                                         [&]() { use.set(arg); });
+          }
         }
       }
 
       rewriter.setInsertionPoint(term);
       rewriter.replaceOpWithNewOp<ConditionOp>(term, term.getCondition(),
                                                condArgs);
 
+      SmallVector<unsigned> indices;
       for (int i = m.size() - 1; i >= 0; i--) {
+        assert(m[i].first.getType() == m[i].second.getType());
         m[i].first.replaceAllUsesWith(m[i].second);
-        afterB->eraseArgument(m[i].first.getArgNumber());
+        indices.push_back(m[i].first.getArgNumber());
       }
+      afterB->eraseArguments(indices);
 
       rewriter.eraseOp(ifOp.thenYield());
       Block *thenB = ifOp.thenBlock();
@@ -1266,7 +1290,10 @@ struct MoveWhileDown3 : public OpRewritePattern<WhileOp> {
 struct WhileLICM : public OpRewritePattern<WhileOp> {
   using OpRewritePattern<WhileOp>::OpRewritePattern;
   static bool canBeHoisted(Operation *op,
-                           function_ref<bool(Value)> definedOutside) {
+                           function_ref<bool(Value)> definedOutside,
+                           bool isSpeculatable) {
+    // TODO consider requirement of isSpeculatable
+
     // Check that dependencies are defined outside of loop.
     if (!llvm::all_of(op->getOperands(), definedOutside))
       return false;
@@ -1293,7 +1320,7 @@ struct WhileLICM : public OpRewritePattern<WhileOp> {
     for (auto &region : op->getRegions()) {
       for (auto &block : region) {
         for (auto &innerOp : block.without_terminator())
-          if (!canBeHoisted(&innerOp, definedOutside))
+          if (!canBeHoisted(&innerOp, definedOutside, isSpeculatable))
             return false;
       }
     }
@@ -1311,10 +1338,14 @@ struct WhileLICM : public OpRewritePattern<WhileOp> {
     // properties.
     auto isDefinedOutsideOfBody = [&](Value value) {
       auto definingOp = value.getDefiningOp();
-      bool definedOutside =
-          (definingOp && !!willBeMovedSet.count(definingOp)) ||
-          !op.getBefore().isAncestor(value.getParentRegion());
-      return definedOutside;
+      if (!definingOp) {
+        if (auto ba = value.dyn_cast<BlockArgument>())
+          definingOp = ba.getOwner()->getParentOp();
+        assert(definingOp);
+      }
+      if (willBeMovedSet.count(definingOp))
+        return true;
+      return op != definingOp && !op->isAncestor(definingOp);
     };
 
     // Do not use walk here, as we do not want to go into nested regions and
@@ -1323,7 +1354,17 @@ struct WhileLICM : public OpRewritePattern<WhileOp> {
     // processed.
     for (auto &block : op.getBefore()) {
       for (auto &op : block.without_terminator()) {
-        bool legal = canBeHoisted(&op, isDefinedOutsideOfBody);
+        bool legal = canBeHoisted(&op, isDefinedOutsideOfBody, false);
+        if (legal) {
+          opsToMove.push_back(&op);
+          willBeMovedSet.insert(&op);
+        }
+      }
+    }
+
+    for (auto &block : op.getAfter()) {
+      for (auto &op : block.without_terminator()) {
+        bool legal = canBeHoisted(&op, isDefinedOutsideOfBody, true);
         if (legal) {
           opsToMove.push_back(&op);
           willBeMovedSet.insert(&op);

lib/polygeist/Passes/LoopRestructure.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/RegionGraphTraits.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
@@ -338,12 +339,8 @@ bool LoopRestructure::removeIfFromRegion(DominanceInfo &domInfo, Region &region,
                                           oldTerm->getOperands());
             oldTerm->erase();
 
-            SmallVector<Value, 4> res;
-            for (size_t i = 1; i < ifOp->getNumResults(); ++i) {
-              res.push_back(ifOp->getResult(i));
-            }
-            builder.create<scf::ConditionOp>(builder.getUnknownLoc(),
-                                             ifOp->getResult(0), res);
+            builder.create<scf::YieldOp>(builder.getUnknownLoc(),
+                                         ifOp->getResults());
             condBr->erase();
 
             pseudoExit->erase();
@@ -386,14 +383,12 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
       mlir::OpBuilder builder(wrapper, wrapper->begin());
 
       // Copy the arguments across
-      SmallVector<Type, 4> headerArgumentTypes;
-      for (auto arg : header->getArguments()) {
-        headerArgumentTypes.push_back(arg.getType());
-      }
+      SmallVector<Type, 4> headerArgumentTypes(header->getArgumentTypes());
       wrapper->addArguments(headerArgumentTypes);
 
-      SmallVector<Value> valsCallingLoop(wrapper->getArguments().begin(),
-                                         wrapper->getArguments().end());
+      SmallVector<Value> valsCallingLoop;
+      for (auto a : wrapper->getArguments())
+        valsCallingLoop.push_back(a);
 
       SmallVector<std::pair<Value, size_t>> preservedVals;
       for (auto B : L->getBlocks()) {
@@ -415,15 +410,9 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
         }
       }
 
-      // TODO values used outside loop should be wrapped.
-
-      SmallVector<Type, 4> combinedTypes(headerArgumentTypes.begin(),
-                                         headerArgumentTypes.end());
-      SmallVector<Type, 4> returns;
-      for (auto arg : target->getArguments()) {
-        returns.push_back(arg.getType());
-        combinedTypes.push_back(arg.getType());
-      }
+      SmallVector<Type, 4> combinedTypes = headerArgumentTypes;
+      SmallVector<Type, 4> returns(target->getArgumentTypes());
+      combinedTypes.append(returns);
 
       auto loop = builder.create<mlir::scf::WhileOp>(
           builder.getUnknownLoc(), combinedTypes, valsCallingLoop);
@@ -451,28 +440,45 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
           Preds.push_back(block);
       }
 
-      loop.getBefore().getBlocks().splice(loop.getBefore().getBlocks().begin(),
-                                          region.getBlocks(), header);
+      Block *loopEntry = new Block();
+      loop.getBefore().push_back(loopEntry);
+      builder.setInsertionPointToEnd(loopEntry);
+      SmallVector<Type, 4> tys = {builder.getI1Type()};
+      for (auto t : combinedTypes)
+        tys.push_back(t);
+      auto exec =
+          builder.create<scf::ExecuteRegionOp>(builder.getUnknownLoc(), tys);
+
+      {
+        SmallVector<Value> yields;
+        for (auto a : exec.getResults())
+          yields.push_back(a);
+        yields.erase(yields.begin());
+        builder.create<scf::ConditionOp>(builder.getUnknownLoc(),
+                                         exec.getResult(0), yields);
+      }
+
+      Region &insertRegion = exec.getRegion();
+
+      insertRegion.getBlocks().splice(insertRegion.getBlocks().begin(),
+                                      region.getBlocks(), header);
+      assert(header->getParent() == &insertRegion);
       for (auto *w : L->getBlocks()) {
         Block *b = &**w;
         if (b != header) {
-          loop.getBefore().getBlocks().splice(
-              loop.getBefore().getBlocks().end(), region.getBlocks(), b);
+          insertRegion.getBlocks().splice(insertRegion.getBlocks().end(),
+                                          region.getBlocks(), b);
         }
       }
 
       Block *pseudoExit = new Block();
-      auto i1Ty = builder.getI1Type();
       {
-        loop.getBefore().push_back(pseudoExit);
-        SmallVector<Type, 4> tys = {i1Ty};
-        for (auto t : combinedTypes)
-          tys.push_back(t);
+        insertRegion.push_back(pseudoExit);
         pseudoExit->addArguments(tys);
         OpBuilder builder(pseudoExit, pseudoExit->begin());
         tys.clear();
-        builder.create<scf::ConditionOp>(builder.getUnknownLoc(), tys,
-                                         pseudoExit->getArguments());
+        builder.create<scf::YieldOp>(builder.getUnknownLoc(), tys,
+                                     pseudoExit->getArguments());
       }
 
       for (auto *w : exitingBlocks) {
@@ -486,7 +492,7 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
             auto vfalse = builder.create<arith::ConstantIntOp>(
                 builder.getUnknownLoc(), false, 1);
 
-            std::vector<Value> args = {vfalse};
+            SmallVector<Value> args = {vfalse};
             for (auto arg : header->getArguments())
               args.push_back(arg);
             for (auto v : preservedVals)
@@ -541,11 +547,10 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
                 builder.getUnknownLoc(), true, 1);
 
             if (auto op = dyn_cast<BranchOp>(terminator)) {
-              std::vector<Value> args(op.getOperands().begin(),
-                                      op.getOperands().end());
+              SmallVector<Value> args(op.getOperands());
               args.insert(args.begin(), vtrue);
-              for (auto pair : preservedVals)
-                args.push_back(pair.first);
+              for (auto p : preservedVals)
+                args.push_back(p.first);
               for (auto ty : returns) {
                 args.push_back(builder.create<mlir::LLVM::UndefOp>(
                     builder.getUnknownLoc(), ty));
@@ -618,41 +623,38 @@ void LoopRestructure::runOnRegion(DominanceInfo &domInfo, Region &region) {
             });
       }
 
+      for (auto pair :
+           llvm::zip(header->getArguments(),
+                     loopEntry->addArguments(header->getArgumentTypes()))) {
+        std::get<0>(pair).replaceAllUsesWith(std::get<1>(pair));
+      }
+      header->eraseArguments([](BlockArgument) { return true; });
+
       builder2.create<scf::YieldOp>(builder.getUnknownLoc(), yieldargs);
-      domInfo.invalidate(&loop.getBefore());
-      runOnRegion(domInfo, loop.getBefore());
-      if (!removeIfFromRegion(domInfo, loop.getBefore(), pseudoExit)) {
+      domInfo.invalidate(&insertRegion);
+
+      assert(header->getParent() == &insertRegion);
+
+      runOnRegion(domInfo, insertRegion);
+
+      if (!removeIfFromRegion(domInfo, insertRegion, pseudoExit)) {
         attemptToFoldIntoPredecessor(pseudoExit);
       }
 
       attemptToFoldIntoPredecessor(wrapper);
       attemptToFoldIntoPredecessor(target);
-      if (loop.getBefore().getBlocks().size() != 1) {
-        Block *blk = new Block();
-        OpBuilder B(loop.getContext());
-        B.setInsertionPointToEnd(blk);
-        auto cop =
-            cast<scf::ConditionOp>(loop.getBefore().getBlocks().back().back());
-        auto er = B.create<scf::ExecuteRegionOp>(loop.getLoc(),
-                                                 cop.getOperandTypes());
-        er.getRegion().getBlocks().splice(er.getRegion().getBlocks().begin(),
-                                          loop.getBefore().getBlocks());
-        loop.getBefore().push_back(blk);
-        SmallVector<Value> yields;
-        for (auto a : er.getResults())
-          yields.push_back(a);
-        yields.erase(yields.begin());
-        B.create<scf::ConditionOp>(cop.getLoc(), er.getResult(0), yields);
-        B.setInsertionPoint(&*cop);
-        for (auto arg : er.getRegion().front().getArguments()) {
-          auto na = blk->addArgument(arg.getType());
-          arg.replaceAllUsesWith(na);
-        }
-        er.getRegion().front().eraseArguments(
-            [](BlockArgument) { return true; });
-        B.create<scf::YieldOp>(cop.getLoc(), cop.getOperands());
-        cop.erase();
+
+      if (llvm::hasSingleElement(insertRegion)) {
+        Block *block = &insertRegion.front();
+        IRRewriter B(exec->getContext());
+        Operation *terminator = block->getTerminator();
+        ValueRange results = terminator->getOperands();
+        terminator->erase();
+        B.mergeBlockBefore(block, exec);
+        exec.replaceAllUsesWith(results);
+        exec.erase();
       }
+
       assert(loop.getBefore().getBlocks().size() == 1);
       runOnRegion(domInfo, loop.getAfter());
       assert(loop.getAfter().getBlocks().size() == 1);

lib/polygeist/Passes/Mem2Reg.cpp

@@ -579,6 +579,17 @@ bool Mem2Reg::forwardStoreToLoad(mlir::Value AI, std::vector<ssize_t> idx,
               }
 
               if (thenVal != nullptr && elseVal != nullptr) {
+                if (ifOp.getElseRegion().getBlocks().size()) {
+                  for (auto tup : llvm::reverse(llvm::zip(
+                           ifOp.getResults(), ifOp.thenYield().getOperands(),
+                           ifOp.elseYield().getOperands()))) {
+                    if (std::get<1>(tup) == thenVal &&
+                        std::get<2>(tup) == elseVal) {
+                      lastVal = std::get<0>(tup);
+                      continue;
+                    }
+                  }
+                }
                 OpBuilder B(ifOp.getContext());
                 B.setInsertionPoint(ifOp);
                 SmallVector<mlir::Type, 4> tys(ifOp.getResultTypes().begin(),

lib/polygeist/Passes/ParallelLoopDistribute.cpp

@@ -754,6 +754,8 @@ struct DistributeAroundBarrier : public OpRewritePattern<scf::ParallelOp> {
           u.set(buf);
         }
       } else if (auto ao = v.getDefiningOp<LLVM::AllocaOp>()) {
+        llvm::errs() << ao->getParentOfType<FuncOp>() << "\n";
+        llvm::errs() << ao << "\n";
         llvm_unreachable("split around llvm alloca unhandled\n");
       } else
         rewriter.create<memref::StoreOp>(v.getLoc(), v, alloc,