[mlir] Add support for DebugCounters using the new DebugAction infrastructure

[rriddle]

Summary

DebugCounters allow for selectively enabling the execution of a debug action based upon a "counter". This counter is comprised of two components that are used in the control of execution of an action, a "skip" value and a "count" value. The "skip" value is used to skip a certain number of initial executions of a debug action. The "count" value is used to prevent a debug action from executing after it has executed for a set number of times (not including any executions that have been skipped). For example, a counter for a debug action with skip=47 and count=2, would skip the first 47 executions, then execute twice, and finally prevent any further executions.

This is effectively the same as the DebugCounter infrastructure in LLVM, but using the DebugAction infrastructure in MLIR. We can't simply reuse the DebugCounter support already present in LLVM due to its heavy reliance on global constructors (which are not allowed in MLIR). The DebugAction infrastructure already nicely supports the debug counter use case, and promotes the separation of policy and mechanism design philosophy.

Depends On D84986

[jpienaar]

Nice, thanks

mlir/docs/DebugActions.md

188

Can these be chained? E.g., skip 47, execute 2, skip 10 again, execute 17?

205

Seems like we alternate between `shell and `sh , but `shell is mostly used.

mlir/lib/Support/DebugCounter.cpp

25

Should these still be registered in opt builds? (I'm assuming they'd have no effect there)

126

Why not fail here? Given this is only used in debug builds

[rriddle]

update

mlir/docs/DebugActions.md

188

No, each counter currently only has one skip and one count, so chaining isn't possible. Added a note.

205

We should align around one, do you have a preference?

mlir/lib/Support/DebugCounter.cpp

25

Good point, wrapped registration with NDEBUG.

[mehdi_amini]

mlir/docs/DebugActions.md

214

Please document the -mlir-print-debug-counter aspect, that seems fairly important in the workflow. I'd expect someone to start from this to get initial values for bisecting for example.

[rriddle]

update

[rriddle]

Update

mlir/docs/DebugActions.md

 *   An external entity registers an [`action handler`](#debug-action-handler)
     with the action manager, and provides the logic to resolve queries on
     actions.
 
 The exact definition of an `external entity` is left opaque, to allow for more
 interesting handlers. The set of possible action queries is detailed in the
 [`action manager`](#debug-action-manager) section below.
 
-(TODO: Add connection to existing handlers when they are added)
-
 ## Debug Action
 
 A `debug action` is essentially a marker for a type of action that may be
 performed within the compiler. There are no constraints on the granularity of an
 “action”, it can be as simple as “perform this fold” and as complex as “run this
 pass pipeline”. An action is comprised of the following:
 
 *   Tag:
struct MyPatternHandler : public DebugActionManager::GenericHandler {
   /// of the action is still somewhat user readable, or at least loggable as
   /// such.
   /// - actionTag: The tag specified by the action type.
   /// - actionDesc: The description specified by the action type.
   virtual FailureOr<bool> shouldExecute(StringRef actionTag,
                                         StringRef actionDesc);
 };
 ```
+
+### Common Action Handlers
+
+MLIR provides several common debug action handlers for immediate use that have
+proven useful in general.
+
+#### DebugCounter
+
+When debugging a compiler issue,
+["bisection"](https://en.wikipedia.org/wiki/Bisection_\(software_engineering\))
+is a useful technique for locating the root cause of the issue. `Debug Counters`
+enable using this technique for debug actions by attaching a counter value to a
+specific debug action and enabling/disabling execution of this action based on
+the value of the counter. The counter controls the execution of the action with
+a "skip" and "count" value. The "skip" value is used to skip a certain number of
+initial executions of a debug action. The "count" value is used to prevent a
+debug action from executing after it has executed for a set number of times (not
+including any executions that have been skipped). If the "skip" value is
+negative, the action will always execute. If the "count" value is negative, the
+action will always execute after the "skip" value has been reached. For example,
+a counter for a debug action with `skip=47` and `count=2`, would skip the first
+47 executions, then execute twice, and finally prevent any further executions.
+With a bit of tooling, the values to use for the counter can be automatically
+selected; allowing for finding the exact execution of a debug action that
+potentially causes the bug being investigated.
+
+Note: The DebugCounter action handler does not support multi-threaded execution,
+and should only be used in MLIRContexts where multi-threading is disabled (e.g.
+via `-mlir-disable-threading`).
+
+##### CommandLine Configuration
+
+The `DebugCounter` handler provides several that allow for configuring counters.
+The main option is `mlir-debug-counter`, which accepts a comma separated list of
+`<count-name>=<counter-value>`. A `<counter-name>` is the debug action tag to
+attach the counter, suffixed with either `-skip` or `-count`. A `-skip` suffix
+will set the "skip" value of the counter. A `-count` suffix will set the "count"
+value of the counter. The `<counter-value>` component is a numeric value to use
+for the counter. An example is shown below using `ApplyPatternAction` defined
+above:
+
+```shell
+$ mlir-opt foo.mlir -mlir-debug-counter=apply-pattern-skip=47,apply-pattern-count=2
+```
+
+The above configuration would skip the first 47 executions of
+`ApplyPatternAction`, then execute twice, and finally prevent any further
+executions.
+
+Note: Each counter currently only has one `skip` and one `count` value, meaning
+that sequences of `skip`/`count` will not be chained.
+
+The `mlir-print-debug-counter` option may be used to print out debug counter
+information after all counters have been accumulated. The information is printed
+in the following format:
+
+```shell
+DebugCounter counters:
+<action-tag>                   : {<current-count>,<skip>,<count>}
+```
+
+For example, using the options above we can see how many times an action is
+executed:
+
+```shell
+$ mlir-opt foo.mlir -mlir-debug-counter=apply-pattern-skip=-1 -mlir-print-debug-counter
+
+DebugCounter counters:
+apply-pattern                   : {370,-1,-1}
+```

mlir/include/mlir/Support/DebugCounter.h

+//===- DebugCounter.h - Debug Counter 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_SUPPORT_DEBUGCOUNTER_H
+#define MLIR_SUPPORT_DEBUGCOUNTER_H
+
+#include "mlir/Support/DebugAction.h"
+#include "llvm/ADT/StringMap.h"
+#include <string>
+
+namespace mlir {
+
+/// This class implements a debug action handler that attaches a counter value
+/// to debug actions and enables/disables execution of these action based on the
+/// value of the counter. The counter controls the execution of the action with
+/// a "skip" and "count" value. The "skip" value is used to skip a certain
+/// number of initial executions of a debug action. The "count" value is used to
+/// prevent a debug action from executing after it has executed for a set number
+/// of times (not including any executions that have been skipped). For example,
+/// a counter for a debug action with `skip=47` and `count=2`, would skip the
+/// first 47 executions, then execute twice, and finally prevent any further
+/// executions.
+class DebugCounter : public DebugActionManager::GenericHandler {
+public:
+  DebugCounter();
+  ~DebugCounter() override;
+
+  /// Add a counter for the given debug action tag. `countToSkip` is the number
+  /// of counter executions to skip before enabling execution of the action.
+  /// `countToStopAfter` is the number of executions of the counter to allow
+  /// before preventing the action from executing any more.
+  void addCounter(StringRef actionTag, int64_t countToSkip,
+                  int64_t countToStopAfter);
+
+  /// Register a counter with the specified name.
+  FailureOr<bool> shouldExecute(StringRef tag, StringRef description) final;
+
+  /// 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();
+
+private:
+  /// Apply the registered CL options to this debug counter instance.
+  void applyCLOptions();
+
+  /// This struct represents a specific counter being tracked.
+  struct Counter {
+    Counter(int64_t countToSkip = 0, int64_t countToStopAfter = -1)
+        : count(0), countToSkip(countToSkip),
+          countToStopAfter(countToStopAfter) {}
+
+    /// The current count of this counter.
+    int64_t count;
+    /// The number of initial executions of this counter to skip.
+    int64_t countToSkip;
+    /// The number of times to execute this counter before stopping.
+    int64_t countToStopAfter;
+  };
+
+  /// A mapping between a given action tag and its counter information.
+  llvm::StringMap<Counter> counters;
+};
+
+} // namespace mlir
+
+#endif

mlir/lib/Support/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
+  DebugCounter.cpp
   FileUtilities.cpp
   IndentedOstream.cpp
   MlirOptMain.cpp
   StorageUniquer.cpp
   ToolUtilities.cpp
 )
 
 add_mlir_library(MLIRSupport
+  DebugCounter.cpp
   FileUtilities.cpp
   StorageUniquer.cpp
   ToolUtilities.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Support
 
   LINK_LIBS PUBLIC

mlir/lib/Support/DebugCounter.cpp

+//===- DebugCounter.cpp - Debug Counter Facilities ------------------------===//
+//
+// 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/Support/DebugCounter.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/ManagedStatic.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// DebugCounter CommandLine Options
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This struct contains command line options that can be used to initialize
+/// various bits of a DebugCounter. This uses a struct wrapper to avoid the need
+/// for global command line options.
+struct DebugCounterOptions {
+  llvm::cl::list<std::string> counters{
+      "mlir-debug-counter",
+      llvm::cl::desc(
+          "Comma separated list of debug counter skip and count arguments"),
+      llvm::cl::CommaSeparated, llvm::cl::ZeroOrMore};
+
+  llvm::cl::opt<bool> printCounterInfo{
+      "mlir-print-debug-counter", llvm::cl::init(false), llvm::cl::Optional,
+      llvm::cl::desc("Print out debug counter information after all counters "
+                     "have been accumulated")};
+};
+} // end anonymous namespace
+
+static llvm::ManagedStatic<DebugCounterOptions> clOptions;
+
+//===----------------------------------------------------------------------===//
+// DebugCounter
+//===----------------------------------------------------------------------===//
+
+DebugCounter::DebugCounter() { applyCLOptions(); }
+
+DebugCounter::~DebugCounter() {
+  // Print information when destroyed, iff command line option is specified.
+  if (clOptions.isConstructed() && clOptions->printCounterInfo)
+    print(llvm::dbgs());
+}
+
+/// Add a counter for the given debug action tag. `countToSkip` is the number
+/// of counter executions to skip before enabling execution of the action.
+/// `countToStopAfter` is the number of executions of the counter to allow
+/// before preventing the action from executing any more.
+void DebugCounter::addCounter(StringRef actionTag, int64_t countToSkip,
+                              int64_t countToStopAfter) {
+  assert(!counters.count(actionTag) &&
+         "a counter for the given action was already registered");
+  counters.try_emplace(actionTag, countToSkip, countToStopAfter);
+}
+
+// Register a counter with the specified name.
+FailureOr<bool> DebugCounter::shouldExecute(StringRef tag,
+                                            StringRef description) {
+  auto counterIt = counters.find(tag);
+  if (counterIt == counters.end())
+    return true;
+
+  ++counterIt->second.count;
+
+  // We only execute while the `countToSkip` is not smaller than `count`, and
+  // `countToStopAfter + countToSkip` is larger than `count`. Negative counters
+  // always execute.
+  if (counterIt->second.countToSkip < 0)
+    return true;
+  if (counterIt->second.countToSkip >= counterIt->second.count)
+    return false;
+  if (counterIt->second.countToStopAfter < 0)
+    return true;
+  return counterIt->second.countToStopAfter + counterIt->second.countToSkip >=
+         counterIt->second.count;
+}
+
+void DebugCounter::print(raw_ostream &os) const {
+  // Order the registered counters by name.
+  SmallVector<const llvm::StringMapEntry<Counter> *, 16> sortedCounters(
+      llvm::make_pointer_range(counters));
+  llvm::sort(sortedCounters, [](const llvm::StringMapEntry<Counter> *lhs,
+                                const llvm::StringMapEntry<Counter> *rhs) {
+    return lhs->getKey() < rhs->getKey();
+  });
+
+  os << "DebugCounter counters:\n";
+  for (const llvm::StringMapEntry<Counter> *counter : sortedCounters) {
+    os << llvm::left_justify(counter->getKey(), 32) << ": {"
+       << counter->second.count << "," << counter->second.countToSkip << ","
+       << counter->second.countToStopAfter << "}\n";
+  }
+}
+
+/// Register a set of useful command-line options that can be used to configure
+/// various flags within the DebugCounter. These flags are used when
+/// constructing a DebugCounter for initialization.
+void DebugCounter::registerCLOptions() {
+#ifndef NDEBUG
+  // Make sure that the options struct has been initialized.
+  *clOptions;
+#endif
+}
+
+// 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())
+      continue;
+
+    // Debug counter arguments are expected to be in the form: `counter=value`.
+    StringRef counterName, counterValueStr;
+    std::tie(counterName, counterValueStr) = arg.split('=');
+    if (counterValueStr.empty()) {
+      llvm::errs() << "error: expected DebugCounter argument to have an `=` "
+                      "separating the counter name and value, but the provided "
+                      "argument was: `"
+                   << arg << "`\n";
+      llvm::report_fatal_error(
+          "Invalid DebugCounter command-line configuration");
+    }
+
+    // Extract the counter value.
+    int64_t counterValue;
+    if (counterValueStr.getAsInteger(0, counterValue)) {
+      llvm::errs() << "error: expected DebugCounter counter value to be "
+                      "numeric, but got `"
+                   << counterValueStr << "`\n";
+      llvm::report_fatal_error(
+          "Invalid DebugCounter command-line configuration");
+    }
+
+    // Now we need to see if this is the skip or the count, remove the suffix,
+    // and add it to the counter values.
+    if (counterName.consume_back("-skip")) {
+      counters[counterName].countToSkip = counterValue;
+
+    } else if (counterName.consume_back("-count")) {
+      counters[counterName].countToStopAfter = counterValue;
+
+    } else {
+      llvm::errs() << "error: expected DebugCounter counter name to end with "
+                      "either `-skip` or `-count`, but got`"
+                   << counterName << "`\n";
+      llvm::report_fatal_error(
+          "Invalid DebugCounter command-line configuration");
+    }
+  }
+}

mlir/lib/Support/MlirOptMain.cpp

 #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.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/DebugCounter.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/ToolUtilities.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
static LogicalResult processBuffer(raw_ostream &os,
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Parse the input file.
   MLIRContext context(registry);
   if (preloadDialectsInContext)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   context.printOpOnDiagnostic(!verifyDiagnostics);
+  context.getDebugActionManager().registerActionHandler<DebugCounter>();
 
   // 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 (!verifyDiagnostics) {
     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,
                           &context, passPipeline);
   }
static cl::opt<bool> runRepro(
       cl::init(false));
 
   InitLLVM y(argc, argv);
 
   // Register any command line options.
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
+  DebugCounter::registerCLOptions();
   PassPipelineCLParser passPipeline("", "Compiler passes to run");
 
   // Build the list of dialects as a header for the --help message.
   std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
   {
     llvm::raw_string_ostream os(helpHeader);
     MLIRContext context;
     interleaveComma(registry.getDialectNames(), os,

mlir/unittests/Support/CMakeLists.txt

 add_mlir_unittest(MLIRSupportTests
   DebugActionTest.cpp
+  DebugCounterTest.cpp
   IndentedOstreamTest.cpp
   MathExtrasTest.cpp
 )
 
 target_link_libraries(MLIRSupportTests
   PRIVATE MLIRSupportIndentedOstream MLIRSupport)

mlir/unittests/Support/DebugCounterTest.cpp

+//===- DebugCounterTest.cpp - Debug Counter Tests -------------------------===//
+//
+// 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/Support/DebugCounter.h"
+#include "gmock/gmock.h"
+
+using namespace mlir;
+
+// DebugActionManager is only enabled in DEBUG mode.
+#ifndef NDEBUG
+
+namespace {
+
+struct CounterAction : public DebugAction<> {
+  static StringRef getTag() { return "counter-action"; }
+  static StringRef getDescription() { return "Test action for debug counters"; }
+};
+
+TEST(DebugCounterTest, CounterTest) {
+  std::unique_ptr<DebugCounter> counter = std::make_unique<DebugCounter>();
+  counter->addCounter(CounterAction::getTag(), /*countToSkip=*/1,
+                      /*countToStopAfter=*/3);
+
+  DebugActionManager manager;
+  manager.registerActionHandler(std::move(counter));
+
+  // The first execution is skipped.
+  EXPECT_FALSE(manager.shouldExecute<CounterAction>());
+
+  // The counter stops after 3 successful executions.
+  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
+  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
+  EXPECT_TRUE(manager.shouldExecute<CounterAction>());
+  EXPECT_FALSE(manager.shouldExecute<CounterAction>());
+}
+
+} // namespace
+
+#endif