From 8b956db8dbef826ae5447e5f4c9909e6ef79c4b5 Mon Sep 17 00:00:00 2001
From: Marko Arandjelovic <Marko.Arandjelovic@amd.com>
Date: Tue, 21 Jan 2025 09:50:37 +0000
Subject: [PATCH] SWDEV-489619 - Added checks for memcpy capture path

Change-Id: I0e156099282f0b6393bcbcee2e9b96c31034a851


[ROCm/clr commit: 269ec5425248720cffa90355654c772f9d77819a]
---
 projects/clr/hipamd/src/hip_graph.cpp | 30 +++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp
index f6bd51450b..e5966a87a4 100644
--- a/projects/clr/hipamd/src/hip_graph.cpp
+++ b/projects/clr/hipamd/src/hip_graph.cpp
@@ -561,9 +561,15 @@ hipError_t capturehipMemcpy2DToArrayAsync(hipStream_t& stream, hipArray_t& dst,
                                           size_t& width, size_t& height, hipMemcpyKind& kind) {
   ClPrint(amd::LOG_INFO, amd::LOG_API,
           "[hipGraph] Current capture node Memcpy2DFromArray on stream : %p", stream);
-  if (src == nullptr || dst == nullptr) {
-    return hipErrorInvalidValue;
+
+  // Skip zero-sized copies
+  if (width == 0 || height == 0) {
+    return hipSuccess;
   }
+
+  HIP_RETURN_ONFAIL(hipMemcpy2DValidateArray(dst, wOffset, hOffset, width, height));
+  HIP_RETURN_ONFAIL(hipMemcpy2DValidateBuffer(src, spitch, width));
+
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
@@ -789,6 +795,16 @@ hipError_t capturehipMemcpyFromSymbolAsync(hipStream_t& stream, void*& dst, cons
                                            size_t& sizeBytes, size_t& offset, hipMemcpyKind& kind) {
   ClPrint(amd::LOG_INFO, amd::LOG_API,
           "[hipGraph] Current capture node MemcpyFromSymbolNode on stream : %p", stream);
+
+  if (kind != hipMemcpyDeviceToHost && kind != hipMemcpyDeviceToDevice &&
+      kind != hipMemcpyDeviceToDeviceNoCU) {
+    return hipErrorInvalidMemcpyDirection;
+  }
+
+  if (dst == nullptr) {
+    return hipErrorInvalidValue;
+  }
+
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
@@ -815,6 +831,16 @@ hipError_t capturehipMemcpyToSymbolAsync(hipStream_t& stream, const void*& symbo
                                          size_t& sizeBytes, size_t& offset, hipMemcpyKind& kind) {
   ClPrint(amd::LOG_INFO, amd::LOG_API,
           "[hipGraph] Current capture node MemcpyToSymbolNode on stream : %p", stream);
+
+  if (kind != hipMemcpyHostToDevice && kind != hipMemcpyDeviceToDevice &&
+      kind != hipMemcpyDeviceToDeviceNoCU) {
+    return hipErrorInvalidMemcpyDirection;
+  }
+
+  if (src == nullptr) {
+    return hipErrorInvalidValue;
+  }
+
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }