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

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