diff --git a/lib/mlir/Transforms/ExtractTopFunc.cc b/lib/mlir/Transforms/ExtractTopFunc.cc
index f800fb6c175..e3283a517c8 100644
--- 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
new file mode 100644
index 00000000000..dca6c62b9d0
--- /dev/null
+++ 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