diff --git a/projects/clr/hipamd/src/hip_graph_internal.cpp b/projects/clr/hipamd/src/hip_graph_internal.cpp index 4b126235cb..5a63202e53 100644 --- a/projects/clr/hipamd/src/hip_graph_internal.cpp +++ b/projects/clr/hipamd/src/hip_graph_internal.cpp @@ -374,20 +374,26 @@ hipError_t GraphExec::Init() { //! Chunk size to add to kern arg pool constexpr uint32_t kKernArgChunkSize = 128 * Ki; // ================================================================================================ -void GraphExec::GetKernelArgSizeForGraph(size_t& kernArgSizeForGraph) { - // GPU packet capture is enabled for kernel nodes. Calculate the kernel - // arg size required for all graph kernel nodes to allocate +void GraphExec::GetKernelArgSizeForGraph(std::unordered_map& kernArgSizeForGraph) { + // Calculate the kernel argument size required for all graph kernel nodes + // when GPU packet capture is enabled for (hip::GraphNode* node : topoOrder_) { if (node->GraphCaptureEnabled()) { - kernArgSizeForGraph += node->GetKerArgSize(); + // Accumulate the kernel argument size for each device + kernArgSizeForGraph[node->dev_id_] += node->GetKerArgSize(); } else if (node->GetType() == hipGraphNodeTypeGraph) { + // Handle child graph nodes auto childNode = reinterpret_cast(node); + // Child graph shares same kernel arg manager GraphKernelArgManager* KernelArgManager = GetKernelArgManager(); - KernelArgManager->retain(); - childNode->SetKernelArgManager(KernelArgManager); - if (childNode->GetChildGraph()->max_streams_ == 1) { - childNode->GetKernelArgSizeForGraph(kernArgSizeForGraph); + if (KernelArgManager != nullptr) { + KernelArgManager->retain(); + childNode->SetKernelArgManager(KernelArgManager); + // Recursively process child graph if it uses single stream + if (childNode->GetChildGraph()->max_streams_ == 1) { + childNode->GetKernelArgSizeForGraph(kernArgSizeForGraph); + } } } } @@ -467,17 +473,32 @@ hipError_t GraphExec::CaptureAndFormPacketsForGraph() { // ================================================================================================ hipError_t GraphExec::CaptureAQLPackets() { hipError_t status = hipSuccess; - size_t kernArgSizeForGraph = 0; + + // Create a map to track kernel argument sizes for each device + std::unordered_map kernArgSizeForGraph; + // Reserve space for all available devices and Initialize to 0 + kernArgSizeForGraph.reserve(g_devices.size()); + for (int devId = 0; devId < g_devices.size(); devId++) { + kernArgSizeForGraph[devId] = 0; + } GetKernelArgSizeForGraph(kernArgSizeForGraph); - // When we support multi device graph lauch we need to allocate the kenel args on respective - // device for each kernel Assume graph has nodes of same device allocate kernel args on the - // device from the first node - auto device = g_devices[topoOrder_[0]->GetDeviceId()]->devices()[0]; - // Add a larger initial pool to accomodate for any updates to kernel args - bool bStatus = - kernArgManager_->AllocGraphKernargPool(kernArgSizeForGraph + kKernArgChunkSize, device); - if (bStatus != true) { - return hipErrorMemoryAllocation; + + // Allocate kernel argument pools on respective devices with extra space for updates + for (const auto& deviceKernArgPair : kernArgSizeForGraph) { + const int deviceId = deviceKernArgPair.first; + const size_t kernArgSize = deviceKernArgPair.second; + + if (kernArgSize == 0) { + continue; + } + + const size_t totalPoolSize = kernArgSize + kKernArgChunkSize; + if (!kernArgManager_->AllocGraphKernargPool(totalPoolSize, g_devices[deviceId]->devices()[0])) { + ClPrint(amd::LOG_ERROR, amd::LOG_CODE, + "[hipGraph] Failed to allocate kernel argument pool of size %zu for device %d", + totalPoolSize, deviceId); + return hipErrorMemoryAllocation; + } } status = CaptureAndFormPacketsForGraph(); @@ -486,13 +507,13 @@ hipError_t GraphExec::CaptureAQLPackets() { } kernArgManager_->ReadBackOrFlush(); - return status; + return hipSuccess;; } // ================================================================================================ hipError_t GraphExec::UpdateAQLPacket(hip::GraphNode* node) { hipError_t status = hipSuccess; - if (max_streams_ == 1) { + if (max_streams_ == 1 && node->GraphCaptureEnabled()) { status = node->CaptureAndFormPacket(kernArgManager_); } return status; @@ -831,9 +852,6 @@ bool GraphKernelArgManager::AllocGraphKernargPool(size_t pool_size, amd::Device* bool bStatus = true; assert(pool_size > 0); address graph_kernarg_base; - // Current device is stored as part of tls. Save current device to destroy kernelArgs from the - // callback thread. - device_ = device; if (device->info().largeBar_) { amd::Device::AllocationFlags flags = {}; flags.executable_ = true; @@ -847,48 +865,76 @@ bool GraphKernelArgManager::AllocGraphKernargPool(size_t pool_size, amd::Device* if (graph_kernarg_base == nullptr) { return false; } - kernarg_graph_.push_back(KernelArgPoolGraph(graph_kernarg_base, pool_size)); + kernarg_graph_[device].push_back(KernelArgPoolGraph(graph_kernarg_base, pool_size)); return true; } -address GraphKernelArgManager::AllocKernArg(size_t size, size_t alignment) { - assert(alignment != 0); - address result = nullptr; - result = amd::alignUp( - kernarg_graph_.back().kernarg_pool_addr_ + kernarg_graph_.back().kernarg_pool_offset_, - alignment); - const size_t pool_new_usage = (result + size) - kernarg_graph_.back().kernarg_pool_addr_; - if (pool_new_usage <= kernarg_graph_.back().kernarg_pool_size_) { - kernarg_graph_.back().kernarg_pool_offset_ = pool_new_usage; - } else { - // If current chunck is full allocate new chunck with same size as current - bool bStatus = AllocGraphKernargPool(kernarg_graph_.back().kernarg_pool_size_, device_); - if (bStatus == false) { - return nullptr; - } else { - // Allocte kernel arg memory from new chunck - return AllocKernArg(size, alignment); - } +address GraphKernelArgManager::AllocKernArg(size_t size, size_t alignment, int devId) { + if (size == 0) { + return nullptr; } - return result; + + amd::Device* device = g_devices[devId]->devices()[0]; + assert(alignment != 0 && "Alignment must be non-zero"); + + // Check if we have any pools allocated for this device + auto& device_pools = kernarg_graph_[device]; + if (device_pools.empty()) { + return nullptr; + } + + auto& current_pool = device_pools.back(); + + // Calculate aligned address for the allocation + address aligned_addr = amd::alignUp(current_pool.kernarg_pool_addr_ + current_pool.kernarg_pool_offset_, alignment); + const size_t new_pool_usage = (aligned_addr + size) - current_pool.kernarg_pool_addr_; + + // Check if allocation fits in current pool + if (new_pool_usage <= current_pool.kernarg_pool_size_) { + current_pool.kernarg_pool_offset_ = new_pool_usage; + return aligned_addr; + } + + // Current pool is full - allocate a new pool with the same size + if (!AllocGraphKernargPool(current_pool.kernarg_pool_size_, device)) { + return nullptr; + } + + // Recursively allocate from the new pool + return AllocKernArg(size, alignment, devId); } void GraphKernelArgManager::ReadBackOrFlush() { - if (device_kernarg_pool_ && device_) { - auto kernArgImpl = device_->settings().kernel_arg_impl_; + if (!device_kernarg_pool_) { + return; + } + + for (const auto& kernarg : kernarg_graph_) { + const auto kernArgImpl = kernarg.first->settings().kernel_arg_impl_; if (kernArgImpl == KernelArgImpl::DeviceKernelArgsHDP) { - *device_->info().hdpMemFlushCntl = 1u; - auto kSentinel = *reinterpret_cast(device_->info().hdpMemFlushCntl); - } else if (kernArgImpl == KernelArgImpl::DeviceKernelArgsReadback && - kernarg_graph_.back().kernarg_pool_addr_ != 0) { - address dev_ptr = - kernarg_graph_.back().kernarg_pool_addr_ + kernarg_graph_.back().kernarg_pool_size_; - auto kSentinel = *reinterpret_cast(dev_ptr - 1); + // Trigger HDP flush + *kernarg.first->info().hdpMemFlushCntl = 1u; + // Read back to ensure flush completion + volatile int kSentinel = *reinterpret_cast(kernarg.first->info().hdpMemFlushCntl); + (void)kSentinel; // Suppress unused variable warning + } else if (kernArgImpl == KernelArgImpl::DeviceKernelArgsReadback) { + const auto& pool = kernarg.second.back(); + if (pool.kernarg_pool_addr_ == 0) { + continue; + } + + // Perform readback operation on the last byte of the pool + address dev_ptr = pool.kernarg_pool_addr_ + pool.kernarg_pool_size_; + volatile unsigned char* sentinel_ptr = reinterpret_cast(dev_ptr - 1); + + // Read-modify-write sequence with memory barriers + volatile unsigned char kSentinel = *sentinel_ptr; _mm_sfence(); - *(dev_ptr - 1) = kSentinel; + *sentinel_ptr = kSentinel; _mm_mfence(); - kSentinel = *reinterpret_cast(dev_ptr - 1); + kSentinel = *sentinel_ptr; + (void)kSentinel; // Suppress unused variable warning } } } diff --git a/projects/clr/hipamd/src/hip_graph_internal.hpp b/projects/clr/hipamd/src/hip_graph_internal.hpp index 69113f81f4..0d212e38e9 100644 --- a/projects/clr/hipamd/src/hip_graph_internal.hpp +++ b/projects/clr/hipamd/src/hip_graph_internal.hpp @@ -153,21 +153,21 @@ class GraphKernelArgManager : public amd::ReferenceCountedObject, public: GraphKernelArgManager() : amd::ReferenceCountedObject() {} ~GraphKernelArgManager() { - //! Release the kernel arg pools - if (device_ != nullptr) { - for (auto& element : kernarg_graph_) { - device_->hostFree(element.kernarg_pool_addr_, element.kernarg_pool_size_); + for (auto kernarg : kernarg_graph_) { + //! Release the kernel arg pools + for (auto& element : kernarg.second) { + kernarg.first->hostFree(element.kernarg_pool_addr_, element.kernarg_pool_size_); } - kernarg_graph_.clear(); + kernarg.second.clear(); } } - // Allocate kernel arg pool on device for the given size. + //! Allocate kernel arg pool on device for the given size. bool AllocGraphKernargPool(size_t pool_size, amd::Device* device); // Allocate kernel args from current chunck for given size and alignment. // If kernel arg pool is full allocate new chunck and alloc kern args from new pool. - address AllocKernArg(size_t size, size_t alignment) override; + address AllocKernArg(size_t size, size_t alignment, int devId) override; // Do HDP flush/When HDP flush register is invalid fallback to Readback void ReadBackOrFlush(); @@ -181,8 +181,8 @@ class GraphKernelArgManager : public amd::ReferenceCountedObject, size_t kernarg_pool_offset_; //! Current offset in the kernel arg alloc }; bool device_kernarg_pool_ = false; //! Indicate if kernel pool in device mem - amd::Device* device_ = nullptr; //! Device from where kernel arguments are allocated - std::vector kernarg_graph_; //! Vector of allocated kernarg pool + std::unordered_map> + kernarg_graph_; //! Vector of allocated kernarg pool per device using KernelArgImpl = device::Settings::KernelArgImpl; }; @@ -870,7 +870,7 @@ class GraphExec : public amd::ReferenceCountedObject, public Graph { GraphKernelArgManager* GetKernelArgManager() { return kernArgManager_; } static void DecrementRefCount(cl_event event, cl_int command_exec_status, void* user_data); hipError_t CaptureAndFormPacketsForGraph(); - void GetKernelArgSizeForGraph(size_t& kernArgSizeForGraph); + void GetKernelArgSizeForGraph(std::unordered_map& kernArgSizeForGraph); hipError_t EnqueueGraphWithSingleList(hip::Stream* hip_stream); bool TopologicalOrder() { return Graph::TopologicalOrder(topoOrder_); } diff --git a/projects/clr/rocclr/device/rocm/rocblit.cpp b/projects/clr/rocclr/device/rocm/rocblit.cpp index 8cefbe93c0..49377f6492 100644 --- a/projects/clr/rocclr/device/rocm/rocblit.cpp +++ b/projects/clr/rocclr/device/rocm/rocblit.cpp @@ -2046,8 +2046,9 @@ bool KernelBlitManager::fillBuffer1D(device::Memory& memory, const void* pattern bool isGraphPktCapturing = gpu().command() != nullptr && gpu().command()->getPktCapturingState(); - auto constBuf = isGraphPktCapturing ? gpu().command()->getGraphKernArg(kCBSize, kCBAlignment) - : gpu().allocKernArg(kCBSize, kCBAlignment); + auto constBuf = isGraphPktCapturing + ? gpu().command()->getGraphKernArg(kCBSize, kCBAlignment, dev().index()) + : gpu().allocKernArg(kCBSize, kCBAlignment); // If pattern has been expanded, use the expanded pattern, otherwise use the default pattern. if (packed_obj.pattern_expanded_) { @@ -2141,8 +2142,9 @@ bool KernelBlitManager::fillBuffer2D(device::Memory& memory, const void* pattern // Get constant buffer to allow multipel fills bool isGraphPktCapturing = gpu().command() != nullptr && gpu().command()->getPktCapturingState(); - auto constBuf = isGraphPktCapturing ? gpu().command()->getGraphKernArg(kCBSize, kCBAlignment) - : gpu().allocKernArg(kCBSize, kCBAlignment); + auto constBuf = isGraphPktCapturing + ? gpu().command()->getGraphKernArg(kCBSize, kCBAlignment, dev().index()) + : gpu().allocKernArg(kCBSize, kCBAlignment); memcpy(constBuf, pattern, patternSize); constexpr bool kDirectVa = true; diff --git a/projects/clr/rocclr/device/rocm/rocvirtual.cpp b/projects/clr/rocclr/device/rocm/rocvirtual.cpp index a633c18b94..2ee6a11793 100644 --- a/projects/clr/rocclr/device/rocm/rocvirtual.cpp +++ b/projects/clr/rocclr/device/rocm/rocvirtual.cpp @@ -3691,7 +3691,7 @@ bool VirtualGPU::submitKernelInternal(const amd::NDRangeContainer& sizes, const // Allocate buffer to hold kernel arguments if (isGraphCapture) { argBuffer = command_->getGraphKernArg(gpuKernel.KernargSegmentByteSize(), - gpuKernel.KernargSegmentAlignment()); + gpuKernel.KernargSegmentAlignment(), dev().index()); command_->SetKernelName(gpuKernel.getDemangledName().c_str()); } else { ClPrint(amd::LOG_DETAIL_DEBUG, amd::LOG_KERN, diff --git a/projects/clr/rocclr/platform/command.hpp b/projects/clr/rocclr/platform/command.hpp index f36636eb48..bd91cd334c 100644 --- a/projects/clr/rocclr/platform/command.hpp +++ b/projects/clr/rocclr/platform/command.hpp @@ -248,7 +248,7 @@ union CopyMetadata { // Interface to callback to allocate kernel args from the graph kernel arg pool. class GraphKernelArgManager { public: - virtual address AllocKernArg(size_t size, size_t alignment) = 0; + virtual address AllocKernArg(size_t size, size_t alignment, int devId) = 0; }; /*! \brief An operation that is submitted to a command queue. @@ -341,8 +341,8 @@ class Command : public Event { return packet; } - address getGraphKernArg(int size, int alignment) { - return graphKernArgMgr_->AllocKernArg(size, alignment); + address getGraphKernArg(int size, int alignment, int devId) { + return graphKernArgMgr_->AllocKernArg(size, alignment, devId); } //! Overload new/delete for fast commands allocation/destruction