[ExtractTopFunc] avoid setting private to public

kumasento · Oct 27, 2021 · be07e3e · be07e3e
1 parent 70ed446
commit be07e3e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/lib/mlir/Transforms/ExtractTopFunc.cc b/lib/mlir/Transforms/ExtractTopFunc.cc
@@ -128,9 +128,10 @@ struct ExtractTopFuncPass
     unionCallee(f, m, keep);
 
     m.walk([&](FuncOp g) {
-      if (!keep.contains(g) && !keepAll)
-        g.erase();
-      else // Enforce every extracted function to be public
+      if (!keep.contains(g)) {
+        if (!keepAll)
+          g.erase();
+      } else // Enforce every extracted function to be public
         g.setVisibility(SymbolTable::Visibility::Public);
     });
   }

diff --git a/test/mlir/Transforms/ExtractTopFunc/keep-all.mlir b/test/mlir/Transforms/ExtractTopFunc/keep-all.mlir
@@ -0,0 +1,7 @@
+// RUN: phism-opt %s -extract-top-func="name=foo keepall=1" | FileCheck %s
+
+func private @bar() { return }
+func @foo() { return }
+
+// CHECK: func private @bar
+// CHECK: func @foo