From 3e8d5599d4cca9c434413227f958a7a95d1bcd6e Mon Sep 17 00:00:00 2001 From: Vladana Stojiljkovic Date: Thu, 31 Oct 2024 17:40:10 +0100 Subject: [PATCH] SWDEV-494612- Add capture support for hipLaunchCooperativeKernel Change-Id: I6b3c6af55c60cffd43ce6f47b75998f750b75703 [ROCm/clr commit: b75b0d9a53a8acf15826b678015a3d74a312351f] --- projects/clr/hipamd/src/hip_graph.cpp | 36 +++++++++++++++++-- projects/clr/hipamd/src/hip_graph_capture.hpp | 4 +++ .../clr/hipamd/src/hip_graph_internal.hpp | 12 ++++--- projects/clr/hipamd/src/hip_module.cpp | 3 ++ 4 files changed, 49 insertions(+), 6 deletions(-) diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp index 76b9a10a7d..d1590cd283 100644 --- a/projects/clr/hipamd/src/hip_graph.cpp +++ b/projects/clr/hipamd/src/hip_graph.cpp @@ -85,7 +85,7 @@ hipError_t ihipGraphAddKernelNode(hip::GraphNode** pGraphNode, hip::Graph* graph hip::GraphNode* const* pDependencies, size_t numDependencies, const hipKernelNodeParams* pNodeParams, const ihipExtKernelEvents* pNodeEvents = nullptr, - bool capture = true) { + bool capture = true, int coopKernel = 0) { if (pGraphNode == nullptr || graph == nullptr || (numDependencies > 0 && pDependencies == nullptr) || pNodeParams == nullptr || pNodeParams->func == nullptr) { @@ -114,7 +114,7 @@ hipError_t ihipGraphAddKernelNode(hip::GraphNode** pGraphNode, hip::Graph* graph return hipErrorInvalidConfiguration; } - *pGraphNode = new hip::GraphKernelNode(pNodeParams, pNodeEvents); + *pGraphNode = new hip::GraphKernelNode(pNodeParams, pNodeEvents, coopKernel); status = ihipGraphAddNode(*pGraphNode, graph, pDependencies, numDependencies, capture); return status; } @@ -377,6 +377,38 @@ hipError_t capturehipLaunchByPtr(hipStream_t& stream, hipFunction_t func, dim3 b return hipSuccess; } +hipError_t capturehipLaunchCooperativeKernel(hipStream_t& stream, const void*& f, dim3& gridDim, + dim3& blockDim, void**& kernelParams, + uint32_t& sharedMemBytes) +{ + ClPrint(amd::LOG_INFO, amd::LOG_API, + "[hipGraph] Current capture node LaunchCooperativeKernel on stream : %p", stream); + if (!hip::isValid(stream)) { + return hipErrorContextIsDestroyed; + } + + hip::Stream* s = reinterpret_cast(stream); + hipKernelNodeParams nodeParams; + nodeParams.func = const_cast(f); + nodeParams.blockDim = blockDim; + nodeParams.gridDim = gridDim; + nodeParams.kernelParams = kernelParams; + nodeParams.sharedMemBytes = sharedMemBytes; + nodeParams.extra = nullptr; + + hip::GraphNode* pGraphNode; + hipError_t status = + ihipGraphAddKernelNode(&pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(), + s->GetLastCapturedNodes().size(), &nodeParams, nullptr, true, + amd::NDRangeKernelCommand::CooperativeGroups); + if (status != hipSuccess) { + return status; + } + s->SetLastCapturedNode(pGraphNode); + + return hipSuccess; +} + hipError_t capturehipMemcpy3DAsync(hipStream_t& stream, const hipMemcpy3DParms*& p) { ClPrint(amd::LOG_INFO, amd::LOG_API, "[hipGraph] Current capture node Memcpy3D on stream : %p", stream); diff --git a/projects/clr/hipamd/src/hip_graph_capture.hpp b/projects/clr/hipamd/src/hip_graph_capture.hpp index 5c02648f99..ac62bee750 100644 --- a/projects/clr/hipamd/src/hip_graph_capture.hpp +++ b/projects/clr/hipamd/src/hip_graph_capture.hpp @@ -45,6 +45,10 @@ hipError_t capturehipModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f, u hipError_t capturehipLaunchByPtr(hipStream_t& stream, hipFunction_t func, dim3 blockDim, dim3 gridDim, unsigned int sharedMemBytes, void** extra); +hipError_t capturehipLaunchCooperativeKernel(hipStream_t& stream, const void*& f, dim3& gridDim, + dim3& blockDim, void**& kernelParams, + uint32_t& sharedMemBytes); + hipError_t capturehipMemcpy2DAsync(hipStream_t& stream, void*& dst, size_t& dpitch, const void*& src, size_t& spitch, size_t& width, size_t& height, hipMemcpyKind& kind); diff --git a/projects/clr/hipamd/src/hip_graph_internal.hpp b/projects/clr/hipamd/src/hip_graph_internal.hpp index 682fb2776c..0d66c0550c 100644 --- a/projects/clr/hipamd/src/hip_graph_internal.hpp +++ b/projects/clr/hipamd/src/hip_graph_internal.hpp @@ -992,6 +992,7 @@ class GraphKernelNode : public GraphNode { unsigned int kernelAttrInUse_; //!< Kernel attributes in use ihipExtKernelEvents kernelEvents_; //!< Events for Ext launch kernel bool hasHiddenHeap_; //!< Kernel has hidden heap(device side allocation) + int coopKernel_; //!< Launch cooperative kernel public: bool HasHiddenHeap() const { return hasHiddenHeap_; } @@ -1183,7 +1184,8 @@ class GraphKernelNode : public GraphNode { return hipSuccess; } - GraphKernelNode(const hipKernelNodeParams* pNodeParams, const ihipExtKernelEvents* pEvents) + GraphKernelNode(const hipKernelNodeParams* pNodeParams, const ihipExtKernelEvents* pEvents, + int coopKernel = 0) : GraphNode(hipGraphNodeTypeKernel, "bold", "octagon", "KERNEL") { kernelParams_ = *pNodeParams; kernelEvents_ = { 0 }; @@ -1196,6 +1198,7 @@ class GraphKernelNode : public GraphNode { memset(&kernelAttr_, 0, sizeof(kernelAttr_)); kernelAttrInUse_ = 0; hasHiddenHeap_ = false; + coopKernel_ = coopKernel; } ~GraphKernelNode() { freeParams(); } @@ -1225,6 +1228,7 @@ class GraphKernelNode : public GraphNode { GraphKernelNode(const GraphKernelNode& rhs) : GraphNode(rhs) { kernelParams_ = rhs.kernelParams_; kernelEvents_ = rhs.kernelEvents_; + coopKernel_ = rhs.coopKernel_; hipError_t status = copyParams(&rhs.kernelParams_); if (status != hipSuccess) { ClPrint(amd::LOG_ERROR, amd::LOG_CODE, "[hipGraph] Failed to allocate memory to copy params"); @@ -1275,9 +1279,9 @@ class GraphKernelNode : public GraphNode { command, func, kernelParams_.gridDim.x * kernelParams_.blockDim.x, kernelParams_.gridDim.y * kernelParams_.blockDim.y, kernelParams_.gridDim.z * kernelParams_.blockDim.z, kernelParams_.blockDim.x, - kernelParams_.blockDim.y, kernelParams_.blockDim.z, kernelParams_.sharedMemBytes, - stream, kernelParams_.kernelParams, kernelParams_.extra, kernelEvents_.startEvent_, - kernelEvents_.stopEvent_, flags, 0, 0, 0, 0, 0, 0); + kernelParams_.blockDim.y, kernelParams_.blockDim.z, kernelParams_.sharedMemBytes, stream, + kernelParams_.kernelParams, kernelParams_.extra, kernelEvents_.startEvent_, + kernelEvents_.stopEvent_, flags, coopKernel_, 0, 0, 0, 0, 0); if (signal_is_required_) { // Optimize the barriers by adding a signal into the dispatch packet directly command->SetProfiling(); diff --git a/projects/clr/hipamd/src/hip_module.cpp b/projects/clr/hipamd/src/hip_module.cpp index 187e08bb86..65fa2c40b6 100644 --- a/projects/clr/hipamd/src/hip_module.cpp +++ b/projects/clr/hipamd/src/hip_module.cpp @@ -773,6 +773,9 @@ hipError_t hipLaunchCooperativeKernel_common(const void* f, dim3 gridDim, dim3 b return hipErrorContextIsDestroyed; } + STREAM_CAPTURE(hipLaunchCooperativeKernel, hStream, f, gridDim, blockDim, kernelParams, + sharedMemBytes); + if (f == nullptr) { return hipErrorInvalidDeviceFunction; }