Cleanup HeapReporter command line options (#2783)

* Cleanup HeapReporter command line options

Signed-off-by: philass <[email protected]>

* Address comments

Signed-off-by: philass <[email protected]>

---------

Signed-off-by: philass <[email protected]>

Author: philass

src/Compiler/CompilerOptions.cpp

@@ -72,8 +72,8 @@ bool disableRecomposeOption;                           // onnx-mlir only
 bool enableSimdDataLayout;                             // onnx-mlir only
 bool verifyInputTensors;                               // onnx-mlir only
 bool allowSorting;                                     // onnx-mlir only
-std::string reportHeapBefore;                          // onnx-mlir only
-std::string reportHeapAfter;                           // onnx-mlir only
+std::vector<std::string> reportHeapBefore;             // onnx-mlir only
+std::vector<std::string> reportHeapAfter;              // onnx-mlir only
 std::string modelTag;                                  // onnx-mlir only
 bool enableConvOptPass;                                // onnx-mlir only
 bool disableConstantProp;                              // onnx-mlir only
@@ -470,20 +470,19 @@ static llvm::cl::opt<bool, true> allowSortingOpt("allowSorting",
     llvm::cl::location(allowSorting), llvm::cl::init(true),
     llvm::cl::cat(OnnxMlirOptions));
 
-static llvm::cl::opt<std::string, true> reportHeapBeforeOpt(
-    "report-heap-before",
-    llvm::cl::desc("Comma separated list of names of passes.\n"
-                   "Before each heap statistics are dumped to "
-                   "<output-files-base-path>.heap.log"),
-    llvm::cl::location(reportHeapBefore), llvm::cl::init(""),
-    llvm::cl::cat(OnnxMlirOptions));
-
-static llvm::cl::opt<std::string, true> reportHeapAfterOpt("report-heap-after",
-    llvm::cl::desc("Comma separated list of names of passes.\n"
+static llvm::cl::list<std::string, std::vector<std::string>>
+    reportHeapBeforeOpt("report-heap-before",
+        llvm::cl::desc("A list of names of passes.\n"
+                       "Before each heap statistics are dumped to "
+                       "<output-files-base-path>.heap.log"),
+        llvm::cl::location(reportHeapBefore), llvm::cl::cat(OnnxMlirOptions));
+
+static llvm::cl::list<std::string, std::vector<std::string>> reportHeapAfterOpt(
+    "report-heap-after",
+    llvm::cl::desc("A list of names of passes.\n"
                    "After each heap statistics are dumped to "
                    "<output-files-base-path>.heap.log"),
-    llvm::cl::location(reportHeapAfter), llvm::cl::init(""),
-    llvm::cl::cat(OnnxMlirOptions));
+    llvm::cl::location(reportHeapAfter), llvm::cl::cat(OnnxMlirOptions));
 
 static llvm::cl::opt<std::string, true> modelTagOpt("tag",
     llvm::cl::desc(

src/Compiler/CompilerOptions.hpp

@@ -115,8 +115,8 @@ extern bool disableRecomposeOption;                           // onnx-mlir only
 extern bool enableSimdDataLayout;                             // onnx-mlir only
 extern bool verifyInputTensors;                               // onnx-mlir only
 extern bool allowSorting;                                     // onnx-mlir only
-extern std::string reportHeapBefore;                          // onnx-mlir only
-extern std::string reportHeapAfter;                           // onnx-mlir only
+extern std::vector<std::string> reportHeapBefore;             // onnx-mlir only
+extern std::vector<std::string> reportHeapAfter;              // onnx-mlir only
 extern std::string modelTag;                                  // onnx-mlir only
 extern bool enableConvOptPass;                                // onnx-mlir only
 extern bool disableConstantProp;                              // onnx-mlir only

src/Compiler/HeapReporter.cpp

@@ -12,11 +12,11 @@
 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/raw_ostream.h"
 
-#include <string>
-
 #if defined(__APPLE__)
 #include <unistd.h> // Unsupported on MSVC.
 #endif
@@ -26,13 +26,6 @@ using namespace mlir;
 namespace onnx_mlir {
 
 namespace {
-void splitToSet(StringRef commaSeparated, llvm::StringSet<> &set) {
-  SmallVector<StringRef> splits;
-  commaSeparated.split(splits, ',', /*MaxSplit=*/-1, /*KeepEmpty=*/false);
-  for (StringRef s : splits)
-    set.insert(s.trim());
-}
-
 #if defined(__APPLE__)
 void logMessage(StringRef logFilename, StringRef msg,
     llvm::sys::fs::OpenFlags extraFlags = llvm::sys::fs::OF_None) {
@@ -48,15 +41,15 @@ void logMessage(StringRef logFilename, StringRef msg,
 #endif
 } // namespace
 
-HeapReporter::HeapReporter(
-    std::string logFilename, StringRef beforePasses, StringRef afterPasses)
-    : logFilename(logFilename) {
-  splitToSet(beforePasses, this->beforePassesSet);
-  splitToSet(afterPasses, this->afterPassesSet);
+HeapReporter::HeapReporter(std::string logFilename,
+    std::vector<std::string> beforePasses, std::vector<std::string> afterPasses)
+    : logFilename(logFilename), beforePassesSet(beforePasses),
+      afterPassesSet(afterPasses) {
+
   reportBegin("onnx-mlir heap report"
               "\n--report-heap-before='" +
-              beforePasses.str() + "'\n--report-heap-after='" +
-              afterPasses.str() + "'");
+              llvm::join(beforePasses, ",") + "'\n--report-heap-after='" +
+              llvm::join(afterPasses, ",") + "'");
 }
 
 HeapReporter::~HeapReporter() {}

src/Compiler/HeapReporter.hpp

@@ -12,14 +12,17 @@
 
 #pragma once
 
+#include <string>
+#include <vector>
+
 #include "mlir/Pass/PassInstrumentation.h"
 #include "llvm/ADT/StringSet.h"
 
 namespace onnx_mlir {
 
 struct HeapReporter : public mlir::PassInstrumentation {
-  HeapReporter(std::string logFilename, llvm::StringRef beforePasses,
-      llvm::StringRef afterPasses);
+  HeapReporter(std::string logFilename, std::vector<std::string> beforePasses,
+      std::vector<std::string> afterPasses);
   ~HeapReporter() override;
 
   void runBeforePass(mlir::Pass *pass, mlir::Operation *op) override;