1 files changed, 7 insertions, 10 deletions
diff --git a/mlir/lib/Transforms/ViewFunctionGraph.cpp b/mlir/lib/Transforms/ViewFunctionGraph.cpp
index 14e21770e25..30fae94139f 100644
--- a/mlir/lib/Transforms/ViewFunctionGraph.cpp
+++ b/mlir/lib/Transforms/ViewFunctionGraph.cpp
@@ -73,18 +73,15 @@ void mlir::Function::viewGraph() const {
 }
 
 namespace {
-struct PrintCFGPass : public FunctionPass {
+struct PrintCFGPass : public FunctionPass<PrintCFGPass> {
   PrintCFGPass(llvm::raw_ostream &os = llvm::errs(), bool shortNames = false,
                const llvm::Twine &title = "")
-      : FunctionPass(&PrintCFGPass::passID), os(os), shortNames(shortNames),
-        title(title) {}
-  PassResult runOnFunction(Function *function) override {
-    mlir::writeGraph(os, function, shortNames, title);
+      : os(os), shortNames(shortNames), title(title) {}
+  PassResult runOnFunction() {
+    mlir::writeGraph(os, &getFunction(), shortNames, title);
     return success();
   }
 
-  constexpr static PassID passID = {};
-
 private:
   llvm::raw_ostream &os;
   bool shortNames;
@@ -92,9 +89,9 @@ private:
 };
 } // namespace
 
-FunctionPass *mlir::createPrintCFGGraphPass(llvm::raw_ostream &os,
-                                            bool shortNames,
-                                            const llvm::Twine &title) {
+FunctionPassBase *mlir::createPrintCFGGraphPass(llvm::raw_ostream &os,
+                                                bool shortNames,
+                                                const llvm::Twine &title) {
   return new PrintCFGPass(os, shortNames, title);
 }