From 664c4627bdc46e406960f654c7fe83d23586610b Mon Sep 17 00:00:00 2001 From: Michal Zientkiewicz Date: Thu, 2 Jun 2022 12:06:01 +0200 Subject: [PATCH] Make copy_to_external in framework iterators non-blocking. Signed-off-by: Michal Zientkiewicz --- dali/python/nvidia/dali/plugin/mxnet.py | 3 ++- dali/python/nvidia/dali/plugin/paddle.py | 4 ++-- dali/python/nvidia/dali/plugin/pytorch.py | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/dali/python/nvidia/dali/plugin/mxnet.py b/dali/python/nvidia/dali/plugin/mxnet.py index 3c94664996a..6c5a0faf367 100644 --- a/dali/python/nvidia/dali/plugin/mxnet.py +++ b/dali/python/nvidia/dali/plugin/mxnet.py @@ -69,7 +69,8 @@ def feed_ndarray(dali_tensor, arr, cuda_stream = None): # Copy data from DALI tensor to ptr if isinstance(dali_tensor, (TensorGPU, TensorListGPU)): - dali_tensor.copy_to_external(ptr, None if cuda_stream is None else ctypes.c_void_p(cuda_stream)) + stream = None if cuda_stream is None else ctypes.c_void_p(cuda_stream) + dali_tensor.copy_to_external(ptr, stream, non_blocking=True) else: dali_tensor.copy_to_external(ptr) diff --git a/dali/python/nvidia/dali/plugin/paddle.py b/dali/python/nvidia/dali/plugin/paddle.py index 22a22d8e85d..fa5a3fc72fe 100644 --- a/dali/python/nvidia/dali/plugin/paddle.py +++ b/dali/python/nvidia/dali/plugin/paddle.py @@ -69,8 +69,8 @@ def feed_ndarray(dali_tensor, ptr, cuda_stream=None): c_type_pointer = ctypes.c_void_p(ptr) if isinstance(dali_tensor, (TensorGPU, TensorListGPU)): - dali_tensor.copy_to_external( - c_type_pointer, None if cuda_stream is None else ctypes.c_void_p(cuda_stream)) + stream = None if cuda_stream is None else ctypes.c_void_p(cuda_stream) + dali_tensor.copy_to_external(c_type_pointer, stream, non_blocking=True) else: dali_tensor.copy_to_external(c_type_pointer) return ptr diff --git a/dali/python/nvidia/dali/plugin/pytorch.py b/dali/python/nvidia/dali/plugin/pytorch.py index 4d903f1c382..f281af248bd 100644 --- a/dali/python/nvidia/dali/plugin/pytorch.py +++ b/dali/python/nvidia/dali/plugin/pytorch.py @@ -62,7 +62,8 @@ def feed_ndarray(dali_tensor, arr, cuda_stream = None): # turn raw int to a c void pointer c_type_pointer = ctypes.c_void_p(arr.data_ptr()) if isinstance(dali_tensor, (TensorGPU, TensorListGPU)): - dali_tensor.copy_to_external(c_type_pointer, None if cuda_stream is None else ctypes.c_void_p(cuda_stream)) + stream = None if cuda_stream is None else ctypes.c_void_p(cuda_stream) + dali_tensor.copy_to_external(c_type_pointer, stream, non_blocking=True) else: dali_tensor.copy_to_external(c_type_pointer) return arr