提交 5e0e455d 编写于 作者: W wanghaoshuang

Add CUDA stream when launching kernel.

上级 0c05ea39
......@@ -88,7 +88,13 @@ void CropCUDAFunctoin(const framework::ExecutionContext& context) {
int d = out_dims[1];
int block = 512;
int grid = (n * d + block - 1) / block;
CropKernel<T, D><<<grid, block>>>(out_count, out_shape_gpu, x_shape_gpu,
auto* device_context =
const_cast<platform::DeviceContext*>(context.device_context_);
CropKernel<T,
D><<<grid, block, 0,
reinterpret_cast<platform::CUDADeviceContext*>(device_context)
->stream()>>>(out_count, out_shape_gpu, x_shape_gpu,
crop_rules_gpu, x_data, out_data);
cudaFree(crop_rules_gpu);
cudaFree(x_shape_gpu);
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册