[FoldSCFIf] match store between branches

Author: kumasento

tools/polymer/lib/Transforms/FoldSCFIf.cc

@@ -8,6 +8,7 @@
 #include "mlir/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
@@ -30,11 +31,110 @@ using namespace llvm;
 
 #define DEBUG_TYPE "fold-scf-if"
 
+static bool hasSingleStore(Block *block) {
+  llvm::SetVector<Value> memrefs;
+
+  for (Operation &op : block->getOperations()) {
+    if (isa<mlir::AffineStoreOp, memref::StoreOp>(op)) {
+      Value memref = op.getOperand(1);
+      if (memrefs.count(memref))
+        return false;
+
+      memrefs.insert(memref);
+    }
+  }
+  return true;
+}
+
+namespace {
+struct MatchIfElsePass : PassWrapper<MatchIfElsePass, OperationPass<FuncOp>> {
+  void runOnOperation() override {
+    FuncOp f = getOperation();
+    OpBuilder b(f.getContext());
+
+    // If there is no store in the target block for a specific memref stored in
+    // the source block, we will create a dummy load.
+    auto matchStore = [&](Block *target, Block *source, Location loc) {
+      llvm::SetVector<Value> memrefs;
+
+      for (Operation &op : target->getOperations())
+        if (isa<memref::StoreOp, mlir::AffineStoreOp>(op))
+          memrefs.insert(op.getOperand(1));
+
+      b.setInsertionPoint(target->getTerminator());
+      for (Operation &op : source->getOperations()) {
+        if (!isa<mlir::AffineStoreOp, memref::StoreOp>(op))
+          continue;
+        Value memref = op.getOperand(1);
+        if (memrefs.count(memref)) // has been stored to
+          continue;
+
+        if (auto storeOp = dyn_cast<AffineStoreOp>(op)) {
+          Value value = b.create<AffineLoadOp>(loc, storeOp.getMemRef(),
+                                               storeOp.getAffineMap(),
+                                               storeOp.getMapOperands());
+          b.create<AffineStoreOp>(loc, value, storeOp.getMemRef(),
+                                  storeOp.getAffineMap(),
+                                  storeOp.getMapOperands());
+        } else if (auto storeOp = dyn_cast<memref::StoreOp>(op)) {
+          Value value = b.create<memref::LoadOp>(loc, storeOp.getMemRef(),
+                                                 storeOp.getIndices());
+          b.create<memref::StoreOp>(loc, value, storeOp.getMemRef(),
+                                    storeOp.indices());
+        }
+      }
+    };
+
+    f.walk([&](scf::IfOp ifOp) {
+      Location loc = ifOp.getLoc();
+      OpBuilder::InsertionGuard g(b);
+
+      // If there is no else block, initialize one with a terminating yield.
+      if (!ifOp.elseBlock()) {
+        ifOp.elseRegion().emplaceBlock();
+
+        b.setInsertionPointToStart(ifOp.elseBlock());
+        b.create<scf::YieldOp>(loc);
+      }
+
+      if (!hasSingleStore(ifOp.thenBlock()) ||
+          !hasSingleStore(ifOp.elseBlock())) {
+        LLVM_DEBUG(
+            dbgs()
+            << "Skipped if:\n"
+            << ifOp
+            << "\ndue to there are duplicated stores on the same memref.");
+        return;
+      }
+
+      matchStore(ifOp.elseBlock(), ifOp.thenBlock(), loc);
+      matchStore(ifOp.thenBlock(), ifOp.elseBlock(), loc);
+
+      LLVM_DEBUG(dbgs() << "Matched else block:\n" << ifOp << "\n\n");
+    });
+  }
+};
+} // namespace
+
+/// ---------------------- LiftStoreOps ------------------------------
+
+namespace {
+struct LiftStoreOps : PassWrapper<LiftStoreOps, OperationPass<FuncOp>> {
+  void runOnOperation() override {}
+};
+} // namespace
+
+/// ---------------------- FoldSCFIf ----------------------------------
+
 static void foldSCFIf(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
   Location loc = ifOp.getLoc();
 
   LLVM_DEBUG(dbgs() << "Working on ifOp: " << ifOp << "\n\n");
 
+  if (!hasSingleStore(ifOp.thenBlock()) ||
+      (ifOp.elseBlock() && !hasSingleStore(ifOp.elseBlock())))
+    return;
+
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPointAfter(ifOp);
 
@@ -95,7 +195,10 @@ struct FoldSCFIfPass : PassWrapper<FoldSCFIfPass, OperationPass<FuncOp>> {
 } // namespace
 
 void polymer::registerFoldSCFIfPass() {
-  PassPipelineRegistration<>(
-      "fold-scf-if", "Fold scf.if into select.",
-      [](OpPassManager &pm) { pm.addPass(std::make_unique<FoldSCFIfPass>()); });
+  PassPipelineRegistration<>("fold-scf-if", "Fold scf.if into select.",
+                             [](OpPassManager &pm) {
+                               pm.addPass(std::make_unique<MatchIfElsePass>());
+                               pm.addPass(std::make_unique<LiftStoreOps>());
+                               pm.addPass(std::make_unique<FoldSCFIfPass>());
+                             });
 }

tools/polymer/test/polymer-opt/FoldSCFIf/match-store.mlir

@@ -0,0 +1,10 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%A: memref<10xf32>, %a: f32, %cond: i1) {
+  scf.if %cond {
+    affine.store %a, %A[0] : memref<10xf32>
+  }
+  return
+}
+
+// CHECK: func @foo