Add indirect allocations to residency at kernel submission time

Change-Id: Idc6ce7ac72de84107990a5c9786c868d4bfa4322
Signed-off-by: Raiyan Latif <raiyan.latif@intel.com>

Author: raiyanla

level_zero/core/source/cmdlist/cmdlist.h

@@ -136,6 +136,14 @@ struct CommandList : _ze_command_list_handle_t {
         return commandListPreemptionMode;
     }
 
+    UnifiedMemoryControls getUnifiedMemoryControls() const {
+        return unifiedMemoryControls;
+    }
+
+    bool hasIndirectAllocationsAllowed() const {
+        return indirectAllocationsAllowed;
+    }
+
     NEO::PreemptionMode obtainFunctionPreemptionMode(Kernel *kernel);
 
     std::vector<Kernel *> &getPrintfFunctionContainer() {
@@ -171,6 +179,8 @@ struct CommandList : _ze_command_list_handle_t {
     uint32_t commandListPerThreadScratchSize = 0u;
     NEO::PreemptionMode commandListPreemptionMode = NEO::PreemptionMode::Initial;
     bool isCopyOnlyCmdList = false;
+    UnifiedMemoryControls unifiedMemoryControls;
+    bool indirectAllocationsAllowed = false;
 };
 
 using CommandListAllocatorFn = CommandList *(*)(uint32_t);

level_zero/core/source/cmdlist/cmdlist_hw_base.inl

@@ -50,10 +50,18 @@ ze_result_t CommandListCoreFamily<gfxCoreFamily>::appendLaunchKernelWithParams(z
 
     if (kernel->hasIndirectAllocationsAllowed()) {
         UnifiedMemoryControls unifiedMemoryControls = kernel->getUnifiedMemoryControls();
-        auto svmAllocsManager = device->getDriverHandle()->getSvmAllocsManager();
-        auto &residencyContainer = commandContainer.getResidencyContainer();
 
-        svmAllocsManager->addInternalAllocationsToResidencyContainer(residencyContainer, unifiedMemoryControls.generateMask());
+        if (unifiedMemoryControls.indirectDeviceAllocationsAllowed) {
+            this->unifiedMemoryControls.indirectDeviceAllocationsAllowed = true;
+        }
+        if (unifiedMemoryControls.indirectHostAllocationsAllowed) {
+            this->unifiedMemoryControls.indirectHostAllocationsAllowed = true;
+        }
+        if (unifiedMemoryControls.indirectSharedAllocationsAllowed) {
+            this->unifiedMemoryControls.indirectSharedAllocationsAllowed = true;
+        }
+
+        this->indirectAllocationsAllowed = true;
     }
 
     NEO::EncodeDispatchKernel<GfxFamily>::encode(commandContainer,

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -20,8 +20,10 @@
 #include "shared/source/helpers/interlocked_max.h"
 #include "shared/source/helpers/preamble.h"
 #include "shared/source/memory_manager/memory_manager.h"
+#include "shared/source/memory_manager/residency_container.h"
 #include "shared/source/os_interface/os_context.h"
 #include "shared/source/page_fault_manager/cpu_page_fault_manager.h"
+#include "shared/source/unified_memory/unified_memory.h"
 
 #include "level_zero/core/source/cmdlist/cmdlist.h"
 #include "level_zero/core/source/cmdlist/cmdlist_hw.h"
@@ -114,6 +116,15 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
     for (auto i = 0u; i < numCommandLists; i++) {
         auto commandList = CommandList::fromHandle(phCommandLists[i]);
 
+        bool indirectAllocationsAllowed = commandList->hasIndirectAllocationsAllowed();
+        if (indirectAllocationsAllowed) {
+            UnifiedMemoryControls unifiedMemoryControls = commandList->getUnifiedMemoryControls();
+
+            auto svmAllocsManager = device->getDriverHandle()->getSvmAllocsManager();
+            svmAllocsManager->addInternalAllocationsToResidencyContainer(commandList->commandContainer.getResidencyContainer(),
+                                                                         unifiedMemoryControls.generateMask());
+        }
+
         totalCmdBuffers += commandList->commandContainer.getCmdBufferAllocations().size();
         spaceForResidency += commandList->commandContainer.getResidencyContainer().size();
         auto commandListPreemption = commandList->getCommandListPreemptionMode();

level_zero/core/test/unit_tests/sources/cmdlist/test_cmdlist_append_launch_kernel.cpp

@@ -22,6 +22,36 @@ namespace ult {
 
 using CommandListAppendLaunchKernel = Test<ModuleFixture>;
 
+HWTEST_F(CommandListAppendLaunchKernel, givenKernelWithIndirectAllocationsAllowedThenCommandListReturnsExpectedIndirectAllocationsAllowed) {
+    createKernel();
+    kernel->unifiedMemoryControls.indirectDeviceAllocationsAllowed = true;
+    kernel->unifiedMemoryControls.indirectSharedAllocationsAllowed = true;
+    kernel->unifiedMemoryControls.indirectHostAllocationsAllowed = true;
+    EXPECT_TRUE(kernel->getUnifiedMemoryControls().indirectDeviceAllocationsAllowed);
+    EXPECT_TRUE(kernel->hasIndirectAllocationsAllowed());
+
+    ze_group_count_t groupCount{1, 1, 1};
+    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily, device, false));
+    auto result = commandList->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr);
+
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+    ASSERT_TRUE(commandList->hasIndirectAllocationsAllowed());
+}
+
+HWTEST_F(CommandListAppendLaunchKernel, givenKernelWithIndirectAllocationsNotAllowedThenCommandListReturnsExpectedIndirectAllocationsAllowed) {
+    createKernel();
+    kernel->unifiedMemoryControls.indirectDeviceAllocationsAllowed = false;
+    kernel->unifiedMemoryControls.indirectSharedAllocationsAllowed = false;
+    kernel->unifiedMemoryControls.indirectHostAllocationsAllowed = false;
+
+    ze_group_count_t groupCount{1, 1, 1};
+    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily, device, false));
+    auto result = commandList->appendLaunchKernel(kernel->toHandle(), &groupCount, nullptr, 0, nullptr);
+
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+    ASSERT_FALSE(commandList->hasIndirectAllocationsAllowed());
+}
+
 HWTEST_F(CommandListAppendLaunchKernel, givenNotEnoughSpaceInCommandStreamWhenAppendingKernelThenBbEndIsAddedAndNewCmdBufferAllocated) {
     using MI_BATCH_BUFFER_END = typename FamilyType::MI_BATCH_BUFFER_END;
     createKernel();

level_zero/core/test/unit_tests/sources/cmdqueue/test_cmdqueue.cpp

@@ -14,6 +14,8 @@
 
 #include "level_zero/core/source/driver/driver_handle_imp.h"
 #include "level_zero/core/test/unit_tests/fixtures/device_fixture.h"
+#include "level_zero/core/test/unit_tests/fixtures/module_fixture.h"
+#include "level_zero/core/test/unit_tests/mocks/mock_cmdlist.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_kernel.h"
 #include "level_zero/core/test/unit_tests/mocks/mock_memory_manager.h"
@@ -284,5 +286,64 @@ HWTEST_F(CommandQueueCommands, givenCommandQueueWhenExecutingCommandListsThenHar
     commandQueue->destroy();
 }
 
+using CommandQueueIndirectAllocations = Test<ModuleFixture>;
+HWTEST_F(CommandQueueIndirectAllocations, givenCommandQueueWhenExecutingCommandListsThenExpectedIndirectAllocationsAddedToResidencyContainer) {
+    const ze_command_queue_desc_t desc = {
+        ZE_COMMAND_QUEUE_DESC_VERSION_CURRENT,
+        ZE_COMMAND_QUEUE_FLAG_NONE,
+        ZE_COMMAND_QUEUE_MODE_DEFAULT,
+        ZE_COMMAND_QUEUE_PRIORITY_NORMAL,
+        0};
+
+    MockCsrHw2<FamilyType> csr(*neoDevice->getExecutionEnvironment(), 0);
+    csr.initializeTagAllocation();
+    csr.setupContext(*neoDevice->getDefaultEngine().osContext);
+
+    L0::CommandQueue *commandQueue = CommandQueue::create(productFamily,
+                                                          device,
+                                                          &csr,
+                                                          &desc,
+                                                          true);
+    ASSERT_NE(nullptr, commandQueue);
+
+    std::unique_ptr<L0::CommandList> commandList(CommandList::create(productFamily, device, true));
+
+    void *deviceAlloc = nullptr;
+    auto result = device->getDriverHandle()->allocDeviceMem(device->toHandle(), ZE_DEVICE_MEM_ALLOC_FLAG_DEFAULT, 16384u, 4096u, &deviceAlloc);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto gpuAlloc = device->getDriverHandle()->getSvmAllocsManager()->getSVMAllocs()->get(deviceAlloc)->gpuAllocation;
+    ASSERT_NE(nullptr, gpuAlloc);
+
+    createKernel();
+    kernel->unifiedMemoryControls.indirectDeviceAllocationsAllowed = true;
+    EXPECT_TRUE(kernel->getUnifiedMemoryControls().indirectDeviceAllocationsAllowed);
+
+    ze_group_count_t groupCount{1, 1, 1};
+    result = commandList->appendLaunchKernel(kernel->toHandle(),
+                                             &groupCount,
+                                             nullptr,
+                                             0,
+                                             nullptr);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    auto itorEvent = std::find(std::begin(commandList->commandContainer.getResidencyContainer()),
+                               std::end(commandList->commandContainer.getResidencyContainer()),
+                               gpuAlloc);
+    EXPECT_EQ(itorEvent, std::end(commandList->commandContainer.getResidencyContainer()));
+
+    auto commandListHandle = commandList->toHandle();
+    result = commandQueue->executeCommandLists(1, &commandListHandle, nullptr, false);
+    ASSERT_EQ(ZE_RESULT_SUCCESS, result);
+
+    itorEvent = std::find(std::begin(commandList->commandContainer.getResidencyContainer()),
+                          std::end(commandList->commandContainer.getResidencyContainer()),
+                          gpuAlloc);
+    EXPECT_NE(itorEvent, std::end(commandList->commandContainer.getResidencyContainer()));
+
+    device->getDriverHandle()->getSvmAllocsManager()->freeSVMAlloc(deviceAlloc);
+    commandQueue->destroy();
+}
+
 } // namespace ult
 } // namespace L0