diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp index fcaa2b2449..2a499a02b1 100644 --- a/hipamd/src/hip_graph.cpp +++ b/hipamd/src/hip_graph.cpp @@ -2179,7 +2179,8 @@ hipError_t hipGraphAddMemAllocNode(hipGraphNode_t* pGraphNode, hipGraph_t graph, // ================================================================================================ hipError_t hipGraphMemAllocNodeGetParams(hipGraphNode_t node, hipMemAllocNodeParams* pNodeParams) { HIP_INIT_API(hipGraphMemAllocNodeGetParams, node, pNodeParams); - if (node == nullptr || pNodeParams == nullptr || !hipGraphNode::isNodeValid(node)) { + if (node == nullptr || pNodeParams == nullptr || !hipGraphNode::isNodeValid(node) + || node->GetType() != hipGraphNodeTypeMemAlloc) { HIP_RETURN(hipErrorInvalidValue); } reinterpret_cast(node)->GetParams(pNodeParams); @@ -2214,7 +2215,8 @@ hipError_t hipGraphAddMemFreeNode(hipGraphNode_t* pGraphNode, hipGraph_t graph, // ================================================================================================ hipError_t hipGraphMemFreeNodeGetParams(hipGraphNode_t node, void* dev_ptr) { HIP_INIT_API(hipGraphMemFreeNodeGetParams, node, dev_ptr); - if (node == nullptr || dev_ptr == nullptr || !hipGraphNode::isNodeValid(node)) { + if (node == nullptr || dev_ptr == nullptr || !hipGraphNode::isNodeValid(node) + || node->GetType() != hipGraphNodeTypeMemFree) { HIP_RETURN(hipErrorInvalidValue); } reinterpret_cast(node)->GetParams(reinterpret_cast(dev_ptr)); diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp index a5799081ec..1dcc0e8f39 100644 --- a/hipamd/src/hip_graph_internal.hpp +++ b/hipamd/src/hip_graph_internal.hpp @@ -1847,7 +1847,7 @@ class hipGraphMemAllocNode : public hipGraphNode { public: hipGraphMemAllocNode(const hipMemAllocNodeParams* node_params) - : hipGraphNode(hipGraphNodeTypeEmpty, "solid", "rectangle", "MEM_ALLOC") { + : hipGraphNode(hipGraphNodeTypeMemAlloc, "solid", "rectangle", "MEM_ALLOC") { node_params_ = *node_params; } ~hipGraphMemAllocNode() {} @@ -1889,7 +1889,7 @@ class hipGraphMemFreeNode : public hipGraphNode { public: hipGraphMemFreeNode(void* dptr) - : hipGraphNode(hipGraphNodeTypeEmpty, "solid", "rectangle", "MEM_FREE") + : hipGraphNode(hipGraphNodeTypeMemFree, "solid", "rectangle", "MEM_FREE") , device_ptr_(dptr) {} ~hipGraphMemFreeNode() {}