Add support for AUB subcapture in TBX mode

Resolves: NEO-3051

Change-Id: If7c2d550227d0c5d09c35a53387f34f968e7c7b7
Signed-off-by: Milczarek, Slawomir <slawomir.milczarek@intel.com>

Author: smilczar

runtime/command_stream/tbx_command_stream_receiver_hw.h

@@ -18,6 +18,7 @@
 
 namespace NEO {
 
+class AubSubCaptureManager;
 class TbxStream;
 
 class TbxMemoryManager : public OsAgnosticMemoryManager {
@@ -57,6 +58,8 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     void writeMemory(uint64_t gpuAddress, void *cpuAddress, size_t size, uint32_t memoryBank, uint64_t entryBits) override;
     bool writeMemory(GraphicsAllocation &gfxAllocation) override;
 
+    AubSubCaptureStatus checkAndActivateAubSubCapture(const MultiDispatchInfo &dispatchInfo) override;
+
     // Family specific version
     MOCKABLE_VIRTUAL void submitBatchBuffer(uint64_t batchBufferGpuAddress, const void *batchBuffer, size_t batchBufferSize, uint32_t memoryBank, uint64_t entryBits);
     void pollForCompletion() override;
@@ -73,7 +76,7 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     }
 
     TbxStream tbxStream;
-
+    std::unique_ptr<AubSubCaptureManager> subCaptureManager;
     uint32_t aubDeviceId;
     bool streamInitialized = false;
 
@@ -88,5 +91,7 @@ class TbxCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
     CommandStreamReceiverType getType() override {
         return CommandStreamReceiverType::CSR_TBX;
     }
+
+    bool dumpTbxNonWritable = false;
 };
 } // namespace NEO

runtime/command_stream/tbx_command_stream_receiver_hw.inl

@@ -16,6 +16,7 @@
 #include "runtime/command_stream/aub_command_stream_receiver.h"
 #include "runtime/command_stream/command_stream_receiver_with_aub_dump.h"
 #include "runtime/execution_environment/execution_environment.h"
+#include "runtime/helpers/dispatch_info.h"
 #include "runtime/helpers/hardware_context_controller.h"
 #include "runtime/helpers/hw_helper.h"
 #include "runtime/memory_manager/memory_banks.h"
@@ -156,9 +157,21 @@ CommandStreamReceiver *TbxCommandStreamReceiverHw<GfxFamily>::create(const std::
         executionEnvironment.initAubCenter(localMemoryEnabled, fullName, CommandStreamReceiverType::CSR_TBX_WITH_AUB);
 
         csr = new CommandStreamReceiverWithAUBDump<TbxCommandStreamReceiverHw<GfxFamily>>(baseName, executionEnvironment);
+
+        auto aubCenter = executionEnvironment.aubCenter.get();
+        UNRECOVERABLE_IF(nullptr == aubCenter);
+
+        auto subCaptureCommon = aubCenter->getSubCaptureCommon();
+        UNRECOVERABLE_IF(nullptr == subCaptureCommon);
+
+        if (subCaptureCommon->subCaptureMode > AubSubCaptureManager::SubCaptureMode::Off) {
+            csr->subCaptureManager = std::make_unique<AubSubCaptureManager>(fullName, *subCaptureCommon);
+        }
+
         if (csr->aubManager) {
             if (!csr->aubManager->isOpen()) {
-                csr->aubManager->open(fullName);
+                MultiDispatchInfo dispatchInfo;
+                csr->aubManager->open(csr->subCaptureManager ? csr->subCaptureManager->getSubCaptureFileName(dispatchInfo) : fullName);
                 UNRECOVERABLE_IF(!csr->aubManager->isOpen());
             }
         }
@@ -179,6 +192,12 @@ CommandStreamReceiver *TbxCommandStreamReceiverHw<GfxFamily>::create(const std::
 
 template <typename GfxFamily>
 FlushStamp TbxCommandStreamReceiverHw<GfxFamily>::flush(BatchBuffer &batchBuffer, ResidencyContainer &allocationsForResidency) {
+    if (subCaptureManager) {
+        if (aubManager) {
+            aubManager->pause(false);
+        }
+    }
+
     initializeEngine();
 
     // Write our batch buffer
@@ -196,7 +215,19 @@ FlushStamp TbxCommandStreamReceiverHw<GfxFamily>::flush(BatchBuffer &batchBuffer
     // Write allocations for residency
     processResidency(allocationsForResidency);
 
+    if (subCaptureManager && !subCaptureManager->isSubCaptureEnabled()) {
+        if (aubManager) {
+            aubManager->pause(true);
+        }
+    }
+
     submitBatchBuffer(batchBufferGpuAddress, pBatchBuffer, sizeBatchBuffer, this->getMemoryBank(batchBuffer.commandBufferAllocation), this->getPPGTTAdditionalBits(batchBuffer.commandBufferAllocation));
+
+    if (subCaptureManager) {
+        pollForCompletion();
+        subCaptureManager->disableSubCapture();
+    }
+
     return 0;
 }
 
@@ -412,12 +443,17 @@ void TbxCommandStreamReceiverHw<GfxFamily>::processEviction() {
 template <typename GfxFamily>
 void TbxCommandStreamReceiverHw<GfxFamily>::processResidency(const ResidencyContainer &allocationsForResidency) {
     for (auto &gfxAllocation : allocationsForResidency) {
+        if (dumpTbxNonWritable) {
+            this->setTbxWritable(true, *gfxAllocation);
+        }
         if (!writeMemory(*gfxAllocation)) {
             DEBUG_BREAK_IF(!((gfxAllocation->getUnderlyingBufferSize() == 0) ||
                              !this->isTbxWritable(*gfxAllocation)));
         }
         gfxAllocation->updateResidencyTaskCount(this->taskCount + 1, this->osContext->getContextId());
     }
+
+    dumpTbxNonWritable = false;
 }
 
 template <typename GfxFamily>
@@ -450,4 +486,17 @@ template <typename GfxFamily>
 bool TbxCommandStreamReceiverHw<GfxFamily>::getpollNotEqualValueForPollForCompletion() const {
     return false;
 }
+
+template <typename GfxFamily>
+AubSubCaptureStatus TbxCommandStreamReceiverHw<GfxFamily>::checkAndActivateAubSubCapture(const MultiDispatchInfo &dispatchInfo) {
+    if (!subCaptureManager) {
+        return {false, false};
+    }
+
+    auto status = subCaptureManager->checkAndActivateSubCapture(dispatchInfo);
+    if (status.isActive && !status.wasActiveInPreviousEnqueue) {
+        dumpTbxNonWritable = true;
+    }
+    return status;
+}
 } // namespace NEO

third_party/aub_stream/headers/aub_manager.h

@@ -23,6 +23,7 @@ class AubManager {
     virtual void close() = 0;
     virtual bool isOpen() = 0;
     virtual const std::string getFileName() = 0;
+    virtual void pause(bool onoff) = 0;
 
     virtual void addComment(const char *message) = 0;
     virtual void writeMemory(uint64_t gfxAddress, const void *memory, size_t size, uint32_t memoryBanks, int hint, size_t pageSize) = 0;

unit_tests/command_stream/aub_subcapture_tests.cpp

@@ -356,6 +356,15 @@ TEST_F(AubSubCaptureTest, givenSubCaptureManagerInToggleModeWhenGetSubCaptureFil
     EXPECT_STREQ(toggleFileName.c_str(), aubSubCaptureManager.getSubCaptureFileName(multiDispatchInfo).c_str());
 }
 
+TEST_F(AubSubCaptureTest, givenSubCaptureManagerInToggleModeWhenGetSubCaptureFileNameIsCalledForEmptyDispatchInfoThenGenerateToggleFileNameWithoutKernelName) {
+    AubSubCaptureManagerMock aubSubCaptureManager("aubfile.aub", subCaptureCommon);
+
+    MultiDispatchInfo dispatchInfo;
+    subCaptureCommon.subCaptureMode = AubSubCaptureManager::SubCaptureMode::Toggle;
+    auto toggleFileName = aubSubCaptureManager.generateToggleFileName(dispatchInfo);
+    EXPECT_STREQ(toggleFileName.c_str(), aubSubCaptureManager.getSubCaptureFileName(dispatchInfo).c_str());
+}
+
 TEST_F(AubSubCaptureTest, givenSubCaptureManagerInFilterModeWhenGetSubCaptureFileNameIsCalledManyTimesAndExternalFileNameIsNotSpecifiedThenItGeneratesFilterFileNameOnceOnly) {
     struct AubSubCaptureManagerMockWithFilterFileNameGenerationCount : AubSubCaptureManager {
         using AubSubCaptureManager::AubSubCaptureManager;