From 12245bc88d60ca3f8ca01e0e4b86601f073df046 Mon Sep 17 00:00:00 2001
From: Maciej Dziuban <maciej.dziuban@intel.com>
Date: Tue, 12 Feb 2019 13:37:39 +0100
Subject: [PATCH] Pass hwInfo to AubHelper::getMemBankSize()

Change-Id: If77775cb5fb10dc82e0c7bef06a71e4292ceb6f9
Signed-off-by: Maciej Dziuban <maciej.dziuban@intel.com>
---
 runtime/aub/aub_center.cpp          | 2 +-
 runtime/aub/aub_helper.cpp          | 8 ++++++--
 runtime/aub/aub_helper.h            | 4 +++-
 unit_tests/aub/aub_helper_tests.cpp | 4 ++++
 unit_tests/aub/aub_helper_tests.inl | 5 -----
 5 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/runtime/aub/aub_center.cpp b/runtime/aub/aub_center.cpp
index af04f3212b..5d318b4127 100644
--- a/runtime/aub/aub_center.cpp
+++ b/runtime/aub/aub_center.cpp
@@ -21,7 +21,7 @@ extern aub_stream::AubManager *createAubManager(uint32_t productFamily, uint32_t
 AubCenter::AubCenter(const HardwareInfo *pHwInfo, bool localMemoryEnabled, const std::string &aubFileName, CommandStreamReceiverType csrType) {
     if (DebugManager.flags.UseAubStream.get()) {
         auto devicesCount = AubHelper::getDevicesCount(pHwInfo);
-        auto memoryBankSize = AubHelper::getMemBankSize();
+        auto memoryBankSize = AubHelper::getMemBankSize(pHwInfo);
         CommandStreamReceiverType type = static_cast<CommandStreamReceiverType>(DebugManager.flags.SetCommandStreamReceiver.get() != CommandStreamReceiverType::CSR_HW
                                                                                     ? DebugManager.flags.SetCommandStreamReceiver.get()
                                                                                     : csrType);
diff --git a/runtime/aub/aub_helper.cpp b/runtime/aub/aub_helper.cpp
index a3b572919f..dc6894d681 100644
--- a/runtime/aub/aub_helper.cpp
+++ b/runtime/aub/aub_helper.cpp
@@ -12,6 +12,10 @@
 
 namespace OCLRT {
 
+uint64_t AubHelper::getTotalMemBankSize() {
+    return 2 * GB;
+}
+
 int AubHelper::getMemTrace(uint64_t pdEntryBits) {
     return AubMemDump::AddressSpaceValues::TraceNonlocal;
 }
@@ -27,8 +31,8 @@ uint32_t AubHelper::getMemType(uint32_t addressSpace) {
     return 0;
 }
 
-uint64_t AubHelper::getMemBankSize() {
-    return 2 * GB;
+uint64_t AubHelper::getMemBankSize(const HardwareInfo *pHwInfo) {
+    return getTotalMemBankSize();
 }
 
 uint32_t AubHelper::getDevicesCount(const HardwareInfo *pHwInfo) {
diff --git a/runtime/aub/aub_helper.h b/runtime/aub/aub_helper.h
index a109f5f70a..76b81a203d 100644
--- a/runtime/aub/aub_helper.h
+++ b/runtime/aub/aub_helper.h
@@ -27,11 +27,13 @@ class AubHelper : public NonCopyableOrMovableClass {
             return false;
         }
     }
+
+    static uint64_t getTotalMemBankSize();
     static int getMemTrace(uint64_t pdEntryBits);
     static uint64_t getPTEntryBits(uint64_t pdEntryBits);
     static void checkPTEAddress(uint64_t address);
     static uint32_t getMemType(uint32_t addressSpace);
-    static uint64_t getMemBankSize();
+    static uint64_t getMemBankSize(const HardwareInfo *pHwInfo);
     static uint32_t getDevicesCount(const HardwareInfo *pHwInfo);
     static MMIOList getAdditionalMmioList();
 
diff --git a/unit_tests/aub/aub_helper_tests.cpp b/unit_tests/aub/aub_helper_tests.cpp
index 6375630dcb..8df39035df 100644
--- a/unit_tests/aub/aub_helper_tests.cpp
+++ b/unit_tests/aub/aub_helper_tests.cpp
@@ -6,3 +6,7 @@
  */
 
 #include "unit_tests/aub/aub_helper_tests.inl"
+
+TEST(AubHelper, GivenHwInfoWhenGetMemBankSizeIsCalledThenItReturnsCorrectValue) {
+    EXPECT_EQ(2 * MemoryConstants::gigaByte, AubHelper::getMemBankSize(platformDevices[0]));
+}
diff --git a/unit_tests/aub/aub_helper_tests.inl b/unit_tests/aub/aub_helper_tests.inl
index c6ea57ed88..ff4ebf8e8a 100644
--- a/unit_tests/aub/aub_helper_tests.inl
+++ b/unit_tests/aub/aub_helper_tests.inl
@@ -49,11 +49,6 @@ TEST(AubHelper, WhenCreateMultipleDevicesIsSetThenGetDevicesCountReturnedCorrect
     EXPECT_EQ(devicesCount, 1u);
 }
 
-TEST(AubHelper, WhenGetMemBankSizeIsCalledThenItReturnsCorrectValue) {
-    auto memBankSize = AubHelper::getMemBankSize();
-    EXPECT_EQ(memBankSize, 2 * GB);
-}
-
 typedef Test<DeviceFixture> AubHelperHwTest;
 
 HWTEST_F(AubHelperHwTest, GivenDisabledLocalMemoryWhenGetDataHintForPml4EntryIsCalledThenTraceNotypeIsReturned) {