From 8a333a989bd907e5e8b5bf740fabd38faf08a591 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Thu, 12 May 2022 08:52:27 +0000
Subject: [PATCH] SWDEV-332606 - SWDEV-333585 - Fix hipMemcpySymbol* and
 hipGetSymbol* APIs

This also fixes the QMCPack staging (SWDEV-335826) test failures related
to hipMemcpyToSymbol

Change-Id: Icc3ff04fb1e17999f5bddf12e1a7266e2289032f
---
 hipamd/src/hip_memory.cpp   | 18 ++++++++++++++++++
 hipamd/src/hip_platform.cpp |  6 ++++++
 2 files changed, 24 insertions(+)

diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index bd1ffe1166..d2866fff49 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -1125,6 +1125,11 @@ inline hipError_t ihipMemcpySymbol_validate(const void* symbol, size_t sizeBytes
 hipError_t hipMemcpyToSymbol_common(const void* symbol, const void* src, size_t sizeBytes,
                              size_t offset, hipMemcpyKind kind, hipStream_t stream=nullptr) {
   CHECK_STREAM_CAPTURING();
+
+  if (kind != hipMemcpyHostToDevice && kind != hipMemcpyDeviceToDevice) {
+    HIP_RETURN(hipErrorInvalidMemcpyDirection);
+  }
+
   size_t sym_size = 0;
   hipDeviceptr_t device_ptr = nullptr;
 
@@ -1153,6 +1158,11 @@ hipError_t hipMemcpyToSymbol_spt(const void* symbol, const void* src, size_t siz
 hipError_t hipMemcpyFromSymbol_common(void* dst, const void* symbol, size_t sizeBytes,
                                size_t offset, hipMemcpyKind kind, hipStream_t stream=nullptr) {
   CHECK_STREAM_CAPTURING();
+
+  if (kind != hipMemcpyDeviceToHost && kind != hipMemcpyDeviceToDevice) {
+    HIP_RETURN(hipErrorInvalidMemcpyDirection);
+  }
+
   size_t sym_size = 0;
   hipDeviceptr_t device_ptr = nullptr;
 
@@ -1184,6 +1194,10 @@ hipError_t hipMemcpyToSymbolAsync(const void* symbol, const void* src, size_t si
 
   STREAM_CAPTURE(hipMemcpyToSymbolAsync, stream, symbol, src, sizeBytes, offset, kind);
 
+  if (kind != hipMemcpyHostToDevice && kind != hipMemcpyDeviceToDevice) {
+    HIP_RETURN(hipErrorInvalidMemcpyDirection);
+  }
+
   size_t sym_size = 0;
   hipDeviceptr_t device_ptr = nullptr;
 
@@ -1201,6 +1215,10 @@ hipError_t hipMemcpyFromSymbolAsync(void* dst, const void* symbol, size_t sizeBy
 
   STREAM_CAPTURE(hipMemcpyFromSymbolAsync, stream, dst, symbol, sizeBytes, offset, kind);
 
+  if (kind != hipMemcpyDeviceToHost && kind != hipMemcpyDeviceToDevice) {
+    HIP_RETURN(hipErrorInvalidMemcpyDirection);
+  }
+
   size_t sym_size = 0;
   hipDeviceptr_t device_ptr = nullptr;
 
diff --git a/hipamd/src/hip_platform.cpp b/hipamd/src/hip_platform.cpp
index 26de428030..9c791f3a10 100644
--- a/hipamd/src/hip_platform.cpp
+++ b/hipamd/src/hip_platform.cpp
@@ -278,6 +278,9 @@ hipError_t hipGetSymbolAddress(void** devPtr, const void* symbol) {
   HIP_INIT_API(hipGetSymbolAddress, devPtr, symbol);
 
   hipError_t hip_error = hipSuccess;
+  if (devPtr == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
   size_t sym_size = 0;
 
   HIP_RETURN_ONFAIL(PlatformState::instance().getStatGlobalVar(symbol, ihipGetDevice(), devPtr, &sym_size));
@@ -288,6 +291,9 @@ hipError_t hipGetSymbolAddress(void** devPtr, const void* symbol) {
 hipError_t hipGetSymbolSize(size_t* sizePtr, const void* symbol) {
   HIP_INIT_API(hipGetSymbolSize, sizePtr, symbol);
 
+  if (sizePtr == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
   hipDeviceptr_t device_ptr = nullptr;
   HIP_RETURN_ONFAIL(PlatformState::instance().getStatGlobalVar(symbol, ihipGetDevice(), &device_ptr, sizePtr));