diff --git a/shared/source/memory_manager/definitions/storage_info.cpp b/shared/source/memory_manager/definitions/storage_info.cpp index e4cfcad69d..c823f954c5 100644 --- a/shared/source/memory_manager/definitions/storage_info.cpp +++ b/shared/source/memory_manager/definitions/storage_info.cpp @@ -132,7 +132,9 @@ StorageInfo MemoryManager::createStorageInfoFromProperties(const AllocationPrope storageInfo.cloningOfPageTables = false; storageInfo.tileInstanced = true; } - if (!releaseHelper || releaseHelper->isLocalOnlyAllowed()) { + if (!releaseHelper || + releaseHelper->isLocalOnlyAllowed() || + properties.flags.isUSMDeviceAllocation) { storageInfo.localOnlyRequired = true; } diff --git a/shared/source/memory_manager/memory_manager.cpp b/shared/source/memory_manager/memory_manager.cpp index a883c0dff1..a339b079ed 100644 --- a/shared/source/memory_manager/memory_manager.cpp +++ b/shared/source/memory_manager/memory_manager.cpp @@ -601,7 +601,9 @@ bool MemoryManager::getAllocationData(AllocationData &allocationData, const Allo case AllocationType::buffer: case AllocationType::svmGpu: case AllocationType::image: - if (false == allocationData.flags.uncacheable && useLocalPreferredForCacheableBuffers) { + if (false == allocationData.flags.uncacheable && + useLocalPreferredForCacheableBuffers && + false == allocationData.flags.isUSMDeviceMemory) { if (!allocationData.flags.preferCompressed) { allocationData.storageInfo.localOnlyRequired = false; } diff --git a/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp b/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp index 7730ce9599..b52d68fbd4 100644 --- a/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp +++ b/shared/test/unit_test/memory_manager/memory_manager_allocate_in_preferred_pool_tests.cpp @@ -76,6 +76,17 @@ HWTEST_F(MemoryManagerGetAlloctionDataTests, givenCommandBufferAllocationTypeWhe EXPECT_TRUE(allocData.flags.useSystemMemory); } +HWTEST_F(MemoryManagerGetAlloctionDataTests, givenUsmDeviceAllocationWhenGetAllocationDataIsCalledThenLocalOnlyRequiredIsSet) { + AllocationData allocData; + AllocationProperties properties(mockRootDeviceIndex, true, 10, AllocationType::buffer, false, mockDeviceBitfield); + properties.flags.isUSMDeviceAllocation = true; + + MockMemoryManager mockMemoryManager; + mockMemoryManager.getAllocationData(allocData, properties, nullptr, mockMemoryManager.createStorageInfoFromProperties(properties)); + + EXPECT_TRUE(allocData.storageInfo.localOnlyRequired); +} + TEST_F(MemoryManagerGetAlloctionDataTests, givenAllocateMemoryFlagTrueWhenHostPtrIsNotNullThenAllocationDataHasHostPtrNulled) { AllocationData allocData; char memory = 0; diff --git a/shared/test/unit_test/memory_manager/storage_info_tests.cpp b/shared/test/unit_test/memory_manager/storage_info_tests.cpp index 3f28719700..d755043bf1 100644 --- a/shared/test/unit_test/memory_manager/storage_info_tests.cpp +++ b/shared/test/unit_test/memory_manager/storage_info_tests.cpp @@ -392,6 +392,16 @@ TEST_F(MultiDeviceStorageInfoTest, whenCreatingStorageInfoForSvmGpuThenLocalOnly EXPECT_TRUE(storageInfo.localOnlyRequired); } +TEST_F(MultiDeviceStorageInfoTest, whenCreatingStorageInfoForDeviceUSMAllocationThenLocalOnlyRequiredIsSet) { + AllocationProperties properties{mockRootDeviceIndex, false, numDevices * MemoryConstants::pageSize64k, AllocationType::buffer, false, singleTileMask}; + properties.flags.isUSMDeviceAllocation = true; + auto releaseHelper = std::make_unique(); + releaseHelper->isLocalOnlyAllowedResult = false; + memoryManager->executionEnvironment.rootDeviceEnvironments[properties.rootDeviceIndex]->releaseHelper.reset(releaseHelper.release()); + auto storageInfo = memoryManager->createStorageInfoFromProperties(properties); + EXPECT_TRUE(storageInfo.localOnlyRequired); +} + TEST_F(MultiDeviceStorageInfoTest, givenReleaseWhichDoesNotAllowLocalOnlyWhenCreatingStorageInfoForSvmGpuThenLocalOnlyFlagIsNotRequired) { AllocationProperties properties{mockRootDeviceIndex, false, numDevices * MemoryConstants::pageSize64k, AllocationType::svmGpu, false, singleTileMask}; auto releaseHelper = std::make_unique();