[AMDGPU][NewInsertWaitcnt] Event Tracker (1/2) See #226970 for context and history regarding this patch. Adds the basic infrastructure needed to track in-flight records for each InstCounterType. This is not a full functional replacement of SIInsertWaitcnts::WaitcntBrackets and it does not aim to be one, as we'd like to break up its functionality into multiple different pieces. This patch only contains the functionality needed for tracking events in one block, without importing events from previous blocks. Assisted-by: Claude Opus 5.5 After requests to split the patch, I used Claude to help me create this patch which is a subset of #226970 with all functionality pertaining to importing incoming events removed.
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.cpp b/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.cpp new file mode 100644 index 0000000..1dd93bd --- /dev/null +++ b/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.cpp
@@ -0,0 +1,206 @@ +//===- AMDGPUEventTracking.cpp ----------------------------------*- 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 +// +//===----------------------------------------------------------------------===// + +#include "AMDGPUEventTracking.h" +#include "AMDGPUHWEvents.h" +#include "AMDGPUWaitcntUtils.h" +#include "llvm/CodeGen/MachineBasicBlock.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "llvm/CodeGen/MachineInstr.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/Debug.h" +#include <algorithm> + +#define DEBUG_TYPE "amdgpu-event-tracking" + +namespace llvm { + +#ifndef NDEBUG +static cl::opt<bool> EventTrackerPrintAll( + "amdgpu-event-tracker-print-all", cl::init(false), + cl::desc("When using -debug, print the updated timeline every time a " + "record is added or removed")); +#endif + +namespace AMDGPU { +namespace eventtracking { + +void EventTrackerRecord::print(raw_ostream &OS, bool PrintMI, + unsigned Indent) const { + OS.indent(Indent) << "Record Kind=" << Kind << ": "; + if (PrintMI && MI) + OS << *MI; + else + OS << "MI@" << (void *)MI << '\n'; +} + +#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) +LLVM_DUMP_METHOD void EventTrackerRecord::dump() const { + dbgs() << '\n'; + print(dbgs(), /*PrintMI=*/true); + dbgs() << '\n'; +} +#endif + +EventTracker::EventTracker(const MachineBasicBlock &MBB, + ArrayRef<CounterInfo> CounterInfos) + : MBB(&MBB) { + Counters.resize(CounterInfos.size()); + for (const CounterInfo &Info : CounterInfos) + Counters[Info.CounterT].CI = &Info; +} + +void EventTracker::record(MachineInstr &MI, SingleHWEvent Event) { + LLVM_DEBUG(dbgs() << "[EventTracker] Recording " << Event << ": " << MI); + + EventTrackerRecord Rec = EventTrackerRecord(&MI, Event); + [[maybe_unused]] bool FoundMatch = false; + for (CounterData &CD : Counters) { + if (!CD.CI->Events.contains(Event)) + continue; + + FoundMatch = true; + ++CD.Count; + CD.Timeline.push_back(Rec); + +#ifndef NDEBUG + LLVM_DEBUG(if (EventTrackerPrintAll) { + dbgs().indent(2) << "Updated Timeline:\n"; + print(dbgs(), CD, /*Indent=*/4); + }); +#endif + } + + assert(FoundMatch && "Event has no matching InstCounterType!"); + +#ifdef EXPENSIVE_CHECKS + verify(); +#endif +} + +void EventTracker::drain(InstCounterType T, unsigned N) { + CounterData &CD = Counters[T]; + LLVM_DEBUG(dbgs() << "[EventTracker] Drain " << getInstCounterName(T) + << " at " << N << '\n'); + + // Fast path for clearing the counter. + if (N == 0) { + CD.Timeline.clear(); + CD.Count = 0; + } else { + CD.Count = std::min(CD.Count, N); + + // Only keep the N youngest records. + if (CD.Timeline.size() > N) { + auto *RmEnd = CD.Timeline.end() - N; + LLVM_DEBUG(for (const EventTrackerRecord &E + : make_range(CD.Timeline.begin(), RmEnd)) { + dbgs() << " | Removing "; + E.print(dbgs()); + }); + CD.Timeline.erase(CD.Timeline.begin(), RmEnd); + } + + LLVM_DEBUG(dbgs() << " | => Updated Count:" << CD.Count << '\n'); + +#ifndef NDEBUG + LLVM_DEBUG(if (EventTrackerPrintAll) { + dbgs().indent(2) << "Updated Timeline:\n"; + print(dbgs(), CD, /*Indent=*/4); + }); +#endif + } + +#ifdef EXPENSIVE_CHECKS + verify(); +#endif +} + +unsigned EventTracker::count(InstCounterType T) const { + return Counters[T].Count; +} + +HWEvents EventTracker::getPendingEvents(InstCounterType T) const { + HWEvents Res; + for (const auto &E : Counters[T].Timeline) + Res |= E.getKind(); + return Res; +} + +ArrayRef<EventTrackerRecord> +EventTracker::getTimeline(InstCounterType T) const { + return Counters[T].Timeline; +} + +void EventTracker::print(raw_ostream &OS, InstCounterType T) const { + print(OS, Counters[T]); +} + +#if !defined(NDEBUG) || defined(EXPENSIVE_CHECKS) +void EventTracker::verify() const { + if (!MBB) + llvm_unreachable("EventTracker has no MBB!"); + + for (const CounterData &C : Counters) { + const auto OnError = [&]() { + dbgs() << "EventTracker verification error\n"; + print(dbgs(), C); + }; + + // The timeline cannot have more records than the value of the counter. + if (C.Timeline.size() > C.Count) { + OnError(); + llvm_unreachable("timeline has more records than the counter value"); + } + } +} +#endif + +void EventTracker::print(raw_ostream &OS, bool IgnoreEmpty, + unsigned Indent) const { + OS.indent(Indent) << "EventTracker for "; + MBB->printAsOperand(OS); + OS << ':'; + if (IgnoreEmpty) { + if (all_of(Counters, [](const CounterData &CD) { return CD.Count == 0; })) { + OS << " (empty)\n"; + return; + } + } + + OS << '\n'; + for (const CounterData &C : Counters) { + if (IgnoreEmpty && C.Count == 0) + continue; + print(OS, C, Indent + 2); + } +} + +#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) +LLVM_DUMP_METHOD void EventTracker::dump() const { + dbgs() << '\n'; + print(dbgs()); + dbgs() << '\n'; +} +#endif + +void EventTracker::print(raw_ostream &OS, const CounterData &CD, + unsigned Indent) { + OS.indent(Indent) << getInstCounterName(CD.CI->CounterT) + << " (Count=" << CD.Count + << ", Timeline=" << CD.Timeline.size() << ")\n"; + for (const EventTrackerRecord &E : CD.Timeline) { + OS.indent(Indent + 2); + E.print(OS); + } +} + +} // namespace eventtracking +} // namespace AMDGPU + +} // namespace llvm
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.h b/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.h new file mode 100644 index 0000000..512d058 --- /dev/null +++ b/llvm/lib/Target/AMDGPU/AMDGPUEventTracking.h
@@ -0,0 +1,194 @@ +//===- 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 +/// Minimal API to track in-flight events that increment hardware instruction +/// counters (\ref InstCounterType, e.g. LOAD_CNT, DS_CNT). Intended as a +/// building block for waitcnt insertion: it tracks which events are +/// outstanding and how far back they are, but does not directly compute waits. +/// +/// - \ref CounterInfo: constant per-counter description (events, limit). +/// - \ref EventTrackerRecord: a (MachineInstr, SingleHWEvent) pair. +/// - \ref EventTracker: per-MBB state; for each counter, a count and a +/// timeline of records ordered from oldest to youngest. Driven by the +/// client while visiting the MBB: \ref EventTracker::record adds events, +/// and \ref EventTracker::drain models waits. +/// +/// This is a low-level API. It can be used directly for trivial waitcnt +/// calculation (e.g. wait on a specific MI), but the typical usage is for +/// clients to build helpers on top of it that combine the timeline with +/// additional information (e.g. alias analysis, subtarget-specific behavior) +/// to compute more precise waits. +// +//===----------------------------------------------------------------------===// + +#ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUEVENTTRACKING_H +#define LLVM_LIB_TARGET_AMDGPU_AMDGPUEVENTTRACKING_H + +#include "AMDGPUHWEvents.h" +#include "AMDGPUWaitcntUtils.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" + +namespace llvm { +class raw_ostream; +class MachineBasicBlock; + +namespace AMDGPU { +namespace eventtracking { + +/// FIXME: Move this to a common file and make it a generic utility. +/// +/// Represents target-specific information about an available +/// \ref InstCounterType. This information is constant and inferred from the +/// subtarget. +struct CounterInfo { + CounterInfo() = default; + + constexpr CounterInfo(InstCounterType T, HWEvents Events, unsigned Limit) + : CounterT(T), Events(Events), Limit(Limit) {} + + /// Type of counter this is. + InstCounterType CounterT; + /// HWEvents for this counter. + HWEvents Events; + /// Hardware limit for this counter. + unsigned Limit; +}; + +/// Represents a record within the \ref EventTracker's timeline +/// for one \ref InstCounterType. A record is a pair of \ref MachineInstr +/// and \ref SingleHWEvent. +class EventTrackerRecord { +public: + EventTrackerRecord(MachineInstr *MI, SingleHWEvent Kind) + : MI(MI), Kind(Kind) {} + + /// \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; } + + 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: + MachineInstr *MI; + SingleHWEvent Kind; +}; + +/// 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 which tracks the current value (count) of each +/// instruction counter and the timeline of \ref EventTrackerRecord +/// for each instruction counter. +/// +/// This class is only responsible for tracking records for every +/// \ref InstCounterType. It does not deal with calculating the waitcnts needed, +/// or doing more advanced reasoning over the timeline for specific queries +/// (e.g. finding an aliasing store). These responsibilities are left to +/// helpers built on top of this class (see the file header). +/// +/// The expected usage pattern is to visit the MBB in instruction order, call +/// \ref EventTracker::record whenever an interesting event occurs, and call +/// \ref EventTracker::drain whenever a wait on a counter is performed. +/// In between those operations, the user of the class can use all other methods +/// to query the current state of a counter, iterate its timeline, etc. +/// +/// Allocation/storage/mapping of EventTrackers to MBBs is left to the user of +/// the class. +class EventTracker { +public: + /// \p MBB is the MachineBasicBlock for this tracker. + /// \p CounterInfos describes the available \ref InstCounterType. It must + /// contain exactly one entry for each counter type in + /// [0, CounterInfos.size()), and must outlive this class as its elements are + /// stored by reference. + EventTracker(const MachineBasicBlock &MBB, + ArrayRef<CounterInfo> CounterInfos); + + /// 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. + /// \pre \p Event belongs to at least one \ref InstCounterType. + void record(MachineInstr &MI, SingleHWEvent Event); + + /// Drains the timeline of \p T such that: + /// - The \ref count of \p T is no higher than \p N. + /// - Only the \p N youngest records remain in the timeline. + void drain(InstCounterType T, unsigned N = 0); + + /// \returns a conservative estimate of the current value of the counter \p T. + unsigned count(InstCounterType T) const; + + /// \returns the set of pending HWEvents for \p T. + HWEvents getPendingEvents(InstCounterType T) const; + + /// \returns the current timeline for \p T, which contains all known in-flight + /// (not drained) records, ordered from oldest to youngest. + /// The returned ArrayRef is invalidated whenever the timeline changes (e.g. + /// \ref record, \ref drain). + ArrayRef<EventTrackerRecord> getTimeline(InstCounterType T) const; + + /// Prints a dump of the internal tracking state of this class for \p T to the + /// stream \p OS. + void print(raw_ostream &OS, InstCounterType T) const; + + /// Prints a dump of all internal tracking state of this class to the stream + /// \p OS. If \p IgnoreEmpty is true, do not print counters with a count of 0. + void print(raw_ostream &OS, bool IgnoreEmpty = false, + unsigned Indent = 0) const; + +#if !defined(NDEBUG) || defined(EXPENSIVE_CHECKS) + /// Verifies that the invariants of this class hold. + void verify() const; +#endif + +#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP) + LLVM_DUMP_METHOD void dump() const; +#endif + +private: + /// FIXME: I do not think it is relevant to expose this class like we expose + /// \ref EventTrackerRecord at this time. However, if this becomes complex + /// enough, we could clean it up and expose it if it leads to a better API. + struct CounterData { + const CounterInfo *CI = nullptr; + + SmallVector<EventTrackerRecord, 16> Timeline; + + /// The current value of the counter. + /// + /// FIXME: We shouldn't need it for correctness so perhaps it should be + /// removed entirely, or be marked as debug. + uint32_t Count = 0; + }; + + static void print(raw_ostream &OS, const CounterData &CD, + unsigned Indent = 0); + + const MachineBasicBlock *MBB; + + // NB: This, combined with the inline storage of Timeline, can lead to this + // class becoming quite big - verify the size of this object whenever a change + // is made. + SmallVector<CounterData, InstCounterType::NUM_INST_CNTS> Counters; +}; + +} // namespace eventtracking +} // namespace AMDGPU + +} // namespace llvm + +#endif
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHWEvents.h b/llvm/lib/Target/AMDGPU/AMDGPUHWEvents.h index eb25206..b4ce2fa 100644 --- a/llvm/lib/Target/AMDGPU/AMDGPUHWEvents.h +++ b/llvm/lib/Target/AMDGPU/AMDGPUHWEvents.h
@@ -165,6 +165,33 @@ value_type Data = NONE; }; +/// Class to store singular HW events in the most compact way possible (u8). +class SingleHWEvent { +public: + using value_type = uint8_t; + + static constexpr SingleHWEvent encode(HWEvents E) { + assert(E.size() == 1 && "expected exactly one event!"); + return SingleHWEvent(countr_zero_constexpr(E.value())); + } + + constexpr operator HWEvents() const { return HWEvents(1 << Data); } + + constexpr value_type rawValue() const { return Data; } + + constexpr bool operator==(const SingleHWEvent &Other) const { + return Data == Other.Data; + } + constexpr bool operator!=(const SingleHWEvent &Other) const { + return Data != Other.Data; + } + +private: + constexpr SingleHWEvent(value_type Data) : Data(Data) {} + + value_type Data; +}; + /// \param Inst A VMEM instruction (as per `SIInstrInfo::isVMEM`). /// \returns the simplified set of events triggered by the VMEM instruction \p /// Inst. The returned mask is not exhaustive, but is guaranteed to be a subset
diff --git a/llvm/lib/Target/AMDGPU/CMakeLists.txt b/llvm/lib/Target/AMDGPU/CMakeLists.txt index b6c5f63..2721ead 100644 --- a/llvm/lib/Target/AMDGPU/CMakeLists.txt +++ b/llvm/lib/Target/AMDGPU/CMakeLists.txt
@@ -54,6 +54,7 @@ AMDGPUCodeGenPrepare.cpp AMDGPUCombinerHelper.cpp AMDGPUCtorDtorLowering.cpp + AMDGPUEventTracking.cpp AMDGPUExportClustering.cpp AMDGPUExportKernelRuntimeHandles.cpp AMDGPUFrameLowering.cpp
diff --git a/llvm/unittests/Target/AMDGPU/AMDGPUEventTrackingTest.cpp b/llvm/unittests/Target/AMDGPU/AMDGPUEventTrackingTest.cpp new file mode 100644 index 0000000..92ce5fa --- /dev/null +++ b/llvm/unittests/Target/AMDGPU/AMDGPUEventTrackingTest.cpp
@@ -0,0 +1,207 @@ +//===- AMDGPUEventTrackingTest.cpp ------------------------------*- 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 +// +//===----------------------------------------------------------------------===// + +#include "AMDGPUEventTracking.h" +#include "AMDGPUUnitTests.h" +#include "AMDGPUWaitcntUtils.h" +#include "MCTargetDesc/AMDGPUMCTargetDesc.h" +#include "llvm/CodeGen/MachineFunction.h" +#include "gtest/gtest.h" + +using namespace llvm; +using namespace llvm::AMDGPU; +using namespace llvm::AMDGPU::eventtracking; + +namespace { + +static constexpr unsigned CounterLimit = 12; + +// These are not accurate, they are simply for testing purposes. +// We do not need to test every single counter accurately, that is the job +// of the IR/MIR tests in tests/CodeGen/AMDGPU. We just need enough here +// to validate that the EventTracker works. +static constexpr std::array<CounterInfo, 4> GFX12CounterInfos = {{ + {LOAD_CNT, HWEvents::VMEM_READ_ACCESS, CounterLimit}, + {DS_CNT, HWEvents::LDS_ACCESS, CounterLimit}, + {EXP_CNT, HWEvents::EXP_GPR_LOCK, CounterLimit}, + {STORE_CNT, HWEvents::VMEM_WRITE_ACCESS | HWEvents::SCRATCH_WRITE_ACCESS, + CounterLimit}, +}}; + +class AMDGPUGFX12EventTrackingTest : public AMDGPUCodeGenTestBase { +public: + void SetUp() override { + setUpImpl("amdgpu12.00-amd-amdhsa", "", ""); + + TrackerGetter = [&](const MachineBasicBlock &MBB) -> EventTracker & { + auto &Entry = Trackers[&MBB]; + if (!Entry) + Entry = std::make_unique<EventTracker>(MBB, GFX12CounterInfos); + return *Entry; + }; + } + + EventTracker & + visitAll(MachineBasicBlock &MBB, + function_ref<void(EventTracker &ET)> AfterVisit = nullptr) { + const GCNSubtarget &ST = MBB.getParent()->getSubtarget<GCNSubtarget>(); + + EventTracker &ET = TrackerGetter(MBB); + for (MachineInstr &MI : MBB) { + HWEvents Events = + getEventsFor(MI, ST, /*IsExpertMode=*/false, /*TgSplit=*/false); + for (HWEvents SingleEv : Events) { + ET.record(MI, SingleHWEvent::encode(SingleEv)); + } + } + if (AfterVisit) + AfterVisit(ET); + return ET; + } + + std::function<EventTracker &(const MachineBasicBlock &)> TrackerGetter; + DenseMap<const MachineBasicBlock *, std::unique_ptr<EventTracker>> Trackers; +}; + +namespace { + +/// Provides some helpers to declaratively check the state of the EventTracker +/// for one counter. This provides helpers to check the general counter state +/// (value, etc) but also allows iterating over the timeline from the oldest to +/// the youngest element. +/// +/// This makes the actual test cases clearer. +struct TrackerRecordsChecker { + TrackerRecordsChecker(EventTracker &ET, InstCounterType T) + : ET(ET), T(T), Records(ET.getTimeline(T)) {} + + unsigned getCount() { return ET.count(T); } + + bool empty() { return Records.empty(); } + + // Has no records and count is zero + bool unused() { return empty() && !getCount(); } + + const EventTrackerRecord &cur() { return Records[CurElt]; } + + /// Move to the next record. + /// \returns true on success, false if the end has been reached. + bool next() { + ++CurElt; + if (CurElt >= Records.size()) + return false; + return true; + } + + EventTracker &ET; + InstCounterType T; + ArrayRef<EventTrackerRecord> Records; + unsigned CurElt = 0; +}; + +} // namespace + +/// Check trivial straight line counting. +TEST_F(AMDGPUGFX12EventTrackingTest, BasicTimeline) { + StringRef MIR = R"( +name: BasicTimeline +body: | + bb.0: + + GLOBAL_STORE_DWORD $vgpr0_vgpr1, $vgpr2, 0, 0, implicit $exec + $vgpr4 = GLOBAL_LOAD_DWORD $vgpr4_vgpr5, 0, 0, implicit $exec + $vgpr0 = DS_READ_B32_gfx9 $vgpr1, 0, 0, implicit $exec + GLOBAL_STORE_DWORDX2 $vgpr1_vgpr2, $vgpr3_vgpr4, 0, 0, implicit $exec + S_ENDPGM 0 +... +)"; + ASSERT_TRUE(parseMIR(MIR)); + MachineFunction &MF = getMF("BasicTimeline"); + MachineBasicBlock &BB0 = *MF.getBlockNumbered(0); + + auto &ET = visitAll(BB0); + + auto LoadCnt = TrackerRecordsChecker(ET, AMDGPU::LOAD_CNT); + EXPECT_EQ(LoadCnt.getCount(), 1u); + EXPECT_FALSE(LoadCnt.empty()); + EXPECT_EQ(LoadCnt.cur().getMI()->getOpcode(), AMDGPU::GLOBAL_LOAD_DWORD); + EXPECT_FALSE(LoadCnt.next()); + + auto DsCnt = TrackerRecordsChecker(ET, AMDGPU::DS_CNT); + EXPECT_EQ(DsCnt.getCount(), 1u); + EXPECT_FALSE(DsCnt.empty()); + EXPECT_EQ(DsCnt.cur().getMI()->getOpcode(), AMDGPU::DS_READ_B32_gfx9); + EXPECT_FALSE(DsCnt.next()); + + auto ExpCnt = TrackerRecordsChecker(ET, AMDGPU::EXP_CNT); + EXPECT_TRUE(ExpCnt.unused()); + + auto StoreCnt = TrackerRecordsChecker(ET, AMDGPU::STORE_CNT); + EXPECT_EQ(StoreCnt.getCount(), 2u); + EXPECT_FALSE(StoreCnt.empty()); + EXPECT_EQ(StoreCnt.cur().getMI()->getOpcode(), AMDGPU::GLOBAL_STORE_DWORD); + EXPECT_TRUE(StoreCnt.next()); + EXPECT_EQ(StoreCnt.cur().getMI()->getOpcode(), AMDGPU::GLOBAL_STORE_DWORDX2); + EXPECT_FALSE(StoreCnt.next()); +} + +/// Check partial and full drains of a counter. +TEST_F(AMDGPUGFX12EventTrackingTest, Drain) { + StringRef MIR = R"( +name: Drain +body: | + bb.0: + + GLOBAL_STORE_DWORD $vgpr0_vgpr1, $vgpr2, 0, 0, implicit $exec + GLOBAL_STORE_DWORDX2 $vgpr1_vgpr2, $vgpr3_vgpr4, 0, 0, implicit $exec + $vgpr4 = GLOBAL_LOAD_DWORD $vgpr4_vgpr5, 0, 0, implicit $exec + GLOBAL_STORE_DWORD $vgpr0_vgpr1, $vgpr2, 0, 0, implicit $exec + S_ENDPGM 0 +... +)"; + ASSERT_TRUE(parseMIR(MIR)); + MachineFunction &MF = getMF("Drain"); + MachineBasicBlock &BB0 = *MF.getBlockNumbered(0); + + auto &ET = visitAll(BB0); + + EXPECT_EQ(ET.getPendingEvents(AMDGPU::STORE_CNT), + HWEvents(HWEvents::VMEM_WRITE_ACCESS)); + EXPECT_EQ(ET.getPendingEvents(AMDGPU::LOAD_CNT), + HWEvents(HWEvents::VMEM_READ_ACCESS)); + EXPECT_FALSE(ET.getPendingEvents(AMDGPU::DS_CNT)); + + // Partial drain: only the two youngest stores remain. + ET.drain(AMDGPU::STORE_CNT, 2); + { + auto StoreCnt = TrackerRecordsChecker(ET, AMDGPU::STORE_CNT); + EXPECT_EQ(StoreCnt.getCount(), 2u); + EXPECT_FALSE(StoreCnt.empty()); + EXPECT_EQ(StoreCnt.cur().getMI()->getOpcode(), + AMDGPU::GLOBAL_STORE_DWORDX2); + EXPECT_TRUE(StoreCnt.next()); + EXPECT_EQ(StoreCnt.cur().getMI()->getOpcode(), AMDGPU::GLOBAL_STORE_DWORD); + EXPECT_FALSE(StoreCnt.next()); + } + + // Draining a counter does not affect other counters. + EXPECT_EQ(ET.count(AMDGPU::LOAD_CNT), 1u); + EXPECT_EQ(ET.getTimeline(AMDGPU::LOAD_CNT).size(), 1u); + + // Draining above the current count is a no-op. + ET.drain(AMDGPU::STORE_CNT, 5); + EXPECT_EQ(ET.count(AMDGPU::STORE_CNT), 2u); + EXPECT_EQ(ET.getTimeline(AMDGPU::STORE_CNT).size(), 2u); + + // Full drain. + ET.drain(AMDGPU::STORE_CNT); + EXPECT_TRUE(TrackerRecordsChecker(ET, AMDGPU::STORE_CNT).unused()); + EXPECT_FALSE(ET.getPendingEvents(AMDGPU::STORE_CNT)); +} + +} // namespace
diff --git a/llvm/unittests/Target/AMDGPU/CMakeLists.txt b/llvm/unittests/Target/AMDGPU/CMakeLists.txt index 2bb9b6d..882ceef 100644 --- a/llvm/unittests/Target/AMDGPU/CMakeLists.txt +++ b/llvm/unittests/Target/AMDGPU/CMakeLists.txt
@@ -23,6 +23,7 @@ ) add_llvm_target_unittest(AMDGPUTests + AMDGPUEventTrackingTest.cpp AMDGPUMCExprTest.cpp AMDGPUUnitTests.cpp CSETest.cpp