From 216f03d72571bc6f446ace6fcb582f80b626b124 Mon Sep 17 00:00:00 2001
From: karkuspeter <karkuspeter@gmail.com>
Date: Mon, 17 Aug 2026 14:48:28 +0200
Subject: [PATCH] Wait for the caller's stream before the nvJPEG hardware
 decode writes

nvJPEG's hardware engine can write its destination buffers before work that
was already queued on the stream it was given has run. Those destinations
come from the caching allocator, which hands out a block as soon as it is
freed on the stream that owns it, so the decoder can be handed memory that
queued kernels are still reading and silently corrupt them.

Host-synchronize the stream before the batched hardware decode, which costs
no wall time because decode_images() already synchronizes it before
returning, and add a regression test.
---
 src/torchcodec/_core/DecodeJpegCuda.cpp | 16 ++++++++++
 test/test_decoders.py                   | 40 +++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/src/torchcodec/_core/DecodeJpegCuda.cpp b/src/torchcodec/_core/DecodeJpegCuda.cpp
index f066ae6d0..58264f2b3 100644
--- a/src/torchcodec/_core/DecodeJpegCuda.cpp
+++ b/src/torchcodec/_core/DecodeJpegCuda.cpp
@@ -373,6 +373,22 @@ void CUDAJpegDecoder::decode_batched_hardware(
     formats.push_back(output_format);
   }
 
+  // nvJPEG's hardware engine does not honour the stream it is given: with work
+  // already queued on `stream`, nvjpegDecodeBatched() can write its
+  // destinations before that work has run. Those destinations are the outputs
+  // allocated above, which come from the caching allocator, so a block that was
+  // just freed on this stream can be handed to the decoder while kernels
+  // reading it are still queued, and an early write silently corrupts them.
+  // Waiting on the device (cudaStreamWaitEvent) is not enough, only a host
+  // barrier is. It costs no wall time here, because decode_images() already
+  // host-synchronizes `stream` before returning. The software path below is
+  // unaffected. Remove once nvJPEG orders the write itself.
+  cudaError_t presync_status = cudaStreamSynchronize(stream);
+  STD_TORCH_CHECK(
+      presync_status == cudaSuccess,
+      "Failed to synchronize CUDA stream: ",
+      presync_status);
+
   // The batch nvjpeg API only support a single output format per call, but we
   // may want both grayscale and RGB images here. So we need to split the input
   // into two groups and decode them separately. To be safe, we add  stream sync
diff --git a/test/test_decoders.py b/test/test_decoders.py
index 3693a501c..5d9ca11eb 100644
--- a/test/test_decoders.py
+++ b/test/test_decoders.py
@@ -4714,6 +4714,46 @@ def test_cuda_jpeg_multithreaded(self):
                 assert got.shape == ref.shape
                 assert_tensor_close_on_at_least(got.cpu(), ref, percentage=99, atol=3)
 
+    @needs_cuda
+    @needs_jpeg
+    def test_cuda_jpeg_waits_for_callers_stream(self):
+        # nvJPEG's hardware engine can write its destination before work that was
+        # already queued on the stream it was given has run. The outputs come from
+        # the caching allocator, which hands out a block as soon as it is freed on
+        # the stream that owns it, so a decode can be handed memory that queued
+        # kernels are still reading.
+        #
+        # Here the reductions are queued behind a long-running kernel and their
+        # inputs are freed while they are still queued, so the decoder is handed
+        # that memory. The sums are exact integers, and are only wrong if the
+        # decode wrote early.
+        num_bytes = 3 * GRADIENT_JPEG.height * GRADIENT_JPEG.width  # the output's size
+        values = (1, 2, 3, 4)
+        expected = float(sum(value * num_bytes for value in values))
+
+        landed_on_freed = 0
+        for _ in range(5):
+            victims = [
+                torch.full((num_bytes,), value, dtype=torch.uint8, device="cuda")
+                for value in values
+            ]
+            pointers = {victim.data_ptr() for victim in victims}
+            torch.cuda.synchronize()  # the victims' contents are resident
+
+            torch.cuda._sleep(200_000_000)  # keep the stream busy for ~100ms
+            total = torch.zeros((), dtype=torch.float64, device="cuda")
+            for victim in victims:
+                total = total + victim.sum(dtype=torch.float64)
+            victims.clear()  # freed while their reductions are still queued
+
+            decoded = decode_jpeg(GRADIENT_JPEG.path, device="cuda")
+            landed_on_freed += decoded.data_ptr() in pointers
+            torch.cuda.synchronize()
+            assert float(total) == expected
+
+        if not landed_on_freed:
+            pytest.skip("the allocator never reused a freed block, nothing was tested")
+
     # ===== PNG =====
 
     @needs_png