adding hsa-interceptor, preliminary

Change-Id: I18f7ba3df223cb16bc6c055792834a3fee7c2373

inc/rocprofiler.h

@@ -70,6 +70,7 @@ typedef struct {
   uint32_t trace_local;
   uint64_t timeout;
   uint32_t timestamp_on;
+  uint32_t hsa_intercepting;
 } rocprofiler_settings_t;
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -459,6 +460,71 @@ hsa_status_t rocprofiler_pool_flush(
   rocprofiler_pool_t* pool);          // profiling pool handle
 
 ////////////////////////////////////////////////////////////////////////////////
+// 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
+};
+
+// HSA callback data type
+struct rocprofiler_hsa_callback_data_t {
+  union {
+    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_memory_pool_global_flag_t global_flag;
+      hsa_device_type_t device_type;
+    } amd_alloc;
+    struct {
+      const void* dst;
+      const void* src;
+      size_t size;
+    } memcopy;
+    struct {
+      const void* packet;
+    } submit;
+  };
+};
+
+// HSA callback function type
+typedef hsa_status_t (*rocprofiler_hsa_callback_fun_t)(
+  rocprofiler_hsa_callback_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 memcopy; // memory copy callback
+  rocprofiler_hsa_callback_fun_t submit; // packet submit callback
+};
+
+// Set callbacks. If the callback is NULL then it is disabled.
+// If callback returns a value that is not HSA_STATUS_SUCCESS the  callback
+// will be unregistered.
+hsa_status_t rocprofiler_set_hsa_callbacks(
+  const rocprofiler_hsa_callbacks_t callbacks, // HSA callback function
+  void* arg); // callback user data
+
 #ifdef __cplusplus
 }  // extern "C" block
 #endif  // __cplusplus

src/core/hsa_interceptor.h

@@ -0,0 +1,171 @@
+/******************************************************************************
+MIT License
+
+Copyright (c) 2018 ROCm Core Technology
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*******************************************************************************/
+
+#ifndef _SRC_CORE_HSA_INTERCEPTOR_H
+#define _SRC_CORE_HSA_INTERCEPTOR_H
+
+#include <amd_hsa_kernel_code.h>
+#include <atomic>
+
+#include "inc/rocprofiler.h"
+
+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_async_copy)* hsa_amd_memory_async_copy_fn;
+
+class HsaInterceptor {
+ public:
+  static void Enable(const bool& enable) { enable_ = enable; }
+
+  static void HsaIntercept(HsaApiTable* table) {
+    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;;
+      table->core_->hsa_memory_copy_fn = MemoryCopy;
+      table->amd_ext_->hsa_amd_memory_async_copy_fn = MemoryAsyncCopy;
+    }
+  }
+
+  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);
+  }
+
+ private:
+  static hsa_status_t HSA_API MemoryAllocate(hsa_region_t region,
+    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)) {
+      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;
+
+      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;
+
+      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);
+    }
+    return status;
+  }
+
+  static hsa_status_t MemoryAssignAgent(
+    void *ptr,
+    hsa_agent_t agent,
+    hsa_access_permission_t access)
+  {
+    rocprofiler_hsa_callback_data_t data{};
+    data.hsa_alloc.addr = 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;
+
+    alloc_callback_fun_(ROCPROFILER_HSA_CB_ID_ASSIGN, &data, alloc_callback_arg_);
+    return hsa_memory_assign_agent(ptr, agent, access);
+  }
+
+  static hsa_status_t MemoryPoolAllocate(
+    hsa_amd_memory_pool_t memory_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)) {
+      rocprofiler_hsa_callback_data_t data{};
+      data.alloc.addr = *ptr;
+      data.alloc.size = size;
+      data.hsa_alloc.device_type = HSA_DEVICE_TYPE_CPU;
+
+      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;
+
+      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);
+    }
+    return status;
+  }
+
+  static hsa_status_t MemoryCopy(
+    void *dst,
+    const void *src,
+    size_t size)
+  {
+    if (memcopy_callback_fun_) {
+      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);
+    }
+    return hsa_memory_copy(dst, src, size);
+  }
+
+  static hsa_status_t MemoryAsyncCopy(
+    void* dst, hsa_agent_t dst_agent, const void* src,
+    hsa_agent_t src_agent, size_t size,
+    uint32_t num_dep_signals,
+    const hsa_signal_t* dep_signals,
+    hsa_signal_t completion_signal)
+  {
+    if (memcopy_callback_fun_) {
+      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);
+    }
+    return hsa_amd_memory_async_copy_fn(
+      dst, dst_agent, src, src_agent, size,
+      num_dep_signals, dep_signals, completion_signal);
+  }
+
+  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_;
+};
+
+}  // namespace rocprofiler
+
+#endif  // _SRC_CORE_HSA_INTERCEPTOR_H

src/core/hsa_listener.h

@@ -0,0 +1,47 @@
+/******************************************************************************
+MIT License
+
+Copyright (c) 2018 ROCm Core Technology
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*******************************************************************************/
+
+#ifndef _SRC_CORE_HSA_LISTENER_H
+#define _SRC_CORE_HSA_LISTENER_H
+
+#include <hsa.h>
+
+#include "util/exception.h"
+
+namespace rocprofiler {
+extern decltype(hsa_queue_destroy)* hsa_queue_destroy_fn;
+extern decltype(hsa_amd_queue_intercept_create)* hsa_amd_queue_intercept_create_fn;
+extern decltype(hsa_amd_queue_intercept_register)* hsa_amd_queue_intercept_register_fn;
+
+class HsaListener {
+ public:
+  hsa_status_t SetCallbacks(rocprofiler_hsa_callbacks_t callbacks, void* data) {
+  }
+
+ private:
+};
+
+}  // namespace rocprofiler
+
+#endif  // _SRC_CORE_HSA_LISTENER_H

src/core/intercept_queue.h

@@ -129,6 +129,15 @@ class InterceptQueue {
     InterceptQueue* obj = reinterpret_cast<InterceptQueue*>(data);
     Queue* proxy = obj->proxy_;
 
+    if (submit_callback_fun_) {
+      for (uint64_t j = 0; j < count; ++j) {
+        const void* packet = reinterpret_cast<const void*>(&packets_arr[j]);
+        rocprofiler_hsa_callback_data_t data{};
+        data.submit.packet = packet;
+        submit_callback_fun_(CB_ID_SUBMIT, &data, submit_callback_arg_);
+      }
+    }
+
     // Travers input packets
     for (uint64_t j = 0; j < count; ++j) {
       const packet_t* packet = &packets_arr[j];
@@ -246,6 +255,11 @@ 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) {
+    submit_callback_fun_ = fun;
+    submit_callback_arg_ = arg;
+  }
+
   static void TrackerOn(bool on) { tracker_on_ = on; }
   static bool IsTrackerOn() { return tracker_on_; }
 
@@ -345,6 +359,9 @@ class InterceptQueue {
   static bool in_create_call_;
   static queue_id_t current_queue_id;
 
+  static rocprofiler_hsa_callback_fun_t submit_callback_fun_;
+  static void* submit_callback_arg_;
+
   hsa_queue_t* const queue_;
   ProxyQueue* const proxy_;
   const util::AgentInfo* agent_info_;

src/core/prof_listener.h

@@ -0,0 +1,183 @@
+/******************************************************************************
+MIT License
+
+Copyright (c) 2018 ROCm Core Technology
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*******************************************************************************/
+
+#ifndef _SRC_CORE_PROF_LISTENER_H
+#define _SRC_CORE_PROF_LISTENER_H
+
+#include <amd_hsa_kernel_code.h>
+#include <cxxabi.h>
+#include <dlfcn.h>
+
+#include <atomic>
+#include <iostream>
+#include <map>
+#include <mutex>
+
+#include "core/context.h"
+#include "core/proxy_queue.h"
+#include "core/tracker.h"
+#include "core/types.h"
+#include "inc/rocprofiler.h"
+#include "util/hsa_rsrc_factory.h"
+
+namespace rocprofiler {
+
+class ProfListener {
+ public:
+  static void Callback(const void* in_packets, uint64_t count, uint64_t user_que_idx, void* data,
+                         hsa_amd_queue_intercept_packet_writer writer) {
+    const packet_t* packets_arr = reinterpret_cast<const packet_t*>(in_packets);
+    InterceptQueue* obj = reinterpret_cast<InterceptQueue*>(data);
+    Queue* proxy = obj->proxy_;
+
+    // Travers input packets
+    for (uint64_t j = 0; j < count; ++j) {
+      const packet_t* packet = &packets_arr[j];
+      bool to_submit = true;
+
+      // Checking for dispatch packet type
+      if ((GetHeaderType(packet) == HSA_PACKET_TYPE_KERNEL_DISPATCH) && (dispatch_callback_ != NULL)) {
+        const hsa_kernel_dispatch_packet_t* dispatch_packet =
+            reinterpret_cast<const hsa_kernel_dispatch_packet_t*>(packet);
+
+        // Adding kernel timing tracker
+        const rocprofiler_dispatch_record_t* record = NULL;
+        if (tracker_ != NULL) {
+          const auto* entry = tracker_->Add(obj->agent_info_->dev_id, dispatch_packet->completion_signal);
+          const_cast<hsa_kernel_dispatch_packet_t*>(dispatch_packet)->completion_signal = entry->signal;
+          record = entry->record;
+        }
+
+        // Prepareing dispatch callback data
+        const char* kernel_name = GetKernelName(dispatch_packet);
+        rocprofiler_callback_data_t data = {obj->agent_info_->dev_id,
+                                            obj->agent_info_->dev_index,
+                                            obj->queue_,
+                                            user_que_idx,
+                                            dispatch_packet,
+                                            kernel_name,
+                                            record};
+
+        // Calling dispatch callback
+        rocprofiler_group_t group = {};
+        hsa_status_t status = dispatch_callback_(&data, callback_data_, &group);
+        free(const_cast<char*>(kernel_name));
+        // Injecting profiling start/stop packets
+        if ((status == HSA_STATUS_SUCCESS) && (group.context != NULL)) {
+          Context* context = reinterpret_cast<Context*>(group.context);
+          const pkt_vector_t& start_vector = context->StartPackets(group.index);
+          const pkt_vector_t& stop_vector = context->StopPackets(group.index);
+
+          pkt_vector_t packets = start_vector;
+          packets.insert(packets.end(), *packet);
+          packets.insert(packets.end(), stop_vector.begin(), stop_vector.end());
+          if (writer != NULL) {
+            writer(&packets[0], packets.size());
+          } else {
+            proxy->Submit(&packets[0], packets.size());
+          }
+          to_submit = false;
+        }
+      }
+
+      // Submitting the original packets if profiling was not enabled
+      if (to_submit) {
+        if (writer != NULL) {
+          writer(packet, 1);
+        } else {
+          proxy->Submit(packet, 1);
+        }
+      }
+    }
+  }
+
+  static void SetProfCallbacks(rocprofiler_callback_t dispatch_callback, queue_callback_t destroy_callback, void* data) {
+    std::lock_guard<mutex_t> lck(mutex_);
+    callback_data_ = data;
+    dispatch_callback_ = dispatch_callback;
+    destroy_callback_ = destroy_callback;
+  }
+
+  void TrackerOn(bool on) { tracker_on_ = on; }
+  bool IsTrackerOn() { return tracker_on_; }
+
+ private:
+  static hsa_packet_type_t GetHeaderType(const packet_t* packet) {
+    const packet_word_t* header = reinterpret_cast<const packet_word_t*>(packet);
+    return static_cast<hsa_packet_type_t>((*header >> HSA_PACKET_HEADER_TYPE) & header_type_mask);
+  }
+
+  static const char* GetKernelName(const hsa_kernel_dispatch_packet_t* dispatch_packet) {
+    const amd_kernel_code_t* kernel_code = NULL;
+    hsa_status_t status =
+        util::HsaRsrcFactory::Instance().LoaderApi()->hsa_ven_amd_loader_query_host_address(
+            reinterpret_cast<const void*>(dispatch_packet->kernel_object),
+            reinterpret_cast<const void**>(&kernel_code));
+    if (HSA_STATUS_SUCCESS != status) {
+      kernel_code = reinterpret_cast<amd_kernel_code_t*>(dispatch_packet->kernel_object);
+    }
+    amd_runtime_loader_debug_info_t* dbg_info = reinterpret_cast<amd_runtime_loader_debug_info_t*>(
+        kernel_code->runtime_loader_kernel_symbol);
+    const char* kernel_name = (dbg_info != NULL) ? dbg_info->kernel_name : NULL;
+
+    // Kernel name is mangled name
+    // apply __cxa_demangle() to demangle it
+    const char* funcname = NULL;
+    if (kernel_name != NULL) {
+      size_t funcnamesize = 0;
+      int status;
+      const char* ret = abi::__cxa_demangle(kernel_name, NULL, &funcnamesize, &status);
+      funcname = (ret != 0) ? ret : strdup(kernel_name);
+    }
+    if (funcname == NULL) funcname = strdup(kernel_none_);
+
+    return funcname;
+  }
+
+  InterceptQueue(const hsa_agent_t& agent, hsa_queue_t* const queue, ProxyQueue* proxy) :
+    queue_(queue),
+    proxy_(proxy)
+  {
+    agent_info_ = util::HsaRsrcFactory::Instance().GetAgentInfo(agent);
+  }
+  ~InterceptQueue() { ProxyQueue::Destroy(proxy_); }
+
+  static mutex_t mutex_;
+  static const packet_word_t header_type_mask = (1ul << HSA_PACKET_HEADER_WIDTH_TYPE) - 1;
+  static rocprofiler_callback_t dispatch_callback_;
+  static queue_callback_t destroy_callback_;
+  static void* callback_data_;
+  static obj_map_t* obj_map_;
+  static const char* kernel_none_;
+  static uint64_t timeout_;
+  static Tracker* tracker_;
+  static bool tracker_on_;
+  static bool in_constr_call_;
+
+  const util::AgentInfo* agent_info_;
+};
+
+}  // namespace rocprofiler
+
+#endif  // _SRC_CORE_PROF_LISTENER_H

src/core/rocprofiler.cpp

@@ -30,6 +30,7 @@ THE SOFTWARE.
 #include "core/context.h"
 #include "core/context_pool.h"
 #include "core/hsa_queue.h"
+#include "core/hsa_interceptor.h"
 #include "core/intercept_queue.h"
 #include "core/proxy_queue.h"
 #include "core/simple_proxy_queue.h"
@@ -84,6 +85,9 @@ decltype(hsa_queue_load_read_index_scacquire)* hsa_queue_load_read_index_scacqui
 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_copy)* hsa_memory_copy_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;
 
@@ -180,6 +184,7 @@ uint32_t LoadTool() {
     settings.trace_local = TraceProfile::IsLocal() ? 1: 0;
     settings.timeout = util::HsaRsrcFactory::GetTimeoutNs();
     settings.timestamp_on = InterceptQueue::IsTrackerOn() ? 1 : 0;
+    settings.hsa_intercepting = 0;
 
     if (handler) handler();
     else if (handler_prop) handler_prop(&settings);
@@ -191,6 +196,7 @@ uint32_t LoadTool() {
     if (settings.intercept_mode != 0) intercept_mode = DISPATCH_INTERCEPT_MODE;
     if (settings.code_obj_tracking) intercept_mode |= CODE_OBJ_TRACKING_MODE;
     if (settings.memcopy_tracking) intercept_mode |= MEMCOPY_INTERCEPT_MODE;
+    HsaInterceptor::Enable(settings.hsa_intercepting != 0);
   }
 
   return intercept_mode;
@@ -456,6 +462,7 @@ PUBLIC_API bool OnLoad(HsaApiTable* table, uint64_t runtime_version, uint64_t fa
     rocprofiler::InterceptQueue::HsaIntercept(table);
   } else {
     rocprofiler::StandaloneIntercept();
+    rocprofiler::HsaInterceptor::HsaIntercept(table);
   }
 
   return true;
@@ -864,4 +871,13 @@ hsa_status_t rocprofiler_get_time(
     return rocprofiler::util::HsaRsrcFactory::Instance().GetTime(time_id, timestamp, value_ns);
 }
 
+// Set new 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);
+  API_METHOD_SUFFIX
+}
+
 }  // extern "C"