Enable task count update from wait on direct submission

Signed-off-by: Lukasz Jobczyk <[email protected]>

Author: LukaszJobczyk

opencl/test/unit_test/command_stream/command_stream_receiver_tests.cpp

@@ -398,6 +398,11 @@ HWTEST_F(CommandStreamReceiverTest, whenClearColorAllocationIsCreatedThenItIsDes
     EXPECT_EQ(nullptr, csr.clearColorAllocation);
 }
 
+HWTEST_F(CommandStreamReceiverTest, givenNoDirectSubmissionWhenCheckTaskCountFromWaitEnabledThenReturnsFalse) {
+    auto &csr = pDevice->getUltCommandStreamReceiver<FamilyType>();
+    EXPECT_FALSE(csr.isUpdateTagFromWaitEnabled());
+}
+
 struct InitDirectSubmissionFixture {
     void SetUp() {
         DebugManager.flags.EnableDirectSubmission.set(1);
@@ -491,6 +496,44 @@ HWTEST_F(InitDirectSubmissionTest, whenDirectSubmissionEnabledOnRcsThenExpectFea
     csr.reset();
 }
 
+HWTEST_F(InitDirectSubmissionTest, givenDirectSubmissionWhenCheckTaskCountFromWaitEnabledThenReturnsTrue) {
+    auto csr = std::make_unique<CommandStreamReceiverHw<FamilyType>>(*device->executionEnvironment, device->getRootDeviceIndex(), device->getDeviceBitfield());
+    std::unique_ptr<OsContext> osContext(OsContext::create(device->getExecutionEnvironment()->rootDeviceEnvironments[0]->osInterface.get(), 0,
+                                                           EngineDescriptorHelper::getDefaultDescriptor({aub_stream::ENGINE_RCS, EngineUsage::Regular},
+                                                                                                        PreemptionMode::ThreadGroup, device->getDeviceBitfield())));
+    osContext->ensureContextInitialized();
+    osContext->setDefaultContext(true);
+    auto hwInfo = device->getRootDeviceEnvironment().getMutableHardwareInfo();
+    hwInfo->capabilityTable.directSubmissionEngines.data[aub_stream::ENGINE_RCS].engineSupported = true;
+    hwInfo->capabilityTable.directSubmissionEngines.data[aub_stream::ENGINE_RCS].submitOnInit = false;
+
+    bool ret = csr->initDirectSubmission(*device, *osContext.get());
+    EXPECT_TRUE(ret);
+
+    EXPECT_TRUE(csr->isUpdateTagFromWaitEnabled());
+
+    csr.reset();
+}
+
+HWTEST_F(InitDirectSubmissionTest, givenBlitterDirectSubmissionWhenCheckTaskCountFromWaitEnabledThenReturnsTrue) {
+    auto csr = std::make_unique<CommandStreamReceiverHw<FamilyType>>(*device->executionEnvironment, device->getRootDeviceIndex(), device->getDeviceBitfield());
+    std::unique_ptr<OsContext> osContext(OsContext::create(device->getExecutionEnvironment()->rootDeviceEnvironments[0]->osInterface.get(), 0,
+                                                           EngineDescriptorHelper::getDefaultDescriptor({aub_stream::ENGINE_BCS, EngineUsage::Regular},
+                                                                                                        PreemptionMode::ThreadGroup, device->getDeviceBitfield())));
+    osContext->ensureContextInitialized();
+    osContext->setDefaultContext(true);
+    auto hwInfo = device->getRootDeviceEnvironment().getMutableHardwareInfo();
+    hwInfo->capabilityTable.directSubmissionEngines.data[aub_stream::ENGINE_BCS].engineSupported = true;
+    hwInfo->capabilityTable.directSubmissionEngines.data[aub_stream::ENGINE_BCS].submitOnInit = false;
+
+    bool ret = csr->initDirectSubmission(*device, *osContext.get());
+    EXPECT_TRUE(ret);
+
+    EXPECT_TRUE(csr->isUpdateTagFromWaitEnabled());
+
+    csr.reset();
+}
+
 template <class Type>
 class CommandStreamReceiverHwDirectSubmissionMock : public CommandStreamReceiverHw<Type> {
   public:

shared/source/command_stream/command_stream_receiver_hw_base.inl

@@ -1134,7 +1134,7 @@ uint32_t CommandStreamReceiverHw<GfxFamily>::blitBuffer(const BlitPropertiesCont
     flush(batchBuffer, getResidencyAllocations());
     makeSurfacePackNonResident(getResidencyAllocations());
 
-    if (!isUpdateTagFromWaitEnabled()) {
+    if (updateTag) {
         latestFlushedTaskCount = newTaskCount;
     }
 
@@ -1341,7 +1341,7 @@ inline void CommandStreamReceiverHw<GfxFamily>::flushHandler(BatchBuffer &batchB
 
 template <typename GfxFamily>
 inline bool CommandStreamReceiverHw<GfxFamily>::isUpdateTagFromWaitEnabled() {
-    bool enabled = false;
+    bool enabled = isDirectSubmissionEnabled() || isBlitterDirectSubmissionEnabled();
 
     if (DebugManager.flags.UpdateTaskCountFromWait.get() != -1) {
         enabled = DebugManager.flags.UpdateTaskCountFromWait.get();