librocdxg: Increase AQL frame size calculation

to prevent PM4 command buffer overflow

Signed-off-by: Flora Cui <flora.cui@amd.com>
Reviewed-by: Longlong Yao <Longlong.Yao@amd.com>
Part-of: <http://10.67.69.192/wsl/rocr-runtime/-/merge_requests/113>

Author: fcui-amd

src/wddm/device.cpp

@@ -409,14 +409,20 @@ void WDDMDevice::InitCmdbufInfo(void) {
     cmdbuf_aql_frame_size_ = 2 * sizeof(gfx10::AcquireMemTemplate);
   }
 
-  if (device_info_.major >= 11)
+  if (device_info_.major >= 11) {
     cmdbuf_aql_frame_size_ += sizeof(SetScratchTemplate);
+    cmdbuf_aql_frame_size_ += sizeof(DispatchProgramResourceRegs); // BuildComputeShaderParams
+  }
 
   cmdbuf_aql_frame_size_ +=
     sizeof(PM4MEC_COPY_DATA) * 2 +
     sizeof(BarrierTemplate) * 2 +
     sizeof(DispatchTemplate) +
     sizeof(AtomicTemplate) * 2;
+
+  // Add safety margin to account for alignment and future additions
+  cmdbuf_aql_frame_size_ += 128;
+
   cmdbuf_aql_frame_size_ = AlignUp(cmdbuf_aql_frame_size_, 0x10);
 
   cmdbuf_size_ = AlignUp(cmdbuf_aql_frame_num_ * cmdbuf_aql_frame_size_, 0x1000);

src/wddm/queue.cpp

@@ -723,6 +723,12 @@ ComputeQueue::KernelDispatchAqlToPm4(char *cpu, hsa_kernel_dispatch_packet_t *pa
   else
     i += cmd_util.BuildWriteData64Command(cpu + i, (uint64_t *)ring_rptr, cmdbuf_aql_frame_write_index + 1);
 
+  // Check if we exceeded the frame size
+  if ((i - ib_size) > cmdbuf_aql_frame_size) {
+    pr_err("PM4 command buffer overflow in KernelDispatch: used %d bytes, limit %d bytes\n", i - ib_size, cmdbuf_aql_frame_size);
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
   ib_size = i;
   cmdbuf_aql_frame_write_index++;
   packet->header = HSA_PACKET_TYPE_INVALID;
@@ -811,6 +817,12 @@ ComputeQueue::BarrierGenericAqlToPm4(char *cpu, hsa_barrier_and_packet_t *packet
   else
     i += cmd_util.BuildWriteData64Command(cpu + i, (uint64_t *)ring_rptr, cmdbuf_aql_frame_write_index + 1);
 
+  // Check if we exceeded the frame size
+  if ((i - ib_size) > cmdbuf_aql_frame_size) {
+    pr_err("PM4 command buffer overflow in BarrierGeneric: used %d bytes, limit %d bytes\n", i - ib_size, cmdbuf_aql_frame_size);
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
   ib_size = i;
   cmdbuf_aql_frame_write_index++;
   packet->header = HSA_PACKET_TYPE_INVALID;
@@ -881,6 +893,12 @@ hsa_status_t ComputeQueue::VendorSpecificAqlToPm4(char *cpu, amd_aql_pm4_ib *pac
   else
     i += cmd_util.BuildWriteData64Command(cpu + i, (uint64_t *)ring_rptr, cmdbuf_aql_frame_write_index + 1);
 
+  // Check if we exceeded the frame size
+  if ((i - ib_size) > cmdbuf_aql_frame_size) {
+    pr_err("PM4 command buffer overflow in VendorSpecific: used %d bytes, limit %d bytes\n", i - ib_size, cmdbuf_aql_frame_size);
+    return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+  }
+
   ib_size = i;
   cmdbuf_aql_frame_write_index++;
   packet->header = HSA_PACKET_TYPE_INVALID;