mlir/lib/Debug/DebugCounter.cpp

   1 //===- DebugCounter.cpp - Debug Counter Facilities ------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "mlir/Debug/Counter.h"
  10 #include "llvm/Support/CommandLine.h"
  11 #include "llvm/Support/Debug.h"
  12 #include "llvm/Support/Format.h"
  13 #include "llvm/Support/ManagedStatic.h"
  14
  15 using namespace mlir;
  16 using namespace mlir::tracing;
  17
  18 //===----------------------------------------------------------------------===//
  19 // DebugCounter CommandLine Options
  20 //===----------------------------------------------------------------------===//
  21
  22 namespace {
  23 /// This struct contains command line options that can be used to initialize
  24 /// various bits of a DebugCounter. This uses a struct wrapper to avoid the need
  25 /// for global command line options.
  26 struct DebugCounterOptions {
  27   llvm::cl::list<std::string> counters{
  28       "mlir-debug-counter",
  29       llvm::cl::desc(
  30           "Comma separated list of debug counter skip and count arguments"),
  31       llvm::cl::CommaSeparated};
  32
  33   llvm::cl::opt<bool> printCounterInfo{
  34       "mlir-print-debug-counter", llvm::cl::init(false), llvm::cl::Optional,
  35       llvm::cl::desc("Print out debug counter information after all counters "
  36                      "have been accumulated")};
  37 };
  38 } // namespace
  39
  40 static llvm::ManagedStatic<DebugCounterOptions> clOptions;
  41
  42 //===----------------------------------------------------------------------===//
  43 // DebugCounter
  44 //===----------------------------------------------------------------------===//
  45
  46 DebugCounter::DebugCounter() { applyCLOptions(); }
  47
  48 DebugCounter::~DebugCounter() {
  49   // Print information when destroyed, iff command line option is specified.
  50   if (clOptions.isConstructed() && clOptions->printCounterInfo)
  51     print(llvm::dbgs());
  52 }
  53
  54 /// Add a counter for the given debug action tag. `countToSkip` is the number
  55 /// of counter executions to skip before enabling execution of the action.
  56 /// `countToStopAfter` is the number of executions of the counter to allow
  57 /// before preventing the action from executing any more.
  58 void DebugCounter::addCounter(StringRef actionTag, int64_t countToSkip,
  59                               int64_t countToStopAfter) {
  60   assert(!counters.count(actionTag) &&
  61          "a counter for the given action was already registered");
  62   counters.try_emplace(actionTag, countToSkip, countToStopAfter);
  63 }
  64
  65 void DebugCounter::operator()(llvm::function_ref<void()> transform,
  66                               const Action &action) {
  67   if (shouldExecute(action.getTag()))
  68     transform();
  69 }
  70
  71 bool DebugCounter::shouldExecute(StringRef tag) {
  72   auto counterIt = counters.find(tag);
  73   if (counterIt == counters.end())
  74     return true;
  75
  76   ++counterIt->second.count;
  77
  78   // We only execute while the `countToSkip` is not smaller than `count`, and
  79   // `countToStopAfter + countToSkip` is larger than `count`. Negative counters
  80   // always execute.
  81   if (counterIt->second.countToSkip < 0)
  82     return true;
  83   if (counterIt->second.countToSkip >= counterIt->second.count)
  84     return false;
  85   if (counterIt->second.countToStopAfter < 0)
  86     return true;
  87   return counterIt->second.countToStopAfter + counterIt->second.countToSkip >=
  88          counterIt->second.count;
  89 }
  90
  91 void DebugCounter::print(raw_ostream &os) const {
  92   // Order the registered counters by name.
  93   SmallVector<const llvm::StringMapEntry<Counter> *, 16> sortedCounters(
  94       llvm::make_pointer_range(counters));
  95   llvm::array_pod_sort(sortedCounters.begin(), sortedCounters.end(),
  96                        [](const decltype(sortedCounters)::value_type *lhs,
  97                           const decltype(sortedCounters)::value_type *rhs) {
  98                          return (*lhs)->getKey().compare((*rhs)->getKey());
  99                        });
 100
 101   os << "DebugCounter counters:\n";
 102   for (const llvm::StringMapEntry<Counter> *counter : sortedCounters) {
 103     os << llvm::left_justify(counter->getKey(), 32) << ": {"
 104        << counter->second.count << "," << counter->second.countToSkip << ","
 105        << counter->second.countToStopAfter << "}\n";
 106   }
 107 }
 108
 109 /// Register a set of useful command-line options that can be used to configure
 110 /// various flags within the DebugCounter. These flags are used when
 111 /// constructing a DebugCounter for initialization.
 112 void DebugCounter::registerCLOptions() {
 113   // Make sure that the options struct has been initialized.
 114   *clOptions;
 115 }
 116
 117 bool DebugCounter::isActivated() {
 118   return clOptions->counters.getNumOccurrences() ||
 119          clOptions->printCounterInfo.getNumOccurrences();
 120 }
 121
 122 // This is called by the command line parser when it sees a value for the
 123 // debug-counter option defined above.
 124 void DebugCounter::applyCLOptions() {
 125   if (!clOptions.isConstructed())
 126     return;
 127
 128   for (StringRef arg : clOptions->counters) {
 129     if (arg.empty())
 130       continue;
 131
 132     // Debug counter arguments are expected to be in the form: `counter=value`.
 133     auto [counterName, counterValueStr] = arg.split('=');
 134     if (counterValueStr.empty()) {
 135       llvm::errs() << "error: expected DebugCounter argument to have an `=` "
 136                       "separating the counter name and value, but the provided "
 137                       "argument was: `"
 138                    << arg << "`\n";
 139       llvm::report_fatal_error(
 140           "Invalid DebugCounter command-line configuration");
 141     }
 142
 143     // Extract the counter value.
 144     int64_t counterValue;
 145     if (counterValueStr.getAsInteger(0, counterValue)) {
 146       llvm::errs() << "error: expected DebugCounter counter value to be "
 147                       "numeric, but got `"
 148                    << counterValueStr << "`\n";
 149       llvm::report_fatal_error(
 150           "Invalid DebugCounter command-line configuration");
 151     }
 152
 153     // Now we need to see if this is the skip or the count, remove the suffix,
 154     // and add it to the counter values.
 155     if (counterName.consume_back("-skip")) {
 156       counters[counterName].countToSkip = counterValue;
 157
 158     } else if (counterName.consume_back("-count")) {
 159       counters[counterName].countToStopAfter = counterValue;
 160
 161     } else {
 162       llvm::errs() << "error: expected DebugCounter counter name to end with "
 163                       "either `-skip` or `-count`, but got`"
 164                    << counterName << "`\n";
 165       llvm::report_fatal_error(
 166           "Invalid DebugCounter command-line configuration");
 167     }
 168   }
 169 }