Add support for USM shared in WSL for dGPU

This patch force KMD allocation path for USM shared

Related-To: NEO-6913
Signed-off-by: Kamil Diedrich kamil.diedrich@intel.com

shared/source/memory_manager/allocation_properties.h

@@ -11,6 +11,7 @@
 
 namespace NEO {
 struct ImageInfo;
+
 struct AllocationProperties {
     union {
         struct {
@@ -40,6 +41,8 @@ struct AllocationProperties {
     AllocationType allocationType = AllocationType::UNKNOWN;
     GraphicsAllocation::UsmInitialPlacement usmInitialPlacement = GraphicsAllocation::UsmInitialPlacement::DEFAULT;
     ImageInfo *imgInfo = nullptr;
+    bool forceKMDAllocation = false;
+    bool makeGPUVaDifferentThanCPUPtr = false;
     bool multiStorageResource = false;
     ColouringPolicy colouringPolicy = ColouringPolicy::DeviceCountBased;
     size_t colouringGranularity = MemoryConstants::pageSize64k;
@@ -122,6 +125,8 @@ struct AllocationData {
     size_t alignment = 0;
     StorageInfo storageInfo = {};
     ImageInfo *imgInfo = nullptr;
+    bool forceKMDAllocation = false;
+    bool makeGPUVaDifferentThanCPUPtr = false;
     uint32_t rootDeviceIndex = 0;
     OsContext *osContext = nullptr;
     bool useMmapObject = true;

shared/source/memory_manager/memory_manager.cpp

@@ -447,6 +447,8 @@ bool MemoryManager::getAllocationData(AllocationData &allocationData, const Allo
     allocationData.storageInfo = storageInfo;
     allocationData.alignment = properties.alignment ? properties.alignment : MemoryConstants::preferredAlignment;
     allocationData.imgInfo = properties.imgInfo;
+    allocationData.forceKMDAllocation = properties.forceKMDAllocation;
+    allocationData.makeGPUVaDifferentThanCPUPtr = properties.makeGPUVaDifferentThanCPUPtr;
 
     if (allocationData.flags.allocateMemory) {
         allocationData.hostPtr = nullptr;

shared/source/memory_manager/unified_memory_manager.cpp

@@ -493,9 +493,12 @@ void *SVMAllocsManager::createUnifiedAllocationWithDeviceStorage(size_t size, co
                                        subDevices};
     cpuProperties.alignment = MemoryConstants::pageSize2Mb;
     cpuProperties.flags.isUSMHostAllocation = useExternalHostPtrForCpu;
+    cpuProperties.forceKMDAllocation = true;
+    cpuProperties.makeGPUVaDifferentThanCPUPtr = true;
     auto cacheRegion = MemoryPropertiesHelper::getCacheRegion(unifiedMemoryProperties.allocationFlags);
     MemoryPropertiesHelper::fillCachePolicyInProperties(cpuProperties, false, svmProperties.readOnly, false, cacheRegion);
     GraphicsAllocation *allocationCpu = memoryManager->allocateGraphicsMemoryWithProperties(cpuProperties, externalPtr);
+
     if (!allocationCpu) {
         return nullptr;
     }

shared/source/os_interface/windows/wddm_memory_manager.cpp

@@ -146,7 +146,7 @@ GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemoryUsingKmdAndMapItToC
 
     [[maybe_unused]] auto status = true;
 
-    if (executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getHardwareInfo()->capabilityTable.gpuAddressSpace >= MemoryConstants::max64BitAppAddress || is32bit) {
+    if ((!(preferredAllocationMethod == GfxMemoryAllocationMethod::AllocateByKmd && allocationData.makeGPUVaDifferentThanCPUPtr) && executionEnvironment.rootDeviceEnvironments[allocationData.rootDeviceIndex]->getHardwareInfo()->capabilityTable.gpuAddressSpace >= MemoryConstants::max64BitAppAddress) || is32bit) {
         status = mapGpuVirtualAddress(wddmAllocation.get(), cpuPtr);
     } else {
         status = mapGpuVirtualAddress(wddmAllocation.get(), nullptr);
@@ -224,7 +224,7 @@ GraphicsAllocation *WddmMemoryManager::allocateUSMHostGraphicsMemory(const Alloc
 GraphicsAllocation *WddmMemoryManager::allocateGraphicsMemoryWithAlignment(const AllocationData &allocationData) {
     auto pageSize = NEO::OSInterface::osEnabled64kbPages ? MemoryConstants::pageSize64k : MemoryConstants::pageSize;
     bool requiresNonStandardAlignment = allocationData.alignment > pageSize;
-    if ((preferredAllocationMethod == GfxMemoryAllocationMethod::UseUmdSystemPtr) || requiresNonStandardAlignment) {
+    if ((preferredAllocationMethod == GfxMemoryAllocationMethod::UseUmdSystemPtr) || (requiresNonStandardAlignment && allocationData.forceKMDAllocation == false)) {
         return allocateSystemMemoryAndCreateGraphicsAllocationFromIt(allocationData);
     } else {
         return allocateGraphicsMemoryUsingKmdAndMapItToCpuVA(allocationData, NEO::OSInterface::osEnabled64kbPages);

shared/source/os_interface/windows/wddm_memory_manager.h

@@ -107,7 +107,7 @@ class WddmMemoryManager : public MemoryManager {
     GraphicsAllocation *createAllocationFromHandle(osHandle handle, bool requireSpecificBitness, bool ntHandle, AllocationType allocationType, uint32_t rootDeviceIndex);
     static bool validateAllocation(WddmAllocation *alloc);
     MOCKABLE_VIRTUAL bool createWddmAllocation(WddmAllocation *allocation, void *requiredGpuPtr);
-    bool mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *requiredGpuPtr);
+    MOCKABLE_VIRTUAL bool mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *requiredGpuPtr);
     bool mapGpuVaForOneHandleAllocation(WddmAllocation *graphicsAllocation, const void *requiredGpuPtr);
     bool mapMultiHandleAllocationWithRetry(WddmAllocation *allocation, const void *requiredGpuPtr);
     bool createGpuAllocationsWithRetry(WddmAllocation *graphicsAllocation);

shared/test/unit_test/os_interface/windows/CMakeLists.txt

@@ -24,6 +24,7 @@ if(WIN32)
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_address_space_tests.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_command_stream_l0_tests.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_mapper_tests.cpp
+                 ${CMAKE_CURRENT_SOURCE_DIR}/wddm_memory_manager_test.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_preemption_tests.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_shared_allocations_test.cpp
                  ${CMAKE_CURRENT_SOURCE_DIR}/wddm_special_heap_test.cpp

shared/test/unit_test/os_interface/windows/wddm_memory_manager_test.cpp

@@ -0,0 +1,144 @@
+/*
+ * Copyright (C) 2022 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/test/common/helpers/execution_environment_helper.h"
+#include "shared/test/common/mocks/mock_device.h"
+#include "shared/test/common/mocks/mock_memory_manager.h"
+#include "shared/test/common/os_interface/windows/mock_wddm_memory_manager.h"
+#include "shared/test/common/os_interface/windows/wddm_fixture.h"
+
+using namespace NEO;
+using namespace ::testing;
+
+class MockAllocateGraphicsMemoryWithAlignmentWddm : public MemoryManagerCreate<WddmMemoryManager> {
+  public:
+    using WddmMemoryManager::allocateGraphicsMemoryWithAlignment;
+    MockAllocateGraphicsMemoryWithAlignmentWddm(ExecutionEnvironment &executionEnvironment) : MemoryManagerCreate(false, false, executionEnvironment) {}
+    bool allocateSystemMemoryAndCreateGraphicsAllocationFromItCalled = false;
+    bool allocateGraphicsMemoryUsingKmdAndMapItToCpuVACalled = false;
+    bool mapGpuVirtualAddressWithCpuPtr = false;
+
+    GraphicsAllocation *allocateSystemMemoryAndCreateGraphicsAllocationFromIt(const AllocationData &allocationData) override {
+        allocateSystemMemoryAndCreateGraphicsAllocationFromItCalled = true;
+
+        return nullptr;
+    }
+    GraphicsAllocation *allocateGraphicsMemoryUsingKmdAndMapItToCpuVA(const AllocationData &allocationData, bool allowLargePages) override {
+        allocateGraphicsMemoryUsingKmdAndMapItToCpuVACalled = true;
+
+        return nullptr;
+    }
+    bool mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *requiredGpuPtr) override {
+        if (requiredGpuPtr != nullptr) {
+            mapGpuVirtualAddressWithCpuPtr = true;
+        } else {
+            mapGpuVirtualAddressWithCpuPtr = false;
+        }
+
+        return true;
+    }
+};
+
+class WddmMemoryManagerTests : public ::testing::Test {
+  public:
+    MockAllocateGraphicsMemoryWithAlignmentWddm *memoryManager = nullptr;
+    WddmMock *wddm = nullptr;
+    ExecutionEnvironment *executionEnvironment = nullptr;
+
+    void SetUp() override {
+        HardwareInfo *hwInfo = nullptr;
+        executionEnvironment = getExecutionEnvironmentImpl(hwInfo, 1);
+
+        memoryManager = new MockAllocateGraphicsMemoryWithAlignmentWddm(*executionEnvironment);
+        executionEnvironment->memoryManager.reset(memoryManager);
+        wddm = static_cast<WddmMock *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Wddm>());
+    }
+
+    void TearDown() override {
+        delete executionEnvironment;
+    }
+};
+
+TEST_F(WddmMemoryManagerTests, GivenAllocDataWithSVMCPUSetWhenAllocateGraphicsMemoryWithAlignmentThenProperFunctionIsUsed) {
+    NEO::AllocationData allocData = {};
+    allocData.type = NEO::AllocationType::SVM_CPU;
+    allocData.forceKMDAllocation = true;
+    allocData.makeGPUVaDifferentThanCPUPtr = true;
+    memoryManager->allocateGraphicsMemoryWithAlignment(allocData);
+
+    if (preferredAllocationMethod == GfxMemoryAllocationMethod::AllocateByKmd) {
+        EXPECT_TRUE(memoryManager->allocateGraphicsMemoryUsingKmdAndMapItToCpuVACalled);
+    } else {
+        EXPECT_TRUE(memoryManager->allocateSystemMemoryAndCreateGraphicsAllocationFromItCalled);
+    }
+}
+
+class MockAllocateGraphicsMemoryUsingKmdAndMapItToCpuVAWddm : public MemoryManagerCreate<WddmMemoryManager> {
+  public:
+    using WddmMemoryManager::allocateGraphicsMemoryUsingKmdAndMapItToCpuVA;
+    using WddmMemoryManager::mapGpuVirtualAddress;
+    MockAllocateGraphicsMemoryUsingKmdAndMapItToCpuVAWddm(ExecutionEnvironment &executionEnvironment) : MemoryManagerCreate(false, false, executionEnvironment) {}
+
+    bool mapGpuVirtualAddressWithCpuPtr = false;
+
+    bool mapGpuVirtualAddress(WddmAllocation *graphicsAllocation, const void *requiredGpuPtr) override {
+        if (requiredGpuPtr != nullptr) {
+            mapGpuVirtualAddressWithCpuPtr = true;
+        } else {
+            mapGpuVirtualAddressWithCpuPtr = false;
+        }
+
+        return true;
+    }
+};
+
+class WddmMemoryManagerAllocPathTests : public ::testing::Test {
+  public:
+    MockAllocateGraphicsMemoryUsingKmdAndMapItToCpuVAWddm *memoryManager = nullptr;
+    WddmMock *wddm = nullptr;
+    ExecutionEnvironment *executionEnvironment = nullptr;
+
+    void SetUp() override {
+        HardwareInfo *hwInfo = nullptr;
+        executionEnvironment = getExecutionEnvironmentImpl(hwInfo, 1);
+
+        memoryManager = new MockAllocateGraphicsMemoryUsingKmdAndMapItToCpuVAWddm(*executionEnvironment);
+        executionEnvironment->memoryManager.reset(memoryManager);
+        wddm = static_cast<WddmMock *>(executionEnvironment->rootDeviceEnvironments[0]->osInterface->getDriverModel()->as<Wddm>());
+    }
+
+    void TearDown() override {
+        delete executionEnvironment;
+    }
+};
+
+TEST_F(WddmMemoryManagerAllocPathTests, givenAllocateGraphicsMemoryUsingKmdAndMapItToCpuVAWhenPreferedAllocationMethodThenProperArgumentsAreSet) {
+    {
+        NEO::AllocationData allocData = {};
+        allocData.type = NEO::AllocationType::SVM_CPU;
+        allocData.forceKMDAllocation = true;
+        allocData.makeGPUVaDifferentThanCPUPtr = true;
+        auto graphicsAllocation = memoryManager->allocateGraphicsMemoryUsingKmdAndMapItToCpuVA(allocData, false);
+
+        if (preferredAllocationMethod == GfxMemoryAllocationMethod::AllocateByKmd) {
+            EXPECT_FALSE(memoryManager->mapGpuVirtualAddressWithCpuPtr);
+        } else {
+            EXPECT_TRUE(memoryManager->mapGpuVirtualAddressWithCpuPtr);
+        }
+
+        memoryManager->freeGraphicsMemory(graphicsAllocation);
+    }
+    {
+        NEO::AllocationData allocData = {};
+        allocData.type = NEO::AllocationType::EXTERNAL_HOST_PTR;
+        auto graphicsAllocation = memoryManager->allocateGraphicsMemoryUsingKmdAndMapItToCpuVA(allocData, false);
+
+        EXPECT_TRUE(memoryManager->mapGpuVirtualAddressWithCpuPtr);
+
+        memoryManager->freeGraphicsMemory(graphicsAllocation);
+    }
+}