From 10d87404b695a2f57483dec90860d4b478e49ecc Mon Sep 17 00:00:00 2001
From: "Milczarek, Slawomir" <slawomir.milczarek@intel.com>
Date: Mon, 6 May 2019 16:22:51 +0200
Subject: [PATCH] AUB with kernel names in case of kernel split

Related-To: NEO-2747

Change-Id: I49d3e4716db4634da6744fe91ecfb0763f67722a
Signed-off-by: Milczarek, Slawomir <slawomir.milczarek@intel.com>
---
 runtime/command_queue/enqueue_common.h        | 14 +++++-----
 .../command_queue/enqueue_handler_tests.cpp   | 26 ++++++++++++++++++-
 .../libult/ult_command_stream_receiver.h      |  4 +--
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/runtime/command_queue/enqueue_common.h b/runtime/command_queue/enqueue_common.h
index fda2e3bd1a..e8fbad530f 100644
--- a/runtime/command_queue/enqueue_common.h
+++ b/runtime/command_queue/enqueue_common.h
@@ -183,17 +183,17 @@ void CommandQueueHw<GfxFamily>::enqueueHandler(Surface **surfacesForResidency,
 
     enqueueHandlerHook(commandType, multiDispatchInfo);
 
-    if (getCommandStreamReceiver().getType() > CommandStreamReceiverType::CSR_HW) {
-        if (!multiDispatchInfo.empty()) {
-            auto kernelName = multiDispatchInfo.peekMainKernel()->getKernelInfo().name;
-            getCommandStreamReceiver().addAubComment(kernelName.c_str());
-        }
-    }
-
     if (DebugManager.flags.AUBDumpSubCaptureMode.get()) {
         getCommandStreamReceiver().activateAubSubCapture(multiDispatchInfo);
     }
 
+    if (getCommandStreamReceiver().getType() > CommandStreamReceiverType::CSR_HW) {
+        for (auto &dispatchInfo : multiDispatchInfo) {
+            auto kernelName = dispatchInfo.getKernel()->getKernelInfo().name;
+            getCommandStreamReceiver().addAubComment(kernelName.c_str());
+        }
+    }
+
     if (DebugManager.flags.MakeEachEnqueueBlocking.get()) {
         blocking = true;
     }
diff --git a/unit_tests/command_queue/enqueue_handler_tests.cpp b/unit_tests/command_queue/enqueue_handler_tests.cpp
index a1053ea787..81e134dce8 100644
--- a/unit_tests/command_queue/enqueue_handler_tests.cpp
+++ b/unit_tests/command_queue/enqueue_handler_tests.cpp
@@ -47,7 +47,31 @@ HWTEST_F(EnqueueHandlerTest, givenEnqueueHandlerWithKernelWhenAubCsrIsActiveThen
     mockCmdQ->enqueueKernel(mockKernel.mockKernel, 1, nullptr, gws, nullptr, 0, nullptr, nullptr);
 
     EXPECT_TRUE(aubCsr->addAubCommentCalled);
-    EXPECT_STREQ("kernel_name", aubCsr->aubCommentMessage.c_str());
+
+    EXPECT_EQ(1u, aubCsr->aubCommentMessages.size());
+    EXPECT_STREQ("kernel_name", aubCsr->aubCommentMessages[0].c_str());
+}
+
+HWTEST_F(EnqueueHandlerTest, givenEnqueueHandlerWithKernelSplitWhenAubCsrIsActiveThenAddCommentWithKernelName) {
+    int32_t tag;
+    auto aubCsr = new MockCsrAub<FamilyType>(tag, *pDevice->executionEnvironment);
+    pDevice->resetCommandStreamReceiver(aubCsr);
+
+    MockKernelWithInternals kernel1(*pDevice);
+    MockKernelWithInternals kernel2(*pDevice);
+    kernel1.kernelInfo.name = "kernel_1";
+    kernel2.kernelInfo.name = "kernel_2";
+
+    auto mockCmdQ = std::unique_ptr<MockCommandQueueHw<FamilyType>>(new MockCommandQueueHw<FamilyType>(context, pDevice, 0));
+    MockMultiDispatchInfo multiDispatchInfo(std::vector<Kernel *>({kernel1.mockKernel, kernel2.mockKernel}));
+
+    mockCmdQ->template enqueueHandler<CL_COMMAND_WRITE_BUFFER>(nullptr, 0, true, multiDispatchInfo, 0, nullptr, nullptr);
+
+    EXPECT_TRUE(aubCsr->addAubCommentCalled);
+
+    EXPECT_EQ(2u, aubCsr->aubCommentMessages.size());
+    EXPECT_STREQ("kernel_1", aubCsr->aubCommentMessages[0].c_str());
+    EXPECT_STREQ("kernel_2", aubCsr->aubCommentMessages[1].c_str());
 }
 
 HWTEST_F(EnqueueHandlerTest, givenEnqueueHandlerWithEmptyDispatchInfoWhenAubCsrIsActiveThenDontAddCommentWithKernelName) {
diff --git a/unit_tests/libult/ult_command_stream_receiver.h b/unit_tests/libult/ult_command_stream_receiver.h
index 74e3a6319c..7bac58de0e 100644
--- a/unit_tests/libult/ult_command_stream_receiver.h
+++ b/unit_tests/libult/ult_command_stream_receiver.h
@@ -147,7 +147,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     }
     void addAubComment(const char *message) override {
         CommandStreamReceiverHw<GfxFamily>::addAubComment(message);
-        aubCommentMessage.assign(message);
+        aubCommentMessages.push_back(message);
         addAubCommentCalled = true;
     }
     void flushBatchedSubmissions() override {
@@ -174,7 +174,7 @@ class UltCommandStreamReceiver : public CommandStreamReceiverHw<GfxFamily>, publ
     bool recordFlusheBatchBuffer = false;
     bool activateAubSubCaptureCalled = false;
     bool addAubCommentCalled = false;
-    std::string aubCommentMessage = "";
+    std::vector<std::string> aubCommentMessages;
     bool flushBatchedSubmissionsCalled = false;
     uint32_t makeSurfacePackNonResidentCalled = false;
     bool initProgrammingFlagsCalled = false;