From 4b65ecbb55e611e91e97456b124e4204ec4d6047 Mon Sep 17 00:00:00 2001 From: TFBunny Date: Thu, 11 Mar 2021 15:33:28 -0500 Subject: [PATCH] attach memcpy to a stream in topk --- .../backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.h | 9 +++++++-- .../backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu | 6 ++---- .../backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh | 2 +- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.h b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.h index 3c23dc828e..4d806561c5 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.h +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.h @@ -42,8 +42,13 @@ class TopKGpuKernel : public GpuKernel { T *output_addr = GetDeviceAddress(outputs, 0); S *indices = GetDeviceAddress(outputs, 1); const T init_k = std::numeric_limits::lowest(); - - FastTopK(outer_size_, inner_size_, input_addr, k, output_addr, indices, init_k, + S k_cut = 0; + CHECK_CUDA_RET_WITH_EXCEPT( + kernel_node_, + cudaMemcpyAsync(&k_cut, k, sizeof(S), cudaMemcpyDeviceToHost, reinterpret_cast(stream_ptr)), + "cudaMemcpyAsync k_cut failed"); + CHECK_CUDA_RET_WITH_EXCEPT(kernel_node_, cudaDeviceSynchronize(), "cudaDeviceSyncFailed - TopK"); + FastTopK(outer_size_, inner_size_, input_addr, k_cut, output_addr, indices, init_k, reinterpret_cast(stream_ptr)); return true; } diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu index 6976c46dd2..5a57405260 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cu @@ -204,11 +204,9 @@ __global__ void TopKBlock(int outer_size, int inner_size, const T *input, T *out } template -void FastTopK(const int outer_size, const int inner_size, const T *input, const S *k, T *output, S *output_index, +void FastTopK(const int outer_size, const int inner_size, const T *input, S k_cut, T *output, S *output_index, const T init_K, cudaStream_t stream) { int block_num_limit = outer_size < 128 ? outer_size : 128; - S k_cut = 0; - cudaMemcpy(&k_cut, k, sizeof(S), cudaMemcpyDeviceToHost); if (k_cut > inner_size) k_cut = inner_size; if (k_cut <= 32) { @@ -223,5 +221,5 @@ void FastTopK(const int outer_size, const int inner_size, const T *input, const } } -template void FastTopK(const int outer_size, const int inner_size, const float *input, const int *k, float *output, +template void FastTopK(const int outer_size, const int inner_size, const float *input, int k_cut, float *output, int *output_index, const float init_K, cudaStream_t stream); diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh index 7ecb392c1b..46210a0337 100644 --- a/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh +++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/cuda_impl/topk_impl.cuh @@ -21,7 +21,7 @@ #include "runtime/device/gpu/cuda_common.h" template -void FastTopK(const int outer, const int inner, const T *input_addr, const S *k, T *output, S *indices, const T initK, +void FastTopK(const int outer, const int inner, const T *input_addr, S k_cut, T *output, S *indices, const T initK, cudaStream_t stream); #endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_GPU_CUDA_IMPL_TOPK_IMPL_CUH_