From c27ec6c94e52b82a01706f11a1a6f68d0fa5eb3e Mon Sep 17 00:00:00 2001
From: Ajay <ajay.gunashekar@amd.com>
Date: Tue, 25 Jan 2022 00:11:58 +0000
Subject: [PATCH] SWDEV-316467 - hipGraphMemcpyNodeSetParamsToSymbol negative
 scenarios

Change-Id: Ic2a6ce1572f35dabe754c9459690448c49bbd4a6


[ROCm/clr commit: 0440d945e8e4a655b97b51c30f6ea43cbd268273]
---
 projects/clr/hipamd/src/hip_graph.cpp         |  7 +++++++
 .../clr/hipamd/src/hip_graph_internal.hpp     | 20 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp
index fa625a0902..4fd87a4791 100644
--- a/projects/clr/hipamd/src/hip_graph.cpp
+++ b/projects/clr/hipamd/src/hip_graph.cpp
@@ -1385,6 +1385,13 @@ hipError_t hipGraphMemcpyNodeSetParamsToSymbol(hipGraphNode_t node, const void*
                                                const void* src, size_t count, size_t offset,
                                                hipMemcpyKind kind) {
   HIP_INIT_API(hipGraphMemcpyNodeSetParamsToSymbol, symbol, src, count, offset, kind);
+  if (node == nullptr || src == nullptr || count == 0 || symbol == src) {
+    return hipErrorInvalidValue;
+  }
+  if (symbol  == nullptr) {
+    return hipErrorInvalidSymbol;
+  }
+
   HIP_RETURN(reinterpret_cast<hipGraphMemcpyNodeToSymbol*>(node)->SetParams(symbol, src, count,
                                                                             offset, kind));
 }
diff --git a/projects/clr/hipamd/src/hip_graph_internal.hpp b/projects/clr/hipamd/src/hip_graph_internal.hpp
index 005144a55e..7b457eab40 100644
--- a/projects/clr/hipamd/src/hip_graph_internal.hpp
+++ b/projects/clr/hipamd/src/hip_graph_internal.hpp
@@ -716,6 +716,26 @@ class hipGraphMemcpyNodeToSymbol : public hipGraphMemcpyNode1D {
 
   hipError_t SetParams(const void* symbol, const void* src, size_t count, size_t offset,
                        hipMemcpyKind kind) {
+
+    size_t zeroOffset = 0;
+    amd::Memory* srcMemory = getMemoryObject(src, zeroOffset);
+    amd::Memory* dstMemory = getMemoryObject(symbol, zeroOffset);
+    hipMemoryType srcMemoryType =
+        amd::MemObjMap::FindMemObj(srcMemory) ? hipMemoryTypeDevice : hipMemoryTypeHost;
+    hipMemoryType dstMemoryType =
+        amd::MemObjMap::FindMemObj(dstMemory) ? hipMemoryTypeDevice : hipMemoryTypeHost;
+
+    // Return error if sizeBytes passed to memcpy is more than the actual size allocated
+    if ((dstMemory && count > (dstMemory->getSize() - offset)) ||
+        (srcMemory && count > (srcMemory->getSize() - offset))) {
+      return hipErrorInvalidValue;
+    }
+    // check the kind with memory types
+    if (std::get<0>(hip::getMemoryType(kind)) != srcMemoryType)
+      return hipErrorInvalidValue;
+    if (std::get<1>(hip::getMemoryType(kind)) != dstMemoryType)
+        return hipErrorInvalidValue;
+
     size_t sym_size = 0;
     hipDeviceptr_t device_ptr = nullptr;