Adjust dev memory available based on impl scaling

Signed-off-by: John Falkowski <john.falkowski@intel.com>

level_zero/core/source/context/context_imp.cpp

@@ -7,6 +7,7 @@
 
 #include "level_zero/core/source/context/context_imp.h"
 
+#include "shared/source/command_container/implicit_scaling.h"
 #include "shared/source/memory_manager/memory_operations_handler.h"
 #include "shared/source/memory_manager/unified_memory_manager.h"
 
@@ -135,18 +136,27 @@ ze_result_t ContextImp::allocDeviceMem(ze_device_handle_t hDevice,
         return ZE_RESULT_SUCCESS;
     }
 
+    neoDevice = this->driverHandle->devices[0]->getNEODevice();
+
     if (lookupTable.relaxedSizeAllowed == false &&
-        (size > this->driverHandle->devices[0]->getNEODevice()->getDeviceInfo().maxMemAllocSize)) {
+        (size > neoDevice->getDeviceInfo().maxMemAllocSize)) {
         *ptr = nullptr;
         return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
     }
 
-    if (lookupTable.relaxedSizeAllowed &&
+    uint64_t globalMemSize = neoDevice->getDeviceInfo().globalMemSize;
-        (size > this->driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize)) {
+
+    uint32_t numSubDevices = neoDevice->getNumSubDevices();
+    if ((!(NEO::ImplicitScalingHelper::isImplicitScalingEnabled(neoDevice->getDeviceBitfield(), true))) && (numSubDevices > 1)) {
+        globalMemSize = globalMemSize / numSubDevices;
+    }
+    if (lookupTable.relaxedSizeAllowed && (size > globalMemSize)) {
         *ptr = nullptr;
         return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
     }
 
+    neoDevice = Device::fromHandle(hDevice)->getNEODevice();
+
     deviceBitfields[rootDeviceIndex] = neoDevice->getDeviceBitfield();
 
     NEO::SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::DEVICE_UNIFIED_MEMORY, this->driverHandle->rootDeviceIndices, deviceBitfields);
@@ -186,20 +196,27 @@ ze_result_t ContextImp::allocSharedMem(ze_device_handle_t hDevice,
         }
     }
 
-    if (relaxedSizeAllowed == false &&
-        (size > this->devices.begin()->second->getNEODevice()->getDeviceInfo().maxMemAllocSize)) {
-        *ptr = nullptr;
-        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
-    }
-
-    if (relaxedSizeAllowed &&
-        (size > this->driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize)) {
-        *ptr = nullptr;
-        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
-    }
-
     auto neoDevice = this->devices.begin()->second->getNEODevice();
+    if (relaxedSizeAllowed == false &&
+        (size > neoDevice->getDeviceInfo().maxMemAllocSize)) {
+        *ptr = nullptr;
+        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+    }
 
+    neoDevice = this->driverHandle->devices[0]->getNEODevice();
+    uint64_t globalMemSize = neoDevice->getDeviceInfo().globalMemSize;
+
+    uint32_t numSubDevices = neoDevice->getNumSubDevices();
+    if ((!(NEO::ImplicitScalingHelper::isImplicitScalingEnabled(neoDevice->getDeviceBitfield(), true))) && (numSubDevices > 1)) {
+        globalMemSize = globalMemSize / numSubDevices;
+    }
+    if (relaxedSizeAllowed &&
+        (size > globalMemSize)) {
+        *ptr = nullptr;
+        return ZE_RESULT_ERROR_UNSUPPORTED_SIZE;
+    }
+
+    neoDevice = this->devices.begin()->second->getNEODevice();
     auto deviceBitfields = this->deviceBitfields;
     NEO::Device *unifiedMemoryPropertiesDevice = nullptr;
     if (hDevice) {

level_zero/core/source/device/device_imp.cpp

@@ -8,6 +8,7 @@
 #include "level_zero/core/source/device/device_imp.h"
 
 #include "shared/source/built_ins/sip.h"
+#include "shared/source/command_container/implicit_scaling.h"
 #include "shared/source/debug_settings/debug_settings_manager.h"
 #include "shared/source/device/device_info.h"
 #include "shared/source/device/sub_device.h"
@@ -275,7 +276,12 @@ ze_result_t DeviceImp::getMemoryProperties(uint32_t *pCount, ze_device_memory_pr
     auto &hwInfoConfig = *NEO::HwInfoConfig::get(hwInfo.platform.eProductFamily);
     pMemProperties->maxClockRate = hwInfoConfig.getDeviceMemoryMaxClkRate(&hwInfo);
     pMemProperties->maxBusWidth = deviceInfo.addressBits;
-    pMemProperties->totalSize = deviceInfo.globalMemSize;
+    if (NEO::ImplicitScalingHelper::isImplicitScalingEnabled(this->getNEODevice()->getDeviceBitfield(), true) ||
+        this->numSubDevices == 0) {
+        pMemProperties->totalSize = deviceInfo.globalMemSize;
+    } else {
+        pMemProperties->totalSize = deviceInfo.globalMemSize / this->numSubDevices;
+    }
 
     return ZE_RESULT_SUCCESS;
 }

level_zero/core/test/unit_tests/sources/device/test_device.cpp

@@ -1201,6 +1201,24 @@ struct MultipleDevicesTest : public ::testing::Test {
     const uint32_t numSubDevices = 2u;
 };
 
+TEST_F(MultipleDevicesTest, whenCallingGetMemoryPropertiesWithSubDevicesThenCorrectSizeReturned) {
+    L0::Device *device0 = driverHandle->devices[0];
+    uint32_t count = 1;
+
+    DebugManager.flags.EnableWalkerPartition.set(0);
+    ze_device_memory_properties_t memProperties = {};
+    ze_result_t res = device0->getMemoryProperties(&count, &memProperties);
+    EXPECT_EQ(res, ZE_RESULT_SUCCESS);
+    EXPECT_EQ(1u, count);
+    EXPECT_EQ(memProperties.totalSize, device0->getNEODevice()->getDeviceInfo().globalMemSize / numSubDevices);
+
+    DebugManager.flags.EnableWalkerPartition.set(1);
+    res = device0->getMemoryProperties(&count, &memProperties);
+    EXPECT_EQ(res, ZE_RESULT_SUCCESS);
+    EXPECT_EQ(1u, count);
+    EXPECT_EQ(memProperties.totalSize, device0->getNEODevice()->getDeviceInfo().globalMemSize);
+}
+
 TEST_F(MultipleDevicesTest, whenRetrievingNumberOfSubdevicesThenCorrectNumberIsReturned) {
     L0::Device *device0 = driverHandle->devices[0];
 

level_zero/core/test/unit_tests/sources/memory/test_memory.cpp

@@ -2794,5 +2794,87 @@ TEST_F(SharedAllocMultiDeviceTests, whenAllocatinSharedMemoryWithNonNullDeviceIn
     EXPECT_EQ(res, ZE_RESULT_SUCCESS);
 }
 
+struct MemAllocMultiSubDeviceTests : public ::testing::Test {
+    void SetUp() override {
+        NEO::MockCompilerEnableGuard mock(true);
+        DebugManager.flags.CreateMultipleSubDevices.set(numSubDevices);
+        auto executionEnvironment = new NEO::ExecutionEnvironment;
+        auto devices = NEO::DeviceFactory::createDevices(*executionEnvironment);
+        driverHandle = std::make_unique<DriverHandleImp>();
+        ze_result_t res = driverHandle->initialize(std::move(devices));
+        EXPECT_EQ(ZE_RESULT_SUCCESS, res);
+        prevSvmAllocsManager = driverHandle->svmAllocsManager;
+        currSvmAllocsManager = new SVMAllocsManagerSharedAllocMultiDeviceMock(driverHandle->memoryManager);
+        driverHandle->svmAllocsManager = currSvmAllocsManager;
+
+        context = std::make_unique<ContextMultiDeviceMock>(driverHandle.get());
+        EXPECT_NE(context, nullptr);
+
+        for (uint32_t i = 0; i < numRootDevices; i++) {
+            auto device = driverHandle->devices[i];
+            context->getDevices().insert(std::make_pair(device->toHandle(), device));
+            auto neoDevice = device->getNEODevice();
+            context->rootDeviceIndices.insert(neoDevice->getRootDeviceIndex());
+            context->deviceBitfields.insert({neoDevice->getRootDeviceIndex(), neoDevice->getDeviceBitfield()});
+        }
+    }
+
+    void TearDown() override {
+        driverHandle->svmAllocsManager = prevSvmAllocsManager;
+        delete currSvmAllocsManager;
+    }
+
+    DebugManagerStateRestore restorer;
+    NEO::SVMAllocsManager *prevSvmAllocsManager;
+    SVMAllocsManagerSharedAllocMultiDeviceMock *currSvmAllocsManager;
+    std::unique_ptr<DriverHandleImp> driverHandle;
+    std::unique_ptr<ContextMultiDeviceMock> context;
+    const uint32_t numSubDevices = 2u;
+    const uint32_t numRootDevices = 1u;
+};
+
+TEST_F(MemAllocMultiSubDeviceTests, whenAllocatingDeviceMemorySubDeviceMemorySizeUsedWhenImplicitScalingDisabled) {
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    void *ptr = nullptr;
+    size_t size = driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize;
+    deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
+    ze_relaxed_allocation_limits_exp_desc_t relaxedSizeDesc = {};
+    relaxedSizeDesc.stype = ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC;
+    relaxedSizeDesc.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE;
+    deviceDesc.pNext = &relaxedSizeDesc;
+
+    DebugManager.flags.EnableWalkerPartition.set(0);
+
+    ze_result_t res = context->allocDeviceMem(driverHandle->devices[0]->toHandle(), &deviceDesc, size, 0u, &ptr);
+    EXPECT_EQ(res, ZE_RESULT_ERROR_UNSUPPORTED_SIZE);
+
+    DebugManager.flags.EnableWalkerPartition.set(1);
+
+    res = context->allocDeviceMem(driverHandle->devices[0]->toHandle(), &deviceDesc, size, 0u, &ptr);
+    EXPECT_EQ(res, ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY);
+}
+
+TEST_F(MemAllocMultiSubDeviceTests, whenAllocatingSharedMemorySubDeviceMemorySizeUsedWhenImplicitScalingDisabled) {
+    ze_device_mem_alloc_desc_t deviceDesc = {};
+    ze_host_mem_alloc_desc_t hostDesc = {};
+    void *ptr = nullptr;
+    size_t size = driverHandle->devices[0]->getNEODevice()->getDeviceInfo().globalMemSize;
+    deviceDesc.stype = ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC;
+    ze_relaxed_allocation_limits_exp_desc_t relaxedSizeDesc = {};
+    relaxedSizeDesc.stype = ZE_STRUCTURE_TYPE_RELAXED_ALLOCATION_LIMITS_EXP_DESC;
+    relaxedSizeDesc.flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE;
+    deviceDesc.pNext = &relaxedSizeDesc;
+
+    DebugManager.flags.EnableWalkerPartition.set(0);
+
+    ze_result_t res = context->allocSharedMem(driverHandle->devices[0]->toHandle(), &deviceDesc, &hostDesc, size, 0u, &ptr);
+    EXPECT_EQ(res, ZE_RESULT_ERROR_UNSUPPORTED_SIZE);
+
+    DebugManager.flags.EnableWalkerPartition.set(1);
+
+    res = context->allocSharedMem(driverHandle->devices[0]->toHandle(), &deviceDesc, &hostDesc, size, 0u, &ptr);
+    EXPECT_EQ(res, ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY);
+}
+
 } // namespace ult
 } // namespace L0