diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp index dca2a3cdd1..0246742758 100644 --- a/hipamd/src/hip_graph.cpp +++ b/hipamd/src/hip_graph.cpp @@ -258,6 +258,36 @@ hipError_t capturehipExtModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f return hipSuccess; } +hipError_t capturehipModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f, uint32_t& gridDimX, + uint32_t& gridDimY, uint32_t& gridDimZ, uint32_t& blockDimX, + uint32_t& blockDimY, uint32_t& blockDimZ, + uint32_t& sharedMemBytes, void**& kernelParams, + void**& extra) { + ClPrint(amd::LOG_INFO, amd::LOG_API, + "[hipGraph] current capture node module launch kernel launch on stream : %p", stream); + if (!hip::isValid(stream)) { + return hipErrorInvalidValue; + } + hip::Stream* s = reinterpret_cast(stream); + hipKernelNodeParams nodeParams; + nodeParams.func = f; + nodeParams.blockDim = {blockDimX, blockDimY, blockDimZ}; + nodeParams.extra = extra; + nodeParams.gridDim = {gridDimX, gridDimY, gridDimZ}; + nodeParams.kernelParams = kernelParams; + nodeParams.sharedMemBytes = sharedMemBytes; + + hipGraphNode_t pGraphNode; + hipError_t status = + ihipGraphAddKernelNode(&pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(), + s->GetLastCapturedNodes().size(), &nodeParams); + 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/hipamd/src/hip_graph_capture.hpp b/hipamd/src/hip_graph_capture.hpp index d09041b26f..6650ec9a93 100644 --- a/hipamd/src/hip_graph_capture.hpp +++ b/hipamd/src/hip_graph_capture.hpp @@ -31,6 +31,12 @@ hipError_t capturehipExtModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f void**& extra, hipEvent_t& startEvent, hipEvent_t& stopEvent, uint32_t& flags); +hipError_t capturehipModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f, uint32_t& gridDimX, + uint32_t& gridDimY, uint32_t& gridDimZ, uint32_t& blockDimX, + uint32_t& blockDimY, uint32_t& blockDimZ, + uint32_t& sharedMemBytes, void**& kernelParams, + void**& extra); + 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/hipamd/src/hip_module.cpp b/hipamd/src/hip_module.cpp index fde15a4a18..69abc42193 100644 --- a/hipamd/src/hip_module.cpp +++ b/hipamd/src/hip_module.cpp @@ -396,6 +396,10 @@ hipError_t hipModuleLaunchKernel(hipFunction_t f, uint32_t gridDimX, uint32_t gr void** kernelParams, void** extra) { HIP_INIT_API(hipModuleLaunchKernel, f, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY, blockDimZ, sharedMemBytes, hStream, kernelParams, extra); + + STREAM_CAPTURE(hipModuleLaunchKernel, hStream, f, gridDimX, gridDimY, gridDimZ, blockDimX, + blockDimY, blockDimZ, sharedMemBytes, kernelParams, extra); + size_t globalWorkSizeX = static_cast(gridDimX) * blockDimX; size_t globalWorkSizeY = static_cast(gridDimY) * blockDimY; size_t globalWorkSizeZ = static_cast(gridDimZ) * blockDimZ;