From 8618b1cd0732ce1bf7fe1db0fd6101889039f572 Mon Sep 17 00:00:00 2001
From: Payam <payam.ghafari@amd.com>
Date: Wed, 23 Mar 2022 15:30:06 -0400
Subject: [PATCH] SWDEV-316819 - graph check ChildGraphNodeGetGraph

Change-Id: Id3a4fb64e6fb5ce37e97f1f3518e09c8bd9a2d0c
---
 hipamd/src/hip_graph.cpp          | 6 +++---
 hipamd/src/hip_graph_internal.hpp | 1 +
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index de79d965f4..6679c3d506 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -1211,11 +1211,11 @@ hipError_t hipGraphExecKernelNodeSetParams(hipGraphExec_t hGraphExec, hipGraphNo
 
 hipError_t hipGraphChildGraphNodeGetGraph(hipGraphNode_t node, hipGraph_t* pGraph) {
   HIP_INIT_API(hipGraphChildGraphNodeGetGraph, node, pGraph);
-  if (node == nullptr || pGraph == nullptr) {
+  if (node == nullptr || pGraph == nullptr || !hipGraphNode::isNodeValid(node)) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  *pGraph = reinterpret_cast<hipChildGraphNode*>(node)->GetChildGraph();
-  if (pGraph == nullptr) {
+  *pGraph = reinterpret_cast<hipGraphNode*>(node)->GetChildGraph();
+  if (*pGraph == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
   HIP_RETURN(hipSuccess);
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 5b6a3e719d..eeb5c05cf5 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -224,6 +224,7 @@ struct hipGraphNode {
     }
   }
   ihipGraph* GetParentGraph() { return parentGraph_; }
+  virtual ihipGraph* GetChildGraph() { return nullptr; }
   void SetParentGraph(ihipGraph* graph) { parentGraph_ = graph; }
   virtual hipError_t SetParams(hipGraphNode* node) { return hipSuccess; }
 };