From a984cbf90caf0618a81196bb7efaf03251bb4cd4 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Mon, 16 Aug 2021 02:51:17 -0400
Subject: [PATCH] SWDEV-298982 - hipMemPrefetchAsync should not allow more
 bytes than the actual size allocated

Change-Id: If8df97a79a16d39c66c508a8927ae1b4c276663e
---
 hipamd/src/hip_hmm.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/hipamd/src/hip_hmm.cpp b/hipamd/src/hip_hmm.cpp
index 49f73ec824..4c4b93efd3 100644
--- a/hipamd/src/hip_hmm.cpp
+++ b/hipamd/src/hip_hmm.cpp
@@ -79,6 +79,14 @@ hipError_t hipMemPrefetchAsync(const void* dev_ptr, size_t count, int device,
   if ((dev_ptr == nullptr) || (count == 0)) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+
+  size_t offset = 0;
+  amd::Memory* memObj = getMemoryObject(dev_ptr, offset);
+  // Return error if count passed is more than the actual size allocated
+  if (memObj && count  > (memObj->getSize() - offset)) {
+    return hipErrorInvalidValue;
+  }
+
   amd::HostQueue* queue = nullptr;
   bool cpu_access = (device == hipCpuDeviceId) ? true : false;