[llvm] [AMDGPU][NewInsertWaitcnt] Add Event Tracker (PR #226970)

[Vasileios Porpodas via llvm-commits]

================
@@ -0,0 +1,344 @@
+//===- AMDGPUEventTracking.h ------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file Tracks a per-InstCounterType event timeline which preserves
+/// information about previously encountered MachineInstr and HWEvents.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_AMDGPU_UTILS_AMDGPUEVENTTRACKING_H
+#define LLVM_LIB_TARGET_AMDGPU_UTILS_AMDGPUEVENTTRACKING_H
+
+#include "AMDGPUHWEvents.h"
+#include "AMDGPUWaitcntUtils.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Twine.h"
+#include <memory>
+
+namespace llvm {
+class raw_ostream;
+class MachineBasicBlock;
+class MachineOperand;
+class MachineFunction;
+
+namespace AMDGPU {
+namespace eventtracking {
+
+class EventTrackingContext;
+class EventTracker;
+
+/// FIXME: Make this a generic util?
+struct CounterInfo {
+  constexpr CounterInfo(InstCounterType T, HWEvents Events, unsigned Limit)
+      : CounterT(T), Events(Events), Limit(Limit) {}
+
+  /// Type of counter this is.
+  /// This must match the index in the container, e.g. CounterT=2 must be at
+  /// index 2.
+  InstCounterType CounterT;
+  /// HWEvents for this counter.
+  HWEvents Events;
+  /// Hardware limit for this counter.
+  unsigned Limit;
+};
+
+/// Represents entries in the \ref EventTracker.
+///
+/// Records are all uniquely identified by a \ref DynamicInstanceID. For each
+/// unique \ref DynamicInstanceID value, all \ref EventTrackerRecord that use
+/// that ID should have the same MI and Kind. This is enforced by exposing these
+/// as read-only, and making the constructor assign a new \ref DynamicInstanceID
+/// every time.
+///
+/// Only the score can change as it may be unique to each instance
+/// of \ref EventTracker that carry it.
+class EventTrackerRecord {
+public:
+  /// An always-increasing counter used to represent a dynamic instance of a
+  /// record.
+  ///
+  /// Whenever we add a new \ref EventTrackerRecord, even if it's one we already
+  /// have seen in a previous dataflow iteration, this counter is increased so
+  /// that the new record has a unique `DynamicInstanceID`.
+  using DynamicInstanceID = uint32_t;
+
+  EventTrackerRecord(EventTrackingContext &Ctx, MachineInstr *MI,
+                     SingleHWEvent Kind, uint32_t Score = 0);
+
+  /// \returns the ID uniquely identifying this record across an entire
+  /// EventTrackingContext. Whenever we revisit an instruction (when iterating
+  /// until a fixpoint is reached), we give it a new ID. This is used to
+  /// represent records carried over from previous iterations of the same basic
+  /// block.
+  DynamicInstanceID getID() const { return ID; }
+
+  /// \returns the MachineInstr that originated this record.
+  MachineInstr *getMI() const { return MI; }
+
+  /// \returns the kind of record this is, as a \ref SingleHWEvent.
+  SingleHWEvent getKind() const { return Kind; }
+
+  /// \returns the score of this record.
+  uint32_t getScore() const { return Score; }
+
+  /// Sets the score of this record to \p NewScore.
+  void setScore(uint32_t NewScore) {
+    Score = NewScore;
+    assert(Score == NewScore && "Score overflow!");
+  }
+
+  void print(raw_ostream &OS, bool PrintMI = true, unsigned Indent = 0) const;
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+
+private:
+  EventTrackerRecord(DynamicInstanceID ID, MachineInstr *MI, SingleHWEvent Kind,
+                     uint32_t Score = 0)
+      : MI(MI), ID(ID), Kind(Kind) {
+    setScore(Score);
+  }
+
+  MachineInstr *MI;
+  DynamicInstanceID ID;
+  SingleHWEvent Kind;
+  // Score should already never exceed uint8_t limit in normal circumstances as
+  // most counter types only use up to 6 bits encoding for the waitcnts. 16 bit
+  // is a very generous limit, we can probably shrink that at some point.
+  uint16_t Score;
+};
+
+/// This assert serves as a reminder to be mindful of the size of the object.
+static_assert(sizeof(EventTrackerRecord) == 16,
+              "EventTrackerRecord should remain small to optimize its layout "
+              "within cache lines, for maximum iteration speed");
+
+/// Per-MBB tracking context.
+///
+/// Tracks data across the following domains:
+///   - Current value (count) of each instruction counter.
+///   - In-flight (alive) \ref EventTrackerRecord of each instruction counter.
+///
+/// This class is only responsible for tracking records for every
+/// InstCounterType. It does not deal with calculating the waitcnts needed, or
+/// doing more advance reasoning over the timeline for specific queries (e.g.
+/// finding an aliasing store). These responsibilities are for
+/// utils/wrappers/users of the class.
+///
+/// The API should be kept as simple and clear as possible.
+class EventTracker {
+public:
+  EventTracker(MachineBasicBlock &MBB, EventTrackingContext &ET);
+
+  /// \defgroup MachineBasicBlock entry and exit
+  /// \{
+
+  /// Notify this EventTracker that we are going to begin recording events.
+  /// In case this is not the first time we are going through this block, this
+  /// clears the internal state of the tracker and re-imports all incoming
+  /// tracking state from the predecessors.
+  void enterBlock();
+
+  /// Notify this EventTracker that we are done recording events.
+  void leaveBlock();
+
+  /// \}
+
+  /// \defgroup InstCounters Tracking Entrypoints
+  /// Methods update the state of the InstCounters by adding/removing events
+  /// or signaling certain special conditions.
+  /// \{
+
+  /// Record an event of type \p Event at a MachineInstr \p MI, which will
+  /// affect all counters that have \p Event in their event set.
+  void record(MachineInstr &MI, SingleHWEvent Event);
+
+  /// Notify that we waited until the counter \p T reached the value \p N before
+  /// continuing execution of the program (and recording more events).
+  ///
+  /// This affects the count of \p T, an removes all records that have a score
+  /// greater than or equal to \p N.
+  ///
+  /// If \p N is zero, then \p T will no longer be in an indeterminate or
+  /// out-of-order state afterwards if it previously was in such a state.
+  void wait(InstCounterType T, unsigned N = 0);
----------------
vporpo wrote:

I would probably call this drainCounter()

https://github.com/llvm/llvm-project/pull/226970