From fb835c6f3e74f0f1097f51bcab13b4e8b4271fc3 Mon Sep 17 00:00:00 2001
From: German Andryeyev <German.Andryeyev@amd.com>
Date: Fri, 1 Apr 2022 11:12:04 -0400
Subject: [PATCH] SWDEV-328670 - Pass the size of allocation

The actual size of allocation can be different from the original
request due to alignments. Pass the size of the original request
for validation during the arena object look-up

Change-Id: I30000046b1ec56abe743e43113046144fd3edbc6
---
 hipamd/src/hip_internal.hpp | 2 +-
 hipamd/src/hip_memory.cpp   | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp
index 67733e15aa..12b61cb3f1 100644
--- a/hipamd/src/hip_internal.hpp
+++ b/hipamd/src/hip_internal.hpp
@@ -477,7 +477,7 @@ extern hipError_t ihipDeviceGetCount(int* count);
 extern int ihipGetDevice();
 
 extern hipError_t ihipMalloc(void** ptr, size_t sizeBytes, unsigned int flags);
-extern amd::Memory* getMemoryObject(const void* ptr, size_t& offset);
+extern amd::Memory* getMemoryObject(const void* ptr, size_t& offset, size_t size = 0);
 extern amd::Memory* getMemoryObjectWithOffset(const void* ptr, const size_t size);
 extern void getStreamPerThread(hipStream_t& stream);
 extern hipError_t ihipUnbindTexture(textureReference* texRef);
diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index 12da05183e..b8b1b2f778 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -28,7 +28,7 @@
 #include "amdocl/cl_vk_amd.hpp"
 
 // ================================================================================================
-amd::Memory* getMemoryObject(const void* ptr, size_t& offset) {
+amd::Memory* getMemoryObject(const void* ptr, size_t& offset, size_t size) {
   amd::Memory *memObj = amd::MemObjMap::FindMemObj(ptr);
   if (memObj != nullptr) {
     const char* hostPtr = reinterpret_cast<const char*>(ptr);
@@ -51,7 +51,8 @@ amd::Memory* getMemoryObject(const void* ptr, size_t& offset) {
     }
   } else {
     // If memObj not found, use arena_mem_obj. arena_mem_obj is null, if HMM and Xnack is disabled.
-    memObj = (hip::getCurrentDevice()->asContext()->svmDevices()[0])->GetArenaMemObj(ptr, offset);
+    memObj = (hip::getCurrentDevice()->asContext()->svmDevices()[0])->GetArenaMemObj(
+        ptr, offset, size);
   }
   return memObj;
 }
@@ -2413,7 +2414,7 @@ hipError_t hipMemsetD32Async(hipDeviceptr_t dst, int value, size_t count,
 hipError_t ihipMemset3D_validate(hipPitchedPtr pitchedDevPtr, int value, hipExtent extent,
                                         size_t sizeBytes) {
   size_t offset = 0;
-  amd::Memory* memory = getMemoryObject(pitchedDevPtr.ptr, offset);
+  amd::Memory* memory = getMemoryObject(pitchedDevPtr.ptr, offset, sizeBytes);
 
   if (memory == nullptr) {
     return hipErrorInvalidValue;