From ca0b50c9ca28a11b04a0999aa95cc0ecb037cf7e Mon Sep 17 00:00:00 2001
From: Anusha GodavarthySurya <anusha.godavarthysurya@amd.com>
Date: Mon, 5 Feb 2024 10:34:58 +0000
Subject: [PATCH] SWDEV-444558 - SWDEV-444418 - Fix capturing of AQL packets
 when kernel arg size is 0

When graph doesn't have kernel nodes.

Change-Id: I6b3b476654d7eedc9ff0cec4b7269168aa115360
---
 hipamd/src/hip_graph_internal.cpp | 44 ++++++++++++++++---------------
 hipamd/src/hip_graph_internal.hpp |  8 +++---
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/hipamd/src/hip_graph_internal.cpp b/hipamd/src/hip_graph_internal.cpp
index 472c94af9a..c164a8c55a 100644
--- a/hipamd/src/hip_graph_internal.cpp
+++ b/hipamd/src/hip_graph_internal.cpp
@@ -365,35 +365,38 @@ hipError_t GraphExec::CaptureAQLPackets() {
         }
       }
     }
-
     auto device = g_devices[ihipGetDevice()]->devices()[0];
-    if (device->info().largeBar_) {
-      // Pad kernel argument buffer with sentinal size bytes to do a readback later
-      kernArgSizeForGraph += sizeof(int);
-      kernarg_pool_graph_ =
-          reinterpret_cast<address>(device->deviceLocalAlloc(kernArgSizeForGraph));
-      device_kernarg_pool_ = true;
-    } else {
-      kernarg_pool_graph_ = reinterpret_cast<address>(
-          device->hostAlloc(kernArgSizeForGraph, 0, amd::Device::MemorySegment::kKernArg));
-    }
+    if (kernArgSizeForGraph != 0) {
+      if (device->info().largeBar_) {
+        // Pad kernel argument buffer with sentinal size bytes to do a readback later
+        kernArgSizeForGraph += sizeof(int);
+        kernarg_pool_graph_ =
+            reinterpret_cast<address>(device->deviceLocalAlloc(kernArgSizeForGraph));
+        device_kernarg_pool_ = true;
+      } else {
+        kernarg_pool_graph_ = reinterpret_cast<address>(
+            device->hostAlloc(kernArgSizeForGraph, 0, amd::Device::MemorySegment::kKernArg));
+      }
 
-    if (kernarg_pool_graph_ == nullptr) {
-      return hipErrorMemoryAllocation;
+      if (kernarg_pool_graph_ == nullptr) {
+        return hipErrorMemoryAllocation;
+      }
+      kernarg_pool_size_graph_ = kernArgSizeForGraph;
     }
-    kernarg_pool_size_graph_ = kernArgSizeForGraph;
-
     for (auto& node : topoOrder_) {
       if (node->GetType() == hipGraphNodeTypeKernel) {
         auto kernelNode = reinterpret_cast<hip::GraphKernelNode*>(node);
         // From the kernel pool allocate the kern arg size required for the current kernel node.
-        address kernArgOffset = allocKernArg(kernelNode->GetKernargSegmentByteSize(),
-                                             kernelNode->GetKernargSegmentAlignment());
-        if (kernArgOffset == nullptr) {
-          return hipErrorMemoryAllocation;
+        address kernArgOffset = nullptr;
+        if (kernelNode->GetKernargSegmentByteSize()) {
+          kernArgOffset = allocKernArg(kernelNode->GetKernargSegmentByteSize(),
+                                       kernelNode->GetKernargSegmentAlignment());
+          if (kernArgOffset == nullptr) {
+            return hipErrorMemoryAllocation;
+          }
         }
         // Form GPU packet capture for the kernel node.
-        kernelNode->CaptureAndFormPacket(capture_stream_, kernArgOffset) ;
+        kernelNode->CaptureAndFormPacket(capture_stream_, kernArgOffset);
       }
     }
 
@@ -419,7 +422,6 @@ hipError_t GraphExec::CaptureAQLPackets() {
         host_val++;
       }
     }
-
     ResetQueueIndex();
   }
   return status;
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 3679f93942..05fb224d5b 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -597,9 +597,11 @@ struct GraphExec {
     // Release the kernel arg memory.
     auto device = g_devices[ihipGetDevice()]->devices()[0];
     if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
-      device->hostFree(kernarg_pool_graph_, kernarg_pool_size_graph_);
-      for (auto& element : kernarg_graph_) {
-        device->hostFree(element, kernarg_graph_size_);
+      if (kernarg_pool_size_graph_ != 0) {
+        device->hostFree(kernarg_pool_graph_, kernarg_pool_size_graph_);
+        for (auto& element : kernarg_graph_) {
+          device->hostFree(element, kernarg_graph_size_);
+        }
       }
     }
     amd::ScopedLock lock(graphExecSetLock_);