[AnnotateScop] new pass annotates scop.ignored

include/polymer/CMakeLists.txt

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

include/polymer/Transforms/CMakeLists.txt

@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Transforms)
+add_public_tablegen_target(PolymerTransformsIncGen)

include/polymer/Transforms/Passes.h

@@ -0,0 +1,17 @@
+//===- Passes.h - Include Tblgen pass defs ------------C++-===//
+#ifndef POLYMER_TRANSFORMS_PASSES_H
+#define POLYMER_TRANSFORMS_PASSES_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace polymer {
+
+std::unique_ptr<mlir::OperationPass<mlir::FuncOp>> createAnnotateScopPass();
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "polymer/Transforms/Passes.h.inc"
+
+} // namespace polymer
+
+#endif

include/polymer/Transforms/Passes.td

@@ -0,0 +1,21 @@
+#ifndef POLYMER_TRANSFORMS_PASSES
+#define POLYMER_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+include "mlir/Rewrite/PassUtil.td"
+
+def AnnotateScop : FunctionPass<"annotate-scop"> {
+  let summary = "Annotate scop.ignore to unselected functions.";
+  let constructor = "polymer::createAnnotateScopPass()";
+
+  let options = [
+    ListOption<
+      "includedFunctions",
+      "functions",
+      "std::string",   
+      "A list of functions to be included.",
+      "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
+  ];
+}
+
+#endif

lib/Transforms/AnnotateScop.cc

@@ -0,0 +1,34 @@
+//===- AnnotateScop.cc --------------------------------------*- C++ -*-===//
+
+#include "PassDetail.h"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Transforms/LoopUtils.h"
+#include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/RegionUtils.h"
+#include "llvm/Support/Debug.h"
+
+using namespace mlir;
+using namespace llvm;
+using namespace polymer;
+
+namespace {
+struct AnnotateScop : public polymer::AnnotateScopBase<AnnotateScop> {
+  void runOnFunction() override {
+    FuncOp f = getOperation();
+    OpBuilder b(f.getContext());
+
+    for (auto &name : includedFunctions)
+      if (name == f.getName())
+        return;
+
+    f->setAttr("scop.ignored", b.getUnitAttr());
+  } // namespace
+};
+} // namespace
+
+std::unique_ptr<OperationPass<FuncOp>> polymer::createAnnotateScopPass() {
+  return std::make_unique<AnnotateScop>();
+}

lib/Transforms/CMakeLists.txt

@@ -6,12 +6,13 @@ add_mlir_conversion_library(PolymerTransforms
   LoopAnnotate.cc
   LoopExtract.cc
   FoldSCFIf.cc
+  AnnotateScop.cc
 
   ADDITIONAL_HEADER_DIRS
   "${POLYMER_MAIN_INCLUDE_DIR}/polymer/Transforms"
 
   DEPENDS
-  ${POLYMER_TRANSFORM_DEPENDS}
+  PolymerTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine

lib/Transforms/ExtractScopStmt.cc

@@ -510,7 +510,11 @@ class ExtractScopStmtPass
     OpBuilder b(m.getContext());
 
     SmallVector<mlir::FuncOp, 4> funcs;
-    m.walk([&](mlir::FuncOp f) { funcs.push_back(f); });
+    m.walk([&](mlir::FuncOp f) {
+      if (f->hasAttr("scop.ignored"))
+        return;
+      funcs.push_back(f);
+    });
 
     unsigned numCallees = 0;
     for (mlir::FuncOp f : funcs) {

lib/Transforms/PassDetail.h

@@ -0,0 +1,14 @@
+//===- PassDetail.h - Transforms Pass class details -------------*- C++ -*-===//
+
+#ifndef POLYMER_TRANSFORMS_PASSDETAIL_H_
+#define POLYMER_TRANSFORMS_PASSDETAIL_H_
+
+#include "mlir/Pass/Pass.h"
+#include "polymer/Transforms/Passes.h"
+
+namespace polymer {
+#define GEN_PASS_CLASSES
+#include "polymer/Transforms/Passes.h.inc"
+} // namespace polymer
+
+#endif

lib/Transforms/PlutoTransform.cc

@@ -156,7 +156,7 @@ class PlutoTransformPass
     llvm::DenseMap<mlir::FuncOp, mlir::FuncOp> funcMap;
 
     m.walk([&](mlir::FuncOp f) {
-      if (!f->getAttr("scop.stmt"))
+      if (!f->getAttr("scop.stmt") && !f->hasAttr("scop.ignored"))
         funcOps.push_back(f);
     });
 

lib/Transforms/Reg2Mem.cc

@@ -315,6 +315,9 @@ class RegToMemPass
     mlir::FuncOp f = getOperation();
     auto builder = OpBuilder(f.getContext());
 
+    if (f->hasAttr("scop.ignored"))
+      return;
+
     separateAffineIfBlocks(f, builder);
     demoteRegisterToMemory(f, builder);
   }

test/polymer-opt/AnnotateScop/annotate-single.mlir

@@ -0,0 +1,12 @@
+// RUN: polymer-opt %s -annotate-scop="functions=foo" | FileCheck %s
+
+func @foo() {
+  return
+} 
+
+func @bar() {
+  return
+} 
+
+// CHECK: func @foo() {
+// CHECK: func @bar() attributes {scop.ignored} {

test/polymer-opt/Reg2Mem/ignored.mlir

@@ -0,0 +1,14 @@
+// RUN: polymer-opt %s -reg2mem | FileCheck %s
+func @foo(%A: memref<?xf32>) attributes {scop.ignored} {
+  %0 = affine.load %A[0] : memref<?xf32>
+  affine.for %i = 1 to 30 {
+    affine.store %0, %A[%i] : memref<?xf32>
+  }
+
+  return
+}
+
+// CHECK: func @foo
+// CHECK-NEXT:   %[[v0:.*]] = affine.load 
+// CHECK-NEXT:   affine.for %{{.*}} = 1 to 30 
+// CHECK-NEXT:     affine.store %[[v0]]

tools/polymer-opt/polymer-opt.cc

@@ -9,6 +9,7 @@
 #include "polymer/Transforms/FoldSCFIf.h"
 #include "polymer/Transforms/LoopAnnotate.h"
 #include "polymer/Transforms/LoopExtract.h"
+#include "polymer/Transforms/Passes.h"
 #include "polymer/Transforms/PlutoTransform.h"
 #include "polymer/Transforms/Reg2Mem.h"
 #include "polymer/Transforms/ScopStmtOpt.h"
@@ -61,6 +62,7 @@ int main(int argc, char *argv[]) {
   registerLoopAnnotatePasses();
   registerLoopExtractPasses();
   registerFoldSCFIfPass();
+  registerAnnotateScopPass();
 
   // Register any pass manager command line options.
   registerMLIRContextCLOptions();