[ROCm] Add rocJPEG support for GPU JPEG decoding (backport of pytorch/vision#9342)

setup.py

@@ -22,6 +22,7 @@
 USE_JPEG = os.getenv("TORCHVISION_USE_JPEG", "1") == "1"
 USE_WEBP = os.getenv("TORCHVISION_USE_WEBP", "1") == "1"
 USE_NVJPEG = os.getenv("TORCHVISION_USE_NVJPEG", "1") == "1"
+USE_ROCJPEG = os.getenv("TORCHVISION_USE_ROCJPEG", "1") == "1"
 NVCC_FLAGS = os.getenv("NVCC_FLAGS", None)
 
 TORCHVISION_INCLUDE = os.environ.get("TORCHVISION_INCLUDE", "")
@@ -44,6 +45,7 @@
 print(f"{USE_JPEG = }")
 print(f"{USE_WEBP = }")
 print(f"{USE_NVJPEG = }")
+print(f"{USE_ROCJPEG = }")
 print(f"{NVCC_FLAGS = }")
 print(f"{TORCHVISION_INCLUDE = }")
 print(f"{TORCHVISION_LIBRARY = }")
@@ -340,7 +342,7 @@ def make_image_extension():
         else:
             warnings.warn("Building torchvision without WEBP support")
 
-    if USE_NVJPEG and (torch.cuda.is_available() or FORCE_CUDA):
+    if USE_NVJPEG and not IS_ROCM and (torch.cuda.is_available() or FORCE_CUDA):
         nvjpeg_found = CUDA_HOME is not None and (Path(CUDA_HOME) / "include/nvjpeg.h").exists()
 
         if nvjpeg_found:
@@ -350,9 +352,19 @@ def make_image_extension():
             Extension = CUDAExtension
         else:
             warnings.warn("Building torchvision without NVJPEG support")
-    elif USE_NVJPEG:
+    elif USE_NVJPEG and not IS_ROCM:
         warnings.warn("Building torchvision without NVJPEG support")
 
+    if USE_ROCJPEG and IS_ROCM and (torch.cuda.is_available() or FORCE_CUDA):
+        rocjpeg_found = ROCM_HOME is not None and (Path(ROCM_HOME) / "include/rocjpeg/rocjpeg.h").exists()
+        if rocjpeg_found:
+            print("Building torchvision with ROCJPEG image support")
+            libraries.append("rocjpeg")
+            define_macros += [("ROCJPEG_FOUND", 1)]
+            Extension = CUDAExtension
+        else:
+            warnings.warn("Building torchvision without ROCJPEG support")
+
     return Extension(
         name="torchvision.image",
         sources=sorted(str(s) for s in sources),

test/test_image.py

@@ -43,6 +43,7 @@
 IS_WINDOWS = sys.platform in ("win32", "cygwin")
 IS_MACOS = sys.platform == "darwin"
 IS_LINUX = sys.platform == "linux"
+IS_ROCM = torch.version.hip is not None
 PILLOW_VERSION = tuple(int(x) for x in PILLOW_VERSION.split("."))
 WEBP_TEST_IMAGES_DIR = os.environ.get("WEBP_TEST_IMAGES_DIR", "")
 # See https://github.com/pytorch/vision/pull/8724#issuecomment-2503964558
@@ -426,12 +427,13 @@ def test_decode_jpegs_cuda(mode, scripted):
         futures = [executor.submit(decode_fn, encoded_images, mode, "cuda") for _ in range(num_workers)]
     decoded_images_threaded = [future.result() for future in futures]
     assert len(decoded_images_threaded) == num_workers
+    tol = 2.5 if IS_ROCM else 2
     for decoded_images in decoded_images_threaded:
         assert len(decoded_images) == len(encoded_images)
         for decoded_image_cuda, decoded_image_cpu in zip(decoded_images, decoded_images_cpu):
             assert decoded_image_cuda.shape == decoded_image_cpu.shape
             assert decoded_image_cuda.dtype == decoded_image_cpu.dtype == torch.uint8
-            assert (decoded_image_cuda.cpu().float() - decoded_image_cpu.cpu().float()).abs().mean() < 2
+            assert (decoded_image_cuda.cpu().float() - decoded_image_cpu.cpu().float()).abs().mean() < tol
 
 
 @needs_cuda
@@ -576,6 +578,7 @@ def test_encode_jpeg(img_path, scripted):
 
 
 @needs_cuda
+@pytest.mark.skipif(IS_ROCM, reason="rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
 def test_encode_jpeg_cuda_device_param():
     path = next(path for path in get_images(IMAGE_ROOT, ".jpg") if "cmyk" not in path)
 
@@ -596,6 +599,7 @@ def test_encode_jpeg_cuda_device_param():
 
 
 @needs_cuda
+@pytest.mark.skipif(IS_ROCM, reason="rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
 @pytest.mark.parametrize(
     "img_path",
     [pytest.param(jpeg_path, id=_get_safe_image_name(jpeg_path)) for jpeg_path in get_images(IMAGE_ROOT, ".jpg")],
@@ -625,6 +629,7 @@ def test_encode_jpeg_cuda(img_path, scripted, contiguous):
 
 
 @needs_cuda
+@pytest.mark.skipif(IS_ROCM, reason="rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
 def test_encode_jpeg_cuda_sync():
     """
     Non-regression test for https://github.com/pytorch/vision/issues/8587.
@@ -666,6 +671,8 @@ def test_encode_jpeg_cuda_sync():
 def test_encode_jpegs_batch(scripted, contiguous, device):
     if device == "cpu" and IS_MACOS:
         pytest.skip("https://github.com/pytorch/vision/issues/8031")
+    if device == "cuda" and IS_ROCM:
+        pytest.skip("rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
     decoded_images_tv = []
     for jpeg_path in get_images(IMAGE_ROOT, ".jpg"):
         if "cmyk" in jpeg_path:
@@ -711,6 +718,7 @@ def test_encode_jpegs_batch(scripted, contiguous, device):
 
 
 @needs_cuda
+@pytest.mark.skipif(IS_ROCM, reason="rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
 def test_single_encode_jpeg_cuda_errors():
     with pytest.raises(RuntimeError, match="Input tensor dtype should be uint8"):
         encode_jpeg(torch.empty((3, 100, 100), dtype=torch.float32, device="cuda"))
@@ -729,6 +737,7 @@ def test_single_encode_jpeg_cuda_errors():
 
 
 @needs_cuda
+@pytest.mark.skipif(IS_ROCM, reason="rocJPEG is decode-only; GPU JPEG encoding is not supported on ROCm")
 def test_batch_encode_jpegs_cuda_errors():
     with pytest.raises(RuntimeError, match="Input tensor dtype should be uint8"):
         encode_jpeg(

torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp

@@ -1,5 +1,201 @@
 #include "decode_jpegs_cuda.h"
-#if !NVJPEG_FOUND
+
+#if ROCJPEG_FOUND
+#include <cstdlib>
+#include <exception>
+#include <memory>
+#include <mutex>
+
+namespace vision {
+namespace image {
+
+namespace {
+uint32_t align_up(uint32_t value) {
+  constexpr uint32_t kRocJpegPitchAlignment = 16;
+  return (value + kRocJpegPitchAlignment - 1) & ~(kRocJpegPitchAlignment - 1);
+}
+
+std::mutex rocDecoderMutex;
+std::unique_ptr<RocJpegDecoder> rocJpegDecoder;
+} // namespace
+
+std::vector<torch::Tensor> decode_jpegs_cuda(
+    const std::vector<torch::Tensor>& encoded_images,
+    vision::image::ImageReadMode mode,
+    torch::Device device) {
+  C10_LOG_API_USAGE_ONCE(
+      "torchvision.csrc.io.image.cuda.decode_jpegs_cuda.decode_jpegs_cuda");
+
+  std::lock_guard<std::mutex> lock(rocDecoderMutex);
+
+  TORCH_CHECK(
+      device.is_cuda(), "Expected the device parameter to be a cuda device");
+
+  std::vector<torch::Tensor> contig_images;
+  contig_images.reserve(encoded_images.size());
+  for (auto& encoded_image : encoded_images) {
+    TORCH_CHECK(
+        encoded_image.dtype() == torch::kU8, "Expected a torch.uint8 tensor");
+    TORCH_CHECK(
+        !encoded_image.is_cuda(),
+        "The input tensor must be on CPU when decoding with rocJPEG");
+    TORCH_CHECK(
+        encoded_image.dim() == 1 && encoded_image.numel() > 0,
+        "Expected a non empty 1-dimensional tensor");
+    // rocJPEG requires contiguous input; contiguous() is a no-op when it
+    // already is.
+    contig_images.push_back(encoded_image.contiguous());
+  }
+
+  torch::Device target_device = device;
+  if (!target_device.has_index()) {
+    int current_device = 0;
+    CHECK_HIP(hipGetDevice(&current_device));
+    target_device = torch::Device(torch::kCUDA, current_device);
+  }
+  CHECK_HIP(hipSetDevice(target_device.index()));
+
+  if (rocJpegDecoder == nullptr ||
+      target_device != rocJpegDecoder->target_device) {
+    if (rocJpegDecoder != nullptr) {
+      rocJpegDecoder.reset(new RocJpegDecoder(target_device));
+    } else {
+      rocJpegDecoder = std::make_unique<RocJpegDecoder>(target_device);
+      std::atexit([]() { rocJpegDecoder.reset(); });
+    }
+  }
+
+  try {
+    return rocJpegDecoder->decode_images(contig_images, mode);
+  } catch (const std::exception& e) {
+    TORCH_CHECK(false, "Error while decoding JPEG images: ", e.what());
+  }
+}
+
+RocJpegDecoder::RocJpegDecoder(const torch::Device& target_device)
+    : target_device{target_device} {
+  CHECK_HIP(hipSetDevice(target_device.index()));
+  CHECK_ROCJPEG(rocJpegCreate(
+      ROCJPEG_BACKEND_HARDWARE, target_device.index(), &rocjpeg_handle_));
+}
+
+RocJpegDecoder::~RocJpegDecoder() {
+  rocJpegDestroy(rocjpeg_handle_);
+  for (auto stream_handle : rocjpeg_stream_handles_) {
+    rocJpegStreamDestroy(stream_handle);
+  }
+}
+
+std::vector<torch::Tensor> RocJpegDecoder::decode_images(
+    const std::vector<torch::Tensor>& encoded_images,
+    vision::image::ImageReadMode mode) {
+  const std::size_t num_images = encoded_images.size();
+  // Reuse existing rocJPEG stream handles and create only the missing ones.
+  while (rocjpeg_stream_handles_.size() < num_images) {
+    RocJpegStreamHandle stream_handle;
+    CHECK_ROCJPEG(rocJpegStreamCreate(&stream_handle));
+    rocjpeg_stream_handles_.push_back(stream_handle);
+  }
+
+  std::vector<RocJpegDecodeParams> decode_params(num_images);
+  std::vector<RocJpegImage> output_images(num_images);
+  std::vector<torch::Tensor> output_tensors;
+  output_tensors.reserve(num_images);
+
+  for (std::size_t i = 0; i < num_images; ++i) {
+    CHECK_ROCJPEG(rocJpegStreamParse(
+        encoded_images[i].data_ptr<uint8_t>(),
+        encoded_images[i].numel(),
+        rocjpeg_stream_handles_[i]));
+
+    uint8_t num_components = 0;
+    RocJpegChromaSubsampling subsampling = ROCJPEG_CSS_UNKNOWN;
+    uint32_t widths[ROCJPEG_MAX_COMPONENT] = {};
+    uint32_t heights[ROCJPEG_MAX_COMPONENT] = {};
+    CHECK_ROCJPEG(rocJpegGetImageInfo(
+        rocjpeg_handle_,
+        rocjpeg_stream_handles_[i],
+        &num_components,
+        &subsampling,
+        widths,
+        heights));
+
+    const uint32_t width = widths[0];
+    const uint32_t height = heights[0];
+    TORCH_CHECK(
+        width >= 64 && height >= 64,
+        "Image resolution ",
+        width,
+        "x",
+        height,
+        " is below the rocJPEG hardware JPEG decoder minimum of 64x64");
+    TORCH_CHECK(
+        subsampling != ROCJPEG_CSS_411 && subsampling != ROCJPEG_CSS_UNKNOWN,
+        "The image chroma subsampling is not supported by the rocJPEG hardware JPEG decoder");
+
+    RocJpegOutputFormat image_output_format;
+    uint32_t num_channels;
+    switch (mode) {
+      case vision::image::IMAGE_READ_MODE_UNCHANGED:
+        // torchvision's UNCHANGED mode is expected to match the CPU/nvJPEG
+        // behavior: grayscale JPEGs return one channel, while color JPEGs
+        // return RGB.
+        if (num_components == 1) {
+          image_output_format = ROCJPEG_OUTPUT_Y;
+          num_channels = 1;
+        } else {
+          image_output_format = ROCJPEG_OUTPUT_RGB_PLANAR;
+          num_channels = 3;
+        }
+        break;
+      case vision::image::IMAGE_READ_MODE_GRAY:
+        image_output_format = ROCJPEG_OUTPUT_Y;
+        num_channels = 1;
+        break;
+      case vision::image::IMAGE_READ_MODE_RGB:
+        image_output_format = ROCJPEG_OUTPUT_RGB_PLANAR;
+        num_channels = 3;
+        break;
+      default:
+        TORCH_CHECK(
+            false,
+            "The provided mode is not supported for JPEG decoding on GPU");
+    }
+
+    // rocJPEG writes rows at a 16-byte-aligned pitch, so allocate a buffer
+    // padded to that alignment and return a view of the valid region.
+    uint32_t pitch = align_up(width);
+    auto buffer = torch::empty(
+        {int64_t(num_channels), int64_t(align_up(height)), int64_t(pitch)},
+        torch::dtype(torch::kU8).device(target_device));
+
+    decode_params[i].output_format = image_output_format;
+    for (uint32_t c = 0; c < num_channels; ++c) {
+      output_images[i].channel[c] = buffer[c].data_ptr<uint8_t>();
+      output_images[i].pitch[c] = pitch;
+    }
+    auto valid_height = buffer.narrow(1, 0, height);
+    output_tensors.push_back(valid_height.narrow(2, 0, width));
+  }
+
+  // Choosing a batch size that is a multiple of the available JPEG cores is
+  // recommended.
+  CHECK_ROCJPEG(rocJpegDecodeBatched(
+      rocjpeg_handle_,
+      rocjpeg_stream_handles_.data(),
+      static_cast<int>(num_images),
+      decode_params.data(),
+      output_images.data()));
+  // rocJpegDecodeBatched synchronizes rocJPEG's internal HIP stream before
+  // returning, so the decoded output tensors are ready for PyTorch streams.
+
+  return output_tensors;
+}
+
+} // namespace image
+} // namespace vision
+
+#elif !NVJPEG_FOUND
 namespace vision {
 namespace image {
 std::vector<torch::Tensor> decode_jpegs_cuda(

torchvision/csrc/io/image/cuda/decode_jpegs_cuda.h

@@ -43,3 +43,54 @@ class CUDAJpegDecoder {
 } // namespace image
 } // namespace vision
 #endif
+
+#if ROCJPEG_FOUND
+#include <hip/hip_runtime.h>
+#include <rocjpeg/rocjpeg.h>
+
+// rocJPEG decode API documentation:
+// https://rocm.docs.amd.com/projects/rocJPEG/en/latest/how-to/rocjpeg-decoding-a-jpeg-stream.html
+
+namespace vision {
+namespace image {
+
+class RocJpegDecoder {
+ public:
+  RocJpegDecoder(const torch::Device& target_device);
+  ~RocJpegDecoder();
+
+  std::vector<torch::Tensor> decode_images(
+      const std::vector<torch::Tensor>& encoded_images,
+      ImageReadMode mode);
+
+  const torch::Device target_device;
+
+ private:
+  std::vector<RocJpegStreamHandle> rocjpeg_stream_handles_;
+  RocJpegHandle rocjpeg_handle_;
+};
+
+} // namespace image
+} // namespace vision
+
+#define CHECK_ROCJPEG(call)                       \
+  {                                               \
+    RocJpegStatus rocjpeg_status = (call);        \
+    TORCH_CHECK(                                   \
+        rocjpeg_status == ROCJPEG_STATUS_SUCCESS, \
+        #call,                                    \
+        " returned ",                             \
+        rocJpegGetErrorName(rocjpeg_status));     \
+  }
+
+#define CHECK_HIP(call)               \
+  {                                   \
+    hipError_t hip_status = (call);   \
+    TORCH_CHECK(                       \
+        hip_status == hipSuccess,     \
+        #call,                        \
+        " failed with status: ",      \
+        hipGetErrorName(hip_status)); \
+  }
+
+#endif

torchvision/io/image.py

@@ -189,14 +189,13 @@ def decode_jpeg(
             for available modes.
         device (str or torch.device): The device on which the decoded image will
             be stored. If a cuda device is specified, the image will be decoded
-            with `nvjpeg <https://developer.nvidia.com/nvjpeg>`_. This is only
-            supported for CUDA version >= 10.1
+            with `nvjpeg <https://developer.nvidia.com/nvjpeg>`_ on NVIDIA GPUs,
+            or with `rocJPEG
+            <https://rocm.docs.amd.com/projects/rocJPEG/en/latest/>`_ on AMD
+            (ROCm) GPUs. On both, pass ``device="cuda"``.
 
             .. betastatus:: device parameter
 
-            .. warning::
-                There is a memory leak in the nvjpeg library for CUDA versions < 11.6.
-                Make sure to rely on CUDA 11.6 or above before using ``device="cuda"``.
         apply_exif_orientation (bool): apply EXIF orientation transformation to the output tensor.
             Default: False. Only implemented for JPEG format on CPU.