[mlir] Add an Observer for profiling actions to a stream. (#67251)

The profile is stored in the Chrome trace event format.

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

Author: antonio-cortes-perez

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;
 };

mlir/include/mlir/Debug/Observers/ActionProfiler.h

@@ -0,0 +1,51 @@
+//===- 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>
+#include <mutex>
+
+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;
+
+  /// A mutex used to guard profiling.
+  std::mutex mutex;
+};
+
+} // namespace tracing
+} // namespace mlir
+
+#endif // MLIR_TRACING_OBSERVERS_ACTIONPROFILER_H

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,

mlir/lib/Debug/Observers/ActionProfiler.cpp

@@ -0,0 +1,63 @@
+//===- 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) {
+  // Create the event.
+  std::string str;
+  llvm::raw_string_ostream event(str);
+  event << "{";
+  event << R"("name": ")" << action->getAction().getTag() << "\", ";
+  event << R"("cat": "PERF", )";
+  event << R"("ph": ")" << phase << "\", ";
+  event << R"("pid": 0, )";
+  event << R"("tid": )" << llvm::get_threadid() << ", ";
+  auto ts = std::chrono::steady_clock::now() - startTime;
+  event << R"("ts": )"
+        << std::chrono::duration_cast<std::chrono::microseconds>(ts).count();
+  if (phase == "B") {
+    event << R"(, "args": {)";
+    event << R"("desc": ")";
+    action->getAction().print(event);
+    event << "\"}";
+  }
+  event << "}";
+
+  // Print the event.
+  std::lock_guard<std::mutex> guard(mutex);
+  if (printComma)
+    os << ",\n";
+  printComma = true;
+  os << event.str();
+  os.flush();
+}

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

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]+}}}