--- a/mlir/lib/Transforms/ViewOpGraph.cpp	Mon May 25 11:55:54 2020 +0900
+++ b/mlir/lib/Transforms/ViewOpGraph.cpp	Tue Jun 08 06:07:14 2021 +0900
@@ -9,8 +9,8 @@
 #include "mlir/Transforms/ViewOpGraph.h"
 #include "PassDetail.h"
 #include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
-#include "mlir/IR/StandardTypes.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace mlir;
@@ -104,10 +104,12 @@
 // PrintOpPass is simple pass to write graph per function.
 // Note: this is a module pass only to avoid interleaving on the same ostream
 // due to multi-threading over functions.
-struct PrintOpPass : public PrintOpBase<PrintOpPass> {
-  explicit PrintOpPass(raw_ostream &os = llvm::errs(), bool short_names = false,
-                       const Twine &title = "")
-      : os(os), title(title.str()), short_names(short_names) {}
+class PrintOpPass : public ViewOpGraphPassBase<PrintOpPass> {
+public:
+  PrintOpPass(raw_ostream &os, bool shortNames, const Twine &title) : os(os) {
+    this->shortNames = shortNames;
+    this->title = title.str();
+  }
 
   std::string getOpName(Operation &op) {
     auto symbolAttr =
@@ -130,10 +132,10 @@
       for (Region &region : op.getRegions()) {
         for (auto indexed_block : llvm::enumerate(region)) {
           // Suffix block number if there are more than 1 block.
-          auto blockName = region.getBlocks().size() == 1
+          auto blockName = llvm::hasSingleElement(region)
                                ? ""
                                : ("__" + llvm::utostr(indexed_block.index()));
-          llvm::WriteGraph(os, &indexed_block.value(), short_names,
+          llvm::WriteGraph(os, &indexed_block.value(), shortNames,
                            Twine(title) + opName + blockName);
         }
       }
@@ -144,9 +146,7 @@
 
 private:
   raw_ostream &os;
-  std::string title;
   int unnamedOpCtr = 0;
-  bool short_names;
 };
 } // namespace