Modernize pass instrument

src/Accelerators/NNPA/Compiler/NNPACompilerUtils.cpp

@@ -216,9 +216,12 @@ void addONNXToZHighPasses(mlir::PassManager &pm) {
   if (hasSignatureInstrumentation(onnx_mlir::InstrumentStages::ZHigh))
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentONNXSignaturePass(
         instrumentSignatures, instrumentOnnxNode));
-  if (hasInstrumentation(onnx_mlir::InstrumentStages::ZHigh))
-    pm.addNestedPass<func::FuncOp>(
-        onnx_mlir::createInstrumentPass(instrumentOps, instrumentActions));
+  if (hasInstrumentation(onnx_mlir::InstrumentStages::ZHigh)) {
+    InstrumentPassOptions options;
+    options.instrumentOps = instrumentOps;
+    options.actions = instrumentActions;
+    pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentPass(options));
+  }
 }
 
 void normalizeMemRefsPasses(mlir::PassManager &pm) {
@@ -326,9 +329,13 @@ void addPassesNNPA(mlir::OwningOpRef<mlir::ModuleOp> &module,
         // Omit printing signatures that late.
         assert(false && "Printing signature information at ZLow instrument "
                         "stage is currently unsupported");
-      if (hasInstrumentation(onnx_mlir::InstrumentStages::ZLow))
-        pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentPass(
-            instrumentOps, instrumentControlBits));
+      if (hasInstrumentation(onnx_mlir::InstrumentStages::ZLow)) {
+        InstrumentPassOptions options;
+        options.instrumentOps = instrumentOps;
+        options.actions = instrumentControlBits;
+        pm.addNestedPass<func::FuncOp>(
+            onnx_mlir::createInstrumentPass(options));
+      }
     }
   }
 

src/Compiler/CompilerPasses.cpp

@@ -216,9 +216,12 @@ void addONNXToMLIRPasses(mlir::PassManager &pm, bool targetCPU,
   if (hasSignatureInstrumentation(onnx_mlir::InstrumentStages::Onnx))
     pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentONNXSignaturePass(
         instrumentSignatures, instrumentOnnxNode));
-  if (hasInstrumentation(onnx_mlir::InstrumentStages::Onnx))
-    pm.addNestedPass<func::FuncOp>(
-        onnx_mlir::createInstrumentPass(instrumentOps, instrumentActions));
+  if (hasInstrumentation(onnx_mlir::InstrumentStages::Onnx)) {
+    InstrumentPassOptions options;
+    options.instrumentOps = instrumentOps;
+    options.actions = instrumentActions;
+    pm.addNestedPass<func::FuncOp>(onnx_mlir::createInstrumentPass(options));
+  }
 
   // Convert ONNX to Linalg if requested (either --use-linalg-path or
   // --linalg-ops is specified)

src/Pass/Passes.hpp

@@ -62,10 +62,6 @@ std::unique_ptr<mlir::Pass> createConstPropONNXToONNXPass();
 /// Pass for instrument the ops in specific stage.
 #define GEN_PASS_DECL_INSTRUMENTPASS
 #include "src/Transform/Passes.h.inc"
-// GEN_PASS_DEF method only adds default constructor only,
-// we add custom constructor with multi positional argument explicitly
-std::unique_ptr<mlir::Pass> createInstrumentPass(
-    const std::string &ops, unsigned actions);
 /// Pass for instrument cleanup.
 #define GEN_PASS_DECL_INSTRUMENTCLEANUPPASS
 #include "src/Transform/Passes.h.inc"

src/Transform/InstrumentCleanupPass.cpp

@@ -122,15 +122,3 @@ class InstrumentCleanupPass
   }
 };
 } // namespace onnx_mlir
-
-/*!
- * Create an instrumentation pass.
- */
-// Below is defined by GEN_PASS_DEF in onnx_mlir namespace
-/*
-   namespace onnx_mlir {
-   std::unique_ptr<mlir::Pass> createInstrumentCleanupPass() {
-   return std::make_unique<InstrumentCleanupPass>();
-   }
-   }//onnx_mlir
- */

src/Transform/InstrumentPass.cpp

@@ -51,36 +51,15 @@ class InstrumentPass : public impl::InstrumentPassBase<InstrumentPass> {
 public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InstrumentPass)
 
-  Option<std::string> instrumentOps{*this, "instrument-ops",
-      llvm::cl::desc("Specify regex for ops to be instrumented:\n"
-                     "\"NONE\" or \"\" for no instrument,\n"
-                     "\"regex1,regex2, ...\" for the specified ops.\n"
-                     "e.g. \"onnx.,zhigh.\" for onnx and zhigh ops.\n"
-                     "e.g. \"onnx.Conv\" for onnx Conv ops.\n"),
-      llvm::cl::init("")};
-
-  Option<bool> instrumentBefore{*this, "instrument-before",
-      llvm::cl::desc("insert instrument before op"), llvm::cl::init(false)};
-
-  Option<bool> instrumentAfter{*this, "instrument-after",
-      llvm::cl::desc("insert instrument after op"), llvm::cl::init(false)};
-
-  Option<bool> reportTime{*this, "report-time",
-      llvm::cl::desc("instrument runtime reports time usage"),
-      llvm::cl::init(false)};
-
-  Option<bool> reportMemory{*this, "report-memory",
-      llvm::cl::desc("instrument runtime reports memory usage"),
-      llvm::cl::init(false)};
-
   InstrumentPass() : allowedOps(/*emptyIsNone*/ true){};
   InstrumentPass(const InstrumentPass &pass)
       : impl::InstrumentPassBase<InstrumentPass>(),
         allowedOps(/*emptyIsNone*/ true) {}
-  InstrumentPass(const std::string &ops, unsigned actions)
-      : allowedOps(/*emptyIsNone*/ true) {
-    this->instrumentOps = ops;
-    unsigned long long tag = actions;
+  InstrumentPass(const InstrumentPassOptions &options)
+      : impl::InstrumentPassBase<InstrumentPass>(options),
+        allowedOps(/*emptyIsNone*/ true) {
+    this->instrumentOps = options.instrumentOps;
+    unsigned long long tag = options.actions;
     this->instrumentBefore = IS_INSTRUMENT_BEFORE_OP(tag);
     this->instrumentAfter = IS_INSTRUMENT_AFTER_OP(tag);
     this->reportTime = IS_INSTRUMENT_REPORT_TIME(tag);
@@ -91,10 +70,6 @@ class InstrumentPass : public impl::InstrumentPassBase<InstrumentPass> {
   EnableByRegexOption allowedOps;
 
 public:
-  StringRef getArgument() const override { return "instrument"; }
-
-  StringRef getDescription() const override { return "instrument on ops."; }
-
   // merge all action options into a bitset
   // used to create tags for instrumentation ops
   uint64_t actions() const {
@@ -180,20 +155,3 @@ class InstrumentPass : public impl::InstrumentPassBase<InstrumentPass> {
   }
 };
 } // namespace onnx_mlir
-
-/*!
- * Create an instrumentation pass.
- */
-namespace onnx_mlir {
-// Below is defined by GEN_PASS_DEF in onnx_mlir namespace
-/*
-std::unique_ptr<mlir::Pass> createInstrument() {
-  return std::make_unique<InstrumentPass>();
-}
-*/
-
-std::unique_ptr<mlir::Pass> createInstrumentPass(
-    const std::string &ops, unsigned actions) {
-  return std::make_unique<InstrumentPass>(ops, actions);
-}
-} // namespace onnx_mlir

src/Transform/Passes.td

@@ -23,6 +23,34 @@ def InstrumentPass : Pass<"instrument", "func::FuncOp"> {
   let description = [{
     instrument on ops.
   }];
+  let options = [
+    Option<"instrumentOps", "instrument-ops", "::std::string",
+            /*default=*/"\"\"", [{Specify regex for ops to be instrumented:\n"
+                                "\"NONE\" or \"\" for no instrument,\n"
+                                "\"regex1,regex2, ...\" for the specified ops.\n"
+                                "e.g. \"onnx.,zhigh.\" for onnx and zhigh ops.\n"
+                                "e.g. \"onnx.Conv\" for onnx Conv ops.\n}]>,
+
+    Option<"instrumentBefore", "instrument-before", "bool",
+            /*default=*/"false", "insert instrument before op">,
+
+    Option<"instrumentAfter", "instrument-after", "bool",
+            /*default=*/"false", "insert instrument after op">,
+
+    Option<"reportTime", "report-time", "bool",
+            /*default=*/"false", "instrument runtime reports time usage">,
+
+    Option<"reportMemory", "report-memory", "bool",
+            /*default=*/"false", "instrument runtime reports memory usage">,
+
+    Option<"actions", "actions", "unsigned",
+            /*default=*/"0", [{Bit-tag for Instrument Actions\n"
+                             "InstrumentBeforeOp = 0x0\n"
+                             "InstrumentAfterOp = 0x1\n"
+                             "InstrumentReportTime = 0x2\n"
+                             "InstrumentReportMemory = 0x3\n"
+                             "InstrumentInit = 0x4\n}]>
+  ];
 }
 
 def InstrumentCleanupPass : Pass<"instrument-cleanup", "func::FuncOp"> {