Make copy_to_external in framework iterators non-blocking.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>
NVIDIA · Jun 2, 2022 · 664c462 · 664c462
1 parent 73203a6
commit 664c462
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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