From 8ced453e49d3cf57f8a713f13fc49bc0b513aa0d Mon Sep 17 00:00:00 2001 From: ethanwee1 Date: Wed, 8 Jul 2026 16:41:23 +0000 Subject: [PATCH] [ROCm] Add rocJPEG support for GPU JPEG decoding Backport of pytorch/vision#9342 adapted to the release/0.26 (classic, non stable-ABI) JPEG decoder layout. Adds a rocJPEG-backed GPU JPEG decode path for ROCm builds, gated by TORCHVISION_USE_ROCJPEG / ROCJPEG_FOUND, alongside the existing nvJPEG CUDA path. This enables decode_jpeg(..., device="cuda") on ROCm for UNCHANGED, GRAY and RGB modes using rocJPEG's hardware backend. rocJPEG uses ROCJPEG_BACKEND_HARDWARE, which rejects images smaller than 64x64 and 4:1:1/unknown chroma subsampling. --- setup.py | 16 +- test/test_image.py | 11 +- .../csrc/io/image/cuda/decode_jpegs_cuda.cpp | 198 +++++++++++++++++- .../csrc/io/image/cuda/decode_jpegs_cuda.h | 51 +++++ torchvision/io/image.py | 9 +- 5 files changed, 276 insertions(+), 9 deletions(-) diff --git a/setup.py b/setup.py index 7bf83fcca73..ba432d86db1 100644 --- a/setup.py +++ b/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), diff --git a/test/test_image.py b/test/test_image.py index b11dd67ca12..7dfa1fe7967 100644 --- a/test/test_image.py +++ b/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( diff --git a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp b/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp index 85aa6c760c1..7b64ccde4a4 100644 --- a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.cpp +++ b/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 +#include +#include +#include + +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; +} // namespace + +std::vector decode_jpegs_cuda( + const std::vector& 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 lock(rocDecoderMutex); + + TORCH_CHECK( + device.is_cuda(), "Expected the device parameter to be a cuda device"); + + std::vector 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(¤t_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(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 RocJpegDecoder::decode_images( + const std::vector& 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 decode_params(num_images); + std::vector output_images(num_images); + std::vector 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(), + 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(); + 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(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 decode_jpegs_cuda( diff --git a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.h b/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.h index 6f72d9e35b2..df1195d0835 100644 --- a/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.h +++ b/torchvision/csrc/io/image/cuda/decode_jpegs_cuda.h @@ -43,3 +43,54 @@ class CUDAJpegDecoder { } // namespace image } // namespace vision #endif + +#if ROCJPEG_FOUND +#include +#include + +// 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 decode_images( + const std::vector& encoded_images, + ImageReadMode mode); + + const torch::Device target_device; + + private: + std::vector 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 diff --git a/torchvision/io/image.py b/torchvision/io/image.py index 69e50d3fd58..2d0067c5d0e 100644 --- a/torchvision/io/image.py +++ b/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 `_. This is only - supported for CUDA version >= 10.1 + with `nvjpeg `_ on NVIDIA GPUs, + or with `rocJPEG + `_ 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.