From 1e1692e6c09b3b04733f2fefadea57bba0d7ab66 Mon Sep 17 00:00:00 2001
From: Saurabh Verma <sauverma@amd.com>
Date: Wed, 18 Oct 2023 18:14:32 -0500
Subject: [PATCH] SWDEV-427921: Fixing rocprof v1 counter collection for PAPI
 on MI300

Change-Id: I8589fe78fca481ed326e584ea1d01c0f90de655e
---
 src/core/context.h            | 13 ++++++++++---
 src/util/hsa_rsrc_factory.cpp |  5 +++++
 src/util/hsa_rsrc_factory.h   |  1 +
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/core/context.h b/src/core/context.h
index 3ce46e080d..9f11eff8f0 100644
--- a/src/core/context.h
+++ b/src/core/context.h
@@ -265,6 +265,8 @@ class Context {
     info_vector_t* info_vector;
     size_t index;
     char* ptr;
+    size_t single_xcc_buff_size;
+    size_t cb_invocation_count; 
   };
 
   void RestoreSignals(const profile_tuple_t& tuple) {
@@ -282,10 +284,14 @@ class Context {
       hsa_rsrc_->SignalWaitRestore(tuple.completion_signal, 1);
       // Restore other signals
       RestoreSignals(tuple);
-      for (rocprofiler_feature_t* rinfo : *(tuple.info_vector))
+      for (rocprofiler_feature_t* rinfo : *(tuple.info_vector)){
         rinfo->data.kind = ROCPROFILER_DATA_KIND_UNINIT;
+        rinfo->data.result_int64 = 0;
+      }
+      uint32_t xcc_count = agent_info_->xcc_num;
+      uint32_t single_xcc_buff_size = tuple.profile->output_buffer.size / (sizeof(uint64_t) * xcc_count);
       callback_data_t callback_data{tuple.profile, tuple.info_vector, tuple.info_vector->size(),
-                                    NULL};
+                                    NULL, single_xcc_buff_size, 0};
       const hsa_status_t status =
           api_->hsa_ven_amd_aqlprofile_iterate_data(tuple.profile, DataCallback, &callback_data);
       if (status != HSA_STATUS_SUCCESS) AQL_EXC_RAISING(status, "context iterate data failed");
@@ -540,6 +546,8 @@ class Context {
     } else {
       if (sample_id == 0) index += 1;
     }
+    if(callback_data->cb_invocation_count++ % callback_data->single_xcc_buff_size == 0)
+      index = 0;
     callback_data->index = index;
 
     if (index < info_vector.size()) {
@@ -547,7 +555,6 @@ class Context {
       rinfo->data.kind = ROCPROFILER_DATA_KIND_UNINIT;
 
       if (ainfo_type == HSA_VEN_AMD_AQLPROFILE_INFO_PMC_DATA) {
-        if (ainfo_data->sample_id == 0) rinfo->data.result_int64 = 0;
         rinfo->data.result_int64 += ainfo_data->pmc_data.result;
         rinfo->data.kind = ROCPROFILER_DATA_KIND_INT64;
       } else if (ainfo_type == HSA_VEN_AMD_AQLPROFILE_INFO_TRACE_DATA) {
diff --git a/src/util/hsa_rsrc_factory.cpp b/src/util/hsa_rsrc_factory.cpp
index 156d52520e..8689b92208 100644
--- a/src/util/hsa_rsrc_factory.cpp
+++ b/src/util/hsa_rsrc_factory.cpp
@@ -390,6 +390,11 @@ const AgentInfo* HsaRsrcFactory::AddAgentInfo(const hsa_agent_t agent) {
     agent_info->sgpr_block_size = 8;
     agent_info->vgpr_block_size = 4;
 
+    if (hsa_api_.hsa_agent_get_info(agent, static_cast<hsa_agent_info_t>(HSA_AMD_AGENT_INFO_NUM_XCC),
+                        &agent_info->xcc_num) != HSA_STATUS_SUCCESS) {
+        agent_info->xcc_num = 1;
+    };
+
     // Set GPU index
     uint32_t driver_node_id;
     status = hsa_api_.hsa_agent_get_info(
diff --git a/src/util/hsa_rsrc_factory.h b/src/util/hsa_rsrc_factory.h
index c9bad5f957..6426cf953d 100644
--- a/src/util/hsa_rsrc_factory.h
+++ b/src/util/hsa_rsrc_factory.h
@@ -172,6 +172,7 @@ struct AgentInfo {
   uint32_t sgpr_block_size;
   uint32_t vgpr_block_size;
   static const uint32_t lds_block_size = 128 * 4;
+  uint32_t xcc_num;
 };
 
 // HSA timer class