Cleanup Mem2Reg and GPU (#204)

* Upgrade mem2reg

* Fix test

* Correct barrier elimination

* Fix affine raising

* Fix affine raise

* Fix affine raising

* Lowering update

* Simplify affine and whilelicm

* Fix affine raise

* Fix infinite loop in memset/cpy to load/store

* Optionally disable loop unroll

* Fix cmpi bug

* Fix allocation location

* Add memcpy behavior

* Properly handle continue

* Do not duplicate already cached values crossing a barrier

* Mem2reg and licm with affine

* Fix unused scf inductive var

* Fix parallel affine raise

* Remove print

* Partially speed up mem2reg

* Reduce max unroll size to 32

* Add rank reduction

* Disable RankReduction, fix for InductiveVarRemoval

* Work around affineexpr bug

* Improve ordering of loop distribute

* Improve if reg2mem

* Make cacheload

* Fix reg2mem for (2)

* Skip recursive load/store

* Now ignoring barriers

* No unnecessary internal store forwarding

* Attempt fix

* Add inner serialize

* Fix set operands

* Fix infinite loop

* Remove unnecessary yield

* Fix inner serialize

* Do not use cacheload when not mincut

* Fix recompute threshold

* Strengthen alias analysis

* No infinite loop

* Add single execution query

* Fix bug

* Fix compile bug

* First pass fix tests

* Bump LLVM

* Cleanup polygeist opt tests

* Fix tests

* Fix format

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

Author: wsmoses

include/polygeist/Ops.h

@@ -18,4 +18,122 @@
 #define GET_OP_CLASSES
 #include "polygeist/PolygeistOps.h.inc"
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/CommandLine.h"
+
+bool getEffectsBefore(
+    mlir::Operation *op,
+    llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects,
+    bool stopAtBarrier);
+
+bool getEffectsAfter(
+    mlir::Operation *op,
+    llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects,
+    bool stopAtBarrier);
+
+bool isReadOnly(mlir::Operation *);
+bool isReadNone(mlir::Operation *);
+
+bool mayReadFrom(mlir::Operation *, mlir::Value);
+bool mayWriteTo(mlir::Operation *, mlir::Value, bool ignoreBarrier = false);
+
+bool mayAlias(mlir::MemoryEffects::EffectInstance a,
+              mlir::MemoryEffects::EffectInstance b);
+
+bool mayAlias(mlir::MemoryEffects::EffectInstance a, mlir::Value b);
+
+extern llvm::cl::opt<bool> BarrierOpt;
+
+template <bool NotTopLevel = false>
+class BarrierElim final
+    : public mlir::OpRewritePattern<mlir::polygeist::BarrierOp> {
+public:
+  using mlir::OpRewritePattern<mlir::polygeist::BarrierOp>::OpRewritePattern;
+  mlir::LogicalResult
+  matchAndRewrite(mlir::polygeist::BarrierOp barrier,
+                  mlir::PatternRewriter &rewriter) const override {
+    using namespace mlir;
+    using namespace polygeist;
+    if (!BarrierOpt)
+      return failure();
+    // Remove if it only sync's constant indices.
+    if (llvm::all_of(barrier.getOperands(), [](mlir::Value v) {
+          IntegerAttr constValue;
+          return matchPattern(v, m_Constant(&constValue));
+        })) {
+      rewriter.eraseOp(barrier);
+      return success();
+    }
+
+    Operation *op = barrier;
+    if (NotTopLevel && isa<mlir::scf::ParallelOp, mlir::AffineParallelOp>(
+                           barrier->getParentOp()))
+      return failure();
+
+    {
+      SmallVector<MemoryEffects::EffectInstance> beforeEffects;
+      getEffectsBefore(op, beforeEffects, /*stopAtBarrier*/ true);
+
+      SmallVector<MemoryEffects::EffectInstance> afterEffects;
+      getEffectsAfter(op, afterEffects, /*stopAtBarrier*/ false);
+
+      bool conflict = false;
+      for (auto before : beforeEffects)
+        for (auto after : afterEffects) {
+          if (mayAlias(before, after)) {
+            // Read, read is okay
+            if (isa<MemoryEffects::Read>(before.getEffect()) &&
+                isa<MemoryEffects::Read>(after.getEffect())) {
+              continue;
+            }
+
+            // Write, write is not okay because may be different offsets and the
+            // later must subsume other conflicts are invalid.
+            conflict = true;
+            break;
+          }
+        }
+
+      if (!conflict) {
+        rewriter.eraseOp(barrier);
+        return success();
+      }
+    }
+
+    {
+      SmallVector<MemoryEffects::EffectInstance> beforeEffects;
+      getEffectsBefore(op, beforeEffects, /*stopAtBarrier*/ false);
+
+      SmallVector<MemoryEffects::EffectInstance> afterEffects;
+      getEffectsAfter(op, afterEffects, /*stopAtBarrier*/ true);
+
+      bool conflict = false;
+      for (auto before : beforeEffects)
+        for (auto after : afterEffects) {
+          if (mayAlias(before, after)) {
+            // Read, read is okay
+            if (isa<MemoryEffects::Read>(before.getEffect()) &&
+                isa<MemoryEffects::Read>(after.getEffect())) {
+              continue;
+            }
+            // Write, write is not okay because may be different offsets and the
+            // later must subsume other conflicts are invalid.
+            conflict = true;
+            break;
+          }
+        }
+
+      if (!conflict) {
+        rewriter.eraseOp(barrier);
+        return success();
+      }
+    }
+
+    return failure();
+  }
+};
+
 #endif // BFV_BFVOPS_H

include/polygeist/Passes/Passes.h

@@ -5,10 +5,13 @@
 #include "mlir/Pass/Pass.h"
 #include <memory>
 namespace mlir {
+class PatternRewriter;
+class DominanceInfo;
 namespace polygeist {
 std::unique_ptr<Pass> createParallelLICMPass();
 std::unique_ptr<Pass> createMem2RegPass();
 std::unique_ptr<Pass> createLoopRestructurePass();
+std::unique_ptr<Pass> createInnerSerializationPass();
 std::unique_ptr<Pass> replaceAffineCFGPass();
 std::unique_ptr<Pass> createOpenMPOptPass();
 std::unique_ptr<Pass> createCanonicalizeForPass();
@@ -26,8 +29,8 @@ std::unique_ptr<Pass> createConvertPolygeistToLLVMPass();
 } // namespace mlir
 
 void fully2ComposeAffineMapAndOperands(
-    mlir::OpBuilder &, mlir::AffineMap *map,
-    llvm::SmallVectorImpl<mlir::Value> *operands);
+    mlir::PatternRewriter &rewriter, mlir::AffineMap *map,
+    llvm::SmallVectorImpl<mlir::Value> *operands, mlir::DominanceInfo &DI);
 bool isValidIndex(mlir::Value val);
 
 namespace mlir {

include/polygeist/Passes/Passes.td

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