[mlir] Add an Observer for profiling actions to a stream.

[antonio-cortes-perez]

The profile is stored in the Chrome trace event format.

Added the --profile-action-to= option to mlir-opt.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Changes

The profile is stored in the Chrome trace event format.

Added the --profile-action-to=<file> option to mlir-opt.

---

Full diff: https://github.com/llvm/llvm-project/pull/67251.diff

6 Files Affected:

• (modified) mlir/include/mlir/Debug/CLOptionsSetup.h (+6)
• (added) mlir/include/mlir/Debug/Observers/ActionProfiler.h (+47)
• (modified) mlir/lib/Debug/CLOptionsSetup.cpp (+28)
• (added) mlir/lib/Debug/Observers/ActionProfiler.cpp (+55)
• (modified) mlir/lib/Debug/Observers/CMakeLists.txt (+1)
• (added) mlir/test/Pass/action-profiler.mlir (+6)

diff --git a/mlir/include/mlir/Debug/CLOptionsSetup.h b/mlir/include/mlir/Debug/CLOptionsSetup.h
index ab8e7b436a1e9f6..a30ffe0ff1af485 100644
--- a/mlir/include/mlir/Debug/CLOptionsSetup.h
+++ b/mlir/include/mlir/Debug/CLOptionsSetup.h
@@ -51,6 +51,9 @@ class DebugConfig {
   /// Get the filename to use for logging actions.
   StringRef getLogActionsTo() const { return logActionsToFlag; }
 
+  /// Get the filename to use for profiling actions.
+  StringRef getProfileActionsTo() const { return profileActionsToFlag; }
+
   /// Set a location breakpoint manager to filter out action logging based on
   /// the attached IR location in the Action context. Ownership stays with the
   /// caller.
@@ -71,6 +74,9 @@ class DebugConfig {
   /// Log action execution to the given file (or "-" for stdout)
   std::string logActionsToFlag;
 
+  /// Profile action execution to the given file (or "-" for stdout)
+  std::string profileActionsToFlag;
+
   /// Location Breakpoints to filter the action logging.
   std::vector<tracing::BreakpointManager *> logActionLocationFilter;
 };
diff --git a/mlir/include/mlir/Debug/Observers/ActionProfiler.h b/mlir/include/mlir/Debug/Observers/ActionProfiler.h
new file mode 100644
index 000000000000000..83a7949f6ac0128
--- /dev/null
+++ b/mlir/include/mlir/Debug/Observers/ActionProfiler.h
@@ -0,0 +1,47 @@
+//===- ActionProfiler.h -  Profiling Actions *- C++ -*-=======================//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRACING_OBSERVERS_ACTIONPROFILER_H
+#define MLIR_TRACING_OBSERVERS_ACTIONPROFILER_H
+
+#include "mlir/Debug/ExecutionContext.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <chrono>
+
+namespace mlir {
+namespace tracing {
+
+/// This class defines an observer that profiles events before and after
+/// execution on the provided stream. The events are stored in the Chrome trace
+/// event format.
+struct ActionProfiler : public ExecutionContext::Observer {
+  ActionProfiler(raw_ostream &os)
+      : os(os), startTime(std::chrono::steady_clock::now()) {
+    os << "[";
+  }
+
+  ~ActionProfiler() override { os << "]"; }
+
+  void beforeExecute(const ActionActiveStack *action, Breakpoint *breakpoint,
+                     bool willExecute) override;
+  void afterExecute(const ActionActiveStack *action) override;
+
+private:
+  void print(const ActionActiveStack *action, llvm::StringRef phase);
+
+  raw_ostream &os;
+  std::chrono::time_point<std::chrono::steady_clock> startTime;
+  bool printComma = false;
+};
+
+} // namespace tracing
+} // namespace mlir
+
+#endif // MLIR_TRACING_OBSERVERS_ACTIONPROFILER_H
diff --git a/mlir/lib/Debug/CLOptionsSetup.cpp b/mlir/lib/Debug/CLOptionsSetup.cpp
index 24d0bd8808261da..76330b2a0a92f1e 100644
--- a/mlir/lib/Debug/CLOptionsSetup.cpp
+++ b/mlir/lib/Debug/CLOptionsSetup.cpp
@@ -12,6 +12,7 @@
 #include "mlir/Debug/DebuggerExecutionContextHook.h"
 #include "mlir/Debug/ExecutionContext.h"
 #include "mlir/Debug/Observers/ActionLogging.h"
+#include "mlir/Debug/Observers/ActionProfiler.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/FileUtilities.h"
 #include "llvm/Support/CommandLine.h"
@@ -30,6 +31,12 @@ struct DebugConfigCLOptions : public DebugConfig {
                  " '-' is passed"),
         cl::location(logActionsToFlag)};
 
+    static cl::opt<std::string, /*ExternalStorage=*/true> profileActionsTo{
+        "profile-actions-to",
+        cl::desc("Profile action execution to a file, or stderr if "
+                 " '-' is passed"),
+        cl::location(profileActionsToFlag)};
+
     static cl::list<std::string> logActionLocationFilter(
         "log-mlir-actions-filter",
         cl::desc(
@@ -71,6 +78,7 @@ class InstallDebugHandler::Impl {
 public:
   Impl(MLIRContext &context, const DebugConfig &config) {
     if (config.getLogActionsTo().empty() &&
+        config.getProfileActionsTo().empty() &&
         !config.isDebuggerActionHookEnabled()) {
       if (tracing::DebugCounter::isActivated())
         context.registerActionHandler(tracing::DebugCounter());
@@ -97,6 +105,24 @@ class InstallDebugHandler::Impl {
         actionLogger->addBreakpointManager(locationBreakpoint);
       executionContext.registerObserver(actionLogger.get());
     }
+
+    if (!config.getProfileActionsTo().empty()) {
+      std::string errorMessage;
+      profileActionsFile =
+          openOutputFile(config.getProfileActionsTo(), &errorMessage);
+      if (!profileActionsFile) {
+        emitError(UnknownLoc::get(&context),
+                  "Opening file for --profile-actions-to failed: ")
+            << errorMessage << "\n";
+        return;
+      }
+      profileActionsFile->keep();
+      raw_fd_ostream &profileActionsStream = profileActionsFile->os();
+      actionProfiler =
+          std::make_unique<tracing::ActionProfiler>(profileActionsStream);
+      executionContext.registerObserver(actionProfiler.get());
+    }
+
     if (config.isDebuggerActionHookEnabled()) {
       errs() << " (with Debugger hook)";
       setupDebuggerExecutionContextHook(executionContext);
@@ -111,6 +137,8 @@ class InstallDebugHandler::Impl {
   std::unique_ptr<tracing::ActionLogger> actionLogger;
   std::vector<std::unique_ptr<tracing::FileLineColLocBreakpoint>>
       locationBreakpoints;
+  std::unique_ptr<ToolOutputFile> profileActionsFile;
+  std::unique_ptr<tracing::ActionProfiler> actionProfiler;
 };
 
 InstallDebugHandler::InstallDebugHandler(MLIRContext &context,
diff --git a/mlir/lib/Debug/Observers/ActionProfiler.cpp b/mlir/lib/Debug/Observers/ActionProfiler.cpp
new file mode 100644
index 000000000000000..8daffb68def4395
--- /dev/null
+++ b/mlir/lib/Debug/Observers/ActionProfiler.cpp
@@ -0,0 +1,55 @@
+//===- ActionProfiler.cpp -  Profiling Actions *- C++ -*-=====================//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Debug/Observers/ActionProfiler.h"
+#include "mlir/Debug/BreakpointManager.h"
+#include "mlir/IR/Action.h"
+#include "mlir/Rewrite/PatternApplicator.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Threading.h"
+#include "llvm/Support/raw_ostream.h"
+#include <chrono>
+
+using namespace mlir;
+using namespace mlir::tracing;
+
+//===----------------------------------------------------------------------===//
+// ActionProfiler
+//===----------------------------------------------------------------------===//
+void ActionProfiler::beforeExecute(const ActionActiveStack *action,
+                                   Breakpoint *breakpoint, bool willExecute) {
+  print(action, "B"); // begin event.
+}
+
+void ActionProfiler::afterExecute(const ActionActiveStack *action) {
+  print(action, "E"); // end event.
+}
+
+// Print an event in JSON format.
+void ActionProfiler::print(const ActionActiveStack *action,
+                           llvm::StringRef phase) {
+  if (printComma)
+    os << ",\n";
+  printComma = true;
+  os << "{";
+  os << R"("name": ")" << action->getAction().getTag() << "\", ";
+  os << R"("cat": "PERF", )";
+  os << R"("ph": ")" << phase << "\", ";
+  os << R"("pid": 0, )";
+  os << R"("tid": )" << llvm::get_threadid() << ", ";
+  auto ts = std::chrono::steady_clock::now() - startTime;
+  os << R"("ts": )"
+     << std::chrono::duration_cast<std::chrono::microseconds>(ts).count();
+  if (phase == "B") {
+    os << R"(, "args": {)";
+    os << R"("desc": ")";
+    action->getAction().print(os);
+    os << "\"}";
+  }
+  os << "}";
+}
diff --git a/mlir/lib/Debug/Observers/CMakeLists.txt b/mlir/lib/Debug/Observers/CMakeLists.txt
index 0a4261b0e45cc1e..78d8038f6f1265d 100644
--- a/mlir/lib/Debug/Observers/CMakeLists.txt
+++ b/mlir/lib/Debug/Observers/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_mlir_library(MLIRObservers
   ActionLogging.cpp
+  ActionProfiler.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Debug/Observers
diff --git a/mlir/test/Pass/action-profiler.mlir b/mlir/test/Pass/action-profiler.mlir
new file mode 100644
index 000000000000000..7735041d03c4a97
--- /dev/null
+++ b/mlir/test/Pass/action-profiler.mlir
@@ -0,0 +1,6 @@
+// RUN: mlir-opt %s --profile-actions-to=- -test-stats-pass -test-module-pass | FileCheck %s
+
+// CHECK: {"name": "pass-execution", "cat": "PERF", "ph": "B", "pid": 0, "tid": {{[0-9]+}}, "ts": {{[0-9]+}}, "args": {"desc": "`pass-execution` running `{{.*}}TestStatisticPass` on Operation `builtin.module`"}},
+// CHECK-NEXT: {"name": "pass-execution", "cat": "PERF", "ph": "E", "pid": 0, "tid": {{[0-9]+}}, "ts": {{[0-9]+}}},
+// CHECK-NEXT: {"name": "pass-execution", "cat": "PERF", "ph": "B", "pid": 0, "tid": {{[0-9]+}}, "ts": {{[0-9]+}}, "args": {"desc": "`pass-execution` running `{{.*}}TestModulePass` on Operation `builtin.module`"}},
+// CHECK-NEXT: {"name": "pass-execution", "cat": "PERF", "ph": "E", "pid": 0, "tid": {{[0-9]+}}, "ts": {{[0-9]+}}}

[antonio-cortes-perez]

Thanks for the review @joker-eph ! If everything looks OK, could you please help me to merge it?