diff --git a/shared/source/os_interface/linux/memory_info.cpp b/shared/source/os_interface/linux/memory_info.cpp
index 25a1f2890c..92c9e79cc2 100644
--- a/shared/source/os_interface/linux/memory_info.cpp
+++ b/shared/source/os_interface/linux/memory_info.cpp
@@ -76,32 +76,29 @@ uint32_t MemoryInfo::getTileIndex(uint32_t memoryBank) {
 MemoryClassInstance MemoryInfo::getMemoryRegionClassAndInstance(uint32_t memoryBank, const HardwareInfo &hwInfo) {
 
     auto &gfxCoreHelper = this->drm.getRootDeviceEnvironment().getHelper<GfxCoreHelper>();
-    if (!gfxCoreHelper.getEnableLocalMemory(hwInfo) || memoryBank == 0) {
-        return systemMemoryRegion.region;
+    if (!gfxCoreHelper.getEnableLocalMemory(hwInfo)) {
+        memoryBank = 0;
+    }
+
+    return getMemoryRegion(memoryBank).region;
+}
+
+const MemoryRegion &MemoryInfo::getMemoryRegion(uint32_t memoryBank) {
+    if (memoryBank == 0) {
+        return systemMemoryRegion;
     }
 
     auto index = getTileIndex(memoryBank);
 
     UNRECOVERABLE_IF(index >= localMemoryRegions.size());
-
-    return localMemoryRegions[index].region;
+    return localMemoryRegions[index];
 }
 
 size_t MemoryInfo::getMemoryRegionSize(uint32_t memoryBank) {
     if (DebugManager.flags.PrintMemoryRegionSizes.get()) {
         printRegionSizes();
     }
-    if (memoryBank == 0) {
-        return systemMemoryRegion.probedSize;
-    }
-
-    auto index = Math::log2(memoryBank);
-
-    if (index < localMemoryRegions.size()) {
-        return localMemoryRegions[index].probedSize;
-    }
-
-    return 0;
+    return getMemoryRegion(memoryBank).probedSize;
 }
 
 void MemoryInfo::printRegionSizes() {
diff --git a/shared/source/os_interface/linux/memory_info.h b/shared/source/os_interface/linux/memory_info.h
index f582319bcd..61a33448f8 100644
--- a/shared/source/os_interface/linux/memory_info.h
+++ b/shared/source/os_interface/linux/memory_info.h
@@ -32,6 +32,8 @@ class MemoryInfo {
 
     MOCKABLE_VIRTUAL size_t getMemoryRegionSize(uint32_t memoryBank);
 
+    const MemoryRegion &getMemoryRegion(uint32_t memoryBank);
+
     void printRegionSizes();
 
     uint32_t getTileIndex(uint32_t memoryBank);
diff --git a/shared/test/common/os_interface/linux/drm_memory_manager_prelim_fixtures.h b/shared/test/common/os_interface/linux/drm_memory_manager_prelim_fixtures.h
index e527a3e4a1..abdbe6107a 100644
--- a/shared/test/common/os_interface/linux/drm_memory_manager_prelim_fixtures.h
+++ b/shared/test/common/os_interface/linux/drm_memory_manager_prelim_fixtures.h
@@ -30,9 +30,9 @@ class DrmMemoryManagerWithSubDevicesPrelimTest : public ::testing::Test {
 
         mock = new DrmQueryMock(*executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]);
         auto memoryInfo = new MockExtendedMemoryInfo(*mock);
-        mock->memoryInfo.reset(memoryInfo);
 
         mock->queryEngineInfo();
+        mock->memoryInfo.reset(memoryInfo);
 
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface = std::make_unique<OSInterface>();
         executionEnvironment->rootDeviceEnvironments[rootDeviceIndex]->osInterface->setDriverModel(std::unique_ptr<DriverModel>(mock));
diff --git a/shared/test/unit_test/os_interface/linux/drm_memory_info_prelim_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_memory_info_prelim_tests.cpp
index f9cdf0ce63..d28e9e4870 100644
--- a/shared/test/unit_test/os_interface/linux/drm_memory_info_prelim_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_memory_info_prelim_tests.cpp
@@ -294,7 +294,7 @@ TEST_F(MultiTileMemoryInfoPrelimTest, givenMemoryInfoWithRegionsWhenGettingMemor
     EXPECT_EQ(regionInfo[2].region.memoryClass, regionClassAndInstance.memoryClass);
     EXPECT_EQ(regionInfo[2].region.memoryInstance, regionClassAndInstance.memoryInstance);
     regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
-    EXPECT_EQ(16 * GB, regionSize);
+    EXPECT_EQ(regionInfo[2].probedSize, regionSize);
 
     // route to tile 0 banks
     DebugManager.flags.OverrideDrmRegion.set(0);
@@ -303,7 +303,7 @@ TEST_F(MultiTileMemoryInfoPrelimTest, givenMemoryInfoWithRegionsWhenGettingMemor
     EXPECT_EQ(regionInfo[1].region.memoryClass, regionClassAndInstance.memoryClass);
     EXPECT_EQ(regionInfo[1].region.memoryInstance, regionClassAndInstance.memoryInstance);
     regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
-    EXPECT_EQ(32 * GB, regionSize);
+    EXPECT_EQ(regionInfo[1].probedSize, regionSize);
 
     // try to force tile 0 banks
     DebugManager.flags.OverrideDrmRegion.set(-1);
@@ -336,9 +336,7 @@ TEST_F(MultiTileMemoryInfoPrelimTest, givenMemoryInfoWithoutRegionsWhenGettingMe
     setupMemoryInfo(regionInfo, 0);
 
     EXPECT_ANY_THROW(memoryInfo->getMemoryRegionClassAndInstance(1, *pHwInfo));
-
-    auto regionSize = memoryInfo->getMemoryRegionSize(1);
-    EXPECT_EQ(0 * GB, regionSize);
+    EXPECT_ANY_THROW(memoryInfo->getMemoryRegionSize(1));
 }
 
 TEST_F(MultiTileMemoryInfoPrelimTest, whenDebugVariablePrintMemoryRegionSizeIsSetAndGetMemoryRegionSizeIsCalledThenMessagePrintedToStdOutput) {
diff --git a/shared/test/unit_test/os_interface/linux/drm_memory_info_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_memory_info_tests.cpp
index a08a2c0b36..26f963270f 100644
--- a/shared/test/unit_test/os_interface/linux/drm_memory_info_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_memory_info_tests.cpp
@@ -115,8 +115,7 @@ TEST(MemoryInfo, givenMemoryInfoWithoutDeviceRegionWhenGettingDeviceRegionSizeTh
     auto drm = std::make_unique<DrmTipMock>(*executionEnvironment->rootDeviceEnvironments[0]);
     auto memoryInfo = std::make_unique<MemoryInfo>(regionInfo, *drm);
     ASSERT_NE(nullptr, memoryInfo);
-    auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
-    EXPECT_EQ(0 * GB, regionSize);
+    EXPECT_ANY_THROW(memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0)));
 }
 
 TEST(MemoryInfo, givenMemoryInfoWithRegionsAndLocalMemoryDisabledWhenGettingMemoryRegionClassAndInstanceThenReturnCorrectValues) {
@@ -190,7 +189,7 @@ TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstan
     EXPECT_EQ(regionInfo[2].region.memoryClass, regionClassAndInstance.memoryClass);
     EXPECT_EQ(regionInfo[2].region.memoryInstance, regionClassAndInstance.memoryInstance);
     auto regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(0));
-    EXPECT_EQ(16 * GB, regionSize);
+    EXPECT_EQ(regionInfo[2].probedSize, regionSize);
 
     DebugManager.flags.OverrideDrmRegion.set(0);
 
@@ -198,7 +197,7 @@ TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstan
     EXPECT_EQ(regionInfo[1].region.memoryClass, regionClassAndInstance.memoryClass);
     EXPECT_EQ(regionInfo[1].region.memoryInstance, regionClassAndInstance.memoryInstance);
     regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
-    EXPECT_EQ(32 * GB, regionSize);
+    EXPECT_EQ(regionInfo[1].probedSize, regionSize);
 
     DebugManager.flags.OverrideDrmRegion.set(-1);
     DebugManager.flags.ForceMemoryBankIndexOverride.set(1);
@@ -207,15 +206,16 @@ TEST(MemoryInfo, givenMemoryInfoWithRegionsWhenGettingMemoryRegionClassAndInstan
     auto &productHelper = executionEnvironment->rootDeviceEnvironments[0]->getHelper<ProductHelper>();
 
     regionClassAndInstance = memoryInfo->getMemoryRegionClassAndInstance(MemoryBanks::getBankForLocalMemory(1), *defaultHwInfo);
+    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
     if (gfxCoreHelper.isBankOverrideRequired(*defaultHwInfo, productHelper)) {
         EXPECT_EQ(regionInfo[1].region.memoryClass, regionClassAndInstance.memoryClass);
         EXPECT_EQ(regionInfo[1].region.memoryInstance, regionClassAndInstance.memoryInstance);
+        EXPECT_EQ(regionInfo[1].probedSize, regionSize);
     } else {
         EXPECT_EQ(regionInfo[2].region.memoryClass, regionClassAndInstance.memoryClass);
         EXPECT_EQ(regionInfo[2].region.memoryInstance, regionClassAndInstance.memoryInstance);
+        EXPECT_EQ(regionInfo[2].probedSize, regionSize);
     }
-    regionSize = memoryInfo->getMemoryRegionSize(MemoryBanks::getBankForLocalMemory(1));
-    EXPECT_EQ(32 * GB, regionSize);
 }
 
 using MemoryInfoTest = ::testing::Test;
diff --git a/shared/test/unit_test/os_interface/linux/drm_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_tests.cpp
index 93b2fe0d60..832eddf85c 100644
--- a/shared/test/unit_test/os_interface/linux/drm_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_tests.cpp
@@ -1574,5 +1574,5 @@ TEST(DistanceInfoTest, givenDistanceInfosWhenAssignRegionsFromDistancesThenCorre
     memoryInfo->assignRegionsFromDistances(distances);
     EXPECT_EQ(1024u, memoryInfo->getMemoryRegionSize(1));
     EXPECT_EQ(1024u, memoryInfo->getMemoryRegionSize(2));
-    EXPECT_EQ(0u, memoryInfo->getMemoryRegionSize(4));
+    EXPECT_ANY_THROW(memoryInfo->getMemoryRegionSize(4));
 }
diff --git a/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp b/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp
index b6322f0e35..a33dd8f9e1 100644
--- a/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp
+++ b/shared/test/unit_test/os_interface/linux/drm_with_prelim_tests.cpp
@@ -372,7 +372,7 @@ TEST_F(IoctlHelperPrelimFixture, givenPrelimWhenQueryEngineInfoWithDeviceMemoryT
 
     EXPECT_EQ(1024u, drm->memoryInfo->getMemoryRegionSize(1));
     EXPECT_EQ(1024u, drm->memoryInfo->getMemoryRegionSize(2));
-    EXPECT_EQ(0u, drm->memoryInfo->getMemoryRegionSize(4));
+    EXPECT_ANY_THROW(drm->memoryInfo->getMemoryRegionSize(4));
 
     std::vector<EngineClassInstance> engines;
     engineInfo->getListOfEnginesOnATile(0u, engines);