From c32d8b66964b4d54040b6e779f9f7a754cf53b59 Mon Sep 17 00:00:00 2001
From: Benjamin Welton <bewelton@amd.com>
Date: Wed, 12 Jun 2024 20:39:18 -0700
Subject: [PATCH] Fix agent profiling for SQ counters (#919)

* Fix agent profiling for SQ counters


---------

Co-authored-by: Benjamin Welton <ben@amd.com>
---
 .../counters/agent_profiling.cpp              |  8 ++--
 .../counters/agent_profiling.hpp              |  3 ++
 .../counters/tests/agent_profiling.cpp        | 48 ++++++++++++++++++-
 source/lib/rocprofiler-sdk/hsa/queue.cpp      | 22 +++++++++
 .../rocprofiler-sdk/hsa/queue_controller.cpp  |  3 +-
 5 files changed, 76 insertions(+), 8 deletions(-)

diff --git a/source/lib/rocprofiler-sdk/counters/agent_profiling.cpp b/source/lib/rocprofiler-sdk/counters/agent_profiling.cpp
index f6c33916dd..4d515c27b8 100644
--- a/source/lib/rocprofiler-sdk/counters/agent_profiling.cpp
+++ b/source/lib/rocprofiler-sdk/counters/agent_profiling.cpp
@@ -198,6 +198,9 @@ init_callback_data(const rocprofiler::context::context& ctx, const hsa::AgentCac
 
     // Set state of the queue to allow profiling (may not be needed since AQL
     // may do this in the future).
+    CHECK(agent.cpu_pool().handle != 0);
+    CHECK(agent.get_hsa_agent().handle != 0);
+
     aql::set_profiler_active_on_queue(
         CHECK_NOTNULL(hsa::get_queue_controller())->get_ext_table(),
         agent.cpu_pool(),
@@ -270,11 +273,6 @@ read_agent_ctx(const context::context*    ctx,
                               agent->get_rocp_agent()->cu_count,
                               agent->get_rocp_agent()->simd_arrays_per_engine);
 
-    // Remove when AQL is updated to not require stop to be called first
-    submitPacket(agent_ctx.callback_data.table,
-                 agent->profile_queue(),
-                 (void*) &agent_ctx.callback_data.packet->stop);
-
     // Submit the read packet to the queue
     submitPacket(agent_ctx.callback_data.table,
                  agent->profile_queue(),
diff --git a/source/lib/rocprofiler-sdk/counters/agent_profiling.hpp b/source/lib/rocprofiler-sdk/counters/agent_profiling.hpp
index 603828a3a1..f97ca166f5 100644
--- a/source/lib/rocprofiler-sdk/counters/agent_profiling.hpp
+++ b/source/lib/rocprofiler-sdk/counters/agent_profiling.hpp
@@ -82,5 +82,8 @@ read_agent_ctx(const context::context*    ctx,
                rocprofiler_user_data_t    user_data,
                rocprofiler_counter_flag_t flags);
 
+uint64_t
+submitPacket(const CoreApiTable& table, hsa_queue_t* queue, const void* packet);
+
 }  // namespace counters
 }  // namespace rocprofiler
\ No newline at end of file
diff --git a/source/lib/rocprofiler-sdk/counters/tests/agent_profiling.cpp b/source/lib/rocprofiler-sdk/counters/tests/agent_profiling.cpp
index eea73a64e4..4fbf308d5f 100644
--- a/source/lib/rocprofiler-sdk/counters/tests/agent_profiling.cpp
+++ b/source/lib/rocprofiler-sdk/counters/tests/agent_profiling.cpp
@@ -93,7 +93,6 @@ findDeviceMetrics(const hsa::AgentCache& agent, const std::unordered_set<std::st
             ret.push_back(counter);
         }
     }
-    ROCP_ERROR << "No counters found for " << std::string(agent.name());
     return ret;
 }
 
@@ -281,8 +280,37 @@ protected:
                                       &queue),
                      HSA_STATUS_SUCCESS);
 
+            // We don't use the queue interceptor, need to enabling profiling manually
+            hsa_amd_profiling_set_profiler_enabled(queue, 1);
+
+            hsa_signal_t completion_signal;
+            hsa_signal_create(1, 0, nullptr, &completion_signal);
+
+            CHECK(agent.cpu_pool().handle != 0);
+            CHECK(agent.get_hsa_agent().handle != 0);
+            // Set state of the queue to allow profiling (may not be needed since AQL
+            // may do this in the future).
+            aql::set_profiler_active_on_queue(
+                get_ext_table(),
+                agent.cpu_pool(),
+                agent.get_hsa_agent(),
+                [&](hsa::rocprofiler_packet pkt) {
+                    pkt.ext_amd_aql_pm4.completion_signal = completion_signal;
+                    submitPacket(queue, (void*) &pkt);
+
+                    if(hsa_signal_wait_relaxed(completion_signal,
+                                               HSA_SIGNAL_CONDITION_EQ,
+                                               0,
+                                               20000000,
+                                               HSA_WAIT_STATE_BLOCKED) != 0)
+                    {
+                        ROCP_FATAL << "Failed to set profiling mode on queue";
+                    }
+                    hsa_signal_store_relaxed(completion_signal, 1);
+                });
+
             rocprofiler::hsa::rocprofiler_packet barrier{};
-            hsa_signal_t                         completion_signal;
+
             hsa_signal_create(1, 0, nullptr, &completion_signal);
             barrier.barrier_and.header            = packet_header(HSA_PACKET_TYPE_BARRIER_AND);
             barrier.barrier_and.completion_signal = completion_signal;
@@ -423,6 +451,22 @@ TEST_F(agent_profile_test, sync_gpu_util_verify)
     test_run(ROCPROFILER_COUNTER_FLAG_NONE, {"GPU_UTIL"}, 50000);
     ROCP_ERROR << global_recs().size();
 
+    for(const auto& val : global_recs())
+    {
+        rocprofiler_counter_id_t id;
+        rocprofiler_query_record_counter_id(val.id, &id);
+        rocprofiler_counter_info_v0_t info;
+        rocprofiler_query_counter_info(id, ROCPROFILER_COUNTER_INFO_VERSION_0, &info);
+        ROCP_ERROR << fmt::format("Name: {} Counter value: {}", info.name, val.counter_value);
+        EXPECT_GT(val.counter_value, 0.0);
+    }
+}
+
+TEST_F(agent_profile_test, sync_sq_waves_verify)
+{
+    test_run(ROCPROFILER_COUNTER_FLAG_NONE, {"SQ_WAVES_sum"}, 50000);
+    ROCP_ERROR << global_recs().size();
+
     for(const auto& val : global_recs())
     {
         rocprofiler_counter_id_t id;
diff --git a/source/lib/rocprofiler-sdk/hsa/queue.cpp b/source/lib/rocprofiler-sdk/hsa/queue.cpp
index d8d0139e22..3d4f6b338b 100644
--- a/source/lib/rocprofiler-sdk/hsa/queue.cpp
+++ b/source/lib/rocprofiler-sdk/hsa/queue.cpp
@@ -495,6 +495,28 @@ Queue::Queue(const AgentCache&  agent,
                HSA_STATUS_SUCCESS)
         << "Could not setup intercept profiler";
 
+    CHECK(_agent.cpu_pool().handle != 0);
+    CHECK(_agent.get_hsa_agent().handle != 0);
+    // Set state of the queue to allow profiling
+    aql::set_profiler_active_on_queue(
+        CHECK_NOTNULL(hsa::get_queue_controller())->get_ext_table(),
+        _agent.cpu_pool(),
+        _agent.get_hsa_agent(),
+        [&](hsa::rocprofiler_packet pkt) {
+            hsa_signal_t completion;
+            create_signal(0, &completion);
+            pkt.ext_amd_aql_pm4.completion_signal = completion;
+            counters::submitPacket(CHECK_NOTNULL(hsa::get_queue_controller())->get_core_table(),
+                                   _intercept_queue,
+                                   (void*) &pkt);
+            if(core_api.hsa_signal_wait_relaxed_fn(
+                   completion, HSA_SIGNAL_CONDITION_EQ, 0, 20000000, HSA_WAIT_STATE_ACTIVE) != 0)
+            {
+                ROCP_FATAL << "Could not set agent to be profiled";
+            }
+            core_api.hsa_signal_destroy_fn(completion);
+        });
+
     LOG_IF(FATAL,
            _ext_api.hsa_amd_queue_intercept_register_fn(_intercept_queue, WriteInterceptor, this))
         << "Could not register interceptor";
diff --git a/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp b/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp
index b966f9e6d6..f81c9ea7be 100644
--- a/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp
+++ b/source/lib/rocprofiler-sdk/hsa/queue_controller.cpp
@@ -267,7 +267,8 @@ QueueController::init(CoreApiTable& core_table, AmdExtTable& ext_table)
             (itr->buffered_tracer &&
              itr->buffered_tracer->domains(ROCPROFILER_BUFFER_TRACING_KERNEL_DISPATCH));
 
-        if(itr->counter_collection || itr->pc_sampler || has_kernel_tracing)
+        if(itr->counter_collection || itr->pc_sampler || has_kernel_tracing ||
+           itr->agent_counter_collection)
         {
             enable_intercepter = true;
             break;