Extract top function pass

Author: kumasento

.gitmodules

@@ -7,3 +7,6 @@
 [submodule "vitis-hls"]
 	path = vitis-hls
 	url = https://github.com/Xilinx/HLS
+[submodule "example/polybench"]
+	path = example/polybench
+	url = git@github.com:MatthiasJReisinger/PolyBenchC-4.2.1.git

CMakeLists.txt

@@ -41,7 +41,7 @@ set(LLVM_LIT_ARGS "${LIT_ARGS_DEFAULT}" CACHE STRING "Default options for lit")
 # ------------------------------------------------- Dependencies
 
 # ------------------------------------------------- This project
-set(PHISM_MAIN_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/include" ) # --src-root
+set(PHISM_MAIN_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/include" ) # --src-root
 set(PHISM_INCLUDE_DIR "${CMAKE_CURRENT_BINARY_DIR}/include" ) # --includedir
 
 set(PHISM_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
@@ -65,3 +65,4 @@ include_directories("${PROJECT_BINARY_DIR}/include")
 
 add_subdirectory(lib)
 add_subdirectory(test)
+add_subdirectory(tools)

include/phism/mlir/Transforms/PhismTransforms.h

@@ -0,0 +1,15 @@
+//===- PhismTransforms.h ----------------------------------------*- C++ -*-===//
+//
+// This file declares all the registration interfaces for Phism passes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PHISM_MLIR_TRANSFORMS_PHISMTRANSFORMS_H
+#define PHISM_MLIR_TRANSFORMS_PHISMTRANSFORMS_H
+
+namespace phism {
+void registerExtractTopFuncPass();
+void registerAllPhismPasses();
+} // namespace phism
+
+#endif

lib/CMakeLists.txt

@@ -1 +1,2 @@
 add_subdirectory(llvm)
+add_subdirectory(mlir)

lib/mlir/CMakeLists.txt

@@ -0,0 +1 @@
+add_subdirectory(Transforms)

lib/mlir/Transforms/CMakeLists.txt

@@ -0,0 +1,18 @@
+add_mlir_library(PhismTransforms
+  ExtractTopFunc.cc
+  PhismTransforms.cc
+
+  ADDITIONAL_HEADER_DIRS
+  "${PHISM_MAIN_INCLUDE_DIR}/phism/mlir/Transforms"  
+
+  LINK_LIBS PUBLIC
+  MLIRAffine
+  MLIRAnalysis
+  MLIRPass
+  MLIRTransforms
+  MLIRTransformUtils
+  MLIRIR
+  MLIRStandard
+  MLIRSupport
+  MLIRAffineToStandard
+)

lib/mlir/Transforms/ExtractTopFunc.cc

@@ -0,0 +1,97 @@
+//===- ExtractTopFunc.cc ----------------------------------------*- C++ -*-===//
+//
+// This file implements a pass that extract the specified top function out of
+// the given module.
+//
+//===----------------------------------------------------------------------===//
+
+#include "phism/mlir/Transforms/PhismTransforms.h"
+
+#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.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/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "mlir/Transforms/Utils.h"
+
+#include "llvm/ADT/SetVector.h"
+
+using namespace mlir;
+using namespace llvm;
+using namespace phism;
+
+/// Union the given function and all the others it calls into 'keep'.
+static void unionCallee(FuncOp f, ModuleOp m,
+                        SmallPtrSetImpl<Operation *> &keep) {
+  std::deque<Operation *> worklist;
+  worklist.push_back(f);
+  keep.insert(f);
+
+  // A breadth first search on the call graph, starting from f.
+  while (!worklist.empty()) {
+    Operation *g = worklist.front();
+    worklist.pop_front();
+
+    g->walk([&](Operation *op) {
+      if (CallOp callOp = dyn_cast<CallOp>(op)) {
+        FuncOp callee =
+            dyn_cast_or_null<FuncOp>(m.lookupSymbol(callOp.getCallee()));
+        assert(callee && "A callee cannot be resolved in the module.");
+
+        if (!keep.contains(callee)) {
+          worklist.push_back(callee);
+          keep.insert(callee);
+        }
+      }
+    });
+  }
+}
+
+namespace {
+
+/// Extract the specified top function out of its module.
+struct ExtractTopFuncPass
+    : public PassWrapper<ExtractTopFuncPass, OperationPass<ModuleOp>> {
+
+  ExtractTopFuncPass() = default;
+  ExtractTopFuncPass(const ExtractTopFuncPass &pass) {}
+
+  Option<std::string> topFuncName{
+      *this, "name",
+      llvm::cl::desc("Name of the top function to be extracted.")};
+
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+    OpBuilder b(m.getContext());
+
+    assert(!topFuncName.empty() && "-name should be specified.");
+
+    FuncOp f = dyn_cast_or_null<FuncOp>(m.lookupSymbol(topFuncName));
+    assert(f && "Given name cannot be found in the module as a FuncOp.");
+
+    SmallPtrSet<Operation *, 4> keep;
+    unionCallee(f, m, keep);
+
+    m.walk([&](FuncOp g) {
+      if (!keep.contains(g))
+        g.erase();
+    });
+  }
+};
+
+} // namespace
+
+void phism::registerExtractTopFuncPass() {
+  PassRegistration<ExtractTopFuncPass>(
+      "extract-top-func", "Extract the top function out of its module.");
+}

lib/mlir/Transforms/PhismTransforms.cc

@@ -0,0 +1,13 @@
+//===- PhismTransforms.cc ---------------------------------------*- C++ -*-===//
+//
+// This file implements the registration interfaces for all Phism passes.
+//
+//===----------------------------------------------------------------------===//
+
+#include "phism/mlir/Transforms/PhismTransforms.h"
+
+namespace phism {
+
+void registerAllPhismPasses() { registerExtractTopFuncPass(); }
+
+} // namespace phism

llvm

@@ -7,7 +7,9 @@ configure_lit_site_cfg(
 
 set(PHISM_TEST_DEPENDS
   FileCheck count not
-  # phism-opt
+  phism-opt
+  # LLVM passes
+  VhlsLLVMRewriter
   )
 
 add_lit_testsuite(check-phism "Running the Phism regression tests"