SWDEV-306122 - [catch2][dtest] Additional tests for hipGraphGetNodes/hipGraphGetRootNodes api (#2860)

Change-Id: I4e9a6eadd92b6de4aded0aa96464bfcc59441ac4

tests/catch/unit/graph/hipGraphGetRootNodes.cc

@@ -26,6 +26,8 @@ Functional ::
  2) Pass nodes as nullptr and verify api returns actual number of root nodes added to graph.
  3) If NumRootNodes passed is greater than the actual number of root nodes, the remaining entries in
  nodes list will be set to NULL, and the number of nodes actually obtained will be returned in NumRootNodes.
+ 4) Create a graph with stream capture done on multiple dependent streams.
+ Verify root nodes of created graph are matching the operations pushed which doesn't have dependencies.
 
 Argument Validation ::
  1) Pass graph as nullptr and verify api returns error code.
@@ -50,6 +52,7 @@ TEST_CASE("Unit_hipGraphGetRootNodes_Functional") {
   constexpr auto addlEntries = 5;
   hipGraph_t graph;
 
+
   hipGraphNode_t memcpyNode, kernelNode;
   hipKernelNodeParams kernelNodeParams{};
   hipStream_t streamForGraph;
@@ -141,6 +144,112 @@ TEST_CASE("Unit_hipGraphGetRootNodes_Functional") {
   free(rootnodes);
 }
 
+
+/**
+ * Create a graph with stream capture done on multiple dependent streams. Verify root nodes
+ * of created graph are matching the operations pushed which doesn't have dependencies.
+ */
+TEST_CASE("Unit_hipGraphGetRootNodes_CapturedStream") {
+  hipStream_t stream1{nullptr}, stream2{nullptr}, mstream{nullptr};
+  hipStream_t streamForGraph{nullptr};
+  hipEvent_t memsetEvent1, memsetEvent2, forkStreamEvent;
+  hipGraph_t graph{nullptr};
+  hipGraphExec_t graphExec{nullptr};
+  constexpr unsigned blocks = 512;
+  constexpr unsigned threadsPerBlock = 256;
+  constexpr size_t N = 1000000;
+  constexpr int numMemsetNodes = 2;
+  size_t Nbytes = N * sizeof(float), numRootNodes{};
+  float *A_d, *C_d;
+  float *A_h, *C_h;
+  A_h = reinterpret_cast<float*>(malloc(Nbytes));
+  C_h = reinterpret_cast<float*>(malloc(Nbytes));
+  REQUIRE(A_h != nullptr);
+  REQUIRE(C_h != nullptr);
+  HIP_CHECK(hipMalloc(&A_d, Nbytes));
+  HIP_CHECK(hipMalloc(&C_d, Nbytes));
+  REQUIRE(A_d != nullptr);
+  REQUIRE(C_d != nullptr);
+
+  HIP_CHECK(hipStreamCreate(&streamForGraph));
+
+  // Initialize input buffer
+  for (size_t i = 0; i < N; ++i) {
+      A_h[i] = 3.146f + i;  // Pi
+  }
+
+  HIP_CHECK(hipStreamCreate(&stream1));
+  HIP_CHECK(hipStreamCreate(&stream2));
+  HIP_CHECK(hipStreamCreate(&mstream));
+  HIP_CHECK(hipEventCreate(&memsetEvent1));
+  HIP_CHECK(hipEventCreate(&memsetEvent2));
+  HIP_CHECK(hipEventCreate(&forkStreamEvent));
+  HIP_CHECK(hipStreamBeginCapture(mstream, hipStreamCaptureModeGlobal));
+  HIP_CHECK(hipEventRecord(forkStreamEvent, mstream));
+  HIP_CHECK(hipStreamWaitEvent(stream1, forkStreamEvent, 0));
+  HIP_CHECK(hipStreamWaitEvent(stream2, forkStreamEvent, 0));
+  HIP_CHECK(hipMemsetAsync(A_d, 0, Nbytes, stream1));
+  HIP_CHECK(hipEventRecord(memsetEvent1, stream1));
+  HIP_CHECK(hipMemsetAsync(C_d, 0, Nbytes, stream2));
+  HIP_CHECK(hipEventRecord(memsetEvent2, stream2));
+  HIP_CHECK(hipStreamWaitEvent(mstream, memsetEvent1, 0));
+  HIP_CHECK(hipStreamWaitEvent(mstream, memsetEvent2, 0));
+  HIP_CHECK(hipMemcpyAsync(A_d, A_h, Nbytes, hipMemcpyHostToDevice, mstream));
+  hipLaunchKernelGGL(HipTest::vector_square, dim3(blocks),
+                              dim3(threadsPerBlock), 0, mstream, A_d, C_d, N);
+  HIP_CHECK(hipMemcpyAsync(C_h, C_d, Nbytes, hipMemcpyDeviceToHost, mstream));
+  HIP_CHECK(hipStreamEndCapture(mstream, &graph));
+
+  // Verify numof root nodes
+  HIP_CHECK(hipGraphGetRootNodes(graph, nullptr, &numRootNodes));
+  REQUIRE(numRootNodes == numMemsetNodes);
+  INFO("Num of nodes returned by GetRootNodes : " << numRootNodes);
+
+  int numBytes = sizeof(hipGraphNode_t) * numRootNodes;
+  hipGraphNode_t* nodes = reinterpret_cast<hipGraphNode_t *>(malloc(numBytes));
+  REQUIRE(nodes != nullptr);
+
+  hipGraphNodeType nodeType;
+  HIP_CHECK(hipGraphGetRootNodes(graph, nodes, &numRootNodes));
+  REQUIRE(numRootNodes == numMemsetNodes);
+
+  // Verify root nodes returned are memset nodes.
+  HIP_CHECK(hipGraphNodeGetType(nodes[0], &nodeType));
+  REQUIRE(nodeType == hipGraphNodeTypeMemset);
+  HIP_CHECK(hipGraphNodeGetType(nodes[1], &nodeType));
+  REQUIRE(nodeType == hipGraphNodeTypeMemset);
+
+  // Instantiate and launch the graph
+  HIP_CHECK(hipGraphInstantiate(&graphExec, graph, NULL, NULL, 0));
+  HIP_CHECK(hipGraphLaunch(graphExec, streamForGraph));
+  HIP_CHECK(hipStreamSynchronize(streamForGraph));
+
+  // Validate the computation
+  for (size_t i = 0; i < N; i++) {
+    if (C_h[i] != A_h[i] * A_h[i]) {
+      INFO("A and C not matching at " << i << " C_h[i] " << C_h[i]
+                                           << " A_h[i] " << A_h[i]);
+      REQUIRE(false);
+    }
+  }
+
+  HIP_CHECK(hipGraphExecDestroy(graphExec));
+  HIP_CHECK(hipGraphDestroy(graph));
+  HIP_CHECK(hipStreamDestroy(streamForGraph));
+  HIP_CHECK(hipStreamDestroy(mstream));
+  HIP_CHECK(hipStreamDestroy(stream1));
+  HIP_CHECK(hipStreamDestroy(stream2));
+  HIP_CHECK(hipEventDestroy(forkStreamEvent));
+  HIP_CHECK(hipEventDestroy(memsetEvent1));
+  HIP_CHECK(hipEventDestroy(memsetEvent2));
+  free(A_h);
+  free(C_h);
+  free(nodes);
+  HIP_CHECK(hipFree(A_d));
+  HIP_CHECK(hipFree(C_d));
+}
+
+
 /**
  * Test performs api parameter validation by passing various values
  * as input and output parameters and validates the behavior.