Make partitioned post sync operations for partitioned workloads

Signed-off-by: Zbigniew Zdanowicz <[email protected]>

Author: zzdanowicz

level_zero/core/source/cmdlist/cmdlist.h

@@ -216,12 +216,6 @@ struct CommandList : _ze_command_list_handle_t {
         TYPE_IMMEDIATE = 1u
     };
 
-    CommandQueue *cmdQImmediate = nullptr;
-    NEO::CommandStreamReceiver *csr = nullptr;
-    uint32_t cmdListType = CommandListType::TYPE_REGULAR;
-    Device *device = nullptr;
-    std::vector<Kernel *> printfFunctionContainer;
-
     virtual ze_result_t executeCommandListImmediate(bool performMigration) = 0;
     virtual ze_result_t initialize(Device *device, NEO::EngineGroupType engineGroupType, ze_command_list_flags_t flags) = 0;
     virtual ~CommandList();
@@ -241,33 +235,41 @@ struct CommandList : _ze_command_list_handle_t {
         return commandsToPatch;
     }
 
-    bool isSyncModeQueue = false;
-    bool commandListSLMEnabled = false;
-    uint32_t commandListPerThreadScratchSize = 0u;
+    void makeResidentAndMigrate(bool);
+    void migrateSharedAllocations();
+
+    std::vector<Kernel *> printfFunctionContainer;
+    CommandQueue *cmdQImmediate = nullptr;
+    NEO::CommandStreamReceiver *csr = nullptr;
+    Device *device = nullptr;
     NEO::PreemptionMode commandListPreemptionMode = NEO::PreemptionMode::Initial;
+    uint32_t cmdListType = CommandListType::TYPE_REGULAR;
+    uint32_t commandListPerThreadScratchSize = 0u;
     uint32_t threadArbitrationPolicy = NEO::ThreadArbitrationPolicy::RoundRobin;
+    uint32_t partitionCount = 1;
     bool isFlushTaskSubmissionEnabled = false;
-
-    void makeResidentAndMigrate(bool);
-    void migrateSharedAllocations();
+    bool isSyncModeQueue = false;
+    bool commandListSLMEnabled = false;
 
   protected:
-    std::map<const void *, NEO::GraphicsAllocation *> hostPtrMap;
-    NEO::EngineGroupType engineGroupType;
-    ze_command_list_flags_t flags = 0u;
-    UnifiedMemoryControls unifiedMemoryControls;
-    bool indirectAllocationsAllowed = false;
-    bool internalUsage = false;
-    bool containsCooperativeKernelsFlag = false;
     NEO::GraphicsAllocation *getAllocationFromHostPtrMap(const void *buffer, uint64_t bufferSize);
     NEO::GraphicsAllocation *getHostPtrAlloc(const void *buffer, uint64_t bufferSize);
-    bool containsStatelessUncachedResource = false;
+
+    std::map<const void *, NEO::GraphicsAllocation *> hostPtrMap;
+    std::vector<NEO::GraphicsAllocation *> ownedPrivateAllocations;
 
     NEO::StreamProperties requiredStreamState{};
     NEO::StreamProperties finalStreamState{};
     CommandsToPatch commandsToPatch{};
 
-    std::vector<NEO::GraphicsAllocation *> ownedPrivateAllocations;
+    ze_command_list_flags_t flags = 0u;
+    UnifiedMemoryControls unifiedMemoryControls;
+
+    NEO::EngineGroupType engineGroupType;
+    bool indirectAllocationsAllowed = false;
+    bool internalUsage = false;
+    bool containsCooperativeKernelsFlag = false;
+    bool containsStatelessUncachedResource = false;
 };
 
 using CommandListAllocatorFn = CommandList *(*)(uint32_t);

level_zero/core/source/cmdlist/cmdlist_hw.inl

@@ -108,7 +108,7 @@ ze_result_t CommandListCoreFamily<gfxCoreFamily>::reset() {
         device->getNEODevice()->getMemoryManager()->freeGraphicsMemory(alloc);
     }
     this->ownedPrivateAllocations.clear();
-
+    partitionCount = 1;
     return ZE_RESULT_SUCCESS;
 }
 

level_zero/core/source/cmdlist/cmdlist_hw_xehp_and_later.inl

@@ -228,6 +228,7 @@ ze_result_t CommandListCoreFamily<gfxCoreFamily>::appendLaunchKernelWithParams(z
                                                  partitionCount,
                                                  internalUsage,
                                                  isCooperative);
+    this->partitionCount = std::max(partitionCount, this->partitionCount);
     if (hEvent) {
         auto event = Event::fromHandle(hEvent);
         if (partitionCount > 1) {

level_zero/core/source/cmdqueue/cmdqueue.cpp

@@ -106,10 +106,17 @@ ze_result_t CommandQueueImp::synchronizeByPollingForTaskCount(uint64_t timeout)
         timeoutMicroseconds = NEO::TimeoutControls::maxTimeout;
     }
 
-    csr->waitForCompletionWithTimeout(enableTimeout, timeoutMicroseconds, this->taskCount);
-
-    if (*csr->getTagAddress() < taskCountToWait) {
-        return ZE_RESULT_NOT_READY;
+    if (partitionCount > 1) {
+        volatile uint32_t *pollAddress = csr->getTagAddress();
+        for (uint32_t i = 0; i < partitionCount; i++) {
+            csr->waitForCompletionWithTimeout(pollAddress, enableTimeout, timeoutMicroseconds, this->taskCount);
+            pollAddress += addressOffsetDwords;
+        }
+    } else {
+        csr->waitForCompletionWithTimeout(enableTimeout, timeoutMicroseconds, this->taskCount);
+        if (*csr->getTagAddress() < taskCountToWait) {
+            return ZE_RESULT_NOT_READY;
+        }
     }
 
     postSyncOperations();

level_zero/core/source/cmdqueue/cmdqueue.h

@@ -57,6 +57,7 @@ struct CommandQueue : _ze_command_queue_handle_t {
 
   protected:
     NEO::PreemptionMode commandQueuePreemptionMode = NEO::PreemptionMode::Initial;
+    uint32_t partitionCount = 1;
     bool preemptionCmdSyncProgramming = true;
     bool commandQueueDebugCmdsProgrammed = false;
     bool isCopyOnlyCommandQueue = false;

level_zero/core/source/cmdqueue/cmdqueue_hw.inl

@@ -10,6 +10,7 @@
 #include "shared/source/built_ins/built_ins.h"
 #include "shared/source/built_ins/sip.h"
 #include "shared/source/command_container/command_encoder.h"
+#include "shared/source/command_container/implicit_scaling.h"
 #include "shared/source/command_stream/command_stream_receiver_hw.h"
 #include "shared/source/command_stream/linear_stream.h"
 #include "shared/source/command_stream/preemption.h"
@@ -72,6 +73,9 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
     using PIPE_CONTROL = typename GfxFamily::PIPE_CONTROL;
     using POST_SYNC_OPERATION = typename PIPE_CONTROL::POST_SYNC_OPERATION;
 
+    using MI_LOAD_REGISTER_MEM = typename GfxFamily::MI_LOAD_REGISTER_MEM;
+    using MI_LOAD_REGISTER_IMM = typename GfxFamily::MI_LOAD_REGISTER_IMM;
+
     auto lockCSR = csr->obtainUniqueOwnership();
 
     auto anyCommandListWithCooperativeKernels = false;
@@ -177,6 +181,8 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
                 heapContainer.push_back(element);
             }
         }
+
+        partitionCount = std::max(partitionCount, commandList->partitionCount);
     }
 
     size_t linearStreamSizeEstimate = totalCmdBuffers * sizeof(MI_BATCH_BUFFER_START);
@@ -240,6 +246,10 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
     }
 
     linearStreamSizeEstimate += isCopyOnlyCommandQueue ? NEO::EncodeMiFlushDW<GfxFamily>::getMiFlushDwCmdSizeForDataWrite() : NEO::MemorySynchronizationCommands<GfxFamily>::getSizeForPipeControlWithPostSyncOperation(hwInfo);
+    if (partitionCount > 1) {
+        linearStreamSizeEstimate += sizeof(MI_LOAD_REGISTER_MEM) + sizeof(MI_LOAD_REGISTER_IMM);
+    }
+
     size_t alignedSize = alignUp<size_t>(linearStreamSizeEstimate, minCmdBufferPtrAlign);
     size_t padding = alignedSize - linearStreamSizeEstimate;
     reserveLinearStreamSize(alignedSize);
@@ -399,6 +409,17 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
 
     commandQueuePreemptionMode = statePreemption;
 
+    if (partitionCount > 1) {
+        uint64_t workPartitionAddress = csr->getWorkPartitionAllocationGpuAddress();
+        NEO::EncodeSetMMIO<GfxFamily>::encodeMEM(child,
+                                                 NEO::PartitionRegisters<GfxFamily>::wparidCCSOffset,
+                                                 workPartitionAddress);
+        NEO::EncodeSetMMIO<GfxFamily>::encodeIMM(child,
+                                                 NEO::PartitionRegisters<GfxFamily>::addressOffsetCCSOffset,
+                                                 addressOffset,
+                                                 true);
+    }
+
     if (hFence) {
         csr->makeResident(fence->getAllocation());
         if (isCopyOnlyCommandQueue) {
@@ -407,6 +428,10 @@ ze_result_t CommandQueueHw<gfxCoreFamily>::executeCommandLists(
             NEO::EncodeMiFlushDW<GfxFamily>::programMiFlushDw(child, fence->getGpuAddress(), Fence::STATE_SIGNALED, args);
         } else {
             NEO::PipeControlArgs args(true);
+            if (partitionCount > 1) {
+                args.workloadPartitionOffset = true;
+                fence->setPartitionCount(partitionCount);
+            }
             NEO::MemorySynchronizationCommands<GfxFamily>::addPipeControlAndProgramPostSyncOperation(
                 child, POST_SYNC_OPERATION::POST_SYNC_OPERATION_WRITE_IMMEDIATE_DATA,
                 fence->getGpuAddress(),
@@ -539,6 +564,9 @@ void CommandQueueHw<gfxCoreFamily>::dispatchTaskCountWrite(NEO::LinearStream &co
         NEO::EncodeMiFlushDW<GfxFamily>::programMiFlushDw(commandStream, gpuAddress, taskCountToWrite, args);
     } else {
         NEO::PipeControlArgs args(true);
+        if (partitionCount > 1) {
+            args.workloadPartitionOffset = true;
+        }
         args.notifyEnable = csr->isUsedNotifyEnableForPostSync();
         NEO::MemorySynchronizationCommands<GfxFamily>::addPipeControlAndProgramPostSyncOperation(
             commandStream,

level_zero/core/source/cmdqueue/cmdqueue_imp.h

@@ -62,6 +62,9 @@ struct CommandQueueImp : public CommandQueue {
         MemoryConstants::cacheLineSize +
         NEO::CSRequirements::csOverfetchSize;
 
+    static constexpr uint32_t addressOffsetDwords = 2u;
+    static constexpr uint32_t addressOffset = sizeof(uint32_t) * addressOffsetDwords;
+
     CommandQueueImp() = delete;
     CommandQueueImp(Device *device, NEO::CommandStreamReceiver *csr, const ze_command_queue_desc_t *desc);
 

level_zero/core/source/fence/fence.cpp

@@ -37,9 +37,15 @@ ze_result_t FenceImp::queryStatus() {
         csr->downloadAllocations();
     }
 
-    uint64_t *hostAddr = static_cast<uint64_t *>(allocation->getUnderlyingBuffer());
+    void *hostAddr = static_cast<uint64_t *>(allocation->getUnderlyingBuffer());
     uint32_t queryVal = Fence::STATE_CLEARED;
-    memcpy_s(static_cast<void *>(&queryVal), sizeof(uint32_t), static_cast<void *>(hostAddr), sizeof(uint32_t));
+    for (uint32_t i = 0; i < partitionCount; i++) {
+        memcpy_s(static_cast<void *>(&queryVal), sizeof(uint32_t), hostAddr, sizeof(uint32_t));
+        if (queryVal == Fence::STATE_CLEARED) {
+            break;
+        }
+        hostAddr = ptrOffset(hostAddr, CommandQueueImp::addressOffset);
+    }
     return queryVal == Fence::STATE_CLEARED ? ZE_RESULT_NOT_READY : ZE_RESULT_SUCCESS;
 }
 

level_zero/core/source/fence/fence.h

@@ -47,8 +47,13 @@ struct Fence : _ze_fence_handle_t {
         return allocation->getGpuAddress();
     }
 
+    void setPartitionCount(uint32_t newPartitionCount) {
+        partitionCount = newPartitionCount;
+    }
+
   protected:
     NEO::GraphicsAllocation *allocation = nullptr;
+    uint32_t partitionCount = 1;
 };
 
 struct FenceImp : public Fence {

level_zero/core/test/unit_tests/mocks/mock_cmdqueue.h

@@ -27,8 +27,10 @@ struct WhiteBox<::L0::CommandQueue> : public ::L0::CommandQueueImp {
     using BaseClass::printfFunctionContainer;
     using BaseClass::submitBatchBuffer;
     using BaseClass::synchronizeByPollingForTaskCount;
+    using BaseClass::taskCount;
     using CommandQueue::commandQueuePreemptionMode;
     using CommandQueue::internalUsage;
+    using CommandQueue::partitionCount;
 
     WhiteBox(Device *device, NEO::CommandStreamReceiver *csr,
              const ze_command_queue_desc_t *desc);
@@ -85,6 +87,7 @@ struct MockCommandQueueHw : public L0::CommandQueueHw<gfxCoreFamily> {
     using BaseClass::commandStream;
     using BaseClass::printfFunctionContainer;
     using L0::CommandQueue::internalUsage;
+    using L0::CommandQueue::partitionCount;
     using L0::CommandQueue::preemptionCmdSyncProgramming;
     using L0::CommandQueueImp::csr;