From 2ae70735e8ef2275f019ba3d74e07657185f027d Mon Sep 17 00:00:00 2001 From: Shweta Khatri Date: Tue, 18 Mar 2025 16:34:46 -0400 Subject: [PATCH] rocr: Fix PcSamplingCreateFromId to pass 32-bit dword count to DmaFill In PcSamplingCreateFromId, convert number of bytes into number of dwords because DmaFill expects a count of 32-bit words, not raw bytes. This prevents OOB writes on large sampling buffers. --- runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp index 4bc7f10b29..5a7563acc5 100644 --- a/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp +++ b/runtime/hsa-runtime/core/runtime/amd_gpu_agent.cpp @@ -2714,7 +2714,7 @@ hsa_status_t GpuAgent::PcSamplingCreateFromId(HsaPcSamplingTraceId ioctlId, device_datahost->buf_watermark1 = 0.8 * device_datahost->buf_size; // Allocate device memory for 2nd level trap handler TMA - size_t deviceAllocSize = sizeof(*pcs_data->device_data) + (2 * trap_buffer_size); + size_t deviceAllocSize = sizeof(pcs_sampling_data_t) + (2 * trap_buffer_size); pcs_data->device_data = (pcs_sampling_data_t*)finegrain_allocator()(deviceAllocSize, 0); if (pcs_data->device_data == nullptr) return HSA_STATUS_ERROR_OUT_OF_RESOURCES; @@ -2730,9 +2730,12 @@ hsa_status_t GpuAgent::PcSamplingCreateFromId(HsaPcSamplingTraceId ioctlId, } uint8_t* device_buf_ptr = - ((uint8_t*)pcs_data->device_data) + sizeof(pcs_sampling_data_t); - if (DmaFill(device_buf_ptr, 0, deviceAllocSize - sizeof(pcs_sampling_data_t)) != - HSA_STATUS_SUCCESS) { + reinterpret_cast(pcs_data->device_data) + sizeof(pcs_sampling_data_t); + size_t count_in_bytes = deviceAllocSize - sizeof(pcs_sampling_data_t); + size_t count_in_dwords = count_in_bytes / sizeof(uint32_t); + + if (DmaFill(device_buf_ptr, 0, count_in_dwords) != + HSA_STATUS_SUCCESS) { debug_print("Failed to dmaFill!\n"); return HSA_STATUS_ERROR; }