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

@@ -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);

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;