Workaround race condition in NVJPEG2k.

Signed-off-by: Michal Zientkiewicz <michalz@nvidia.com>

Author: mzient

dali/imgcodec/decoders/nvjpeg2k/nvjpeg2k.cc

@@ -158,19 +158,6 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
     return check_status(ret, in);
   } else {
     auto &image_info = ctx.image_info;
-    int num_tiles = image_info.num_tiles_y * image_info.num_tiles_x;
-    int tile_res_used = std::min<int>(ctx.tile_dec_res.size(), num_tiles);
-    SmallVector<nvjpeg2kStream_t, kNumParallelTiles> tile_jp2_stream;
-    SmallVector<NvJpeg2kStream, kNumParallelTiles> tile_stream_storage;
-    tile_jp2_stream.resize(tile_res_used);
-    tile_stream_storage.resize(tile_res_used);
-    tile_jp2_stream[0] = ctx.nvjpeg2k_stream;
-    for (int i = 1; i < tile_res_used; i++) {
-      tile_stream_storage[i] = NvJpeg2kStream::Create();
-      tile_jp2_stream[i] = tile_stream_storage[i];
-      CUDA_CALL(nvjpeg2kStreamParse(nvjpeg2k_handle_, in->RawData<uint8_t>(), in->Size(),
-                                    0, 0, tile_jp2_stream[i]));
-    }
 
     // Decode tile by tile: nvjpeg2kDecodeImage seems to be bugged
     auto &roi = ctx.roi;
@@ -201,7 +188,6 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
           const TileDecodingResources &per_tile_ctx = ctx.tile_dec_res[state_idx];
 
           CUDA_CALL(cudaEventSynchronize(per_tile_ctx.decode_event));
-          //cudaStreamSynchronize(ctx.cuda_stream);
 
           auto &params = per_tile_ctx.params;
           CUDA_CALL(nvjpeg2kDecodeParamsSetDecodeArea(params, begin_x, end_x, begin_y, end_y));
@@ -211,7 +197,7 @@ bool NvJpeg2000DecoderInstance::DecodeJpeg2000(ImageSource *in, void *out, const
 
           auto ret = nvjpeg2kDecodeTile(nvjpeg2k_handle_,
                                         per_tile_ctx.state,
-                                        tile_jp2_stream[state_idx],
+                                        ctx.nvjpeg2k_stream,
                                         params,
                                         tile_x + tile_y * image_info.num_tiles_x,
                                         0,

dali/imgcodec/decoders/nvjpeg2k/nvjpeg2k.h

@@ -81,14 +81,13 @@ class DLL_PUBLIC NvJpeg2000DecoderInstance : public BatchParallelDecoderImpl {
     CUDAEvent decode_event;
     NvJpeg2kDecodeParams params;
 
-    explicit TileDecodingResources(const NvJpeg2kHandle &nvjpeg2k_handle, int device_id,
-                                   cudaStream_t cuda_stream)
+    explicit TileDecodingResources(const NvJpeg2kHandle &nvjpeg2k_handle, int device_id)
         : state(nvjpeg2k_handle), decode_event(CUDAEvent::Create(device_id)) {
-      CUDA_CALL(cudaEventRecord(decode_event, cuda_stream));
     }
   };
 
-  static constexpr int kNumParallelTiles = 10;
+  static constexpr int kNumParallelTiles = 1;   // TODO(michalz): Use a different memory resource
+  // to allow parallel processing of tiles - existing one isn't stream-safe
   struct PerThreadResources {
     PerThreadResources() = default;
     PerThreadResources(const NvJpeg2kHandle &nvjpeg2k_handle,
@@ -99,11 +98,10 @@ class DLL_PUBLIC NvJpeg2000DecoderInstance : public BatchParallelDecoderImpl {
     , decode_event(CUDAEvent::Create(device_id))
     , cuda_stream(CUDAStreamPool::instance().Get(device_id)) {
       intermediate_buffer.resize(device_memory_padding / 8);
-      CUDA_CALL(cudaEventRecord(decode_event, cuda_stream));
 
       tile_dec_res.reserve(kNumParallelTiles);
       for (int i = 0; i < kNumParallelTiles; i++) {
-        tile_dec_res.emplace_back(nvjpeg2k_handle, device_id, cuda_stream);
+        tile_dec_res.emplace_back(nvjpeg2k_handle, device_id);
       }
     }
 

dali/test/python/test_dali_variable_batch_size.py

@@ -994,27 +994,29 @@ def peek_image_shape_pipe(module, max_batch_size, input_data, device):
         pipe.set_outputs(shape)
         return pipe
 
-    image_decoder_extensions = ['.jp2'] * 1000
+    image_decoder_extensions = ['.jpg', '.bmp', '.png', '.pnm', '.jp2']
     image_decoder_pipes = [
-        #image_decoder_pipe,
+        image_decoder_pipe,
         image_decoder_crop_pipe,
         image_decoder_slice_pipe,
     ]
 
     data_path = os.path.join(test_utils.get_dali_extra_path(), 'db', 'single')
     for ext in image_decoder_extensions:
         for pipe_template in image_decoder_pipes:
+            pipe = partial(pipe_template, fn.decoders)
+            yield test_decoders_check, pipe, data_path, ext, ['cpu', 'mixed']
             pipe = partial(pipe_template, fn.experimental.decoders)
-            yield test_decoders_check, pipe, data_path + '/jpeg2k/2', ext, ['mixed']
-        #pipe = partial(image_decoder_rcrop_pipe, fn.decoders)
-        #yield test_decoders_run, pipe, data_path, ext, ['cpu', 'mixed']
-        #pipe = partial(image_decoder_rcrop_pipe, fn.experimental.decoders)
-        #yield test_decoders_run, pipe, data_path, ext, ['cpu', 'mixed']
-
-    #pipe = partial(peek_image_shape_pipe, fn)
-    #yield test_decoders_check, pipe, data_path, '.jpg', ['cpu']
-    #pipe = partial(peek_image_shape_pipe, fn.experimental)
-    #yield test_decoders_check, pipe, data_path, '.jpg', ['cpu']
+            yield test_decoders_check, pipe, data_path, ext, ['cpu', 'mixed']
+        pipe = partial(image_decoder_rcrop_pipe, fn.decoders)
+        yield test_decoders_run, pipe, data_path, ext, ['cpu', 'mixed']
+        pipe = partial(image_decoder_rcrop_pipe, fn.experimental.decoders)
+        yield test_decoders_run, pipe, data_path, ext, ['cpu', 'mixed']
+
+    pipe = partial(peek_image_shape_pipe, fn)
+    yield test_decoders_check, pipe, data_path, '.jpg', ['cpu']
+    pipe = partial(peek_image_shape_pipe, fn.experimental)
+    yield test_decoders_check, pipe, data_path, '.jpg', ['cpu']
 
 
 def test_python_function():