diff --git a/hipamd/src/hip_device.cpp b/hipamd/src/hip_device.cpp index 0b9e8f4ebf..cc8cfdacf4 100644 --- a/hipamd/src/hip_device.cpp +++ b/hipamd/src/hip_device.cpp @@ -74,6 +74,13 @@ bool Device::Create() { return true; } +// ================================================================================================ +bool Device::IsMemoryPoolValid(MemoryPool* pool) { + amd::ScopedLock lock(lock_); + bool result = (mem_pools_.find(pool) != mem_pools_.end()) ? true : false; + return result; +} + // ================================================================================================ void Device::AddMemoryPool(MemoryPool* pool) { amd::ScopedLock lock(lock_); diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp index 992d030932..ec20e47873 100644 --- a/hipamd/src/hip_internal.hpp +++ b/hipamd/src/hip_internal.hpp @@ -535,6 +535,8 @@ public: /// Removes a destroyed stream from the safe list of memory pools void RemoveStreamFromPools(Stream* stream); + /// Returns true if memory pool is valid on this device + bool IsMemoryPoolValid(MemoryPool* pool); }; /// Thread Local Storage Variables Aggregator Class diff --git a/hipamd/src/hip_mempool.cpp b/hipamd/src/hip_mempool.cpp index 4acb57eeb6..3b259e316c 100644 --- a/hipamd/src/hip_mempool.cpp +++ b/hipamd/src/hip_mempool.cpp @@ -26,6 +26,20 @@ namespace hip { */ extern hipError_t ihipFree(void* ptr); +// Static declarations +namespace { +// ================================================================================================ +inline bool IsMemPoolValid(MemoryPool* mem_pool) { + bool result = false; + for (auto it : g_devices) { + if (result = it->IsMemoryPoolValid(mem_pool) == true) { + break; + } + } + return result; +} +} // namespace + // ================================================================================================ hipError_t hipDeviceGetDefaultMemPool(hipMemPool_t* mem_pool, int device) { HIP_INIT_API(hipDeviceGetDefaultMemPool, mem_pool, device); @@ -221,11 +235,17 @@ hipError_t hipMemPoolSetAccess( if ((mem_pool == nullptr) || (desc_list == nullptr)) { HIP_RETURN(hipErrorInvalidValue); } + if (count > g_devices.size()) { + HIP_RETURN(hipErrorInvalidDevice); + } auto hip_mem_pool = reinterpret_cast(mem_pool); for (int i = 0; i < count; ++i) { if (desc_list[i].location.type == hipMemLocationTypeDevice) { if (desc_list[i].location.id >= g_devices.size()) { - HIP_RETURN(hipErrorInvalidValue); + HIP_RETURN(hipErrorInvalidDevice); + } + if (desc_list[i].flags == hipMemAccessFlagsProtNone) { + HIP_RETURN(hipErrorInvalidDevice); } if (desc_list[i].flags > hipMemAccessFlagsProtReadWrite) { HIP_RETURN(hipErrorInvalidValue); @@ -264,7 +284,7 @@ hipError_t hipMemPoolGetAccess( // ================================================================================================ hipError_t hipMemPoolCreate(hipMemPool_t* mem_pool, const hipMemPoolProps* pool_props) { HIP_INIT_API(hipMemPoolCreate, mem_pool, pool_props); - if (mem_pool == nullptr) { + if ((mem_pool == nullptr) || (pool_props == nullptr)) { HIP_RETURN(hipErrorInvalidValue); } // validate hipMemAllocationType value @@ -293,6 +313,10 @@ hipError_t hipMemPoolDestroy(hipMemPool_t mem_pool) { } hip::MemoryPool* hip_mem_pool = reinterpret_cast(mem_pool); + if (!IsMemPoolValid(hip_mem_pool)) { + HIP_RETURN(hipErrorInvalidValue); + } + auto device = hip_mem_pool->Device(); if (hip_mem_pool == device->GetDefaultMemoryPool()) { HIP_RETURN(hipErrorInvalidValue);