Refactor memory info

Related-To: NEO-6149


Signed-off-by: Szymon Morek <szymon.morek@intel.com>

level_zero/tools/source/sysman/memory/linux/os_memory_imp_dg1.cpp

@@ -49,7 +49,7 @@ ze_result_t LinuxMemoryImp::getState(zes_mem_state_t *pState) {
     if (!memoryInfo) {
         return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
     }
-    for (auto region : memoryInfo->regions) {
+    for (auto region : memoryInfo->getDrmRegionInfos()) {
         if (region.region.memory_class == I915_MEMORY_CLASS_DEVICE) {
             deviceRegions.push_back(region);
         }

opencl/test/unit_test/os_interface/linux/drm_memory_info_tests_dg1.cpp

@@ -8,6 +8,8 @@
 #include "shared/source/execution_environment/execution_environment.h"
 #include "shared/source/memory_manager/memory_banks.h"
 #include "shared/source/os_interface/linux/memory_info_impl.h"
+#include "shared/test/common/helpers/debug_manager_state_restore.h"
+#include "shared/test/common/helpers/default_hw_info.h"
 
 #include "opencl/test/unit_test/os_interface/linux/drm_mock_dg1.h"
 
@@ -27,7 +29,7 @@ TEST(MemoryInfo, givenMemoryRegionQuerySupportedWhenQueryingMemoryInfoThenMemory
     auto memoryInfo = static_cast<MemoryInfoImpl *>(drm->getMemoryInfo());
 
     ASSERT_NE(nullptr, memoryInfo);
-    EXPECT_EQ(2u, memoryInfo->regions.size());
+    EXPECT_EQ(2u, memoryInfo->getDrmRegionInfos().size());
 }
 
 TEST(MemoryInfo, givenMemoryRegionQueryNotSupportedWhenQueryingMemoryInfoThenMemoryInfoIsNotCreated) {
@@ -69,7 +71,9 @@ TEST(MemoryInfo, givenMemoryRegionQueryWhenQueryingFailsThenMemoryInfoIsNotCreat
     EXPECT_EQ(2u, drm->ioctlCallsCount);
 }
 
-TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstanceThenReturnCorrectValues) {
+TEST(MemoryInfo, givenMemoryInfoWithRegionsAndLocalMemoryEnabledWhenGettingMemoryRegionClassAndInstanceThenReturnCorrectValues) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(1);
     drm_i915_memory_region_info regionInfo[2] = {};
     regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
     regionInfo[0].probed_size = 8 * GB;
@@ -79,28 +83,141 @@ TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstan
     auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 2);
     ASSERT_NE(nullptr, memoryInfo);
 
-    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::MainBank);
+    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::MainBank, *defaultHwInfo);
     EXPECT_EQ(regionInfo[0].region.memory_class, regionClassAndInstance.memory_class);
     EXPECT_EQ(regionInfo[0].region.memory_instance, regionClassAndInstance.memory_instance);
     auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::MainBank);
     EXPECT_EQ(8 * GB, regionSize);
 
-    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(0));
+    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(0), *defaultHwInfo);
     EXPECT_EQ(regionInfo[1].region.memory_class, regionClassAndInstance.memory_class);
     EXPECT_EQ(regionInfo[1].region.memory_instance, regionClassAndInstance.memory_instance);
     regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
     EXPECT_EQ(16 * GB, regionSize);
 }
 
-TEST(MemoryInfo, givenMemoryInfoWithoutRegionsWhenGettingMemoryRegionClassAndInstanceThenReturnInvalidMemoryRegion) {
+TEST(MemoryInfo, givenMemoryInfoWithRegionsAndLocalMemoryEnabledWhenAssignRegionsFromDistancesThenRegionsNotChanged) {
-    drm_i915_memory_region_info regionInfo = {};
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(1);
+    drm_i915_memory_region_info regionInfo[2] = {};
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
+    regionInfo[0].probed_size = 8 * GB;
+    regionInfo[1].region = {I915_MEMORY_CLASS_DEVICE, 0};
+    regionInfo[1].probed_size = 16 * GB;
 
-    auto memoryInfo = std::make_unique<MemoryInfoImpl>(&regionInfo, 0);
+    auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 2);
     ASSERT_NE(nullptr, memoryInfo);
-
+    memoryInfo->assignRegionsFromDistances(&regionInfo, 2);
-    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::MainBank);
+    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::MainBank, *defaultHwInfo);
-    EXPECT_EQ(MemoryInfoImpl::invalidMemoryRegion(), regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[0].region.memory_class, regionClassAndInstance.memory_class);
-    EXPECT_EQ(MemoryInfoImpl::invalidMemoryRegion(), regionClassAndInstance.memory_instance);
+    EXPECT_EQ(regionInfo[0].region.memory_instance, regionClassAndInstance.memory_instance);
     auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::MainBank);
+    EXPECT_EQ(8 * GB, regionSize);
+
+    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(0), *defaultHwInfo);
+    EXPECT_EQ(regionInfo[1].region.memory_class, regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[1].region.memory_instance, regionClassAndInstance.memory_instance);
+    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
+    EXPECT_EQ(16 * GB, regionSize);
+}
+
+TEST(MemoryInfo, givenMemoryInfoWithoutDeviceRegionWhenGettingDeviceRegionSizeThenReturnCorrectSize) {
+    drm_i915_memory_region_info regionInfo[1] = {};
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
+    regionInfo[0].probed_size = 8 * GB;
+    auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 1);
+    ASSERT_NE(nullptr, memoryInfo);
+    auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
     EXPECT_EQ(0 * GB, regionSize);
 }
+
+TEST(MemoryInfo, givenMemoryInfoWithRegionsAndLocalMemoryDisabledWhenGettingMemoryRegionClassAndInstanceThenReturnCorrectValues) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(0);
+    drm_i915_memory_region_info regionInfo[2] = {};
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
+    regionInfo[0].probed_size = 8 * GB;
+    regionInfo[1].region = {I915_MEMORY_CLASS_DEVICE, 0};
+    regionInfo[1].probed_size = 16 * GB;
+
+    auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 2);
+    ASSERT_NE(nullptr, memoryInfo);
+
+    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::MainBank, *defaultHwInfo);
+    EXPECT_EQ(regionInfo[0].region.memory_class, regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[0].region.memory_instance, regionClassAndInstance.memory_instance);
+    auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::MainBank);
+    EXPECT_EQ(8 * GB, regionSize);
+
+    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(0), *defaultHwInfo);
+    EXPECT_EQ(regionInfo[0].region.memory_class, regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[0].region.memory_instance, regionClassAndInstance.memory_instance);
+    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
+    EXPECT_EQ(16 * GB, regionSize);
+}
+
+TEST(MemoryInfo, whenDebugVariablePrintMemoryRegionSizeIsSetAndGetMemoryRegionSizeIsCalledThenMessagePrintedToStdOutput) {
+    DebugManagerStateRestore restore;
+    DebugManager.flags.PrintMemoryRegionSizes.set(true);
+
+    drm_i915_memory_region_info regionInfo[1] = {};
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 1};
+    regionInfo[0].probed_size = 16 * GB;
+
+    auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 1);
+    ASSERT_NE(nullptr, memoryInfo);
+
+    testing::internal::CaptureStdout();
+    auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::MainBank);
+    EXPECT_EQ(16 * GB, regionSize);
+
+    std::string output = testing::internal::GetCapturedStdout();
+    std::string expectedOutput("Memory type: 0, memory instance: 1, region size: 17179869184\n");
+    EXPECT_EQ(expectedOutput, output);
+}
+
+TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstanceWhileDebugFlagIsActiveThenReturnCorrectValues) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(1);
+    drm_i915_memory_region_info regionInfo[3] = {};
+    regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
+    regionInfo[0].probed_size = 8 * GB;
+    regionInfo[1].region = {I915_MEMORY_CLASS_DEVICE, 0};
+    regionInfo[1].probed_size = 16 * GB;
+    regionInfo[2].region = {I915_MEMORY_CLASS_DEVICE, 1};
+    regionInfo[2].probed_size = 32 * GB;
+
+    auto memoryInfo = std::make_unique<MemoryInfoImpl>(regionInfo, 3);
+    ASSERT_NE(nullptr, memoryInfo);
+
+    DebugManager.flags.OverrideDrmRegion.set(1);
+
+    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(0), *defaultHwInfo);
+    EXPECT_EQ(regionInfo[2].region.memory_class, regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[2].region.memory_instance, regionClassAndInstance.memory_instance);
+    auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
+    EXPECT_EQ(16 * GB, regionSize);
+
+    DebugManager.flags.OverrideDrmRegion.set(0);
+
+    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(1), *defaultHwInfo);
+    EXPECT_EQ(regionInfo[1].region.memory_class, regionClassAndInstance.memory_class);
+    EXPECT_EQ(regionInfo[1].region.memory_instance, regionClassAndInstance.memory_instance);
+    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
+    EXPECT_EQ(32 * GB, regionSize);
+
+    DebugManager.flags.OverrideDrmRegion.set(-1);
+    DebugManager.flags.ForceMemoryBankIndexOverride.set(1);
+
+    auto &helper = HwHelper::get(defaultHwInfo->platform.eRenderCoreFamily);
+    regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(1), *defaultHwInfo);
+    if (helper.isBankOverrideRequired(*defaultHwInfo)) {
+        EXPECT_EQ(regionInfo[1].region.memory_class, regionClassAndInstance.memory_class);
+        EXPECT_EQ(regionInfo[1].region.memory_instance, regionClassAndInstance.memory_instance);
+    } else {
+        EXPECT_EQ(regionInfo[2].region.memory_class, regionClassAndInstance.memory_class);
+        EXPECT_EQ(regionInfo[2].region.memory_instance, regionClassAndInstance.memory_instance);
+    }
+    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
+    EXPECT_EQ(32 * GB, regionSize);
+}

opencl/test/unit_test/os_interface/linux/drm_memory_manager_local_memory_tests_dg1.cpp

@@ -171,6 +171,8 @@ class DrmMemoryManagerLocalMemoryWithCustomMockTest : public ::testing::Test {
 };
 
 TEST_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWhenCreateBufferObjectInMemoryRegionIsCalledThenBufferObjectWithAGivenGpuAddressAndSizeIsCreatedAndAllocatedInASpecifiedMemoryRegion) {
+    DebugManagerStateRestore restorer;
+    DebugManager.flags.EnableLocalMemory.set(1);
     drm_i915_memory_region_info regionInfo[2] = {};
     regionInfo[0].region = {I915_MEMORY_CLASS_SYSTEM, 0};
     regionInfo[1].region = {I915_MEMORY_CLASS_DEVICE, 0};
@@ -222,22 +224,6 @@ TEST_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWhenCreateBufferObj
     EXPECT_EQ(nullptr, bo);
 }
 
-TEST_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWhenCreateBufferObjectInMemoryRegionIsCalledWithIncorrectRegionInfoThenNullBufferObjectIsReturned) {
-    drm_i915_memory_region_info regionInfo = {};
-
-    mock->memoryInfo.reset(new MemoryInfoImpl(&regionInfo, 0));
-
-    auto gpuAddress = 0x1234u;
-    auto size = MemoryConstants::pageSize;
-
-    auto bo = std::unique_ptr<BufferObject>(memoryManager->createBufferObjectInMemoryRegion(&memoryManager->getDrm(0),
-                                                                                            gpuAddress,
-                                                                                            size,
-                                                                                            MemoryBanks::MainBank,
-                                                                                            1));
-    EXPECT_EQ(nullptr, bo);
-}
-
 TEST_F(DrmMemoryManagerLocalMemoryTest, givenDrmMemoryManagerWhenCreateBufferObjectInMemoryRegionIsCalledWithZeroSizeThenNullBufferObjectIsReturned) {
     mock->memoryInfo.reset(new MockMemoryInfo());
 

opencl/test/unit_test/test_files/igdrcl.config

@@ -329,4 +329,6 @@ WparidRegisterProgramming = -1
 UsePipeControlAfterPartitionedWalker = -1
 OverrideBufferSuitableForRenderCompression = -1
 AllowMixingRegularAndCooperativeKernels = 0
-AllowPatchingVfeStateInCommandLists = 0
+AllowPatchingVfeStateInCommandLists = 0
+PrintMemoryRegionSizes = 0
+OverrideDrmRegion = -1

shared/source/debug_settings/debug_variables_base.inl

@@ -73,6 +73,7 @@ DECLARE_DEBUG_VARIABLE(bool, EnableFlushTaskSubmission, false, "true: driver use
 DECLARE_DEBUG_VARIABLE(bool, DoNotFreeResources, false, "true: driver stops freeing resources")
 DECLARE_DEBUG_VARIABLE(bool, AllowMixingRegularAndCooperativeKernels, false, "true: driver allows mixing regular and cooperative kernels in a single command list and in a single execute")
 DECLARE_DEBUG_VARIABLE(bool, AllowPatchingVfeStateInCommandLists, false, "true: MEDIA_VFE_STATE may be programmed in a command list")
+DECLARE_DEBUG_VARIABLE(bool, PrintMemoryRegionSizes, false, "print memory bank type, instance and it's size")
 DECLARE_DEBUG_VARIABLE(std::string, ForceDeviceId, std::string("unk"), "DeviceId selected for testing")
 DECLARE_DEBUG_VARIABLE(std::string, LoadBinarySipFromFile, std::string("unk"), "Select binary file to load SIP kernel raw binary")
 DECLARE_DEBUG_VARIABLE(int64_t, OverrideMultiStoragePlacement, -1, "-1: disable, 0+: tile mask, each bit corresponds to tile")
@@ -348,3 +349,4 @@ DECLARE_DEBUG_VARIABLE(bool, UseUmKmDataTranslator, false, "Use helper library f
 DECLARE_DEBUG_VARIABLE(bool, SkipFlushingEventsOnGetStatusCalls, false, "When set to 1, events are not causing internal flush when querying for CL_EVENT_COMMAND_EXECUTION_STATUS")
 DECLARE_DEBUG_VARIABLE(bool, AllowUnrestrictedSize, false, "Allow allocating memory with greater size than MAX_MEM_ALLOC_SIZE")
 DECLARE_DEBUG_VARIABLE(bool, ProgramAdditionalPipeControlBeforeStateComputeModeCommand, false, "Program additional PIPE CONTROL command before STATE_COMPUTE_MODE command")
+DECLARE_DEBUG_VARIABLE(int32_t, OverrideDrmRegion, -1, "-1: disable, 0+: override to given memory region for all allocations")

shared/source/os_interface/linux/CMakeLists.txt

@@ -85,6 +85,7 @@ if(SUPPORT_DG1 AND "${BRANCH_TYPE}" STREQUAL "")
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_local_memory_dg1.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_memory_manager_create_multi_host_allocation_dg1.cpp
        ${CMAKE_CURRENT_SOURCE_DIR}/drm_query_dg1.cpp
+       ${CMAKE_CURRENT_SOURCE_DIR}/memory_info_impl.cpp
   )
 else()
   list(APPEND NEO_CORE_OS_INTERFACE_LINUX

shared/source/os_interface/linux/drm_memory_manager_local_memory_dg1.cpp

@@ -27,11 +27,8 @@ BufferObject *DrmMemoryManager::createBufferObjectInMemoryRegion(Drm *drm,
     if (!memoryInfo) {
         return nullptr;
     }
-
+    auto pHwInfo = drm->getRootDeviceEnvironment().getHardwareInfo();
-    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(memoryBanks);
+    auto regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(memoryBanks, *pHwInfo);
-    if (regionClassAndInstance.memory_class == MemoryInfoImpl::invalidMemoryRegion()) {
-        return nullptr;
-    }
 
     drm_i915_gem_memory_class_instance memRegions{};
     memRegions.memory_class = regionClassAndInstance.memory_class;

shared/source/os_interface/linux/memory_info_impl.cpp

@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2020-2021 Intel Corporation
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ */
+
+#include "shared/source/os_interface/linux/memory_info_impl.h"
+
+#include "shared/source/helpers/basic_math.h"
+#include "shared/source/helpers/debug_helpers.h"
+#include "shared/source/os_interface/linux/drm_neo.h"
+
+#include "drm/i915_drm.h"
+
+#include <cstddef>
+#include <vector>
+
+namespace NEO {
+
+MemoryInfoImpl::MemoryInfoImpl(const drm_i915_memory_region_info *regionInfo, size_t count)
+    : drmQueryRegions(regionInfo, regionInfo + count), systemMemoryRegion(drmQueryRegions[0]) {
+    UNRECOVERABLE_IF(systemMemoryRegion.region.memory_class != I915_MEMORY_CLASS_SYSTEM);
+    std::copy_if(drmQueryRegions.begin(), drmQueryRegions.end(), std::back_inserter(localMemoryRegions),
+                 [](const drm_i915_memory_region_info &memoryRegionInfo) {
+                     return (memoryRegionInfo.region.memory_class == I915_MEMORY_CLASS_DEVICE);
+                 });
+}
+
+void MemoryInfoImpl::assignRegionsFromDistances(const void *distanceInfosPtr, size_t size) {
+}
+
+} // namespace NEO

shared/source/os_interface/linux/memory_info_impl.h

@@ -8,6 +8,7 @@
 #pragma once
 #include "shared/source/helpers/basic_math.h"
 #include "shared/source/helpers/debug_helpers.h"
+#include "shared/source/helpers/hw_helper.h"
 #include "shared/source/os_interface/linux/drm_neo.h"
 #include "shared/source/os_interface/linux/memory_info.h"
 
@@ -18,33 +19,68 @@
 
 namespace NEO {
 
-struct MemoryInfoImpl : public MemoryInfo {
+class MemoryInfoImpl : public MemoryInfo {
+  public:
+    using RegionContainer = std::vector<drm_i915_memory_region_info>;
+
     ~MemoryInfoImpl() override = default;
 
-    MemoryInfoImpl(const drm_i915_memory_region_info *regionInfo, size_t count) : regions(regionInfo, regionInfo + count) {
+    MemoryInfoImpl(const drm_i915_memory_region_info *regionInfo, size_t count);
-    }
 
-    drm_i915_gem_memory_class_instance getMemoryRegionClassAndInstance(uint32_t memoryBank) {
+    void assignRegionsFromDistances(const void *distanceInfosPtr, size_t size);
-        auto index = (memoryBank > 0) ? Math::log2(memoryBank) + 1 : 0;
+
-        if (index < regions.size()) {
+    drm_i915_gem_memory_class_instance getMemoryRegionClassAndInstance(uint32_t memoryBank, const HardwareInfo &hwInfo) {
-            return regions[index].region;
+        auto &hwHelper = HwHelper::get(hwInfo.platform.eRenderCoreFamily);
+        if (!hwHelper.getEnableLocalMemory(hwInfo) || memoryBank == 0) {
+            return systemMemoryRegion.region;
         }
-        return {invalidMemoryRegion(), invalidMemoryRegion()};
+
+        auto index = Math::log2(memoryBank);
+
+        index = hwHelper.isBankOverrideRequired(hwInfo) ? 0 : index;
+
+        if (DebugManager.flags.OverrideDrmRegion.get() != -1) {
+            index = DebugManager.flags.OverrideDrmRegion.get();
+        }
+
+        UNRECOVERABLE_IF(index >= localMemoryRegions.size());
+
+        return localMemoryRegions[index].region;
     }
 
     size_t getMemoryRegionSize(uint32_t memoryBank) override {
-        auto index = (memoryBank > 0) ? Math::log2(memoryBank) + 1 : 0;
+        if (DebugManager.flags.PrintMemoryRegionSizes.get()) {
-        if (index < regions.size()) {
+            printRegionSizes();
-            return regions[index].probed_size;
         }
+        if (memoryBank == 0) {
+            return systemMemoryRegion.probed_size;
+        }
+
+        auto index = Math::log2(memoryBank);
+
+        if (index < localMemoryRegions.size()) {
+            return localMemoryRegions[index].probed_size;
+        }
+
         return 0;
     }
 
-    static constexpr uint16_t invalidMemoryRegion() {
+    void printRegionSizes() {
-        return static_cast<uint16_t>(-1);
+        for (auto region : drmQueryRegions) {
+            std::cout << "Memory type: " << region.region.memory_class
+                      << ", memory instance: " << region.region.memory_instance
+                      << ", region size: " << region.probed_size << std::endl;
+        }
     }
 
-    std::vector<drm_i915_memory_region_info> regions;
+    const RegionContainer &getDrmRegionInfos() const { return drmQueryRegions; }
+
+  protected:
+    const RegionContainer drmQueryRegions;
+
+    const drm_i915_memory_region_info &systemMemoryRegion;
+
+    RegionContainer localMemoryRegions;
 };
 
 } // namespace NEO