Add an Observer for logging actions application to a stream

[mehdi_amini]

Summary

Integrate the tracing::ExecutionContext() into mlir-opt with a new
--log-action-to=<file> option to demonstrate the feature.

[mehdi_amini]

rebase

[Mogball]

mlir/include/mlir/Debug/Observers/ActionLogging.h

21

fill the doc? :P

41

newline

mlir/lib/Debug/Observers/ActionLogging.cpp

21

Is this the canonical way to do something like this? LLVM has no utilities?

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

149

should this emitError?

[rriddle]

mlir/include/mlir/Debug/Counter.h

50

mlir/lib/Debug/Observers/ActionLogging.cpp

20

The use of static atomic here is kind of sketchy, why is this necessary?

21

llvm::get_threadid()

mlir/lib/Pass/PassDetail.h

19

Can you shift this to the bottom of the struct, with the other fields?

23–25

I wonder if this would be cleaner using llvm::formatv

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

42

Are these two not already included?

254

[mehdi_amini]

Address comments

mlir/lib/Debug/Observers/ActionLogging.cpp

20

It was necessary because of thread_local int64_t tid = thread_counter++; is a threaded initialization I believe.

21

Replaced with llvm::get_threadid, but note that it makes the output less deterministic, I kind of liked having threads numbered from 0 to X

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

42

VSCode auto-added includes...

[jpienaar]

mlir/include/mlir/Debug/Counter.h

51

OOC why not just isActivated ? (the class is already DebugCounter and above the form is addCounter())

mlir/include/mlir/Debug/Observers/ActionLogging.h

2

Outdated header?

mlir/lib/Debug/Observers/ActionLogging.cpp

2

Same

mlir/lib/Debug/Observers/CMakeLists.txt

2

So this library will contain multiple kind of observers?

mlir/lib/Pass/PassDetail.h

22

Does this have to be inline in header?

22

Mmm, this is quoting using and " , could we just use throughout?

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

144

Error form should be sentence fragments (https://llvm.org/docs/CodingStandards.html#error-and-warning-messages)

164

Are these correct destruction order to avoid dangling references ?

[rriddle]

LGTM with Jacques comments addressed

mlir/lib/Debug/Observers/ActionLogging.cpp

21

You can also use llvm::get_thread_name, which can have a more user friendly name for a thread (e.g. the llvm worker threads usually have llvm-worker-<index>.

mlir/lib/Pass/Pass.cpp

466–476

passFailed won't be initialized here if the action doesn't run, this should be set to false.

[mehdi_amini]

mlir/lib/Debug/Observers/ActionLogging.cpp

21

Ah good point! That was added recently to the LLVM thread pool I believe, I'll do that :)

[mehdi_amini]

Address comment

[mehdi_amini]

minor cleanup

[mehdi_amini]

clang-format

[mehdi_amini]

Minor fix to PAssExecutionAction

mlir/include/mlir/Debug/Counter.h

public:
   void operator()(llvm::function_ref<void()> transform, const Action &action);
 
   /// Print the counters that have been registered with this instance to the
   /// provided output stream.
   void print(raw_ostream &os) const;
 
   /// Register the command line options for debug counters.
   static void registerCLOptions();
+  /// Returns true if any of the CL options is activated.
+  static bool isDebugCounterActivated();
 
 private:
   // Returns true if the next action matching this tag should be executed.
   bool shouldExecute(StringRef tag);
 
   /// Apply the registered CL options to this debug counter instance.
   void applyCLOptions();
 

mlir/include/mlir/Debug/Observers/ActionLogging.h

+//===- ExecutionContext.h -  Execution Context Support *- 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_ACTIONLOGGING_H
+#define MLIR_TRACING_OBSERVERS_ACTIONLOGGING_H
+
+#include "mlir/Debug/BreakpointManager.h"
+#include "mlir/Debug/ExecutionContext.h"
+#include "mlir/IR/Action.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace mlir {
+namespace tracing {
+
+/// This class...
+struct ActionLogger : public ExecutionContext::Observer {
+  ActionLogger(raw_ostream &os, bool printActions = true,
+               bool printBreakpoints = true)
+      : os(os), printActions(printActions), printBreakpoints(printBreakpoints) {
+  }
+
+  void beforeExecute(const ActionActiveStack *action, Breakpoint *breakpoint,
+                     bool willExecute) override;
+  void afterExecute(const ActionActiveStack *action) override;
+
+private:
+  raw_ostream &os;
+  bool printActions;
+  bool printBreakpoints;
+};
+
+} // namespace tracing
+} // namespace mlir
+
+#endif // MLIR_TRACING_OBSERVERS_ACTIONLOGGING_H
 No newline at end of file

mlir/include/mlir/IR/MLIRContext.h

void dispatch(function_ref<void()> actionFn, const tracing::Action &action) {
     else
       actionFn();
   }
 
   /// Dispatch the provided action to the handler if any, or just execute it.
   template <typename ActionTy, typename... Args>
   void dispatch(function_ref<void()> actionFn, Args &&...args) {
     if (LLVM_UNLIKELY(hasActionHandler()))
-      dispatchInteral(actionFn, ActionTy(std::forward<Args>(args)...));
+      dispatchInternal(actionFn, ActionTy(std::forward<Args>(args)...));
     else
       actionFn();
   }
 
   /// These APIs are tracking whether the context will be used in a
   /// multithreading environment: this has no effect other than enabling
   /// assertions on misuses of some APIs.
   void enterMultiThreadedExecution();

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

public:
 
   /// Set the output format to bytecode instead of textual IR.
   MlirOptMainConfig &emitBytecode(bool emit) {
     emitBytecodeFlag = emit;
     return *this;
   }
   bool shouldEmitBytecode() const { return emitBytecodeFlag; }
 
+  /// Set the filename to use for logging actions, use "-" for stdout.
+  MlirOptMainConfig &logActionsTo(StringRef filename) {
+    logActionsToFlag = filename;
+    return *this;
+  }
+  /// Get the filename to use for logging actions.
+  StringRef getLogActionsTo() const { return logActionsToFlag; }
+
   /// Set the callback to populate the pass manager.
   MlirOptMainConfig &
   setPassPipelineSetupFn(std::function<LogicalResult(PassManager &)> callback) {
     passPipelineCallback = std::move(callback);
     return *this;
   }
 
   /// Set the parser to use to populate the pass manager.
protected:
   bool allowUnregisteredDialectsFlag = false;
 
   /// Print the pipeline that will be run.
   bool dumpPassPipelineFlag = false;
 
   /// Emit bytecode instead of textual assembly when generating output.
   bool emitBytecodeFlag = false;
 
+  /// Log action execution to the given file (or "-" for stdout)
+  std::string logActionsToFlag;
+
   /// The callback to populate the pass manager.
   std::function<LogicalResult(PassManager &)> passPipelineCallback;
 
   /// Deprecated.
   bool preloadDialectsInContextFlag = false;
 
   /// Show the registered dialects before trying to load the input file.
   bool showDialectsFlag = false;

mlir/lib/Debug/CMakeLists.txt

+add_subdirectory(Observers)
+
 add_mlir_library(MLIRDebug
   DebugCounter.cpp
   ExecutionContext.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Debug
 
   LINK_LIBS PUBLIC
   ${LLVM_PTHREAD_LIB}
   MLIRIR
   MLIRSupport
 )
+

mlir/lib/Debug/DebugCounter.cpp

 /// constructing a DebugCounter for initialization.
 void DebugCounter::registerCLOptions() {
 #ifndef NDEBUG
   // Make sure that the options struct has been initialized.
   *clOptions;
 #endif
 }
 
+bool DebugCounter::isDebugCounterActivated() {
+  return clOptions->counters.getNumOccurrences() ||
+         clOptions->printCounterInfo.getNumOccurrences();
+}
+
 // This is called by the command line parser when it sees a value for the
 // debug-counter option defined above.
 void DebugCounter::applyCLOptions() {
   if (!clOptions.isConstructed())
     return;
 
   for (StringRef arg : clOptions->counters) {
     if (arg.empty())

mlir/lib/Debug/Observers/ActionLogging.cpp

+//===- ExecutionContext.cpp - Debug Execution Context Support -------------===//
+//
+// 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/ActionLogging.h"
+#include <sstream>
+#include <thread>
+
+using namespace mlir;
+using namespace mlir::tracing;
+
+//===----------------------------------------------------------------------===//
+// ActionLogger
+//===----------------------------------------------------------------------===//
+
+static std::atomic<int64_t> thread_counter;
+int64_t getThreadId() {
+  thread_local int64_t tid = thread_counter++;
+  return tid;
+}
+
+void ActionLogger::beforeExecute(const ActionActiveStack *action,
+                                 Breakpoint *breakpoint, bool willExecute) {
+  os << "[thread " << getThreadId() << "] ";
+  if (willExecute)
+    os << "begins ";
+  else
+    os << "skipping ";
+  if (printBreakpoints) {
+    if (breakpoint)
+      os << " (on breakpoint: " << *breakpoint << ") ";
+    else
+      os << " (no breakpoint) ";
+  }
+  os << "Action ";
+  if (printActions)
+    action->getAction().print(os);
+  else
+    os << action->getAction().getTag();
+}
+
+void ActionLogger::afterExecute(const ActionActiveStack *action) {
+  os << "[thread " << getThreadId() << "] completed `"
+     << action->getAction().getTag() << "`\n";
+}

mlir/lib/Debug/Observers/CMakeLists.txt

+add_mlir_library(MLIRObservers
+  ActionLogging.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Debug/Observers
+
+  LINK_LIBS PUBLIC
+  ${LLVM_PTHREAD_LIB}
+  MLIRSupport
+)

mlir/lib/Pass/Pass.cpp

return OpToOpPassAdaptor::runPipeline(pipeline, root, nestedAm,
                                           pi, &parentInfo);
   };
   pass->passState.emplace(op, am, dynamicPipelineCallback);
 
   // Instrument before the pass has run.
   if (pi)
     pi->runBeforePass(pass, op);
 
+  bool passFailed;
+  op->getContext()->dispatch<PassExecutionAction>(
+      [&]() {
         // Invoke the virtual runOnOperation method.
         if (auto *adaptor = dyn_cast<OpToOpPassAdaptor>(pass))
           adaptor->runOnOperation(verifyPasses);
         else
           pass->runOnOperation();
-  bool passFailed = pass->passState->irAndPassFailed.getInt();
+        passFailed = pass->passState->irAndPassFailed.getInt();
+      },
+      *pass, op);
 
   // Invalidate any non preserved analyses.
   am.invalidate(pass->passState->preservedAnalyses);
 
   // When verifyPasses is specified, we run the verifier (unless the pass
   // failed).
   if (!passFailed && verifyPasses) {
     bool runVerifierNow = true;

mlir/lib/Pass/PassDetail.h

 //===- PassDetail.h - MLIR Pass details -------------------------*- 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_PASS_PASSDETAIL_H_
 #define MLIR_PASS_PASSDETAIL_H_
 
+#include "mlir/IR/Action.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 
 namespace mlir {
+/// Encapsulate the "action" of executing a single pass, used for the MLIR
+/// tracing infrastructure.
+struct PassExecutionAction : public tracing::ActionImpl<PassExecutionAction> {
+  const Pass &pass;
+  PassExecutionAction(const Pass &pass, Operation *op) : pass(pass), op(op) {}
+  static constexpr StringLiteral tag = "pass-execution-action";
+  void print(raw_ostream &os) const override {
+    os << "`" << getTag() << "` "
+       << " running \"" << pass.getName() << "\" on Operation \""
+       << op->getName() << "\"\n";
+  }
+  Operation *op;
+};
+
 namespace detail {
 
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
 /// An adaptor pass used to run operation passes over nested operations.
 class OpToOpPassAdaptor

mlir/lib/Tools/mlir-opt/CMakeLists.txt

 add_mlir_library(MLIROptLib
   MlirOptMain.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-opt
 
   LINK_LIBS PUBLIC
   MLIRBytecodeWriter
   MLIRDebug
+  MLIRObservers
   MLIRPass
   MLIRParser
   MLIRSupport
   )

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

 //===- MlirOptMain.cpp - MLIR Optimizer Driver ----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This is a utility that runs an optimization pass and prints the result back
 // out. It is designed to support unit testing.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/Debug/Counter.h"
+#include "mlir/Debug/ExecutionContext.h"
+#include "mlir/Debug/Observers/ActionLogging.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/Timing.h"
 #include "mlir/Support/ToolUtilities.h"
 #include "mlir/Tools/ParseUtilities.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/ToolOutputFile.h"
+#include "llvm/Support/raw_ostream.h"
+#include <memory>
 
 using namespace mlir;
 using namespace llvm;
 
 namespace {
 /// This class is intended to manage the handling of command line options for
 /// creating a *-opt config. This is a singleton.
 struct MlirOptMainConfigCLOptions : public MlirOptMainConfig {
static cl::opt<bool, /*ExternalStorage=*/true> emitBytecode(
         cl::location(emitBytecodeFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> explicitModule(
         "no-implicit-module",
         cl::desc("Disable implicit addition of a top-level module op during "
                  "parsing"),
         cl::location(useExplicitModuleFlag), cl::init(false));
 
+    static cl::opt<std::string, /*ExternalStorage=*/true> logActionsTo{
+        "log-actions-to",
+        cl::desc("Log action execution to a file, or stderr if "
+                 " '-' is passed"),
+        cl::location(logActionsToFlag)};
+
     static cl::opt<bool, /*ExternalStorage=*/true> showDialects(
         "show-dialects",
         cl::desc("Print the list of registered dialects and exit"),
         cl::location(showDialectsFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> splitInputFile(
         "split-input-file",
         cl::desc("Split the input file into pieces and process each "
if (this->shouldDumpPassPipeline()) {
       pm.dump();
       llvm::errs() << "\n";
     }
     return success();
   };
   return *this;
 }
 
+/// Set the ExecutionContext on the context and handle the observers.
+class InstallDebugHandler {
+public:
+  InstallDebugHandler(MLIRContext &context, const MlirOptMainConfig &config) {
+    if (config.getLogActionsTo().empty()) {
+      if (tracing::DebugCounter::isDebugCounterActivated())
+        context.registerActionHandler(tracing::DebugCounter{});
+      return;
+    }
+    if (tracing::DebugCounter::isDebugCounterActivated())
+      emitError(UnknownLoc::get(&context),
+                "Debug counters are incompatible with --log-actions-to option "
+                "and are disabled");
+    std::string errorMessage;
+    logActionsFile = openOutputFile(config.getLogActionsTo(), &errorMessage);
+    if (!logActionsFile) {
+      llvm::errs() << errorMessage << "\n";
+      return;
+    }
+    logActionsFile->keep();
+    raw_fd_ostream &logActionsStream = logActionsFile->os();
+    actionLogger = std::make_unique<tracing::ActionLogger>(logActionsStream);
+
+    executionContext.registerObserver(actionLogger.get());
+    context.registerActionHandler(executionContext);
+  }
+
+private:
+  std::unique_ptr<llvm::ToolOutputFile> logActionsFile;
+  tracing::ExecutionContext executionContext;
+  std::unique_ptr<tracing::ActionLogger> actionLogger;
+};
+
 /// Perform the actions on the input file indicated by the command line flags
 /// within the specified context.
 ///
 /// This typically parses the main source file, runs zero or more optimization
 /// passes, then prints the output.
 ///
 static LogicalResult
 performActions(raw_ostream &os,
if (threadPool)
     context.setThreadPool(*threadPool);
 
   // Parse the input file.
   if (config.shouldPreloadDialectsInContext())
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
   if (config.shouldVerifyDiagnostics())
     context.printOpOnDiagnostic(false);
-  context.registerActionHandler(tracing::DebugCounter{});
+
+  InstallDebugHandler installDebugHandler{context, config};
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
   if (!config.shouldVerifyDiagnostics()) {
     SourceMgrDiagnosticHandler sourceMgrHandler(*sourceMgr, &context);
     return performActions(os, sourceMgr, &context, config);
   }
 

mlir/test/Pass/action-logging.mlir

+// RUN: mlir-opt %s --log-actions-to=- -canonicalize -test-module-pass | FileCheck %s
+
+// CHECK: [thread 0] begins (no breakpoint) Action `pass-execution-action` running "Canonicalizer" on Operation "builtin.module"
+// CHECK: [thread 0] completed `pass-execution-action`
+// CHECK: [thread 0] begins (no breakpoint) Action `pass-execution-action` running "(anonymous namespace)::TestModulePass" on Operation "builtin.module"
+// CHECK: [thread 0] completed `pass-execution-action`