From 10763f0e7a8426cd2dc7a868f43196e94b2acba0 Mon Sep 17 00:00:00 2001
From: Jaydeep <106300970+jaydeeppatel1111@users.noreply.github.com>
Date: Fri, 31 Oct 2025 22:49:14 +0530
Subject: [PATCH] SWDEV-559505 - Enable back memset optimization and handle the
 cases when setParam can change the number of AQL packets for memset graph
 node. (#1320)

Co-authored-by: jaydeeppatel1111 <jaypatel@amd.com>
---
 .../clr/hipamd/src/hip_graph_internal.cpp     | 42 +++++++++++++++++--
 projects/clr/rocclr/device/blit.cpp           |  9 ----
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_graph_internal.cpp b/projects/clr/hipamd/src/hip_graph_internal.cpp
index 19e3a50dd5..8ccab6a53e 100644
--- a/projects/clr/hipamd/src/hip_graph_internal.cpp
+++ b/projects/clr/hipamd/src/hip_graph_internal.cpp
@@ -653,15 +653,48 @@ hipError_t GraphExec::UpdateAQLPacket(hip::GraphNode* node) {
       if (it != batch.nodeToRangeIndex.end()) {
         // Found the batch containing this node - update packets
         PacketBatch::NodeRange& range = batch.nodeRanges[it->second];
-
         // Capture new packets for this node
         std::vector<uint8_t*> newPackets;
         std::vector<std::string> newKernelNames;
-        hipError_t status =
-            node->CaptureAndFormPacket(kernArgManager_, &newPackets, &newKernelNames);
+        hipError_t status = node->CaptureAndFormPacket(kernArgManager_, &newPackets,
+                                                                        &newKernelNames);
         if (status != hipSuccess) {
           return status;
         }
+        // Number of packets per node can change
+        if (newPackets.size() != range.packetCount) {
+          size_t rangeIdx = it->second;
+          size_t oldPacketCount = range.packetCount;
+          size_t newPacketCount = newPackets.size();
+          int64_t packetDelta = static_cast<int64_t>(newPacketCount) -
+                                static_cast<int64_t>(oldPacketCount);
+          ClPrint(amd::LOG_INFO, amd::LOG_CODE,
+                  "[hipGraph] Handling packet count change for node (type=%d): %zu -> %zu packets",
+                  node->GetType(), oldPacketCount, newPacketCount);
+          if (packetDelta > 0) {
+            // Need to insert additional packets
+            // Insert new slots at the end of this node's range
+            size_t insertPos = range.startIndex + oldPacketCount;
+            batch.dispatchPackets.insert(batch.dispatchPackets.begin() + insertPos,
+                                        packetDelta, nullptr);
+            batch.dispatchKernelNames.insert(batch.dispatchKernelNames.begin() + insertPos,
+                                            packetDelta, std::string());
+          } else if (packetDelta < 0) {
+            // Need to remove packets
+            size_t removePos = range.startIndex + newPacketCount;
+            size_t removeCount = std::abs(packetDelta);
+            batch.dispatchPackets.erase(batch.dispatchPackets.begin() + removePos,
+                                        batch.dispatchPackets.begin() + removePos + removeCount);
+            batch.dispatchKernelNames.erase(batch.dispatchKernelNames.begin() + removePos,
+                                    batch.dispatchKernelNames.begin() + removePos + removeCount);
+          }
+          // Update this node's packet count
+          range.packetCount = newPacketCount;
+          // Adjust startIndex for all subsequent nodes in this batch
+          for (size_t i = rangeIdx + 1; i < batch.nodeRanges.size(); ++i) {
+            batch.nodeRanges[i].startIndex += packetDelta;
+          }
+        }
         // Update dispatch packets (always update regardless of enabled state)
         // The enabled/disabled check happens during dispatch, not here
         for (size_t i = 0; i < range.packetCount && i < newPackets.size(); ++i) {
@@ -679,7 +712,8 @@ hipError_t GraphExec::UpdateAQLPacket(hip::GraphNode* node) {
 }
 
 // ================================================================================================
-hipError_t GraphExec::UpdatePacketBatchesForNodeEnableDisable(hip::GraphNode* node, bool isEnabled) {
+hipError_t GraphExec::UpdatePacketBatchesForNodeEnableDisable(hip::GraphNode* node,
+                                                              bool isEnabled) {
   if (max_streams_ != 1 && max_streams_dev_.size() == 1 && !node->GraphCaptureEnabled()) {
     // Only handle single stream and single device case with captured nodes
     return hipSuccess;
diff --git a/projects/clr/rocclr/device/blit.cpp b/projects/clr/rocclr/device/blit.cpp
index 749c895c4d..d58ae48395 100644
--- a/projects/clr/rocclr/device/blit.cpp
+++ b/projects/clr/rocclr/device/blit.cpp
@@ -735,15 +735,6 @@ void HostBlitManager::FillBufferInfo::PackInfo(const device::Memory& memory, siz
   guarantee(fill_size >= pattern_size, "Pattern Size: %u cannot be greater than fill size: %u \n",
             pattern_size, fill_size);
 
-  constexpr bool kDisablePackingOptimization = true;
-  // Check if packing optimization is disabled
-  if (kDisablePackingOptimization) {
-    // Simple case: create a single FillBufferInfo without alignment optimization
-    FillBufferInfo fill_info(fill_size);
-    packed_info.push_back(fill_info);
-    return;
-  }
-
   // 2. Calculate the next closest dword aligned address for faster processing
   size_t dst_addr = memory.virtualAddress() + fill_origin;
   size_t aligned_dst_addr = amd::alignUp(dst_addr, kExtendedSize);