[FoldSCFIf] added -fold-scf-if pass

Author: kumasento

tools/polymer/include/polymer/Transforms/FoldSCFIf.h

@@ -0,0 +1,10 @@
+//===- FoldSCFIf.h - Fold scf.if into select --------------C++-===//
+
+#ifndef POLYMER_TRANSFORMS_FOLDSCFIF_H
+#define POLYMER_TRANSFORMS_FOLDSCFIF_H
+
+namespace polymer {
+void registerFoldSCFIfPass();
+}
+
+#endif

tools/polymer/lib/Transforms/CMakeLists.txt

@@ -5,6 +5,7 @@ add_mlir_conversion_library(PolymerTransforms
   ScopStmtOpt.cc
   LoopAnnotate.cc
   LoopExtract.cc
+  FoldSCFIf.cc
 
   ADDITIONAL_HEADER_DIRS
   "${POLYMER_MAIN_INCLUDE_DIR}/polymer/Transforms"

tools/polymer/lib/Transforms/FoldSCFIf.cc

@@ -0,0 +1,97 @@
+//===- FoldSCFIf.cc - Fold scf.if into select --------------C++-===//
+
+#include "polymer/Transforms/FoldSCFIf.h"
+
+#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/AffineValueMap.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "mlir/Transforms/Utils.h"
+
+#include "llvm/Support/Debug.h"
+
+using namespace mlir;
+using namespace llvm;
+
+#define DEBUG_TYPE "fold-scf-if"
+
+static void foldSCFIf(scf::IfOp ifOp, FuncOp f, OpBuilder &b) {
+  Location loc = ifOp.getLoc();
+
+  LLVM_DEBUG(dbgs() << "Working on ifOp: " << ifOp << "\n\n");
+
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPointAfter(ifOp);
+
+  SmallVector<Value> thenResults, elseResults;
+
+  auto cloneAfter = [&](Block *block, SmallVectorImpl<Value> &results) {
+    BlockAndValueMapping vmap;
+    for (Operation &op : block->getOperations()) {
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(op))
+        for (Value result : yieldOp.getOperands())
+          results.push_back(vmap.contains(result) ? vmap.lookup(result)
+                                                  : result);
+      else
+        b.clone(op, vmap);
+    }
+  };
+
+  cloneAfter(ifOp.thenBlock(), thenResults);
+  cloneAfter(ifOp.elseBlock(), elseResults);
+
+  for (auto ifResult : enumerate(ifOp.getResults())) {
+    Value newResult =
+        b.create<SelectOp>(loc, ifOp.condition(), thenResults[ifResult.index()],
+                           elseResults[ifResult.index()]);
+    ifResult.value().replaceAllUsesWith(newResult);
+  }
+
+  ifOp.erase();
+}
+
+/// Return true if anything changed.
+static bool process(FuncOp f, OpBuilder &b) {
+  bool changed = false;
+
+  f.walk([&](scf::IfOp ifOp) {
+    /// TODO: add verification.
+    foldSCFIf(ifOp, f, b);
+    changed = true;
+  });
+
+  return changed;
+}
+
+namespace {
+struct FoldSCFIfPass : PassWrapper<FoldSCFIfPass, OperationPass<FuncOp>> {
+  void runOnOperation() override {
+    FuncOp f = getOperation();
+    OpBuilder b(f.getContext());
+
+    while (process(f, b))
+      ;
+  }
+};
+} // namespace
+
+void polymer::registerFoldSCFIfPass() {
+  PassPipelineRegistration<>(
+      "fold-scf-if", "Fold scf.if into select.",
+      [](OpPassManager &pm) { pm.addPass(std::make_unique<FoldSCFIfPass>()); });
+}

tools/polymer/test/polymer-opt/FoldSCFIf/pure-arith-multi-if.mlir

@@ -0,0 +1,24 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%a: f32, %b: f32, %c: i1, %d: i1) -> f32 {
+  %0 = scf.if %c -> f32 {
+    %0 = scf.if %d -> f32 {
+      scf.yield %a : f32
+    } else {
+      scf.yield %b : f32
+    }
+    %1 = arith.addf %0, %b : f32
+    scf.yield %1 : f32
+  } else {
+    %1 = arith.mulf %a, %b : f32
+    scf.yield %1 : f32
+  }
+  return %0 : f32
+}
+
+// CHECK: func @foo(%[[a:.*]]: f32, %[[b:.*]]: f32, %[[c:.*]]: i1, %[[d:.*]]: i1) -> f32 
+// CHECK-NEXT:   %[[v0:.*]] = select %[[d]], %[[a]], %[[b]] : f32
+// CHECK-NEXT:   %[[v1:.*]] = arith.addf %[[v0]], %[[b]] : f32
+// CHECK-NEXT:   %[[v2:.*]] = arith.mulf %[[a]], %[[b]] : f32
+// CHECK-NEXT:   %[[v3:.*]] = select %[[c]], %[[v1]], %[[v2]] : f32
+// CHECK-NEXT:   return %[[v3]] : f32

tools/polymer/test/polymer-opt/FoldSCFIf/pure-arith-single-if.mlir

@@ -0,0 +1,18 @@
+// RUN: polymer-opt %s -fold-scf-if | FileCheck %s
+
+func @foo(%a: f32, %b: f32, %c: i1) -> f32 {
+  %0 = scf.if %c -> f32 {
+    %1 = arith.addf %a, %b : f32
+    scf.yield %1 : f32
+  } else {
+    %1 = arith.mulf %a, %b : f32
+    scf.yield %1 : f32
+  }
+  return %0 : f32
+}
+
+// CHECK: func @foo(%[[a:.*]]: f32, %[[b:.*]]: f32, %[[c:.*]]: i1) -> f32 
+// CHECK-NEXT:   %[[v0:.*]] = arith.addf %[[a]], %[[b]] : f32
+// CHECK-NEXT:   %[[v1:.*]] = arith.mulf %[[a]], %[[b]] : f32
+// CHECK-NEXT:   %[[v2:.*]] = select %[[c]], %[[v0]], %[[v1]] : f32
+// CHECK-NEXT:   return %[[v2]] : f32

tools/polymer/tools/polymer-opt/polymer-opt.cc

@@ -6,6 +6,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "polymer/Transforms/ExtractScopStmt.h"
+#include "polymer/Transforms/FoldSCFIf.h"
 #include "polymer/Transforms/LoopAnnotate.h"
 #include "polymer/Transforms/LoopExtract.h"
 #include "polymer/Transforms/PlutoTransform.h"
@@ -59,6 +60,7 @@ int main(int argc, char *argv[]) {
   registerScopStmtOptPasses();
   registerLoopAnnotatePasses();
   registerLoopExtractPasses();
+  registerFoldSCFIfPass();
 
   // Register any pass manager command line options.
   registerMLIRContextCLOptions();