Stop accessing task count from out-of-bound in opencl and shared ult

Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>

opencl/test/unit_test/command_stream/command_stream_receiver_flush_task_4_tests.cpp

@@ -647,6 +647,9 @@ HWTEST_F(CommandStreamReceiverFlushTaskTests, givenStaticPartitioningEnabledWhen
     mockCsr.storeMakeResidentAllocations = true;
 
     DispatchFlags dispatchFlags = DispatchFlagsHelper::createDefaultDispatchFlags();
+
+    cleanupHeaps();
+    initHeaps();
     mockCsr.flushTask(commandStream,
                       0,
                       &dsh,

opencl/test/unit_test/fixtures/ult_command_stream_receiver_fixture.h

@@ -30,8 +30,6 @@ struct UltCommandStreamReceiverTest
         ClDeviceFixture::SetUp();
         ClHardwareParse::SetUp();
 
-        size_t sizeStream = 512;
-        size_t alignmentStream = 0x1000;
         cmdBuffer = alignedMalloc(sizeStream, alignmentStream);
         dshBuffer = alignedMalloc(sizeStream, alignmentStream);
         iohBuffer = alignedMalloc(sizeStream, alignmentStream);
@@ -42,6 +40,14 @@ struct UltCommandStreamReceiverTest
         ASSERT_NE(nullptr, iohBuffer);
         ASSERT_NE(nullptr, sshBuffer);
 
+        initHeaps();
+
+        flushTaskFlags.threadArbitrationPolicy = NEO::HwHelper::get(hardwareInfo.platform.eRenderCoreFamily).getDefaultThreadArbitrationPolicy();
+
+        pDevice->getGpgpuCommandStreamReceiver().setupContext(*pDevice->getDefaultEngine().osContext);
+    }
+
+    void initHeaps() {
         commandStream.replaceBuffer(cmdBuffer, sizeStream);
         auto graphicsAllocation = new MockGraphicsAllocation(cmdBuffer, sizeStream);
         commandStream.replaceGraphicsAllocation(graphicsAllocation);
@@ -58,18 +64,18 @@ struct UltCommandStreamReceiverTest
         ssh.replaceBuffer(sshBuffer, sizeStream);
         graphicsAllocation = new MockGraphicsAllocation(sshBuffer, sizeStream);
         ssh.replaceGraphicsAllocation(graphicsAllocation);
-
-        flushTaskFlags.threadArbitrationPolicy = NEO::HwHelper::get(hardwareInfo.platform.eRenderCoreFamily).getDefaultThreadArbitrationPolicy();
-
-        pDevice->getGpgpuCommandStreamReceiver().setupContext(*pDevice->getDefaultEngine().osContext);
     }
 
-    void TearDown() override {
+    void cleanupHeaps() {
-        pDevice->getGpgpuCommandStreamReceiver().flushBatchedSubmissions();
         delete dsh.getGraphicsAllocation();
         delete ioh.getGraphicsAllocation();
         delete ssh.getGraphicsAllocation();
         delete commandStream.getGraphicsAllocation();
+    }
+
+    void TearDown() override {
+        pDevice->getGpgpuCommandStreamReceiver().flushBatchedSubmissions();
+        cleanupHeaps();
 
         alignedFree(sshBuffer);
         alignedFree(iohBuffer);
@@ -167,5 +173,8 @@ struct UltCommandStreamReceiverTest
     uint32_t latestSentDcFlushTaskCount;
     uint32_t latestSentNonDcFlushTaskCount;
     uint32_t dcFlushRequiredTaskCount;
+
+    const size_t sizeStream = 512;
+    const size_t alignmentStream = 0x1000;
 };
 } // namespace NEO

opencl/test/unit_test/mocks/mock_buffer.h

@@ -20,9 +20,16 @@ class MockBufferStorage {
   public:
     MockBufferStorage() : mockGfxAllocation(data, sizeof(data) / 2),
                           multiGfxAllocation(GraphicsAllocationHelper::toMultiGraphicsAllocation(&mockGfxAllocation)) {
+        initDevice();
     }
+
     MockBufferStorage(bool unaligned) : mockGfxAllocation(unaligned ? alignUp(&data, 4) : alignUp(&data, 64), sizeof(data) / 2),
                                         multiGfxAllocation(GraphicsAllocationHelper::toMultiGraphicsAllocation(&mockGfxAllocation)) {
+        initDevice();
+    }
+    void initDevice() {
+        VariableBackup<uint32_t> maxOsContextCountBackup(&MemoryManager::maxOsContextCount);
+        device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
     }
     ~MockBufferStorage() {
         if (mockGfxAllocation.getDefaultGmm()) {
@@ -31,7 +38,7 @@ class MockBufferStorage {
     }
     char data[128];
     MockGraphicsAllocation mockGfxAllocation;
-    std::unique_ptr<MockDevice> device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(nullptr));
+    std::unique_ptr<MockDevice> device;
     MultiGraphicsAllocation multiGfxAllocation;
 };
 

shared/test/unit_test/command_container/command_container_tests.cpp

@@ -616,10 +616,12 @@ TEST_P(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenHe
 
     auto executionEnvironment = new NEO::ExecutionEnvironment();
     const size_t numDevices = 2;
+
     executionEnvironment->prepareRootDeviceEnvironments(numDevices);
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));
 
@@ -643,10 +645,12 @@ TEST_P(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenHe
 TEST_F(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenCmdBufferAllocationIsCreatedWithCorrectRootDeviceIndex) {
     auto executionEnvironment = new NEO::ExecutionEnvironment();
     const size_t numDevices = 2;
+
     executionEnvironment->prepareRootDeviceEnvironments(numDevices);
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));
 
@@ -670,6 +674,7 @@ TEST_F(CommandContainerHeaps, givenCommandContainerForDifferentRootDevicesThenIn
     for (auto i = 0u; i < numDevices; i++) {
         executionEnvironment->rootDeviceEnvironments[i]->setHwInfo(defaultHwInfo.get());
     }
+    executionEnvironment->calculateMaxOsContextCount();
     auto device0 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 0u));
     auto device1 = std::unique_ptr<MockDevice>(Device::create<MockDevice>(executionEnvironment, 1u));