From 7cb6d665a543e069bbe37b5d67073bc5494cfeff Mon Sep 17 00:00:00 2001
From: Mateusz Jablonski <mateusz.jablonski@intel.com>
Date: Mon, 23 Sep 2019 11:10:48 +0200
Subject: [PATCH] Use device index from os context when allocating heaps and
 command buffers

Related-To: NEO-3691

Change-Id: I64015d606bba289d250920899ad620171e8303b7
Signed-off-by: Mateusz Jablonski <mateusz.jablonski@intel.com>
---
 .../command_stream_receiver.cpp               |  8 +++++--
 .../command_stream/command_stream_receiver.h  |  1 +
 .../command_stream_receiver_tests.cpp         | 22 +++++++++++++++++++
 unit_tests/mocks/mock_csr.h                   |  2 ++
 4 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/runtime/command_stream/command_stream_receiver.cpp b/runtime/command_stream/command_stream_receiver.cpp
index cb8b434a34..042e9ce9e8 100644
--- a/runtime/command_stream/command_stream_receiver.cpp
+++ b/runtime/command_stream/command_stream_receiver.cpp
@@ -126,7 +126,7 @@ void CommandStreamReceiver::ensureCommandBufferAllocation(LinearStream &commandS
     auto allocation = this->getInternalAllocationStorage()->obtainReusableAllocation(allocationSize, allocationType).release();
     if (allocation == nullptr) {
         const AllocationProperties commandStreamAllocationProperties{true, allocationSize, allocationType,
-                                                                     isMultiOsContextCapable(), deviceIndex};
+                                                                     isMultiOsContextCapable(), getDeviceIndexForAllocation()};
         allocation = this->getMemoryManager()->allocateGraphicsMemoryWithProperties(commandStreamAllocationProperties);
     }
     DEBUG_BREAK_IF(allocation == nullptr);
@@ -297,6 +297,10 @@ IndirectHeap &CommandStreamReceiver::getIndirectHeap(IndirectHeap::Type heapType
     return *heap;
 }
 
+uint32_t CommandStreamReceiver::getDeviceIndexForAllocation() const {
+    return osContext->getDeviceBitfield().any() ? static_cast<uint32_t>(Math::log2(static_cast<uint32_t>(osContext->getDeviceBitfield().to_ulong()))) : 0u;
+}
+
 void CommandStreamReceiver::allocateHeapMemory(IndirectHeap::Type heapType,
                                                size_t minRequiredSize, IndirectHeap *&indirectHeap) {
     size_t reservedSize = 0;
@@ -321,7 +325,7 @@ void CommandStreamReceiver::allocateHeapMemory(IndirectHeap::Type heapType,
 
     if (!heapMemory) {
         heapMemory = getMemoryManager()->allocateGraphicsMemoryWithProperties({true, finalHeapSize, allocationType,
-                                                                               isMultiOsContextCapable(), deviceIndex});
+                                                                               isMultiOsContextCapable(), getDeviceIndexForAllocation()});
     } else {
         finalHeapSize = std::max(heapMemory->getUnderlyingBufferSize(), finalHeapSize);
     }
diff --git a/runtime/command_stream/command_stream_receiver.h b/runtime/command_stream/command_stream_receiver.h
index 66c1691c61..ed8c2ddfff 100644
--- a/runtime/command_stream/command_stream_receiver.h
+++ b/runtime/command_stream/command_stream_receiver.h
@@ -191,6 +191,7 @@ class CommandStreamReceiver {
 
   protected:
     void cleanupResources();
+    uint32_t getDeviceIndexForAllocation() const;
 
     std::unique_ptr<FlushStampTracker> flushStamp;
     std::unique_ptr<SubmissionAggregator> submissionAggregator;
diff --git a/unit_tests/command_stream/command_stream_receiver_tests.cpp b/unit_tests/command_stream/command_stream_receiver_tests.cpp
index ef5bee9446..4bf6c55606 100644
--- a/unit_tests/command_stream/command_stream_receiver_tests.cpp
+++ b/unit_tests/command_stream/command_stream_receiver_tests.cpp
@@ -624,3 +624,25 @@ INSTANTIATE_TEST_CASE_P(
     CommandStreamReceiverWithAubSubCaptureTest_program,
     CommandStreamReceiverWithAubSubCaptureTest,
     testing::ValuesIn(aubSubCaptureStatus));
+
+TEST(CommandStreamReceiverDeviceIndexTest, givenCommandStreamReceiverWithDeviceIndexDifferentThanDeviceBitfieldInOsContextWhenGetDeviceIndexForAllocationThenUseHighestBitFromOsContext) {
+    ExecutionEnvironment executioneEnvironment;
+    executioneEnvironment.initializeMemoryManager();
+    MockCommandStreamReceiver csr(executioneEnvironment);
+    csr.deviceIndex = 1;
+    auto osContext = executioneEnvironment.memoryManager->createAndRegisterOsContext(&csr, aub_stream::EngineType::ENGINE_RCS, 0b01, PreemptionMode::Disabled, false);
+
+    csr.setupContext(*osContext);
+    EXPECT_EQ(0u, csr.getDeviceIndexForAllocation());
+}
+
+TEST(CommandStreamReceiverDeviceIndexTest, givenCommandStreamReceiverWithOsContextWithoutDeviceBitfieldsWhenGetDeviceIndexForAllocationThenZeroIsReturned) {
+    ExecutionEnvironment executioneEnvironment;
+    executioneEnvironment.initializeMemoryManager();
+    MockCommandStreamReceiver csr(executioneEnvironment);
+    csr.deviceIndex = 1;
+    auto osContext = executioneEnvironment.memoryManager->createAndRegisterOsContext(&csr, aub_stream::EngineType::ENGINE_RCS, 0b00, PreemptionMode::Disabled, false);
+
+    csr.setupContext(*osContext);
+    EXPECT_EQ(0u, csr.getDeviceIndexForAllocation());
+}
diff --git a/unit_tests/mocks/mock_csr.h b/unit_tests/mocks/mock_csr.h
index a109793dd8..8d2991d651 100644
--- a/unit_tests/mocks/mock_csr.h
+++ b/unit_tests/mocks/mock_csr.h
@@ -245,6 +245,8 @@ class MockFlatBatchBufferHelper : public FlatBatchBufferHelperHw<GfxFamily> {
 class MockCommandStreamReceiver : public CommandStreamReceiver {
   public:
     using CommandStreamReceiver::CommandStreamReceiver;
+    using CommandStreamReceiver::deviceIndex;
+    using CommandStreamReceiver::getDeviceIndexForAllocation;
     using CommandStreamReceiver::internalAllocationStorage;
     using CommandStreamReceiver::latestFlushedTaskCount;
     using CommandStreamReceiver::latestSentTaskCount;