rocr/driver: add SetTrapHandler API to driver interface

This commit introduces a new SetTrapHandler API to the driver interface

- Implemented SetTrapHandler in KfdDriver to set trap handlers using
  hsaKmtSetTrapHandler.
- Added a stub implementation of SetTrapHandler in XdnaDriver that returns
  HSA_STATUS_ERROR.
- Updated the driver interface in driver.h to include the new SetTrapHandler
  method.
- Modified GpuAgent to use driver().SetTrapHandler instead of directly calling
  hsaKmtSetTrapHandler.

Signed-off-by: Honglei Huang <Honglei1.Huang@amd.com>


[ROCm/ROCR-Runtime commit: d874b8003a]

projects/rocr-runtime/runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp

@@ -559,6 +559,16 @@ bool KfdDriver::BindXnackMode() {
   return (mode != Flag::XNACK_DISABLE);
 }
 
+hsa_status_t KfdDriver::SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                                       const void* buffer_base, uint64_t buffer_base_size) const {
+  if (HSAKMT_CALL(hsaKmtSetTrapHandler(node_id, const_cast<void*>(base), base_size,
+                                       const_cast<void*>(buffer_base), buffer_base_size)) !=
+      HSAKMT_STATUS_SUCCESS)
+    return HSA_STATUS_ERROR;
+
+  return HSA_STATUS_SUCCESS;
+}
+
 hsa_status_t KfdDriver::IsModelEnabled(bool* enable) const {
   // AIE does not support streaming performance monitor.
   HSAKMT_STATUS status = HSAKMT_STATUS_ERROR;

projects/rocr-runtime/runtime/hsa-runtime/core/driver/xdna/amd_xdna_driver.cpp

@@ -872,5 +872,10 @@ hsa_status_t XdnaDriver::ConfigHwCtx(const PDICache& pdi_bo_handles, HSA_QUEUEID
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t XdnaDriver::SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                                        const void* buffer_base, uint64_t buffer_base_size) const {
+  return HSA_STATUS_ERROR;
+}
+
 } // namespace AMD
 } // namespace rocr

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_kfd_driver.h

@@ -122,6 +122,8 @@ public:
   hsa_status_t SPMSetDestBuffer(uint32_t preferred_node_id, uint32_t size_bytes, uint32_t* timeout,
                                 uint32_t* size_copied, void* dest_mem_addr,
                                 bool* is_spm_data_loss) const override;
+  hsa_status_t SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                              const void* buffer_base, uint64_t buffer_base_size) const override;
 
   hsa_status_t OpenSMI(uint32_t node_id, int* fd) const override;
 

projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_xdna_driver.h

@@ -236,6 +236,8 @@ public:
   hsa_status_t SPMSetDestBuffer(uint32_t preferred_node_id, uint32_t size_bytes, uint32_t* timeout,
                                 uint32_t* size_copied, void* dest_mem_addr,
                                 bool* is_spm_data_loss) const override;
+  hsa_status_t SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                              const void* buffer_base, uint64_t buffer_base_size) const override;
 
   hsa_status_t IsModelEnabled(bool* enable) const override;
 

projects/rocr-runtime/runtime/hsa-runtime/core/inc/driver.h

@@ -258,6 +258,17 @@ public:
     return HSA_STATUS_ERROR_INVALID_AGENT;
   }
 
+  /// @brief Sets trap handler and trap buffer to be used for all queues associated
+  /// with the specified NodeId within this process context
+  /// @param[in] node_id Node ID of the agent
+  /// @param[in] base Trap handler base address
+  /// @param[in] base_size Trap handler base size
+  /// @param[in] buffer_base Trap buffer base address
+  /// @param[in] buffer_base_size Trap buffer size
+  /// @return HSA_STATUS_SUCCESS if the driver successfully sets the trap handler.
+  virtual hsa_status_t SetTrapHandler(uint32_t node_id, const void* base, uint64_t base_size,
+                                      const void* buffer_base, uint64_t buffer_base_size) const = 0;
+
   /// @brief Check if the HSA KMT Model is enabled
   /// @param[out] enable True if the model is enabled, false otherwise
   virtual hsa_status_t IsModelEnabled(bool* enable) const = 0;

projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp

@@ -2249,10 +2249,8 @@ hsa_status_t GpuAgent::UpdateTrapHandlerWithPCS(pcs_sampling_data_t* pcs_hosttra
   }
 
   // Bind the trap handler to this node.
-  HSAKMT_STATUS retKmt =
-      HSAKMT_CALL(hsaKmtSetTrapHandler(node_id(), trap_code_buf_, trap_code_buf_size_, tma_addr, tma_size));
-
-  return (retKmt != HSAKMT_STATUS_SUCCESS) ? HSA_STATUS_ERROR : HSA_STATUS_SUCCESS;
+  return driver().SetTrapHandler(node_id(), trap_code_buf_, trap_code_buf_size_, tma_addr,
+                                 tma_size);
 }
 
 void GpuAgent::BindTrapHandler() {
@@ -2292,9 +2290,9 @@ void GpuAgent::BindTrapHandler() {
   }
 
   // Bind the trap handler to this node.
-  HSAKMT_STATUS err = HSAKMT_CALL(hsaKmtSetTrapHandler(node_id(), trap_code_buf_, trap_code_buf_size_,
-                                           tma_addr, tma_size));
-  assert(err == HSAKMT_STATUS_SUCCESS && "hsaKmtSetTrapHandler() failed");
+  hsa_status_t err =
+      driver().SetTrapHandler(node_id(), trap_code_buf_, trap_code_buf_size_, tma_addr, tma_size);
+  assert(err == HSA_STATUS_SUCCESS && "SetTrapHandler() failed");
 }
 
 void GpuAgent::InvalidateCodeCaches(void *ptr, size_t size) {