[FoldSCFIf] lift store out of scf.if

Author: kumasento

tools/polymer/lib/Transforms/FoldSCFIf.cc

@@ -40,6 +40,19 @@ static bool hasSingleStore(Block *block) {
       if (memrefs.count(memref))
         return false;
 
+      // The indices should be defined above the current block.
+      if (auto storeOp = dyn_cast<mlir::AffineStoreOp>(op)) {
+        if (any_of(storeOp.getMapOperands(), [&](Value operand) {
+              return operand.getParentBlock() == block;
+            }))
+          return false;
+      } else if (auto storeOp = dyn_cast<memref::StoreOp>(op)) {
+        if (any_of(storeOp.getIndices(), [&](Value operand) {
+              return operand.getParentBlock() == block;
+            }))
+          return false;
+      }
+
       memrefs.insert(memref);
     }
   }
@@ -118,9 +131,166 @@ struct MatchIfElsePass : PassWrapper<MatchIfElsePass, OperationPass<FuncOp>> {
 
 /// ---------------------- LiftStoreOps ------------------------------
 
+static bool hasMatchingStores(ArrayRef<Block *> blocks) {
+  if (blocks.size() <= 1)
+    return true;
+
+  llvm::SetVector<Value> setUnion;
+
+  for (Block *block : blocks.drop_front()) {
+    llvm::SetVector<Value> memrefs;
+
+    for (Operation &op : block->getOperations())
+      if (isa<memref::StoreOp, mlir::AffineStoreOp>(op)) {
+        Value memref = op.getOperand(1);
+        assert(!memrefs.count(memref) &&
+               "Should only apply on blocks that contain single store to each "
+               "memref.");
+
+        memrefs.insert(op.getOperand(1));
+      }
+
+    bool wasEmpty = setUnion.empty();
+    if (!wasEmpty && setUnion.set_union(memrefs))
+      return false;
+  }
+
+  return true;
+}
+
+namespace {
+struct MemRefStoreInfo {
+  unsigned index;
+  Type type;
+  Operation *source;
+  SmallVector<Value> operands;
+};
+} // namespace
+
+static void
+getMemRefStoreInfo(Block *block,
+                   SmallDenseMap<Value, MemRefStoreInfo> &storeInfo) {
+  unsigned ord = 0;
+  for (Operation &op : block->getOperations())
+    if (isa<memref::StoreOp, mlir::AffineStoreOp>(op)) {
+      MemRefStoreInfo info;
+      info.index = ord++;
+      info.type = op.getOperand(0).getType();
+      info.source = &op;
+
+      if (auto storeOp = dyn_cast<memref::StoreOp>(op))
+        info.operands = storeOp.getIndices();
+      else if (auto storeOp = dyn_cast<mlir::AffineStoreOp>(op))
+        info.operands = storeOp.getMapOperands();
+
+      storeInfo[op.getOperand(1)] = info;
+    }
+}
+
+static LogicalResult liftStoreOps(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
+  Location loc = ifOp.getLoc();
+
+  if (!hasMatchingStores({ifOp.thenBlock(), ifOp.elseBlock()}))
+    return failure();
+
+  SmallDenseMap<Value, MemRefStoreInfo> storeInfo;
+  getMemRefStoreInfo(ifOp.thenBlock(), storeInfo);
+
+  SmallVector<Type> storeTypes(storeInfo.size());
+  for (auto &info : storeInfo)
+    storeTypes[info.second.index] = info.second.type;
+
+  // No need to process further.
+  if (storeInfo.empty())
+    return failure();
+
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPointAfter(ifOp);
+
+  SmallVector<Type> resultTypes(ifOp.getResultTypes());
+  resultTypes.append(storeTypes);
+
+  scf::IfOp newIfOp = b.create<scf::IfOp>(loc, resultTypes, ifOp.condition(),
+                                          /*withElseRegion=*/true);
+
+  auto cloneBlock = [&](Block *target, Block *source) {
+    BlockAndValueMapping vmap;
+
+    scf::YieldOp yieldOp = cast<scf::YieldOp>(source->getTerminator());
+    unsigned numExistingResults = yieldOp.getNumOperands();
+    SmallVector<Value> results(numExistingResults + storeInfo.size());
+
+    OpBuilder::InsertionGuard g(b);
+    b.setInsertionPointToStart(target);
+
+    for (Operation &op : source->getOperations()) {
+      if (isa<memref::StoreOp, mlir::AffineStoreOp>(op)) {
+        Value memref = op.getOperand(1);
+        Value toStore = op.getOperand(0);
+        results[storeInfo[memref].index + numExistingResults] =
+            vmap.lookupOrDefault(toStore);
+      } else if (!isa<scf::YieldOp>(op)) {
+        b.clone(op, vmap);
+      }
+    }
+
+    for (auto operand : enumerate(yieldOp.getOperands()))
+      results[operand.index()] = vmap.lookupOrDefault(operand.value());
+
+    b.create<scf::YieldOp>(loc, results);
+  };
+
+  cloneBlock(newIfOp.thenBlock(), ifOp.thenBlock());
+  cloneBlock(newIfOp.elseBlock(), ifOp.elseBlock());
+
+  b.setInsertionPointAfter(newIfOp);
+  for (auto &p : storeInfo) {
+    Value memref;
+    MemRefStoreInfo info;
+    std::tie(memref, info) = p;
+
+    if (auto storeOp = dyn_cast<mlir::AffineStoreOp>(info.source))
+      b.create<mlir::AffineStoreOp>(
+          loc, newIfOp.getResult(ifOp.getNumResults() + info.index), memref,
+          storeOp.getAffineMap(), info.operands);
+    else if (auto storeOp = dyn_cast<memref::StoreOp>(info.source))
+      b.create<memref::StoreOp>(
+          loc, newIfOp.getResult(ifOp.getNumResults() + info.index), memref,
+          info.operands);
+  }
+
+  ifOp.erase();
+
+  return success();
+}
+
+static bool processLiftStoreOps(FuncOp f, OpBuilder &b) {
+  bool changed = false;
+
+  f.walk([&](scf::IfOp ifOp) {
+    if (!hasSingleStore(ifOp.thenBlock()) ||
+        (ifOp.elseBlock() && !hasSingleStore(ifOp.elseBlock())))
+      return;
+    if (failed(liftStoreOps(ifOp, f, b)))
+      return;
+
+    changed = true;
+  });
+
+  return changed;
+}
+
 namespace {
 struct LiftStoreOps : PassWrapper<LiftStoreOps, OperationPass<FuncOp>> {
-  void runOnOperation() override {}
+  void runOnOperation() override {
+    FuncOp f = getOperation();
+    OpBuilder b(f.getContext());
+
+    // For each scf.if, see if it has single store for each memref on each
+    // branch.
+    while (processLiftStoreOps(f, b))
+      ;
+  }
 };
 } // namespace
 

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

@@ -7,4 +7,7 @@ func @foo(%A: memref<10xf32>, %a: f32, %cond: i1) {
   return
 }
 
-// CHECK: func @foo
+// CHECK: func @foo(%[[A:.*]]: memref<10xf32>, %[[a:.*]]: f32, %[[cond:.*]]: i1) 
+// CHECK-NEXT:   %[[v0:.*]] = affine.load %[[A]][0] : memref<10xf32>
+// CHECK-NEXT:   %[[v1:.*]] = select %[[cond]], %[[a]], %[[v0]] : f32
+// CHECK-NEXT:   affine.store %[[v1]], %[[A]][0] : memref<10xf32>