From 33f02abc419ee6ccc3cdadc848ae69809252e12b Mon Sep 17 00:00:00 2001
From: sdashmiz <shadi.dashmiz@amd.com>
Date: Mon, 19 Sep 2022 09:55:03 -0400
Subject: [PATCH] SWDEV-356569 - check for mempool device

- mem pool should be from same device

Signed-off-by: sdashmiz <shadi.dashmiz@amd.com>
Change-Id: Id06a7f070c019548462c235bff3c0cd972e6eb8a
---
 hipamd/src/hip_mempool.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/hipamd/src/hip_mempool.cpp b/hipamd/src/hip_mempool.cpp
index 8309d19e7c..29f8af99c6 100644
--- a/hipamd/src/hip_mempool.cpp
+++ b/hipamd/src/hip_mempool.cpp
@@ -44,6 +44,12 @@ hipError_t hipDeviceSetMemPool(int device, hipMemPool_t mem_pool) {
   if ((mem_pool == nullptr) || (device >= g_devices.size())) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+
+  auto poolDevice = reinterpret_cast<hip::MemoryPool*>(mem_pool)->Device();
+  if (poolDevice->deviceId() != device) {
+    HIP_RETURN(hipErrorInvalidDevice);
+  }
+
   g_devices[device]->SetCurrentMemoryPool(reinterpret_cast<hip::MemoryPool*>(mem_pool));
   HIP_RETURN(hipSuccess);
 }