adding allocate/device/memcopy/submit callbacks

Change-Id: Ie83ae3cd57cdf5038e83db70891b103439c37d55

inc/rocprofiler.h

@@ -41,8 +41,9 @@ THE SOFTWARE.
 #ifndef INC_ROCPROFILER_H_
 #define INC_ROCPROFILER_H_
 
-#include <hsa.h>
 #include <amd_hsa_kernel_code.h>
+#include <hsa.h>
+#include <hsa_ext_amd.h>
 #include <hsa_ven_amd_aqlprofile.h>
 #include <stdint.h>
 
@@ -460,21 +461,14 @@ hsa_status_t rocprofiler_pool_flush(
   rocprofiler_pool_t* pool);          // profiling pool handle
 
 ////////////////////////////////////////////////////////////////////////////////
+// HSA intercepting API
+
 // HSA callbacks ID enumeration
 enum rocprofiler_hsa_cb_id_t {
-  ROCPROFILER_HSA_CB_ID_ALLOC = 0,
-  ROCPROFILER_HSA_CB_ID_ASSIGN = 0,
-  ROCPROFILER_HSA_CB_ID_MEMCOPY = 1,
-  ROCPROFILER_HSA_CB_ID_SUBMIT = 2
-};
-
-// HSA memory type enumeration
-enum rocprofiler_hsa_mtype_t {
-  ROCPROFILER_HSA_MTYPE_UNASSIGNED = 0,
-  ROCPROFILER_HSA_MTYPE_KERN_ARG = 1,
-  ROCPROFILER_HSA_MTYPE_SYSTEM = 2,
-  ROCPROFILER_HSA_MTYPE_CPU = 3
-  ROCPROFILER_HSA_MTYPE_GPU = 4
+  ROCPROFILER_HSA_CB_ID_ALLOCATE = 0,
+  ROCPROFILER_HSA_CB_ID_DEVICE = 1,
+  ROCPROFILER_HSA_CB_ID_MEMCOPY = 2,
+  ROCPROFILER_HSA_CB_ID_SUBMIT = 3
 };
 
 // HSA callback data type
@@ -483,17 +477,14 @@ struct rocprofiler_hsa_callback_data_t {
     struct {
       const void* addr;
       size_t size;
-      hsa_segment_t sement;
-      hsa_global_flag_t global_flag;
-      hsa_device_type_t device_type;
-    } hsa_alloc;
-    struct {
-      const void* addr;
-      size_t size;
-      hsa_amd_segment_t sement;
+      hsa_amd_segment_t segment;
       hsa_amd_memory_pool_global_flag_t global_flag;
-      hsa_device_type_t device_type;
-    } amd_alloc;
+    } allocate;
+    struct {
+      hsa_device_type_t type;
+      uint32_t id;
+      const void* mem;
+    } device;
     struct {
       const void* dst;
       const void* src;
@@ -501,19 +492,21 @@ struct rocprofiler_hsa_callback_data_t {
     } memcopy;
     struct {
       const void* packet;
+      const amd_kernel_code_t* kernel_code;
     } submit;
   };
 };
 
 // HSA callback function type
 typedef hsa_status_t (*rocprofiler_hsa_callback_fun_t)(
-  rocprofiler_hsa_callback_id_t id, // callback id
+  rocprofiler_hsa_cb_id_t id, // callback id
   const rocprofiler_hsa_callback_data_t* data, // [in] callback data
   void* arg); // [in/out] user passed data
 
 // HSA callbacks structure
 struct rocprofiler_hsa_callbacks_t {
-  rocprofiler_hsa_callback_fun_t alloc; // memory allocate callback
+  rocprofiler_hsa_callback_fun_t allocate; // memory allocate callback
+  rocprofiler_hsa_callback_fun_t device; // agent assign callback
   rocprofiler_hsa_callback_fun_t memcopy; // memory copy callback
   rocprofiler_hsa_callback_fun_t submit; // packet submit callback
 };

src/core/hsa_interceptor.h

@@ -25,40 +25,79 @@ SOFTWARE.
 #ifndef _SRC_CORE_HSA_INTERCEPTOR_H
 #define _SRC_CORE_HSA_INTERCEPTOR_H
 
-#include <amd_hsa_kernel_code.h>
+#include <hsa.h>
+#include <hsa_ext_amd.h>
+
 #include <atomic>
+#include <mutex>
 
 #include "inc/rocprofiler.h"
+#include "util/exception.h"
+#include "util/hsa_rsrc_factory.h"
+
+#define HSA_RT(call) \
+  do { \
+    const hsa_status_t status = call; \
+    if (status != HSA_STATUS_SUCCESS) EXC_ABORT(status, #call); \
+  } while(0)
+
+#define IS_HSA_CALLBACK(ID) \
+  const auto __id = ID; (void)__id; \
+  void *__arg = arg_.load(); (void)__arg; \
+  rocprofiler_hsa_callback_fun_t __callback = \
+    (ID == ROCPROFILER_HSA_CB_ID_ALLOCATE) ? callbacks_.allocate: \
+    (ID == ROCPROFILER_HSA_CB_ID_DEVICE) ? callbacks_.device: \
+    (ID == ROCPROFILER_HSA_CB_ID_MEMCOPY) ? callbacks_.memcopy: \
+                                            callbacks_.submit; \
+  if (__callback != NULL)
+
+#define DO_HSA_CALLBACK \
+  do { __callback(__id, &data, __arg); } while (0)
+
+#define ISSUE_HSA_CALLBACK(ID) \
+  IS_HSA_CALLBACK(ID) { DO_HSA_CALLBACK; }
 
 namespace rocprofiler {
 extern decltype(hsa_memory_allocate)* hsa_memory_allocate_fn;
 extern decltype(hsa_memory_assign_agent)* hsa_memory_assign_agent_fn;
-extern decltype(hsa_amd_memory_pool_allocate)* hsa_amd_memory_pool_allocate_fn;
 extern decltype(hsa_memory_copy)* hsa_memory_copy_fn;
+extern decltype(hsa_amd_memory_pool_allocate)* hsa_amd_memory_pool_allocate_fn;
+extern decltype(hsa_amd_agents_allow_access)* hsa_amd_agents_allow_access_fn;
 extern decltype(hsa_amd_memory_async_copy)* hsa_amd_memory_async_copy_fn;
 
 class HsaInterceptor {
  public:
+  typedef std::atomic<void*> arg_t;
+  typedef std::mutex mutex_t;
+
   static void Enable(const bool& enable) { enable_ = enable; }
 
   static void HsaIntercept(HsaApiTable* table) {
+    fprintf(stderr, "HsaInterceptor ...\n");
     if (enable_) {
-      table->amd_ext_->hsa_memory_allocate_fn = MemoryAllocate;
-      table->amd_ext_->hsa_memory_assign_agent_fn = MemoryAssignAgent;
-      table->amd_ext_->hsa_amd_memory_pool_allocate_fn = MemoryPoolAllocate;;
+      fprintf(stderr, "HsaInterceptor enabled\n");
+      // saving original API functions
+      hsa_memory_allocate_fn = table->core_->hsa_memory_allocate_fn;
+      hsa_memory_assign_agent_fn = table->core_->hsa_memory_assign_agent_fn;
+      hsa_memory_copy_fn = table->core_->hsa_memory_copy_fn;
+      hsa_amd_memory_pool_allocate_fn = table->amd_ext_->hsa_amd_memory_pool_allocate_fn;
+      hsa_amd_agents_allow_access_fn = table->amd_ext_->hsa_amd_agents_allow_access_fn;
+      hsa_amd_memory_async_copy_fn = table->amd_ext_->hsa_amd_memory_async_copy_fn;
+      // intercepting API
+      table->core_->hsa_memory_allocate_fn = MemoryAllocate;
+      table->core_->hsa_memory_assign_agent_fn = MemoryAssignAgent;
       table->core_->hsa_memory_copy_fn = MemoryCopy;
+      table->amd_ext_->hsa_amd_memory_pool_allocate_fn = MemoryPoolAllocate;
+      table->amd_ext_->hsa_amd_agents_allow_access_fn = AgentsAllowAccess;
       table->amd_ext_->hsa_amd_memory_async_copy_fn = MemoryAsyncCopy;
     }
+    fprintf(stderr, "HsaInterceptor done\n");
   }
 
-  static void SetHsaAllocCallback(rocprofiler_hsa_callback_fun_t fun, void* arg) {
-    alloc_callback_arg_ = arg;
-    alloc_callback_fun_.store(fun);
-  }
-
-  static void SetHsaMemcopyCallback(rocprofiler_hsa_callback_fun_t fun, void* arg) {
-    memcopy_callback_arg_ = arg;
-    memcopy_callback_fun_.store(fun);
+  static void SetCallbacks(rocprofiler_hsa_callbacks_t callbacks, void* arg) {
+    std::lock_guard<mutex_t> lck(mutex_);
+    callbacks_ = callbacks;
+    arg_.store(arg);
   }
 
  private:
@@ -66,20 +105,18 @@ class HsaInterceptor {
     size_t size,
     void** ptr)
   {
-    const hsa_status_t status = hsa_memory_allocate_fn(region, size, ptr);
-    if ((status == HSA_STATUS_SUCCESS) && (alloc_callback_fun_ != NULL)) {
+    EXC_ABORT(HSA_STATUS_ERROR, "Deprecated API");
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_memory_allocate_fn(region, size, ptr));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_ALLOCATE) {
       rocprofiler_hsa_callback_data_t data{};
-      data.hsa_alloc.addr = *ptr;
-      data.hsa_alloc.size = size;
-      data.hsa_alloc.device_type = HSA_DEVICE_TYPE_CPU;
+      data.allocate.addr = *ptr;
+      data.allocate.size = size;
 
-      hsa_status_t err = hsa_region_get_info(region, HSA_REGION_INFO_SEGMENT, &data.hsa_alloc.segment);
-      if (err != HSA_STATUS_SUCCESS) data.hsa_alloc.addr = NULL;
-      err = hsa_region_get_info(region, HSA_REGION_INFO_GLOBAL_FLAGS, &data.hsa_alloc.global_flag);
-      if (err != HSA_STATUS_SUCCESS) data.hsa_alloc.addr = NULL;
+      HSA_RT(hsa_region_get_info(region, HSA_REGION_INFO_SEGMENT, &data.allocate.segment));
+      HSA_RT(hsa_region_get_info(region, HSA_REGION_INFO_GLOBAL_FLAGS, &data.allocate.global_flag));
 
-      const hsa_status_t ret = alloc_callback_fun_(ROCPROFILER_HSA_CB_ID_ALLOC, &data, alloc_callback_arg_);
-      if (ret != HSA_STATUS_SUCCESS) memcopy_callback_fun_.store(NULL);
+      DO_HSA_CALLBACK;
     }
     return status;
   }
@@ -89,36 +126,95 @@ class HsaInterceptor {
     hsa_agent_t agent,
     hsa_access_permission_t access)
   {
-    rocprofiler_hsa_callback_data_t data{};
-    data.hsa_alloc.addr = ptr;
+    EXC_ABORT(HSA_STATUS_ERROR, "Deprecated API");
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_memory_assign_agent_fn(ptr, agent, access));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_DEVICE) {
+      rocprofiler_hsa_callback_data_t data{};
+      data.device.mem = ptr;
 
-    hsa_status_t err = hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &data.hsa_alloc.device_type);
-    if (err != HSA_STATUS_SUCCESS) data.hsa_alloc.addr = NULL;
+      HSA_RT(hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &data.device.type));
 
-    alloc_callback_fun_(ROCPROFILER_HSA_CB_ID_ASSIGN, &data, alloc_callback_arg_);
-    return hsa_memory_assign_agent(ptr, agent, access);
+      DO_HSA_CALLBACK;
+    }
+    return status;
+  }
+
+  // Spawn device allow access callback
+  static void DeviceCallback(
+    uint32_t num_agents,
+    const hsa_agent_t* agents,
+    const void* ptr)
+  {
+    for (const hsa_agent_t* agent_p = agents; agent_p < (agents + num_agents); ++agent_p) {
+      hsa_agent_t agent = *agent_p;
+      rocprofiler_hsa_callback_data_t data{};
+      data.device.id = util::HsaRsrcFactory::Instance().GetAgentInfo(agent)->dev_index;
+      data.device.mem = ptr;
+
+      HSA_RT(hsa_agent_get_info(agent, HSA_AGENT_INFO_DEVICE, &data.device.type));
+
+      ISSUE_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_DEVICE);
+    }
+  }
+
+  // Agent allow access callback 'hsa_amd_agents_allow_access'
+  static hsa_status_t AgentsAllowAccess(
+    uint32_t num_agents,
+    const hsa_agent_t* agents,
+    const uint32_t* flags,
+    const void* ptr)
+  {
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_amd_agents_allow_access_fn(num_agents, agents, flags, ptr));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_DEVICE) {
+      DeviceCallback(num_agents, agents, ptr);
+    }
+    return status;
+  }
+
+  // Callback function to get available in the system agents
+  struct agent_callback_data_t {
+    hsa_amd_memory_pool_t pool;
+    void* addr;
+  };
+  static hsa_status_t AgentCallback(hsa_agent_t agent, void* data) {
+    agent_callback_data_t* callback_data = reinterpret_cast<agent_callback_data_t*>(data);
+    hsa_amd_agent_memory_pool_info_t attribute = HSA_AMD_AGENT_MEMORY_POOL_INFO_ACCESS;
+    hsa_amd_memory_pool_access_t value;
+    HSA_RT(hsa_amd_agent_memory_pool_get_info(agent, callback_data->pool, attribute, &value));
+    if (value == HSA_AMD_MEMORY_POOL_ACCESS_ALLOWED_BY_DEFAULT) {
+      DeviceCallback(1, &agent, callback_data->addr);
+    }
+    return HSA_STATUS_SUCCESS;
   }
 
   static hsa_status_t MemoryPoolAllocate(
-    hsa_amd_memory_pool_t memory_pool, size_t size,
-    uint32_t flags, void** ptr)
+    hsa_amd_memory_pool_t pool,
+    size_t size,
+    uint32_t flags,
+    void** ptr)
   {
-    const hsa_status_t status = hsa_amd_memory_pool_allocate_fn(memory_pool, size, flags, ptr);
-    if ((status == HSA_STATUS_SUCCESS) && (alloc_callback_fun_ != NULL)) {
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_amd_memory_pool_allocate_fn(pool, size, flags, ptr));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_ALLOCATE) {
       rocprofiler_hsa_callback_data_t data{};
-      data.alloc.addr = *ptr;
-      data.alloc.size = size;
-      data.hsa_alloc.device_type = HSA_DEVICE_TYPE_CPU;
+      data.allocate.addr = *ptr;
+      data.allocate.size = size;
 
-      hsa_status_t err = hsa_amd_memory_pool_get_info(memory_pool, HSA_AMD_MEMORY_POOL_INFO_SEGMENT, &data.hsa_alloc.segment);
-      if (err != HSA_STATUS_SUCCESS) data.pool_alloc.addr = NULL;
-      hsa_status_t err = hsa_amd_memory_pool_get_info(memory_pool, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &data.hsa_alloc.global_flag);
-      if (err != HSA_STATUS_SUCCESS) data.pool_alloc.addr = NULL;
-      hsa_status_t err = hsa_amd_memory_pool_get_info(memory_pool, HSA_AMD_MEMORY_POOL_INFO_ACCESSIBLE_BY_ALL, &data.hsa_alloc.global_mem);
-      if (err != HSA_STATUS_SUCCESS) data.pool_alloc.addr = NULL;
+      HSA_RT(hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_SEGMENT, &data.allocate.segment));
+      HSA_RT(hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_GLOBAL_FLAGS, &data.allocate.global_flag));
+#if 0
+      HSA_RT(hsa_amd_memory_pool_get_info(pool, HSA_AMD_MEMORY_POOL_INFO_ACCESSIBLE_BY_ALL, &data.allocate.global_mem));
+#endif
 
-      const hsa_status_t ret = alloc_callback_fun_(ROCPROFILER_HSA_CB_ID_POOL_ALLOC, &data, memcopy_callback_arg_);
-      if (ret != HSA_STATUS_SUCCESS) memcopy_callback_fun_.store(NULL);
+      DO_HSA_CALLBACK;
+
+      IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_DEVICE) {
+        // Scan the pool assigned devices
+        agent_callback_data_t callback_data{pool, *ptr};
+        hsa_iterate_agents(AgentCallback, &callback_data);
+      }
     }
     return status;
   }
@@ -128,15 +224,16 @@ class HsaInterceptor {
     const void *src,
     size_t size)
   {
-    if (memcopy_callback_fun_) {
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_memory_copy_fn(dst, src, size));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_MEMCOPY) {
       rocprofiler_hsa_callback_data_t data{};
       data.memcopy.dst = dst;
       data.memcopy.src = src;
       data.memcopy.size = size;
-      const hsa_status_t ret = memcopy_callback_fun_(ROCPROFILER_HSA_CB_ID_MEMCOPY, &data, memcopy_callback_arg_);
-      if (ret != HSA_STATUS_SUCCESS) memcopy_callback_fun_.store(NULL);
+      DO_HSA_CALLBACK;
     }
-    return hsa_memory_copy(dst, src, size);
+    return status;
   }
 
   static hsa_status_t MemoryAsyncCopy(
@@ -146,24 +243,24 @@ class HsaInterceptor {
     const hsa_signal_t* dep_signals,
     hsa_signal_t completion_signal)
   {
-    if (memcopy_callback_fun_) {
+    hsa_status_t status = HSA_STATUS_SUCCESS;
+    HSA_RT(hsa_amd_memory_async_copy_fn(
+      dst, dst_agent, src, src_agent, size,
+      num_dep_signals, dep_signals, completion_signal));
+    IS_HSA_CALLBACK(ROCPROFILER_HSA_CB_ID_MEMCOPY) {
       rocprofiler_hsa_callback_data_t data{};
       data.memcopy.dst = dst;
       data.memcopy.src = src;
       data.memcopy.size = size;
-      const hsa_status_t ret = memcopy_callback_fun_(ROCPROFILER_HSA_CB_ID_MEMCOPY, &data, memcopy_callback_arg_);
-      if (ret != HSA_STATUS_SUCCESS) memcopy_callback_fun_.store(NULL);
+      DO_HSA_CALLBACK;
     }
-    return hsa_amd_memory_async_copy_fn(
-      dst, dst_agent, src, src_agent, size,
-      num_dep_signals, dep_signals, completion_signal);
+    return status;
   }
 
   static bool enable_;
-  static std::atomic<rocprofiler_hsa_callback_fun_t> alloc_callback_fun_;
-  static void* alloc_callback_arg_;
-  static std::atomic<rocprofiler_hsa_callback_fun_t> memcopy_callback_fun_;
-  static void* memcopy_callback_arg_;
+  static rocprofiler_hsa_callbacks_t callbacks_;
+  static arg_t arg_;
+  static mutex_t mutex_;
 };
 
 }  // namespace rocprofiler

src/core/intercept_queue.cpp

@@ -39,4 +39,7 @@ bool InterceptQueue::tracker_on_ = false;
 bool InterceptQueue::in_create_call_ = false;
 InterceptQueue::queue_id_t InterceptQueue::current_queue_id = 0;
 
+rocprofiler_hsa_callback_fun_t InterceptQueue::submit_callback_fun_ = NULL;
+void* InterceptQueue::submit_callback_arg_ = NULL;
+
 }  // namespace rocprofiler

src/core/intercept_queue.h

@@ -131,10 +131,14 @@ class InterceptQueue {
 
     if (submit_callback_fun_) {
       for (uint64_t j = 0; j < count; ++j) {
-        const void* packet = reinterpret_cast<const void*>(&packets_arr[j]);
+        const packet_t* packet = &packets_arr[j];
+        const hsa_kernel_dispatch_packet_t* dispatch_packet =
+            reinterpret_cast<const hsa_kernel_dispatch_packet_t*>(packet);
         rocprofiler_hsa_callback_data_t data{};
-        data.submit.packet = packet;
-        submit_callback_fun_(CB_ID_SUBMIT, &data, submit_callback_arg_);
+        data.submit.packet = (void*)packet;
+        data.submit.kernel_code =
+          (GetHeaderType(packet) == HSA_PACKET_TYPE_KERNEL_DISPATCH) ? GetKernelCode(dispatch_packet) : NULL;
+        submit_callback_fun_(ROCPROFILER_HSA_CB_ID_SUBMIT, &data, submit_callback_arg_);
       }
     }
 
@@ -255,7 +259,8 @@ class InterceptQueue {
   static inline void Start() { dispatch_callback_.store(callbacks_.dispatch, std::memory_order_release); }
   static inline void Stop() { dispatch_callback_.store(NULL, std::memory_order_relaxed); }
 
-  static void SetHsaSubmitCallback(rocprofiler_hsa_callback_fun_t fun, void* arg) {
+  static void SetSubmitCallback(rocprofiler_hsa_callback_fun_t fun, void* arg) {
+    std::lock_guard<mutex_t> lck(mutex_);
     submit_callback_fun_ = fun;
     submit_callback_arg_ = arg;
   }

src/core/rocprofiler.cpp

@@ -86,8 +86,10 @@ decltype(hsa_amd_queue_intercept_create)* hsa_amd_queue_intercept_create_fn;
 decltype(hsa_amd_queue_intercept_register)* hsa_amd_queue_intercept_register_fn;
 
 decltype(hsa_memory_allocate)* hsa_memory_allocate_fn;
-decltype(hsa_amd_memory_pool_allocate)* hsa_amd_memory_pool_allocate_fn;
+decltype(hsa_memory_assign_agent)* hsa_memory_assign_agent_fn;
 decltype(hsa_memory_copy)* hsa_memory_copy_fn;
+decltype(hsa_amd_memory_pool_allocate)* hsa_amd_memory_pool_allocate_fn;
+decltype(hsa_amd_agents_allow_access)* hsa_amd_agents_allow_access_fn;
 decltype(hsa_amd_memory_async_copy)* hsa_amd_memory_async_copy_fn;
 decltype(hsa_amd_memory_async_copy_rect)* hsa_amd_memory_async_copy_rect_fn;
 
@@ -154,6 +156,7 @@ enum {
 uint32_t LoadTool() {
   uint32_t intercept_mode = 0;
   const char* tool_lib = getenv("ROCP_TOOL_LIB");
+  fprintf(stderr, "ROCProfiler: load tool library \"%s\"\n", tool_lib); fflush(stderr);
 
   if (tool_lib) {
     intercept_mode = DISPATCH_INTERCEPT_MODE;
@@ -409,6 +412,7 @@ extern "C" {
 // HSA-runtime tool on-load method
 PUBLIC_API bool OnLoad(HsaApiTable* table, uint64_t runtime_version, uint64_t failed_tool_count,
                        const char* const* failed_tool_names) {
+  fprintf(stderr, "rocprof OnLoad\n"); fflush(stderr);
   rocprofiler::SaveHsaApi(table);
   rocprofiler::ProxyQueue::InitFactory();
   bool intercept_mode = false;
@@ -872,12 +876,23 @@ hsa_status_t rocprofiler_get_time(
 }
 
 // Set new callbacks. If a callback is NULL then it is disabled
+}  // extern "C"
+
+// HSA API callbacks routines
+
+// Static fields
+bool rocprofiler::HsaInterceptor::enable_ = false;
+rocprofiler_hsa_callbacks_t rocprofiler::HsaInterceptor::callbacks_{};
+rocprofiler::HsaInterceptor::arg_t rocprofiler::HsaInterceptor::arg_{};
+rocprofiler::HsaInterceptor::mutex_t rocprofiler::HsaInterceptor::mutex_;
+
+extern "C" {
+// Set HSA callbacks. If a callback is NULL then it is disabled
 PUBLIC_API hsa_status_t rocprofiler_set_hsa_callbacks(const rocprofiler_hsa_callbacks_t callbacks, void* arg) {
   API_METHOD_PREFIX
-  rocprofiler::InterceptQueue::SetHsaSubmitCallback(callbacks.submit, arg);
-  rocprofiler::HsaInterceptor::SetHsaAllocCallback(callbacks.alloc, arg);
-  rocprofiler::HsaInterceptor::SetHsaMemcopyCallback(callbacks.memcopy, arg);
+  rocprofiler::HsaInterceptor::SetCallbacks(callbacks, arg);
+  rocprofiler::InterceptQueue::SetSubmitCallback(callbacks.submit, arg);
   API_METHOD_SUFFIX
 }
-
 }  // extern "C"
+

test/run.sh

@@ -52,18 +52,22 @@ eval_test() {
   test_number=$((test_number + 1))
 }
 
-# enable tools load failure reporting
-export HSA_TOOLS_REPORT_LOAD_FAILURE=1
 # paths to ROC profiler and oher libraries
 export LD_LIBRARY_PATH=$PWD
+
+# enable tools load failure reporting
+export HSA_TOOLS_REPORT_LOAD_FAILURE=1
 # enable error messages logging to '/tmp/rocprofiler_log.txt'
 export ROCPROFILER_LOG=1
-# ROC profiler metrics config file
+# enable error messages logging to '/tmp/aql_profile_log.txt'
+export HSA_VEN_AMD_AQLPROFILE_LOG=1
+# test trace
+export ROC_TEST_TRACE=1
+
+# Disabple profiler own proxy queue
 unset ROCP_PROXY_QUEUE
 # ROC profiler metrics config file
 export ROCP_METRICS=metrics.xml
-# test trace
-export ROC_TEST_TRACE=1
 
 ## C test
 eval_test "C test" ./test/c_test
@@ -93,6 +97,9 @@ export ROCP_TIMESTAMP_ON=1
 # and SQTT trace files 'thread_trace.se<n>.out'
 export ROCP_OUTPUT_DIR=./RESULTS
 
+# enable HSA intercepting
+export ROCP_HSA_INTERC=1
+
 if [ ! -e $ROCP_TOOL_LIB ] ; then
   export ROCP_TOOL_LIB=test/libtool.so
 fi

test/tool/tool.cpp

@@ -826,6 +826,37 @@ static inline void check_env_var(const char* var_name, uint64_t& val) {
   if (str != NULL ) val = atoll(str);
 }
 
+// HSA intercepting routines
+
+// HSA unified callback function
+hsa_status_t hsa_unified_callback(
+  rocprofiler_hsa_cb_id_t id,
+  const rocprofiler_hsa_callback_data_t* data,
+  void* arg)
+{
+  printf("hsa_unified_callback(%d, %p, %p\n", (int)id, data, arg);
+
+  if (id == ROCPROFILER_HSA_CB_ID_SUBMIT) {
+    if (data->submit.kernel_code != NULL) {
+      printf("  submit:\n");
+      printf("    kernel_code_entry_byte_offset %lx\n", data->submit.kernel_code->kernel_code_entry_byte_offset);
+      printf("    kernel_code_prefetch_byte_offset %lx\n", data->submit.kernel_code->kernel_code_prefetch_byte_offset);
+      printf("    kernel_code_prefetch_byte_size %lx\n", data->submit.kernel_code->kernel_code_prefetch_byte_size);
+    }
+  }
+
+  fflush(stdout);
+  return HSA_STATUS_SUCCESS;
+}
+
+// HSA callbacks structure
+rocprofiler_hsa_callbacks_t hsa_callbacks {
+  hsa_unified_callback,
+  hsa_unified_callback,
+  hsa_unified_callback,
+  hsa_unified_callback
+};
+
 // Tool constructor
 extern "C" PUBLIC_API void OnLoadToolProp(rocprofiler_settings_t* settings)
 {
@@ -908,6 +939,8 @@ extern "C" PUBLIC_API void OnLoadToolProp(rocprofiler_settings_t* settings)
   check_env_var("ROCP_OBJ_TRACKING", settings->code_obj_tracking);
   // Set memcopies tracking
   check_env_var("ROCP_MCOPY_TRACKING", settings->memcopy_tracking);
+  // Set HSA intercepting
+  check_env_var("ROCP_HSA_INTERC", settings->hsa_intercepting);
 
   is_trace_local = settings->trace_local;