SWDEV-461072 - Extend AQL Optimization for child graph nodes

Change-Id: I6baf906add7240b29ea653020a9a0b56206ee2a7

hipamd/src/hip_graph.cpp

@@ -1760,7 +1760,24 @@ hipError_t hipGraphExecChildGraphNodeSetParams(hipGraphExec_t hGraphExec, hipGra
   if (clonedNode == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  HIP_RETURN(reinterpret_cast<hip::ChildGraphNode*>(clonedNode)->SetParams(cg));
+  hipError_t status = reinterpret_cast<hip::ChildGraphNode*>(clonedNode)->SetParams(cg);
+  if (status != hipSuccess) {
+    return status;
+  }
+  if (reinterpret_cast<hip::ChildGraphNode*>(clonedNode)->GetGraphCaptureStatus()) {
+    std::vector<hip::GraphNode*> childGraphNodes;
+    reinterpret_cast<hip::ChildGraphNode*>(clonedNode)->TopologicalOrder(childGraphNodes);
+    for (std::vector<hip::GraphNode*>::size_type i = 0; i != childGraphNodes.size(); i++) {
+      if (childGraphNodes[i]->GetType() == hipGraphNodeTypeKernel) {
+        status = reinterpret_cast<hip::GraphExec*>(hGraphExec)
+                     ->UpdateAQLPacket(reinterpret_cast<hip::GraphKernelNode*>(childGraphNodes[i]));
+        if (status != hipSuccess) {
+          return status;
+        }
+      }
+    }
+  }
+  return status;
 }
 
 hipError_t hipStreamGetCaptureInfo_common(hipStream_t stream,

hipamd/src/hip_graph_internal.cpp

@@ -350,20 +350,64 @@ hipError_t GraphExec::Init() {
   return status;
 }
 
+void GetKernelArgSizeForGraph(std::vector<std::vector<Node>>& parallelLists,
+                              size_t& kernArgSizeForGraph) {
+  // GPU packet capture is enabled for kernel nodes. Calculate the kernel
+  // arg size required for all graph kernel nodes to allocate
+  for (const auto& list : parallelLists) {
+    for (auto& node : list) {
+      if (node->GetType() == hipGraphNodeTypeKernel &&
+          !reinterpret_cast<hip::GraphKernelNode*>(node)->HasHiddenHeap()) {
+        kernArgSizeForGraph += reinterpret_cast<hip::GraphKernelNode*>(node)->GetKerArgSize();
+      } else if (node->GetType() == hipGraphNodeTypeGraph) {
+        auto& childParallelLists = reinterpret_cast<hip::ChildGraphNode*>(node)->GetParallelLists();
+        if (childParallelLists.size() == 1) {
+          GetKernelArgSizeForGraph(childParallelLists, kernArgSizeForGraph);
+        }
+      }
+    }
+  }
+}
+
+hipError_t AllocKernelArgForGraph(std::vector<hip::Node>& topoOrder, hip::Stream* capture_stream,
+                                  hip::GraphExec* graphExec) {
+  hipError_t status = hipSuccess;
+  for (auto& node : topoOrder) {
+    if (node->GetType() == hipGraphNodeTypeKernel &&
+        !reinterpret_cast<hip::GraphKernelNode*>(node)->HasHiddenHeap()) {
+      auto kernelNode = reinterpret_cast<hip::GraphKernelNode*>(node);
+      // From the kernel pool allocate the kern arg size required for the current kernel node.
+      address kernArgOffset = nullptr;
+      if (kernelNode->GetKernargSegmentByteSize()) {
+        kernArgOffset = graphExec->allocKernArg(kernelNode->GetKernargSegmentByteSize(),
+                                                kernelNode->GetKernargSegmentAlignment());
+        if (kernArgOffset == nullptr) {
+          return hipErrorMemoryAllocation;
+        }
+      }
+      // Form GPU packet capture for the kernel node.
+      kernelNode->CaptureAndFormPacket(capture_stream, kernArgOffset);
+    } else if (node->GetType() == hipGraphNodeTypeGraph) {
+      auto childNode = reinterpret_cast<hip::ChildGraphNode*>(node);
+      auto& childParallelLists = childNode->GetParallelLists();
+      if (childParallelLists.size() == 1) {
+        childNode->SetGraphCaptureStatus(true);
+        status =
+            AllocKernelArgForGraph(childNode->GetChildGraphNodeOrder(), capture_stream, graphExec);
+        if (status != hipSuccess) {
+          return status;
+        }
+      }
+    }
+  }
+  return status;
+}
+
 hipError_t GraphExec::CaptureAQLPackets() {
   hipError_t status = hipSuccess;
   if (parallelLists_.size() == 1) {
     size_t kernArgSizeForGraph = 0;
-    // GPU packet capture is enabled for kernel nodes. Calculate the kernel
-    // arg size required for all graph kernel nodes to allocate
-    for (const auto& list : parallelLists_) {
-      for (auto& node : list) {
-        if (node->GetType() == hipGraphNodeTypeKernel &&
-            !reinterpret_cast<hip::GraphKernelNode*>(node)->HasHiddenHeap()) {
-          kernArgSizeForGraph += reinterpret_cast<hip::GraphKernelNode*>(node)->GetKerArgSize();
-        }
-      }
-    }
+    GetKernelArgSizeForGraph(parallelLists_, kernArgSizeForGraph);
     auto device = g_devices[ihipGetDevice()]->devices()[0];
     if (kernArgSizeForGraph != 0) {
       if (device->info().largeBar_) {
@@ -380,22 +424,9 @@ hipError_t GraphExec::CaptureAQLPackets() {
       }
       kernarg_pool_size_graph_ = kernArgSizeForGraph;
     }
-    for (auto& node : topoOrder_) {
-      if (node->GetType() == hipGraphNodeTypeKernel &&
-          !reinterpret_cast<hip::GraphKernelNode*>(node)->HasHiddenHeap()) {
-        auto kernelNode = reinterpret_cast<hip::GraphKernelNode*>(node);
-        // From the kernel pool allocate the kern arg size required for the current kernel node.
-        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);
-      }
+    status = AllocKernelArgForGraph(topoOrder_, capture_stream_, this);
+    if (status != hipSuccess) {
+      return status;
     }
 
     if (device_kernarg_pool_) {
@@ -414,8 +445,6 @@ hipError_t GraphExec::CaptureAQLPackets() {
         kSentinel = *reinterpret_cast<volatile int*>(dev_ptr - 1);
       }
     }
-
-    ResetQueueIndex();
   }
   return status;
 }
@@ -550,6 +579,36 @@ void UpdateStream(std::vector<std::vector<Node>>& parallelLists, hip::Stream* st
   }
 }
 
+hipError_t EnqueueGraphWithSingleList(std::vector<hip::Node> topoOrder, hip::Stream* hip_stream,
+                                      hip::GraphExec* graphExec) {
+  // Accumulate command tracks all the AQL packet batch that we submit to the HW. For now
+  // we track only kernel nodes.
+  amd::AccumulateCommand* accumulate = nullptr;
+  hipError_t status = hipSuccess;
+  if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
+    accumulate = new amd::AccumulateCommand(*hip_stream, {}, nullptr);
+  }
+  for (int i = 0; i < topoOrder.size(); i++) {
+    if (DEBUG_CLR_GRAPH_PACKET_CAPTURE && topoOrder[i]->GetType() == hipGraphNodeTypeKernel &&
+        !reinterpret_cast<hip::GraphKernelNode*>(topoOrder[i])->HasHiddenHeap()) {
+      if (topoOrder[i]->GetEnabled()) {
+        hip_stream->vdev()->dispatchAqlPacket(topoOrder[i]->GetAqlPacket(), accumulate);
+        accumulate->addKernelName(topoOrder[i]->GetKernelName());
+      }
+    } else {
+      topoOrder[i]->SetStream(hip_stream, graphExec);
+      status = topoOrder[i]->CreateCommand(topoOrder[i]->GetQueue());
+      topoOrder[i]->EnqueueCommands(reinterpret_cast<hipStream_t>(hip_stream));
+    }
+  }
+
+  if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
+    accumulate->enqueue();
+    accumulate->release();
+  }
+  return status;
+}
+
 hipError_t GraphExec::Run(hipStream_t stream) {
   hipError_t status = hipSuccess;
 
@@ -576,33 +635,9 @@ hipError_t GraphExec::Run(hipStream_t stream) {
     repeatLaunch_ = true;
   }
 
-  if (parallelLists_.size() == 1 && instantiateDeviceId_ == hip_stream->DeviceId()) {
-    // Accumulate command tracks all the AQL packet batch that we submit to the HW. For now
-    // we track only kernel nodes.
-    amd::AccumulateCommand* accumulate = nullptr;
-
-    if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
-      accumulate = new amd::AccumulateCommand(*hip_stream, {}, nullptr);
-    }
-    for (int i = 0; i < topoOrder_.size(); i++) {
-      if (DEBUG_CLR_GRAPH_PACKET_CAPTURE &&
-          topoOrder_[i]->GetType() == hipGraphNodeTypeKernel &&
-          !reinterpret_cast<hip::GraphKernelNode*>(topoOrder_[i])->HasHiddenHeap()) {
-        if (topoOrder_[i]->GetEnabled()) {
-          accumulate->addKernelName(topoOrder_[i]->GetKernelName());
-          hip_stream->vdev()->dispatchAqlPacket(topoOrder_[i]->GetAqlPacket(), accumulate);
-        }
-      } else {
-        topoOrder_[i]->SetStream(hip_stream, this);
-        status = topoOrder_[i]->CreateCommand(topoOrder_[i]->GetQueue());
-        topoOrder_[i]->EnqueueCommands(stream);
-      }
-    }
-
-    if (DEBUG_CLR_GRAPH_PACKET_CAPTURE) {
-      accumulate->enqueue();
-      accumulate->release();
-    }
+  if (parallelLists_.size() == 1 &&
+      instantiateDeviceId_ == hip_stream->DeviceId()) {
+    status = EnqueueGraphWithSingleList(topoOrder_, hip_stream, this);
   } else if (parallelLists_.size() == 1 &&
              instantiateDeviceId_ != hip_stream->DeviceId()) {
     for (int i = 0; i < topoOrder_.size(); i++) {

hipamd/src/hip_graph_internal.hpp

@@ -52,7 +52,8 @@ hipError_t FillCommands(std::vector<std::vector<Node>>& parallelLists,
                         amd::Command*& graphEnd, hip::Stream* stream);
 void UpdateStream(std::vector<std::vector<Node>>& parallelLists, hip::Stream* stream,
                   GraphExec* ptr);
-
+hipError_t EnqueueGraphWithSingleList(std::vector<hip::Node> topoOrder, hip::Stream* hip_stream,
+                                      hip::GraphExec* graphExec = nullptr);
 struct UserObject : public amd::ReferenceCountedObject {
   typedef void (*UserCallbackDestructor)(void* data);
   static std::unordered_set<UserObject*> ObjectSet_;
@@ -661,18 +662,21 @@ struct ChildGraphNode : public GraphNode {
   amd::Command* lastEnqueuedCommand_;
   amd::Command* startCommand_;
   amd::Command* endCommand_;
+  bool graphCaptureStatus_;
  public:
   ChildGraphNode(Graph* g) : GraphNode(hipGraphNodeTypeGraph, "solid", "rectangle") {
     childGraph_ = g->clone();
     lastEnqueuedCommand_ = nullptr;
     startCommand_ = nullptr;
     endCommand_ = nullptr;
+    graphCaptureStatus_ = false;
   }
 
   ~ChildGraphNode() { delete childGraph_; }
 
   ChildGraphNode(const ChildGraphNode& rhs) : GraphNode(rhs) {
     childGraph_ = rhs.childGraph_->clone();
+    graphCaptureStatus_ = rhs.graphCaptureStatus_;
   }
 
   GraphNode* clone() const override {
@@ -681,6 +685,19 @@ struct ChildGraphNode : public GraphNode {
 
   Graph* GetChildGraph() override { return childGraph_; }
 
+  void SetGraphCaptureStatus(bool status) { graphCaptureStatus_ = status; }
+
+  bool GetGraphCaptureStatus() { return graphCaptureStatus_; }
+
+  std::vector<Node>& GetChildGraphNodeOrder() {
+    return childGraphNodeOrder_;
+  }
+
+  std::vector<std::vector<Node>>& GetParallelLists() {
+    return parallelLists_;
+  }
+
+
   hipError_t GetNumParallelStreams(size_t &num) override {
     if (false == TopologicalOrder(childGraphNodeOrder_)) {
       return hipErrorInvalidValue;
@@ -715,8 +732,10 @@ struct ChildGraphNode : public GraphNode {
     }
     startCommand_ = nullptr;
     endCommand_ = nullptr;
-    status = FillCommands(parallelLists_, nodeWaitLists_, childGraphNodeOrder_, childGraph_,
-                          startCommand_, endCommand_, stream);
+    if (!graphCaptureStatus_) {
+      status = FillCommands(parallelLists_, nodeWaitLists_, childGraphNodeOrder_, childGraph_,
+                            startCommand_, endCommand_, stream);
+    }
     return status;
   }
 
@@ -735,19 +754,24 @@ struct ChildGraphNode : public GraphNode {
     return childGraph_->TopologicalOrder(TopoOrder);
   }
   void EnqueueCommands(hipStream_t stream) override {
-    // enqueue child graph start command
-    if (startCommand_ != nullptr) {
-      startCommand_->enqueue();
-      startCommand_->release();
-    }
-    // enqueue nodes in child graph in level order
-    for (auto& node : childGraphNodeOrder_) {
-      node->EnqueueCommands(stream);
-    }
-    // enqueue child graph end command
-    if (endCommand_ != nullptr) {
-      endCommand_->enqueue();
-      endCommand_->release();
+    if (graphCaptureStatus_) {
+      hipError_t status =
+          EnqueueGraphWithSingleList(childGraphNodeOrder_, reinterpret_cast<hip::Stream*>(stream));
+    } else {
+      // enqueue child graph start command
+      if (startCommand_ != nullptr) {
+        startCommand_->enqueue();
+        startCommand_->release();
+      }
+      // enqueue nodes in child graph in level order
+      for (auto& node : childGraphNodeOrder_) {
+        node->EnqueueCommands(stream);
+      }
+      // enqueue child graph end command
+      if (endCommand_ != nullptr) {
+        endCommand_->enqueue();
+        endCommand_->release();
+      }
     }
   }