From dbf268df2fe6851d83517517a628f70413697fea Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Mon, 16 Jan 2023 15:42:00 +0000
Subject: [PATCH] L0: Capability to implicitly assign multi regular context to
 cmd list.

Related-To: NEO-7618

Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 level_zero/core/source/device/device_imp.cpp  | 11 ++++-
 .../sources/cmdqueue/test_cmdqueue_1.cpp      | 49 +++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/level_zero/core/source/device/device_imp.cpp b/level_zero/core/source/device/device_imp.cpp
index c12bd0e54f..46228154d4 100644
--- a/level_zero/core/source/device/device_imp.cpp
+++ b/level_zero/core/source/device/device_imp.cpp
@@ -1408,10 +1408,17 @@ ze_result_t DeviceImp::getCsrForOrdinalAndIndex(NEO::CommandStreamReceiver **csr
     }
 
     if (ordinal < numEngineGroups) {
-        if (index >= engineGroups[ordinal].engines.size()) {
+        auto &engines = engineGroups[ordinal].engines;
+        if (index >= engines.size()) {
             return ZE_RESULT_ERROR_INVALID_ARGUMENT;
         }
-        *csr = engineGroups[ordinal].engines[index].commandStreamReceiver;
+        *csr = engines[index].commandStreamReceiver;
+
+        auto &osContext = (*csr)->getOsContext();
+
+        if (neoDevice->getNumberOfRegularContextsPerEngine() > 1 && !osContext.isRootDevice() && NEO::EngineHelpers::isCcs(osContext.getEngineType())) {
+            *csr = neoDevice->getNextEngineForMultiRegularContextMode().commandStreamReceiver;
+        }
     } else {
         auto subDeviceOrdinal = ordinal - numEngineGroups;
         if (index >= this->subDeviceCopyEngineGroups[subDeviceOrdinal].engines.size()) {
diff --git a/level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue_1.cpp b/level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue_1.cpp
index 9669ef2f13..1fd84ae3b4 100644
--- a/level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue_1.cpp
+++ b/level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue_1.cpp
@@ -8,6 +8,7 @@
 #include "shared/source/command_stream/scratch_space_controller.h"
 #include "shared/source/helpers/state_base_address.h"
 #include "shared/test/common/cmd_parse/gen_cmd_parse.h"
+#include "shared/test/common/helpers/variable_backup.h"
 #include "shared/test/common/libult/ult_command_stream_receiver.h"
 #include "shared/test/common/mocks/mock_bindless_heaps_helper.h"
 #include "shared/test/common/mocks/mock_command_stream_receiver.h"
@@ -939,6 +940,54 @@ TEST_F(DeviceCreateCommandQueueTest, givenNormalPriorityDescWhenCreateCommandQue
     commandQueue->destroy();
 }
 
+struct CommandQueueCreateWithMultipleRegularContextsTests : public DeviceCreateCommandQueueTest {
+    void SetUp() override {
+        DebugManager.flags.NumberOfRegularContextsPerEngine.set(numberOfRegularContextsPerEngine);
+        DebugManager.flags.NodeOrdinal.set(static_cast<int32_t>(aub_stream::EngineType::ENGINE_CCS));
+
+        backupHwInfo = std::make_unique<VariableBackup<HardwareInfo>>(defaultHwInfo.get());
+        defaultHwInfo->featureTable.flags.ftrCCSNode = true;
+
+        DeviceCreateCommandQueueTest::SetUp();
+
+        if (device->getHwInfo().gtSystemInfo.CCSInfo.NumberOfCCSEnabled == 0) {
+            GTEST_SKIP();
+        }
+
+        auto &engineGroups = device->getNEODevice()->getRegularEngineGroups();
+
+        for (uint32_t i = 0; i < engineGroups.size(); i++) {
+            if (engineGroups[i].engineGroupType == EngineGroupType::Compute) {
+                computeOrdinal = i;
+                break;
+            }
+        }
+    }
+
+    std::unique_ptr<VariableBackup<HardwareInfo>> backupHwInfo;
+    const uint32_t numberOfRegularContextsPerEngine = 5;
+    uint32_t computeOrdinal = 0;
+    DebugManagerStateRestore restore;
+};
+
+HWTEST_F(CommandQueueCreateWithMultipleRegularContextsTests, givenSupportedRequestWhenCreatingCommandQueueThenAssignNextAvailableContext) {
+    uint32_t expectedIndex = 0;
+    constexpr uint32_t iterationCount = 3;
+
+    for (uint32_t i = 0; i < (numberOfRegularContextsPerEngine * iterationCount); i++) {
+        NEO::CommandStreamReceiver *csr = nullptr;
+        device->getCsrForOrdinalAndIndex(&csr, computeOrdinal, 0u);
+        ASSERT_NE(nullptr, csr);
+
+        EXPECT_EQ(csr, device->getNEODevice()->getAllEngines()[expectedIndex].commandStreamReceiver);
+
+        expectedIndex++;
+        if (expectedIndex == (numberOfRegularContextsPerEngine - 1)) {
+            expectedIndex = 0;
+        }
+    }
+}
+
 TEST_F(DeviceCreateCommandQueueTest,
        whenCallingGetCsrForOrdinalAndIndexWithInvalidOrdinalThenInvalidArgumentIsReturned) {
     ze_command_queue_desc_t desc{};