[Support] Optimize DebugCounter (#170305)

Currently, DebugCounters work by creating a unique counter ID during
registration, and then using that ID to look up the counter information
in the global registry.

However, this means that anything working with counters has to always go
through the global instance. This includes the fast path that checks
whether any counters are enabled.

Instead, we can drop the counter IDs, and make the counter variables use
CounterInfo themselves. We can then directly check whether the specific
counter is active without going through the global registry. This is
both faster for the fast-path where all counters are disabled, and also
faster for the case where only one counter is active (as the fast-path
can now still be used for all the disabled counters).

After this change, disabled counters become essentially free at runtime,
and we should be able to enable them in non-assert builds as well.

Author: nikic

llvm/include/llvm/Support/DebugCounter.h

@@ -44,8 +44,8 @@
 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringRef.h"
-#include "llvm/ADT/UniqueVector.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include <string>
@@ -63,6 +63,29 @@ class DebugCounter {
     bool contains(int64_t Idx) const { return Idx >= Begin && Idx <= End; }
   };
 
+  /// Struct to store counter info.
+  class CounterInfo {
+    friend class DebugCounter;
+
+    /// Whether counting should be enabled, either due to -debug-counter or
+    /// -print-debug-counter.
+    bool Active = false;
+    /// Whether chunks for the counter are set (differs from Active in that
+    /// -print-debug-counter uses Active=true, IsSet=false).
+    bool IsSet = false;
+
+    int64_t Count = 0;
+    uint64_t CurrChunkIdx = 0;
+    StringRef Name;
+    StringRef Desc;
+    SmallVector<Chunk> Chunks;
+
+  public:
+    CounterInfo(StringRef Name, StringRef Desc) : Name(Name), Desc(Desc) {
+      DebugCounter::registerCounter(this);
+    }
+  };
+
   LLVM_ABI static void printChunks(raw_ostream &OS, ArrayRef<Chunk>);
 
   /// Return true on parsing error and print the error message on the
@@ -75,48 +98,46 @@ class DebugCounter {
   // Used by the command line option parser to push a new value it parsed.
   LLVM_ABI void push_back(const std::string &);
 
-  // Register a counter with the specified name.
+  // Register a counter with the specified counter information.
   //
   // FIXME: Currently, counter registration is required to happen before command
   // line option parsing. The main reason to register counters is to produce a
   // nice list of them on the command line, but i'm not sure this is worth it.
-  static unsigned registerCounter(StringRef Name, StringRef Desc) {
-    return instance().addCounter(std::string(Name), std::string(Desc));
+  static void registerCounter(CounterInfo *Info) {
+    instance().addCounter(Info);
   }
-  LLVM_ABI static bool shouldExecuteImpl(unsigned CounterName);
+  LLVM_ABI static bool shouldExecuteImpl(CounterInfo &Counter);
 
-  inline static bool shouldExecute(unsigned CounterName) {
-    if (!isCountingEnabled())
+  inline static bool shouldExecute(CounterInfo &Counter) {
+    // Compile to nothing when debugging is off
+#ifdef NDEBUG
+    return true;
+#else
+    if (!Counter.Active)
       return true;
-    return shouldExecuteImpl(CounterName);
+    return shouldExecuteImpl(Counter);
+#endif
   }
 
   // Return true if a given counter had values set (either programatically or on
   // the command line).  This will return true even if those values are
   // currently in a state where the counter will always execute.
-  static bool isCounterSet(unsigned ID) {
-    return instance().Counters[ID].IsSet;
-  }
+  static bool isCounterSet(CounterInfo &Info) { return Info.IsSet; }
 
   struct CounterState {
     int64_t Count;
     uint64_t ChunkIdx;
   };
 
   // Return the state of a counter. This only works for set counters.
-  static CounterState getCounterState(unsigned ID) {
-    auto &Us = instance();
-    auto Result = Us.Counters.find(ID);
-    assert(Result != Us.Counters.end() && "Asking about a non-set counter");
-    return {Result->second.Count, Result->second.CurrChunkIdx};
+  static CounterState getCounterState(CounterInfo &Info) {
+    return {Info.Count, Info.CurrChunkIdx};
   }
 
   // Set a registered counter to a given state.
-  static void setCounterState(unsigned ID, CounterState State) {
-    auto &Us = instance();
-    auto &Counter = Us.Counters[ID];
-    Counter.Count = State.Count;
-    Counter.CurrChunkIdx = State.ChunkIdx;
+  static void setCounterState(CounterInfo &Info, CounterState State) {
+    Info.Count = State.Count;
+    Info.CurrChunkIdx = State.ChunkIdx;
   }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
@@ -126,66 +147,38 @@ class DebugCounter {
 
   LLVM_ABI void print(raw_ostream &OS) const;
 
-  // Get the counter ID for a given named counter, or return 0 if none is found.
-  unsigned getCounterId(const std::string &Name) const {
-    return RegisteredCounters.idFor(Name);
+  // Get the counter info for a given named counter,
+  // or return null if none is found.
+  CounterInfo *getCounterInfo(StringRef Name) const {
+    return Counters.lookup(Name);
   }
 
   // Return the number of registered counters.
-  unsigned int getNumCounters() const { return RegisteredCounters.size(); }
+  unsigned int getNumCounters() const { return Counters.size(); }
 
-  // Return the name and description of the counter with the given ID.
-  std::pair<std::string, std::string> getCounterInfo(unsigned ID) const {
-    return {RegisteredCounters[ID], Counters.lookup(ID).Desc};
+  // Return the name and description of the counter with the given info.
+  std::pair<StringRef, StringRef> getCounterDesc(CounterInfo *Info) const {
+    return {Info->Name, Info->Desc};
   }
 
   // Iterate through the registered counters
-  using CounterVector = UniqueVector<std::string>;
-  CounterVector::const_iterator begin() const {
-    return RegisteredCounters.begin();
+  MapVector<StringRef, CounterInfo *>::const_iterator begin() const {
+    return Counters.begin();
+  }
+  MapVector<StringRef, CounterInfo *>::const_iterator end() const {
+    return Counters.end();
   }
-  CounterVector::const_iterator end() const { return RegisteredCounters.end(); }
-
-  // Force-enables counting all DebugCounters.
-  //
-  // Since DebugCounters are incompatible with threading (not only do they not
-  // make sense, but we'll also see data races), this should only be used in
-  // contexts where we're certain we won't spawn threads.
-  static void enableAllCounters() { instance().Enabled = true; }
 
-  static bool isCountingEnabled() {
-// Compile to nothing when debugging is off
-#ifdef NDEBUG
-    return false;
-#else
-    return instance().Enabled;
-#endif
+  void activateAllCounters() {
+    for (auto &[_, Counter] : Counters)
+      Counter->Active = true;
   }
 
 protected:
-  unsigned addCounter(const std::string &Name, const std::string &Desc) {
-    unsigned Result = RegisteredCounters.insert(Name);
-    auto &C = Counters[Result];
-    C = {};
-    C.Desc = Desc;
-    return Result;
-  }
-  // Struct to store counter info.
-  struct CounterInfo {
-    int64_t Count = 0;
-    uint64_t CurrChunkIdx = 0;
-    bool IsSet = false;
-    std::string Desc;
-    SmallVector<Chunk> Chunks;
-  };
+  void addCounter(CounterInfo *Info) { Counters[Info->Name] = Info; }
   bool handleCounterIncrement(CounterInfo &Info);
 
-  DenseMap<unsigned, CounterInfo> Counters;
-  CounterVector RegisteredCounters;
-
-  // Whether we should do DebugCounting at all. DebugCounters aren't
-  // thread-safe, so this should always be false in multithreaded scenarios.
-  bool Enabled = false;
+  MapVector<StringRef, CounterInfo *> Counters;
 
   bool ShouldPrintCounter = false;
 
@@ -195,8 +188,7 @@ class DebugCounter {
 };
 
 #define DEBUG_COUNTER(VARNAME, COUNTERNAME, DESC)                              \
-  static const unsigned VARNAME =                                              \
-      DebugCounter::registerCounter(COUNTERNAME, DESC)
+  static DebugCounter::CounterInfo VARNAME(COUNTERNAME, DESC)
 
 } // namespace llvm
 #endif

llvm/lib/Support/DebugCounter.cpp

@@ -4,6 +4,7 @@
 
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/ManagedStatic.h"
 
 using namespace llvm;
 
@@ -110,12 +111,11 @@ class DebugCounterList : public cl::list<std::string, DebugCounter> {
     // width, so we do the same.
     Option::printHelpStr(HelpStr, GlobalWidth, ArgStr.size() + 6);
     const auto &CounterInstance = DebugCounter::instance();
-    for (const auto &Name : CounterInstance) {
-      const auto Info =
-          CounterInstance.getCounterInfo(CounterInstance.getCounterId(Name));
-      size_t NumSpaces = GlobalWidth - Info.first.size() - 8;
-      outs() << "    =" << Info.first;
-      outs().indent(NumSpaces) << " -   " << Info.second << '\n';
+    for (const auto &Entry : CounterInstance) {
+      const auto &[Name, Desc] = CounterInstance.getCounterDesc(Entry.second);
+      size_t NumSpaces = GlobalWidth - Name.size() - 8;
+      outs() << "    =" << Name;
+      outs().indent(NumSpaces) << " -   " << Desc << '\n';
     }
   }
 };
@@ -138,7 +138,7 @@ struct DebugCounterOwner : DebugCounter {
       cl::desc("Print out debug counter info after all counters accumulated"),
       cl::callback([&](const bool &Value) {
         if (Value)
-          enableAllCounters();
+          activateAllCounters();
       })};
   cl::opt<bool, true> PrintDebugCounterQueries{
       "print-debug-counter-queries",
@@ -171,12 +171,14 @@ struct DebugCounterOwner : DebugCounter {
 
 } // anonymous namespace
 
+// Use ManagedStatic instead of function-local static variable to ensure
+// the destructor (which accesses counters and streams) runs during
+// llvm_shutdown() rather than at some unspecified point.
+static ManagedStatic<DebugCounterOwner> Owner;
+
 void llvm::initDebugCounterOptions() { (void)DebugCounter::instance(); }
 
-DebugCounter &DebugCounter::instance() {
-  static DebugCounterOwner O;
-  return O;
-}
+DebugCounter &DebugCounter::instance() { return *Owner; }
 
 // This is called by the command line parser when it sees a value for the
 // debug-counter option defined above.
@@ -202,32 +204,26 @@ void DebugCounter::push_back(const std::string &Val) {
     return;
   }
 
-  unsigned CounterID = getCounterId(std::string(CounterName));
-  if (!CounterID) {
+  CounterInfo *Counter = getCounterInfo(CounterName);
+  if (!Counter) {
     errs() << "DebugCounter Error: " << CounterName
            << " is not a registered counter\n";
     return;
   }
-  enableAllCounters();
 
-  CounterInfo &Counter = Counters[CounterID];
-  Counter.IsSet = true;
-  Counter.Chunks = std::move(Chunks);
+  Counter->Active = Counter->IsSet = true;
+  Counter->Chunks = std::move(Chunks);
 }
 
 void DebugCounter::print(raw_ostream &OS) const {
-  SmallVector<StringRef, 16> CounterNames(RegisteredCounters.begin(),
-                                          RegisteredCounters.end());
+  SmallVector<StringRef, 16> CounterNames(Counters.keys());
   sort(CounterNames);
 
-  auto &Us = instance();
   OS << "Counters and values:\n";
-  for (auto &CounterName : CounterNames) {
-    unsigned CounterID = getCounterId(std::string(CounterName));
-    const CounterInfo &C = Us.Counters[CounterID];
-    OS << left_justify(RegisteredCounters[CounterID], 32) << ": {" << C.Count
-       << ",";
-    printChunks(OS, C.Chunks);
+  for (StringRef CounterName : CounterNames) {
+    const CounterInfo *C = getCounterInfo(CounterName);
+    OS << left_justify(C->Name, 32) << ": {" << C->Count << ",";
+    printChunks(OS, C->Chunks);
     OS << "}\n";
   }
 }
@@ -257,20 +253,14 @@ bool DebugCounter::handleCounterIncrement(CounterInfo &Info) {
   return Res;
 }
 
-bool DebugCounter::shouldExecuteImpl(unsigned CounterName) {
+bool DebugCounter::shouldExecuteImpl(CounterInfo &Counter) {
   auto &Us = instance();
-  auto Result = Us.Counters.find(CounterName);
-  if (Result != Us.Counters.end()) {
-    auto &CounterInfo = Result->second;
-    bool Res = Us.handleCounterIncrement(CounterInfo);
-    if (Us.ShouldPrintCounterQueries && CounterInfo.IsSet) {
-      dbgs() << "DebugCounter " << Us.RegisteredCounters[CounterName] << "="
-             << (CounterInfo.Count - 1) << (Res ? " execute" : " skip") << "\n";
-    }
-    return Res;
+  bool Res = Us.handleCounterIncrement(Counter);
+  if (Us.ShouldPrintCounterQueries && Counter.IsSet) {
+    dbgs() << "DebugCounter " << Counter.Name << "=" << (Counter.Count - 1)
+           << (Res ? " execute" : " skip") << "\n";
   }
-  // Didn't find the counter, should we warn?
-  return true;
+  return Res;
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -95,6 +95,7 @@
 #include <cassert>
 #include <cstdint>
 #include <iterator>
+#include <map>
 #include <memory>
 #include <optional>
 #include <set>