feature: prepare for l0 usm device pooling

Related-To: NEO-6893

Signed-off-by: Dominik Dabek <[email protected]>

Author: ddabek-i

level_zero/core/source/context/context_imp.cpp

@@ -274,12 +274,17 @@ ze_result_t ContextImp::allocDeviceMem(ze_device_handle_t hDevice,
         unifiedMemoryProperties.allocationFlags.flags.resource48Bit = productHelper.is48bResourceNeededForRayTracing();
     }
 
-    if (false == lookupTable.exportMemory &&
-        neoDevice->getUsmMemAllocPoolsManager()) {
-        neoDevice->getUsmMemAllocPoolsManager()->ensureInitialized(this->driverHandle->svmAllocsManager);
-        if (auto usmPtrFromPool = neoDevice->getUsmMemAllocPoolsManager()->createUnifiedMemoryAllocation(size, unifiedMemoryProperties)) {
-            *ptr = usmPtrFromPool;
-            return ZE_RESULT_SUCCESS;
+    if (false == lookupTable.exportMemory) {
+        if (neoDevice->getUsmMemAllocPoolsManager()) {
+            if (auto usmPtrFromPool = neoDevice->getUsmMemAllocPoolsManager()->createUnifiedMemoryAllocation(size, unifiedMemoryProperties)) {
+                *ptr = usmPtrFromPool;
+                return ZE_RESULT_SUCCESS;
+            }
+        } else if (neoDevice->getUsmMemAllocPool()) {
+            if (auto usmPtrFromPool = neoDevice->getUsmMemAllocPool()->createUnifiedMemoryAllocation(size, unifiedMemoryProperties)) {
+                *ptr = usmPtrFromPool;
+                return ZE_RESULT_SUCCESS;
+            }
         }
     }
 
@@ -464,6 +469,10 @@ ze_result_t ContextImp::freeMem(const void *ptr, bool blocking) {
                                                        blocking)) {
                 return ZE_RESULT_SUCCESS;
             }
+        } else if (auto deviceUsmPool = allocation->device->getUsmMemAllocPool()) {
+            if (deviceUsmPool->freeSVMAlloc(ptr, blocking)) {
+                return ZE_RESULT_SUCCESS;
+            }
         }
     }
     this->driverHandle->svmAllocsManager->freeSVMAlloc(const_cast<void *>(ptr), blocking);
@@ -487,6 +496,18 @@ ze_result_t ContextImp::freeMemExt(const ze_memory_free_ext_desc_t *pMemFreeDesc
             this->freePeerAllocations(ptr, false, Device::fromHandle(pairDevice.second));
         }
 
+        if (InternalMemoryType::hostUnifiedMemory == allocation->memoryType) {
+            if (this->driverHandle->usmHostMemAllocPool.freeSVMAlloc(ptr, false)) {
+                return ZE_RESULT_SUCCESS;
+            }
+        } else if (InternalMemoryType::deviceUnifiedMemory == allocation->memoryType) {
+            if (auto deviceUsmPool = allocation->device->getUsmMemAllocPool()) {
+                if (deviceUsmPool->freeSVMAlloc(ptr, false)) {
+                    return ZE_RESULT_SUCCESS;
+                }
+            }
+        }
+
         this->driverHandle->svmAllocsManager->freeSVMAllocDefer(const_cast<void *>(ptr));
         return ZE_RESULT_SUCCESS;
     }
@@ -641,6 +662,11 @@ void ContextImp::setIPCHandleData(NEO::GraphicsAllocation *graphicsAllocation, u
                     ipcData.poolOffset = poolOffset;
                     break;
                 }
+            } else if (auto deviceUsmMemAllocPool = neoDevice->getUsmMemAllocPool()) {
+                if (auto poolOffset = deviceUsmMemAllocPool->getOffsetInPool(addrToPtr(ptrAddress))) {
+                    ipcData.poolOffset = poolOffset;
+                    break;
+                }
             }
         }
     }

level_zero/core/source/device/device_imp.cpp

@@ -1625,6 +1625,8 @@ void DeviceImp::releaseResources() {
 
     getNEODevice()->getMemoryManager()->freeGraphicsMemory(syncDispatchTokenAllocation);
 
+    getNEODevice()->cleanupUsmAllocationPool();
+
     this->bcsSplit.releaseResources();
 
     if (neoDevice->getExecutionEnvironment()->rootDeviceEnvironments[neoDevice->getRootDeviceIndex()]->debugger.get()) {

level_zero/core/source/driver/driver_handle_imp.cpp

@@ -300,7 +300,15 @@ ze_result_t DriverHandleImp::initialize(std::vector<std::unique_ptr<NEO::Device>
     }
     this->svmAllocsManager->initUsmAllocationsCaches(*this->devices[0]->getNEODevice());
     this->initHostUsmAllocPool();
-
+    for (auto &device : this->devices) {
+        this->initDeviceUsmAllocPool(*device->getNEODevice());
+        if (auto deviceUsmAllocPool = device->getNEODevice()->getUsmMemAllocPool()) {
+            deviceUsmAllocPool->ensureInitialized(this->svmAllocsManager);
+        }
+        if (auto deviceUsmAllocPoolsManager = device->getNEODevice()->getUsmMemAllocPoolsManager()) {
+            deviceUsmAllocPoolsManager->ensureInitialized(this->svmAllocsManager);
+        }
+    }
     this->numDevices = static_cast<uint32_t>(this->devices.size());
 
     uuidTimestamp = static_cast<uint64_t>(std::chrono::system_clock::now().time_since_epoch().count());
@@ -352,6 +360,23 @@ void DriverHandleImp::initHostUsmAllocPool() {
     }
 }
 
+void DriverHandleImp::initDeviceUsmAllocPool(NEO::Device &device) {
+    const uint64_t minServicedSize = 0u;
+    const uint64_t maxServicedSize = 1 * MemoryConstants::megaByte;
+    bool enabled = NEO::ApiSpecificConfig::isDeviceUsmPoolingEnabled() && device.getProductHelper().isDeviceUsmPoolAllocatorSupported();
+    uint64_t poolSize = 2 * MemoryConstants::megaByte;
+
+    if (NEO::debugManager.flags.EnableDeviceUsmAllocationPool.get() != -1) {
+        enabled = NEO::debugManager.flags.EnableDeviceUsmAllocationPool.get() > 0;
+        poolSize = NEO::debugManager.flags.EnableDeviceUsmAllocationPool.get() * MemoryConstants::megaByte;
+    }
+
+    if (enabled) {
+        device.resetUsmAllocationPool(new NEO::UsmMemAllocPool(rootDeviceIndices, deviceBitfields, &device, InternalMemoryType::deviceUnifiedMemory,
+                                                               poolSize, minServicedSize, maxServicedSize));
+    }
+}
+
 ze_result_t DriverHandleImp::getDevice(uint32_t *pCount, ze_device_handle_t *phDevices) {
 
     // If the user has requested FLAT or COMBINED device hierarchy model, then report all the sub devices as devices.

level_zero/core/source/driver/driver_handle_imp.h

@@ -124,6 +124,7 @@ struct DriverHandleImp : public DriverHandle {
     std::map<uint64_t, IpcHandleTracking *> &getIPCHandleMap() { return this->ipcHandles; };
     [[nodiscard]] std::unique_lock<std::mutex> lockIPCHandleMap() { return std::unique_lock<std::mutex>(this->ipcHandleMapMutex); };
     void initHostUsmAllocPool();
+    void initDeviceUsmAllocPool(NEO::Device &device);
 
     std::unique_ptr<HostPointerManager> hostPointerManager;
 

level_zero/core/source/helpers/api_specific_config_l0.cpp

@@ -60,7 +60,7 @@ bool ApiSpecificConfig::isHostAllocationCacheEnabled() {
 }
 
 bool ApiSpecificConfig::isDeviceUsmPoolingEnabled() {
-    return true;
+    return false;
 }
 
 bool ApiSpecificConfig::isHostUsmPoolingEnabled() {

level_zero/core/test/unit_tests/fixtures/device_fixture.h

@@ -59,6 +59,7 @@ struct DeviceFixture {
 template <typename T>
 struct DeviceFixtureWithCustomMemoryManager : public DeviceFixture {
     void setUp() {
+        debugManager.flags.EnableDeviceUsmAllocationPool.set(0);
         auto executionEnvironment = NEO::MockDevice::prepareExecutionEnvironment(defaultHwInfo.get(), 0u);
         memoryManager = new T(*executionEnvironment);
         executionEnvironment->memoryManager.reset(memoryManager);
@@ -68,7 +69,7 @@ struct DeviceFixtureWithCustomMemoryManager : public DeviceFixture {
     void tearDown() {
         DeviceFixture::tearDown();
     }
-
+    DebugManagerStateRestore restorer;
     T *memoryManager = nullptr;
 };
 

level_zero/core/test/unit_tests/fixtures/memory_ipc_fixture.cpp

@@ -536,6 +536,11 @@ void MemoryExportImportImplicitScalingTest::SetUp() {
 }
 
 void MemoryExportImportImplicitScalingTest::TearDown() {
+    // cleanup pool before restoring svm manager
+    for (auto device : driverHandle->devices) {
+        device->getNEODevice()->cleanupUsmAllocationPool();
+        device->getNEODevice()->resetUsmAllocationPool(nullptr);
+    }
     driverHandle->svmAllocsManager = prevSvmAllocsManager;
     delete currSvmAllocsManager;
     driverHandle->setMemoryManager(prevMemoryManager);

level_zero/core/test/unit_tests/sources/cmdlist/test_in_order_cmdlist_1.cpp

@@ -5488,8 +5488,9 @@ HWTEST_F(InOrderCmdListTests, givenExternalSyncStorageWhenCreatingCounterBasedEv
     EXPECT_EQ(counterValue, inOrderExecInfo->getCounterValue());
     EXPECT_EQ(castToUint64(externalStorageAllocProperties.deviceAddress), inOrderExecInfo->getBaseDeviceAddress());
     EXPECT_NE(nullptr, inOrderExecInfo->getDeviceCounterAllocation());
-
-    auto lockedPtr = reinterpret_cast<uint64_t *>(ptrOffset(inOrderExecInfo->getDeviceCounterAllocation()->getLockedPtr(), sizeof(uint64_t)));
+    SvmAllocationData *deviceAlloc = context->getDriverHandle()->getSvmAllocsManager()->getSVMAlloc(reinterpret_cast<void *>(devAddress));
+    auto offset = ptrDiff(devAddress, deviceAlloc->gpuAllocations.getDefaultGraphicsAllocation()->getGpuAddress());
+    auto lockedPtr = reinterpret_cast<uint64_t *>(ptrOffset(inOrderExecInfo->getDeviceCounterAllocation()->getLockedPtr(), sizeof(uint64_t) + offset));
 
     EXPECT_EQ(inOrderExecInfo->getBaseHostAddress(), lockedPtr);
     EXPECT_EQ(inOrderExecInfo->getExternalHostAllocation(), inOrderExecInfo->getDeviceCounterAllocation());

level_zero/core/test/unit_tests/sources/debugger/windows/test_l0_debugger_windows.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020-2024 Intel Corporation
+ * Copyright (C) 2020-2025 Intel Corporation
  *
  * SPDX-License-Identifier: MIT
  *
@@ -34,6 +34,7 @@ namespace ult {
 struct L0DebuggerWindowsFixture {
     void setUp() {
         debugManager.flags.ForcePreferredAllocationMethod.set(static_cast<int32_t>(GfxMemoryAllocationMethod::useUmdSystemPtr));
+        debugManager.flags.EnableDeviceUsmAllocationPool.set(0);
         executionEnvironment = new NEO::ExecutionEnvironment;
         executionEnvironment->setDebuggingMode(NEO::DebuggingMode::online);
         executionEnvironment->prepareRootDeviceEnvironments(1);

level_zero/core/test/unit_tests/sources/driver/test_driver.cpp

@@ -198,6 +198,7 @@ TEST_F(DriverVersionTest, givenExternalAllocatorWhenCallingGetExtensionPropertie
     DebugManagerStateRestore restorer;
     NEO::debugManager.flags.UseBindlessMode.set(1);
     NEO::debugManager.flags.UseExternalAllocatorForSshAndDsh.set(1);
+    NEO::debugManager.flags.EnableDeviceUsmAllocationPool.set(0);
 
     auto hwInfo = *NEO::defaultHwInfo;
     NEO::MockDevice *neoDevice = NEO::MockDevice::createWithNewExecutionEnvironment<NEO::MockDevice>(&hwInfo);