From 7e3ea0c58eda884adc862a6ee0b71a9742a67b5b Mon Sep 17 00:00:00 2001
From: "Welton, Benjamin" <Benjamin.Welton@amd.com>
Date: Fri, 11 Jul 2025 12:46:19 -0700
Subject: [PATCH] [SWDEV-540753] Reduce memory allocations in device profiling
 (#507)

Cache packet creation in all cases to reduce the number of allocations/
destruction operations made down to KFD. There is a bug that we
encounter after a period of runtime in KFD where allocations fail to be
visable to the GPU (suspect this is a FW issue, similar to other FW
issues they have had along the same lines). This sidesteps that issue in
rocprof (and likely should be done regardless)

Co-authored-by: Benjamin Welton <bewelton@amd.com>
---
 source/lib/common/utility.hpp                 | 21 +++++++++++++++
 .../counters/device_counting.cpp              | 27 +++++++++++++++++--
 .../counters/device_counting.hpp              |  2 +-
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/source/lib/common/utility.hpp b/source/lib/common/utility.hpp
index 4f09de8c8d..4c3591187a 100644
--- a/source/lib/common/utility.hpp
+++ b/source/lib/common/utility.hpp
@@ -28,6 +28,7 @@
 #include <sys/syscall.h>
 #include <sys/utsname.h>
 #include <unistd.h>
+#include <atomic>
 #include <chrono>
 #include <condition_variable>
 #include <cstddef>
@@ -269,6 +270,26 @@ yield(PredicateT&&                      predicate,
     // return the result of the last predicate query
     return result;
 }
+
+class assert_single_threaded
+{
+public:
+    assert_single_threaded(std::atomic<bool>& lock)
+    : m_is_initialized(lock)
+    {
+        bool expected = false;
+        if(!m_is_initialized.compare_exchange_strong(expected, true))
+        {
+            ROCP_FATAL << "This code must be run in a single thread!!!";
+        }
+    }
+
+    ~assert_single_threaded() { m_is_initialized.store(false, std::memory_order_release); }
+
+private:
+    std::atomic<bool>& m_is_initialized;
+};
+
 }  // namespace common
 }  // namespace rocprofiler
 
diff --git a/source/lib/rocprofiler-sdk/counters/device_counting.cpp b/source/lib/rocprofiler-sdk/counters/device_counting.cpp
index 60f74814f5..af473986a7 100644
--- a/source/lib/rocprofiler-sdk/counters/device_counting.cpp
+++ b/source/lib/rocprofiler-sdk/counters/device_counting.cpp
@@ -22,6 +22,7 @@
 
 #include "lib/rocprofiler-sdk/counters/device_counting.hpp"
 #include "lib/common/logging.hpp"
+#include "lib/common/utility.hpp"
 #include "lib/rocprofiler-sdk/buffer.hpp"
 #include "lib/rocprofiler-sdk/context/context.hpp"
 #include "lib/rocprofiler-sdk/counters/controller.hpp"
@@ -102,9 +103,29 @@ header_pkt(hsa_packet_type_t type)
     return header;
 }
 
-std::unique_ptr<hsa::CounterAQLPacket>
+/**
+ * Construct the packet or grab it from the cache.
+ * Note this function is not thread safe and is only called from
+ * init_callback_data which can only be called when the context is in the LOCKED state
+ * with only a single thread active.
+ */
+std::shared_ptr<hsa::CounterAQLPacket>
 construct_aql_pkt(std::shared_ptr<counter_config>& profile)
 {
+    static std::atomic<bool> has_thread{false};
+    static std::unordered_map<rocprofiler_profile_config_id_t,
+                              std::shared_ptr<hsa::CounterAQLPacket>>
+        pkt_cache;
+    // Asserts if there are two threads in this function at the same time.
+    auto _ = common::assert_single_threaded(has_thread);
+
+    // If we have a packet in the cache, return it.
+    if(pkt_cache.find(profile->id) != pkt_cache.end())
+    {
+        return pkt_cache[profile->id];
+    }
+
+    // If we do not have a packet in the cache, create it.
     if(counter_callback_info::setup_counter_config(profile) != ROCPROFILER_STATUS_SUCCESS)
     {
         return nullptr;
@@ -119,7 +140,9 @@ construct_aql_pkt(std::shared_ptr<counter_config>& profile)
     pkts->packets.read_packet.header  = header_pkt(HSA_PACKET_TYPE_VENDOR_SPECIFIC);
 
     pkts->packets.start_packet.completion_signal.handle = 0;
-    return pkts;
+
+    pkt_cache[profile->id] = std::move(pkts);
+    return pkt_cache[profile->id];
 }
 
 bool
diff --git a/source/lib/rocprofiler-sdk/counters/device_counting.hpp b/source/lib/rocprofiler-sdk/counters/device_counting.hpp
index b891d55ce2..b168c36724 100644
--- a/source/lib/rocprofiler-sdk/counters/device_counting.hpp
+++ b/source/lib/rocprofiler-sdk/counters/device_counting.hpp
@@ -45,7 +45,7 @@ struct agent_callback_data
 {
     uint64_t                               context_idx = 0;
     hsa_queue_t*                           queue       = nullptr;
-    std::unique_ptr<hsa::CounterAQLPacket> packet      = {};
+    std::shared_ptr<hsa::CounterAQLPacket> packet      = {};
 
     // Tri-state signal used to know what the current state of processing
     // a sample is. The states are: