From 58e05898d1671671fed4cf1c332218ec4cc3d665 Mon Sep 17 00:00:00 2001 From: Ioannis Assiouras Date: Tue, 7 Mar 2023 13:53:44 +0000 Subject: [PATCH] SWDEV-386926 - Added check for correct node type in hipGraphMem[Alloc/Free]NodeGetParams Change-Id: Ic8446210a683d1329391288d0ed99d835c5aa140 --- hipamd/src/hip_graph.cpp | 6 ++++-- hipamd/src/hip_graph_internal.hpp | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) 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() {}