SWDEV-448586 - Added implementation for new API hipStreamBeginCaptureToGraph

Change-Id: I1ce802102cef2b66c92d3375f769983841de793f

hipamd/src/amdhip.def

@@ -460,4 +460,5 @@ hipGraphExternalSemaphoresWaitNodeGetParams
 hipGraphExecExternalSemaphoresSignalNodeSetParams
 hipGraphExecExternalSemaphoresWaitNodeSetParams
 hipGraphAddNode
-hipGraphInstantiateWithParams
+hipGraphInstantiateWithParams
+hipStreamBeginCaptureToGraph

hipamd/src/hip_api_trace.cpp

@@ -763,6 +763,10 @@ hipError_t hipTexRefGetBorderColor(float* pBorderColor, const textureReference*
 hipError_t hipTexRefGetArray(hipArray_t* pArray, const textureReference* texRef);
 hipError_t hipGetProcAddress(const char* symbol, void** pfn, int hipVersion, uint64_t flags,
                              hipDriverProcAddressQueryResult* symbolStatus = NULL);
+hipError_t hipStreamBeginCaptureToGraph(hipStream_t stream, hipGraph_t graph,
+                                        const hipGraphNode_t* dependencies,
+                                        const hipGraphEdgeData* dependencyData,
+                                        size_t numDependencies, hipStreamCaptureMode mode);
 }  // namespace hip
 
 namespace hip {
@@ -1237,6 +1241,7 @@ void UpdateDispatchTable(HipDispatchTable* ptrDispatchTable) {
   ptrDispatchTable->hipTexRefGetBorderColor_fn =  hip::hipTexRefGetBorderColor;
   ptrDispatchTable->hipTexRefGetArray_fn = hip::hipTexRefGetArray;
   ptrDispatchTable->hipGetProcAddress_fn = hip::hipGetProcAddress;
+  ptrDispatchTable->hipStreamBeginCaptureToGraph_fn = hip::hipStreamBeginCaptureToGraph;
 }
 
 #if HIP_ROCPROFILER_REGISTER > 0
@@ -1783,20 +1788,22 @@ HIP_ENFORCE_ABI(HipDispatchTable, hipStreamGetCaptureInfo_v2_spt_fn, 425)
 HIP_ENFORCE_ABI(HipDispatchTable, hipLaunchHostFunc_spt_fn, 426)
 HIP_ENFORCE_ABI(HipDispatchTable, hipGetStreamDeviceId_fn, 427)
 HIP_ENFORCE_ABI(HipDispatchTable, hipDrvGraphAddMemsetNode_fn, 428)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddExternalSemaphoresWaitNode_fn, 429)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddExternalSemaphoresSignalNode_fn, 430)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresSignalNodeSetParams_fn, 431)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresWaitNodeSetParams_fn, 432)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresSignalNodeGetParams_fn, 433)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresWaitNodeGetParams_fn, 434)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExecExternalSemaphoresSignalNodeSetParams_fn, 435)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExecExternalSemaphoresWaitNodeSetParams_fn, 436)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddNode_fn, 437)
-HIP_ENFORCE_ABI(HipDispatchTable, hipGraphInstantiateWithParams_fn, 438)
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddExternalSemaphoresWaitNode_fn, 429);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddExternalSemaphoresSignalNode_fn, 430);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresSignalNodeSetParams_fn, 431);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresWaitNodeSetParams_fn, 432);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresSignalNodeGetParams_fn, 433);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExternalSemaphoresWaitNodeGetParams_fn, 434);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExecExternalSemaphoresSignalNodeSetParams_fn, 435);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphExecExternalSemaphoresWaitNodeSetParams_fn, 436);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphAddNode_fn, 437);
+HIP_ENFORCE_ABI(HipDispatchTable, hipGraphInstantiateWithParams_fn, 438);
 HIP_ENFORCE_ABI(HipDispatchTable, hipExtGetLastError_fn, 439)
 HIP_ENFORCE_ABI(HipDispatchTable, hipTexRefGetBorderColor_fn, 440)
 HIP_ENFORCE_ABI(HipDispatchTable, hipTexRefGetArray_fn, 441)
 HIP_ENFORCE_ABI(HipDispatchTable, hipGetProcAddress_fn, 442)
+HIP_ENFORCE_ABI(HipDispatchTable, hipStreamBeginCaptureToGraph_fn, 443);
+
 
 // if HIP_ENFORCE_ABI entries are added for each new function pointer in the table, the number below
 // will be +1 of the number in the last HIP_ENFORCE_ABI line. E.g.:
@@ -1804,9 +1811,9 @@ HIP_ENFORCE_ABI(HipDispatchTable, hipGetProcAddress_fn, 442)
 //  HIP_ENFORCE_ABI(<table>, <functor>, 8)
 //
 //  HIP_ENFORCE_ABI_VERSIONING(<table>, 9) <- 8 + 1 = 9
-HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 443)
+HIP_ENFORCE_ABI_VERSIONING(HipDispatchTable, 444)
 
-static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 1,
+static_assert(HIP_RUNTIME_API_TABLE_MAJOR_VERSION == 0 && HIP_RUNTIME_API_TABLE_STEP_VERSION == 2,
               "If you get this error, add new HIP_ENFORCE_ABI(...) code for the new function "
               "pointers and then update this check so it is true");
 #endif

hipamd/src/hip_graph.cpp

@@ -33,7 +33,8 @@ amd::Monitor g_captureStreamsLock{"StreamCaptureGlobalList"};
 amd::Monitor g_streamSetLock{"StreamCaptureset"};
 std::unordered_set<hip::Stream*> g_allCapturingStreams;
 hipError_t ihipGraphDebugDotPrint(hipGraph_t graph, const char* path, unsigned int flags);
-
+hipError_t ihipStreamUpdateCaptureDependencies(hipStream_t stream, hipGraphNode_t* dependencies,
+                                               size_t numDependencies, unsigned int flags);
 
 inline hipError_t ihipGraphUpload(hipGraphExec_t graphExec, hipStream_t stream) {
   if (graphExec == nullptr) {
@@ -863,7 +864,8 @@ hipError_t hipThreadExchangeStreamCaptureMode(hipStreamCaptureMode* mode) {
   HIP_RETURN_DURATION(hipSuccess);
 }
 
-hipError_t hipStreamBeginCapture_common(hipStream_t stream, hipStreamCaptureMode mode) {
+hipError_t hipStreamBeginCapture_common(hipStream_t stream, hipStreamCaptureMode mode,
+                                        hipGraph_t graph = nullptr) {
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
@@ -879,8 +881,11 @@ hipError_t hipStreamBeginCapture_common(hipStream_t stream, hipStreamCaptureMode
   if (s->GetCaptureStatus() == hipStreamCaptureStatusActive) {
     return hipErrorIllegalState;
   }
-
-  s->SetCaptureGraph(new hip::Graph(s->GetDevice()));
+  if(graph == nullptr) {
+    s->SetCaptureGraph(new hip::Graph(s->GetDevice()));
+  } else {
+    s->SetCaptureGraph(reinterpret_cast<hip::Graph*>(graph));
+  }
   s->SetCaptureId();
   s->SetCaptureMode(mode);
   s->SetOriginStream();
@@ -903,6 +908,37 @@ hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode mode)
   HIP_RETURN_DURATION(hipStreamBeginCapture_common(stream, mode));
 }
 
+hipError_t hipStreamBeginCaptureToGraph(hipStream_t stream, hipGraph_t graph,
+                                        const hipGraphNode_t* dependencies,
+                                        const hipGraphEdgeData* dependencyData,
+                                        size_t numDependencies, hipStreamCaptureMode mode) {
+  HIP_INIT_API(hipStreamBeginCapture, stream, graph, dependencies, dependencyData, numDependencies,
+               mode);
+  if (dependencyData != nullptr) {
+    return hipErrorNotSupported;
+  } else if (graph == nullptr) {
+    return hipErrorInvalidValue;
+  } else if (dependencies == nullptr && numDependencies != 0) {
+    return hipErrorInvalidValue;
+  } else if (dependencies != nullptr && numDependencies == 0) {
+    return hipErrorInvalidValue;
+  }
+  hip::Graph* g = reinterpret_cast<hip::Graph*>(graph);
+  const std::vector<Node> nodes = g->GetNodes();
+  for (auto& node : nodes) {
+    g->AddManualNodeDuringCapture(node);
+  }
+  hipError_t status = hipStreamBeginCapture_common(stream, mode, graph);
+  if (status != hipSuccess) {
+    HIP_RETURN_DURATION(status);
+  }
+  if (dependencies != nullptr) {
+    status = ihipStreamUpdateCaptureDependencies(stream, const_cast<hipGraphNode_t*>(dependencies),
+                                                 numDependencies, hipStreamSetCaptureDependencies);
+  }
+  HIP_RETURN_DURATION(status);
+}
+
 hipError_t hipStreamBeginCapture_spt(hipStream_t stream, hipStreamCaptureMode mode) {
   HIP_INIT_API(hipStreamBeginCapture, stream, mode);
   PER_THREAD_DEFAULT_STREAM(stream);
@@ -1774,7 +1810,7 @@ hipError_t hipStreamGetCaptureInfo_v2_spt(hipStream_t stream,
                                                dependencies_out, numDependencies_out));
 }
 
-hipError_t hipStreamUpdateCaptureDependencies(hipStream_t stream, hipGraphNode_t* dependencies,
+hipError_t ihipStreamUpdateCaptureDependencies(hipStream_t stream, hipGraphNode_t* dependencies,
                                               size_t numDependencies, unsigned int flags) {
   HIP_INIT_API(hipStreamUpdateCaptureDependencies, stream, dependencies, numDependencies, flags);
   if (!hip::isValid(stream)) {
@@ -1809,6 +1845,12 @@ hipError_t hipStreamUpdateCaptureDependencies(hipStream_t stream, hipGraphNode_t
   HIP_RETURN(hipSuccess);
 }
 
+hipError_t hipStreamUpdateCaptureDependencies(hipStream_t stream, hipGraphNode_t* dependencies,
+                                              size_t numDependencies, unsigned int flags) {
+  HIP_INIT_API(hipStreamUpdateCaptureDependencies, stream, dependencies, numDependencies, flags);
+  HIP_RETURN(ihipStreamUpdateCaptureDependencies(stream, dependencies, numDependencies, flags));
+}
+
 hipError_t hipGraphRemoveDependencies(hipGraph_t graph, const hipGraphNode_t* from,
                                       const hipGraphNode_t* to, size_t numDependencies) {
   HIP_INIT_API(hipGraphRemoveDependencies, graph, from, to, numDependencies);

hipamd/src/hip_hcc.map.in

@@ -552,6 +552,7 @@ hip_6.1 {
 global:
     hipGraphInstantiateWithParams;
     hipGetProcAddress;
+    hipStreamBeginCaptureToGraph;
 local:
     *;
 } hip_6.0;

hipamd/src/hip_table_interface.cpp

@@ -1732,4 +1732,11 @@ extern "C" hipError_t hipGetProcAddress(const char* symbol, void** pfn, int  hip
                                         hipDriverProcAddressQueryResult* symbolStatus = nullptr) {
   return hip::GetHipDispatchTable()->hipGetProcAddress_fn(symbol, pfn, hipVersion, flags,
                                                           symbolStatus);
-}
+}
+hipError_t hipStreamBeginCaptureToGraph(hipStream_t stream, hipGraph_t graph,
+                                        const hipGraphNode_t* dependencies,
+                                        const hipGraphEdgeData* dependencyData,
+                                        size_t numDependencies, hipStreamCaptureMode mode) {
+  return hip::GetHipDispatchTable()->hipStreamBeginCaptureToGraph_fn(
+      stream, graph, dependencies, dependencyData, numDependencies, mode);
+}