Add support for AUB subcapture in TBX mode

Resolves: NEO-3051

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

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;

unit_tests/command_stream/tbx_command_stream_tests.cpp

@@ -9,6 +9,7 @@
 #include "core/unit_tests/helpers/debug_manager_state_restore.h"
 #include "runtime/command_stream/aub_command_stream_receiver.h"
 #include "runtime/command_stream/command_stream_receiver_hw.h"
+#include "runtime/command_stream/command_stream_receiver_with_aub_dump.h"
 #include "runtime/command_stream/tbx_command_stream_receiver_hw.h"
 #include "runtime/helpers/hardware_context_controller.h"
 #include "runtime/helpers/hw_helper.h"
@@ -25,8 +26,11 @@
 #include "unit_tests/helpers/variable_backup.h"
 #include "unit_tests/mocks/mock_aub_center.h"
 #include "unit_tests/mocks/mock_aub_manager.h"
+#include "unit_tests/mocks/mock_aub_subcapture_manager.h"
 #include "unit_tests/mocks/mock_execution_environment.h"
 #include "unit_tests/mocks/mock_graphics_allocation.h"
+#include "unit_tests/mocks/mock_kernel.h"
+#include "unit_tests/mocks/mock_mdi.h"
 #include "unit_tests/mocks/mock_os_context.h"
 #include "unit_tests/mocks/mock_tbx_csr.h"
 
@@ -69,6 +73,16 @@ struct TbxFixture : public TbxCommandStreamFixture,
 using TbxCommandStreamTests = Test<TbxFixture>;
 using TbxCommandSteamSimpleTest = TbxCommandStreamTests;
 
+template <typename GfxFamily>
+struct MockTbxCsrToTestDumpTbxNonWritable : public TbxCommandStreamReceiverHw<GfxFamily> {
+    using TbxCommandStreamReceiverHw<GfxFamily>::TbxCommandStreamReceiverHw;
+    using TbxCommandStreamReceiverHw<GfxFamily>::dumpTbxNonWritable;
+
+    bool writeMemory(GraphicsAllocation &gfxAllocation) override {
+        return true;
+    }
+};
+
 TEST_F(TbxCommandStreamTests, DISABLED_testFactory) {
 }
 
@@ -589,6 +603,25 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenCreatedWithAubDumpThenOpenIsCalle
     EXPECT_TRUE(tbxCsrWithAubDump->aubManager->isOpen());
 }
 
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenCreatedWithAubDumpInSubCaptureModeThenCreateSubCaptureManagerAndGenerateSubCaptureFileName) {
+    DebugManagerStateRestore dbgRestore;
+    DebugManager.flags.AUBDumpSubCaptureMode.set(static_cast<int32_t>(AubSubCaptureManager::SubCaptureMode::Filter));
+
+    MockExecutionEnvironment executionEnvironment;
+    executionEnvironment.setHwInfo(*platformDevices);
+    executionEnvironment.initializeMemoryManager();
+
+    std::unique_ptr<TbxCommandStreamReceiverHw<FamilyType>> tbxCsrWithAubDump(static_cast<TbxCommandStreamReceiverHw<FamilyType> *>(
+        TbxCommandStreamReceiver::create("aubfile", true, executionEnvironment)));
+    EXPECT_TRUE(tbxCsrWithAubDump->aubManager->isOpen());
+
+    auto subCaptureManager = tbxCsrWithAubDump->subCaptureManager.get();
+    EXPECT_NE(nullptr, subCaptureManager);
+
+    MultiDispatchInfo dispatchInfo;
+    EXPECT_STREQ(subCaptureManager->getSubCaptureFileName(dispatchInfo).c_str(), tbxCsrWithAubDump->aubManager->getFileName().c_str());
+}
+
 HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenCreatedWithAubDumpSeveralTimesThenOpenIsCalledOnAubManagerOnceOnly) {
     MockExecutionEnvironment executionEnvironment(*platformDevices, true);
     executionEnvironment.setHwInfo(*platformDevices);
@@ -603,3 +636,199 @@ HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenCreatedWithAubDumpSeveralTimesThe
     auto mockManager = reinterpret_cast<MockAubManager *>(executionEnvironment.aubCenter->getAubManager());
     EXPECT_EQ(1u, mockManager->openCalledCnt);
 }
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenFlushIsCalledAndSubCaptureIsDisabledThenPauseShouldBeTurnedOn) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+    EXPECT_FALSE(tbxCsr.subCaptureManager->isSubCaptureEnabled());
+
+    auto commandBuffer = pDevice->executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize});
+    LinearStream cs(commandBuffer);
+    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, QueueSliceCount::defaultSliceCount, cs.getUsed(), &cs};
+    ResidencyContainer allocationsForResidency = {};
+
+    tbxCsr.flush(batchBuffer, allocationsForResidency);
+
+    auto mockAubManager = reinterpret_cast<MockAubManager *>(pDevice->executionEnvironment->aubCenter->getAubManager());
+    EXPECT_TRUE(mockAubManager->isPaused);
+
+    pDevice->executionEnvironment->memoryManager->freeGraphicsMemory(commandBuffer);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenFlushIsCalledAndSubCaptureIsEnabledThenPauseShouldBeTurnedOff) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    aubSubCaptureManagerMock->setSubCaptureIsActive(true);
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+    EXPECT_TRUE(tbxCsr.subCaptureManager->isSubCaptureEnabled());
+
+    auto commandBuffer = pDevice->executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize});
+    LinearStream cs(commandBuffer);
+    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, QueueSliceCount::defaultSliceCount, cs.getUsed(), &cs};
+    ResidencyContainer allocationsForResidency = {};
+
+    tbxCsr.flush(batchBuffer, allocationsForResidency);
+
+    auto mockAubManager = reinterpret_cast<MockAubManager *>(pDevice->executionEnvironment->aubCenter->getAubManager());
+    EXPECT_FALSE(mockAubManager->isPaused);
+
+    pDevice->executionEnvironment->memoryManager->freeGraphicsMemory(commandBuffer);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenFlushIsCalledAndSubCaptureIsEnabledThenCallPollForCompletionAndDisableSubCapture) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    aubSubCaptureManagerMock->setSubCaptureIsActive(true);
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+    EXPECT_TRUE(tbxCsr.subCaptureManager->isSubCaptureEnabled());
+
+    auto commandBuffer = pDevice->executionEnvironment->memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize});
+    LinearStream cs(commandBuffer);
+    BatchBuffer batchBuffer{cs.getGraphicsAllocation(), 0, 0, nullptr, false, false, QueueThrottle::MEDIUM, QueueSliceCount::defaultSliceCount, cs.getUsed(), &cs};
+
+    ResidencyContainer allocationsForResidency = {};
+
+    tbxCsr.flush(batchBuffer, allocationsForResidency);
+
+    EXPECT_TRUE(tbxCsr.pollForCompletionCalled);
+    EXPECT_FALSE(tbxCsr.subCaptureManager->isSubCaptureEnabled());
+
+    pDevice->executionEnvironment->memoryManager->freeGraphicsMemory(commandBuffer);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenProcessResidencyIsCalledWithDumpTbxNonWritableFlagThenAllocationsForResidencyShouldBeMadeTbxWritable) {
+    std::unique_ptr<MemoryManager> memoryManager(nullptr);
+    std::unique_ptr<MockTbxCsrToTestDumpTbxNonWritable<FamilyType>> tbxCsr(new MockTbxCsrToTestDumpTbxNonWritable<FamilyType>(*pDevice->executionEnvironment));
+    memoryManager.reset(new OsAgnosticMemoryManager(*pDevice->executionEnvironment));
+    tbxCsr->setupContext(*pDevice->getDefaultEngine().osContext);
+
+    auto gfxAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize, GraphicsAllocation::AllocationType::BUFFER});
+    tbxCsr->setTbxWritable(false, *gfxAllocation);
+
+    tbxCsr->dumpTbxNonWritable = true;
+
+    ResidencyContainer allocationsForResidency = {gfxAllocation};
+    tbxCsr->processResidency(allocationsForResidency);
+
+    EXPECT_TRUE(tbxCsr->isTbxWritable(*gfxAllocation));
+    EXPECT_FALSE(tbxCsr->dumpTbxNonWritable);
+
+    memoryManager->freeGraphicsMemory(gfxAllocation);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrWhenProcessResidencyIsCalledWithoutDumpTbxWritableFlagThenAllocationsForResidencyShouldBeKeptNonTbxWritable) {
+    std::unique_ptr<MemoryManager> memoryManager(nullptr);
+    std::unique_ptr<MockTbxCsrToTestDumpTbxNonWritable<FamilyType>> tbxCsr(new MockTbxCsrToTestDumpTbxNonWritable<FamilyType>(*pDevice->executionEnvironment));
+    memoryManager.reset(new OsAgnosticMemoryManager(*pDevice->executionEnvironment));
+    tbxCsr->setupContext(*pDevice->getDefaultEngine().osContext);
+
+    auto gfxAllocation = memoryManager->allocateGraphicsMemoryWithProperties(MockAllocationProperties{MemoryConstants::pageSize, GraphicsAllocation::AllocationType::BUFFER});
+    tbxCsr->setTbxWritable(false, *gfxAllocation);
+
+    EXPECT_FALSE(tbxCsr->dumpTbxNonWritable);
+
+    ResidencyContainer allocationsForResidency = {gfxAllocation};
+    tbxCsr->processResidency(allocationsForResidency);
+
+    EXPECT_FALSE(tbxCsr->isTbxWritable(*gfxAllocation));
+    EXPECT_FALSE(tbxCsr->dumpTbxNonWritable);
+
+    memoryManager->freeGraphicsMemory(gfxAllocation);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenCheckAndActivateAubSubCaptureIsCalledAndSubCaptureIsInactiveThenDontForceDumpingAllocationsTbxNonWritable) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    aubSubCaptureCommon.subCaptureMode = AubSubCaptureManager::SubCaptureMode::Toggle;
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+
+    MockKernelWithInternals kernelInternals(*pDevice);
+    Kernel *kernel = kernelInternals.mockKernel;
+    MockMultiDispatchInfo multiDispatchInfo(kernel);
+
+    EXPECT_FALSE(tbxCsr.dumpTbxNonWritable);
+
+    auto status = tbxCsr.checkAndActivateAubSubCapture(multiDispatchInfo);
+    EXPECT_FALSE(status.isActive);
+    EXPECT_FALSE(status.wasActiveInPreviousEnqueue);
+
+    EXPECT_FALSE(tbxCsr.dumpTbxNonWritable);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenCheckAndActivateAubSubCaptureIsCalledAndSubCaptureGetsActivatedThenForceDumpingAllocationsTbxNonWritable) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    aubSubCaptureCommon.subCaptureMode = AubSubCaptureManager::SubCaptureMode::Toggle;
+    aubSubCaptureManagerMock->setSubCaptureIsActive(false);
+    aubSubCaptureManagerMock->setSubCaptureToggleActive(true);
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+
+    MockKernelWithInternals kernelInternals(*pDevice);
+    Kernel *kernel = kernelInternals.mockKernel;
+    MockMultiDispatchInfo multiDispatchInfo(kernel);
+
+    EXPECT_FALSE(tbxCsr.dumpTbxNonWritable);
+
+    auto status = tbxCsr.checkAndActivateAubSubCapture(multiDispatchInfo);
+    EXPECT_TRUE(status.isActive);
+    EXPECT_FALSE(status.wasActiveInPreviousEnqueue);
+
+    EXPECT_TRUE(tbxCsr.dumpTbxNonWritable);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInSubCaptureModeWhenCheckAndActivateAubSubCaptureIsCalledAndSubCaptureRemainsActivatedThenDontForceDumpingAllocationsTbxNonWritable) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    AubSubCaptureCommon aubSubCaptureCommon;
+    auto aubSubCaptureManagerMock = new AubSubCaptureManagerMock("", aubSubCaptureCommon);
+    aubSubCaptureCommon.subCaptureMode = AubSubCaptureManager::SubCaptureMode::Toggle;
+    aubSubCaptureManagerMock->setSubCaptureIsActive(true);
+    aubSubCaptureManagerMock->setSubCaptureToggleActive(true);
+    tbxCsr.subCaptureManager = std::unique_ptr<AubSubCaptureManagerMock>(aubSubCaptureManagerMock);
+
+    MockKernelWithInternals kernelInternals(*pDevice);
+    Kernel *kernel = kernelInternals.mockKernel;
+    MockMultiDispatchInfo multiDispatchInfo(kernel);
+
+    EXPECT_FALSE(tbxCsr.dumpTbxNonWritable);
+
+    auto status = tbxCsr.checkAndActivateAubSubCapture(multiDispatchInfo);
+    EXPECT_TRUE(status.isActive);
+    EXPECT_TRUE(status.wasActiveInPreviousEnqueue);
+
+    EXPECT_FALSE(tbxCsr.dumpTbxNonWritable);
+}
+
+HWTEST_F(TbxCommandStreamTests, givenTbxCsrInNonSubCaptureModeWhenCheckAndActivateAubSubCaptureIsCalledThenReturnStatusInactive) {
+    MockTbxCsr<FamilyType> tbxCsr{*pDevice->executionEnvironment};
+    MockOsContext osContext(0, 1, aub_stream::ENGINE_RCS, PreemptionMode::Disabled, false);
+    tbxCsr.setupContext(osContext);
+
+    MultiDispatchInfo dispatchInfo;
+    auto status = tbxCsr.checkAndActivateAubSubCapture(dispatchInfo);
+    EXPECT_FALSE(status.isActive);
+    EXPECT_FALSE(status.wasActiveInPreviousEnqueue);
+}

unit_tests/mocks/mock_aub_manager.h

@@ -85,6 +85,9 @@ class MockAubManager : public aub_stream::AubManager {
         getFileNameCalled = true;
         return fileName;
     }
+    void pause(bool onoff) override {
+        isPaused = onoff;
+    }
 
     void addComment(const char *message) override {
         receivedComment.assign(message);
@@ -106,6 +109,7 @@ class MockAubManager : public aub_stream::AubManager {
     bool closeCalled = false;
     bool isOpenCalled = false;
     bool getFileNameCalled = false;
+    bool isPaused = false;
     bool addCommentCalled = false;
     std::string receivedComment = "";
     bool writeMemoryCalled = false;