From 72aad33bde135635bf4f34672a504025a5ab8936 Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Mon, 17 Feb 2020 13:25:57 +0100
Subject: [PATCH] Add aub comments in AUB+HW mode

Change-Id: Id9227153ef3b43f2bb8c57be4804560ea223cbf8
Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 .../command_stream_receiver_with_aub_dump.h   |  2 ++
 .../command_stream_receiver_with_aub_dump.inl |  8 +++++
 ...nd_stream_receiver_with_aub_dump_tests.cpp | 29 +++++++++++++++++++
 3 files changed, 39 insertions(+)

diff --git a/runtime/command_stream/command_stream_receiver_with_aub_dump.h b/runtime/command_stream/command_stream_receiver_with_aub_dump.h
index dddeaeec5b..d1a445a28f 100644
--- a/runtime/command_stream/command_stream_receiver_with_aub_dump.h
+++ b/runtime/command_stream/command_stream_receiver_with_aub_dump.h
@@ -41,6 +41,8 @@ class CommandStreamReceiverWithAUBDump : public BaseCSR {
 
     size_t getPreferredTagPoolSize() const override { return 1; }
 
+    void addAubComment(const char *comment) override;
+
     std::unique_ptr<CommandStreamReceiver> aubCSR;
 };
 
diff --git a/runtime/command_stream/command_stream_receiver_with_aub_dump.inl b/runtime/command_stream/command_stream_receiver_with_aub_dump.inl
index 00ccf6bd6c..19067e87dd 100644
--- a/runtime/command_stream/command_stream_receiver_with_aub_dump.inl
+++ b/runtime/command_stream/command_stream_receiver_with_aub_dump.inl
@@ -74,4 +74,12 @@ void CommandStreamReceiverWithAUBDump<BaseCSR>::waitForTaskCountWithKmdNotifyFal
 
     BaseCSR::waitForTaskCountWithKmdNotifyFallback(taskCountToWait, flushStampToWait, useQuickKmdSleep, forcePowerSavingMode);
 }
+
+template <typename BaseCSR>
+void CommandStreamReceiverWithAUBDump<BaseCSR>::addAubComment(const char *comment) {
+    if (aubCSR) {
+        aubCSR->addAubComment(comment);
+    }
+    BaseCSR::addAubComment(comment);
+}
 } // namespace NEO
diff --git a/unit_tests/command_stream/command_stream_receiver_with_aub_dump_tests.cpp b/unit_tests/command_stream/command_stream_receiver_with_aub_dump_tests.cpp
index f617312782..5c75d78b5a 100644
--- a/unit_tests/command_stream/command_stream_receiver_with_aub_dump_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_with_aub_dump_tests.cpp
@@ -220,6 +220,35 @@ HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenCsrWithAubDumpWhenCrea
     EXPECT_EQ(std::numeric_limits<uint32_t>::max(), *csrWithAubDump.aubCSR->getTagAddress());
 }
 
+HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenAubCsrWithHwWhenAddingCommentThenAddCommentToAubManager) {
+    MockAubCenter *mockAubCenter = new MockAubCenter(*platformDevices, false, "file_name.aub", CommandStreamReceiverType::CSR_HW_WITH_AUB);
+    auto mockAubManager = new MockAubManager();
+    mockAubCenter->aubManager.reset(mockAubManager);
+
+    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->initializeMemoryManager();
+    executionEnvironment->rootDeviceEnvironments[0]->aubCenter = std::unique_ptr<MockAubCenter>(mockAubCenter);
+
+    EXPECT_FALSE(mockAubManager->addCommentCalled);
+    CommandStreamReceiverWithAUBDump<UltCommandStreamReceiver<FamilyType>> csrWithAubDump("file_name.aub", *executionEnvironment, 0);
+    csrWithAubDump.addAubComment("test");
+    EXPECT_TRUE(mockAubManager->addCommentCalled);
+}
+
+HWTEST_F(CommandStreamReceiverWithAubDumpSimpleTest, givenAubCsrWithTbxWhenAddingCommentThenDontAddCommentToAubManager) {
+    MockAubCenter *mockAubCenter = new MockAubCenter(*platformDevices, false, "file_name.aub", CommandStreamReceiverType::CSR_TBX_WITH_AUB);
+    auto mockAubManager = new MockAubManager();
+    mockAubCenter->aubManager.reset(mockAubManager);
+
+    auto executionEnvironment = platform()->peekExecutionEnvironment();
+    executionEnvironment->initializeMemoryManager();
+    executionEnvironment->rootDeviceEnvironments[0]->aubCenter = std::unique_ptr<MockAubCenter>(mockAubCenter);
+
+    CommandStreamReceiverWithAUBDump<TbxCommandStreamReceiverHw<FamilyType>> csrWithAubDump("file_name.aub", *executionEnvironment, 0);
+    csrWithAubDump.addAubComment("test");
+    EXPECT_FALSE(mockAubManager->addCommentCalled);
+}
+
 struct CommandStreamReceiverTagTests : public ::testing::Test {
     template <typename FamilyType>
     using AubWithHw = CommandStreamReceiverWithAUBDump<UltCommandStreamReceiver<FamilyType>>;