Enable early inner serialization pass and more aggressive serial LICM (#252)

* Early inner serialization

* Add loweraffine

* Add serial licm improvement

* Fix yield bug

* Simplify

* Fix bug and simplify

* Add fallback serialization

* fix typo

* add ser pass

* Consider speculation branch

* Correct isSpec

* Expand if sink

* Fix bug where BarrierElim removes barrier required for interchange

* Fix format

Co-authored-by: Ivan Radanov Ivanov <[email protected]>

Author: wsmoses

include/polygeist/Passes/Passes.h

@@ -13,6 +13,7 @@ std::unique_ptr<Pass> createParallelLICMPass();
 std::unique_ptr<Pass> createMem2RegPass();
 std::unique_ptr<Pass> createLoopRestructurePass();
 std::unique_ptr<Pass> createInnerSerializationPass();
+std::unique_ptr<Pass> createSerializationPass();
 std::unique_ptr<Pass> replaceAffineCFGPass();
 std::unique_ptr<Pass> createOpenMPOptPass();
 std::unique_ptr<Pass> createCanonicalizeForPass();

include/polygeist/Passes/Passes.td

@@ -41,6 +41,12 @@ def InnerSerialization : Pass<"inner-serialize"> {
   let dependentDialects =
       ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect"];
 }
+def Serialization : Pass<"serialize"> {
+  let summary = "remove scf.barrier";
+  let constructor = "mlir::polygeist::createSerializationPass()";
+  let dependentDialects =
+      ["memref::MemRefDialect", "func::FuncDialect", "LLVM::LLVMDialect"];
+}
 
 def SCFBarrierRemovalContinuation : InterfacePass<"barrier-removal-continuation", "FunctionOpInterface"> {
   let summary = "Remove scf.barrier using continuations";

lib/polygeist/Passes/InnerSerialization.cpp

@@ -21,16 +21,18 @@ namespace {
 struct InnerSerialization : public InnerSerializationBase<InnerSerialization> {
   void runOnOperation() override;
 };
+struct Serialization : public SerializationBase<Serialization> {
+  void runOnOperation() override;
+};
 } // namespace
 
 struct ParSerialize : public OpRewritePattern<scf::ParallelOp> {
   using OpRewritePattern<scf::ParallelOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(scf::ParallelOp nextParallel,
                                 PatternRewriter &rewriter) const override {
-    if (!(nextParallel->getParentOfType<scf::ParallelOp>()
-          // || nextParallel->getParentOfType<AffineParallelOp>()
-          ))
+    if (!(nextParallel->getParentOfType<scf::ParallelOp>() ||
+          nextParallel->getParentOfType<AffineParallelOp>()))
       return failure();
 
     SmallVector<Value> inds;
@@ -53,6 +55,31 @@ struct ParSerialize : public OpRewritePattern<scf::ParallelOp> {
   }
 };
 
+struct Serialize : public OpRewritePattern<scf::ParallelOp> {
+  using OpRewritePattern<scf::ParallelOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(scf::ParallelOp nextParallel,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value> inds;
+    scf::ForOp last = nullptr;
+    for (auto tup :
+         llvm::zip(nextParallel.getLowerBound(), nextParallel.getUpperBound(),
+                   nextParallel.getStep(), nextParallel.getInductionVars())) {
+      last =
+          rewriter.create<scf::ForOp>(nextParallel.getLoc(), std::get<0>(tup),
+                                      std::get<1>(tup), std::get<2>(tup));
+      inds.push_back(last.getInductionVar());
+      rewriter.setInsertionPointToStart(last.getBody());
+    }
+    rewriter.eraseOp(last.getBody()->getTerminator());
+    rewriter.mergeBlocks(&nextParallel.getRegion().front(), last.getBody(),
+                         inds);
+
+    rewriter.eraseOp(nextParallel);
+    return success();
+  }
+};
+
 void InnerSerialization::runOnOperation() {
   mlir::RewritePatternSet rpl(getOperation()->getContext());
   rpl.add<ParSerialize>(getOperation()->getContext());
@@ -61,6 +88,17 @@ void InnerSerialization::runOnOperation() {
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(rpl), config);
 }
 
+void Serialization::runOnOperation() {
+  mlir::RewritePatternSet rpl(getOperation()->getContext());
+  rpl.add<Serialize>(getOperation()->getContext());
+  GreedyRewriteConfig config;
+  config.maxIterations = 47;
+  (void)applyPatternsAndFoldGreedily(getOperation(), std::move(rpl), config);
+}
+
 std::unique_ptr<Pass> mlir::polygeist::createInnerSerializationPass() {
   return std::make_unique<InnerSerialization>();
 }
+std::unique_ptr<Pass> mlir::polygeist::createSerializationPass() {
+  return std::make_unique<Serialization>();
+}

lib/polygeist/Passes/OpenMPOpt.cpp

@@ -100,6 +100,11 @@ bool isReadNone(Operation *op) {
   return false;
 }
 
+Value getBase(Value v);
+bool isStackAlloca(Value v);
+bool isCaptured(Value v, Operation *potentialUser = nullptr,
+                bool *seenuse = nullptr);
+
 bool mayReadFrom(Operation *op, Value val) {
   bool hasRecursiveEffects = op->hasTrait<OpTrait::HasRecursiveSideEffects>();
   if (hasRecursiveEffects) {
@@ -128,14 +133,16 @@ bool mayReadFrom(Operation *op, Value val) {
     }
     return false;
   }
+  if (isa<LLVM::CallOp, func::CallOp>(op)) {
+    auto base = getBase(val);
+    bool seenuse = false;
+    if (isStackAlloca(base) && !isCaptured(base, op, &seenuse) && !seenuse) {
+      return false;
+    }
+  }
   return true;
 }
 
-Value getBase(Value v);
-bool isStackAlloca(Value v);
-bool isCaptured(Value v, Operation *potentialUser = nullptr,
-                bool *seenuse = nullptr);
-
 bool mayWriteTo(Operation *op, Value val, bool ignoreBarrier) {
   bool hasRecursiveEffects = op->hasTrait<OpTrait::HasRecursiveSideEffects>();
   if (hasRecursiveEffects) {