Add ROCJPEG decoder

Author: NicolasHug

.github/workflows/linux_rocm.yaml

@@ -56,10 +56,10 @@ jobs:
       package-name: torchcodec
       trigger-event: ${{ github.event_name }}
       build-platform: "python-build-package"
-      # CPU-only compile: we build against torch-ROCm, but no ROCm/GPU decode
-      # sources exist yet (rocJPEG is a follow-up), so ENABLE_CUDA/ENABLE_ROCM
-      # are intentionally unset and this produces a CPU-only wheel.
-      build-command: "BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 python -m build --wheel -vvv --no-isolation"
+      # ENABLE_ROCM=1 compiles the rocJPEG GPU JPEG decoder into
+      # libtorchcodec_image.so (the ROCm counterpart of nvJPEG). Video decoding
+      # stays on CPU. The FFmpeg core libs are still built against all FFmpeg.
+      build-command: "BUILD_AGAINST_ALL_FFMPEG_FROM_S3=1 ENABLE_ROCM=1 python -m build --wheel -vvv --no-isolation"
 
   install-and-test:
     # ROCm runners are not on EC2 and authenticate to AWS (ECR) via GitHub OIDC.

packaging/repair_wheel.py

@@ -168,6 +168,13 @@ def repair_linux(wheels):
         "libnvshmem*",
         "libnvfatbin*",
         "libnvcuvid*",
+        # rocJPEG: the GPU JPEG decoder our image lib links on ROCm. Unlike
+        # nvJPEG (which we bundle), rocJPEG is not shipped by the torch-ROCm
+        # wheel, and bundling it would drag in torch's ROCm libs under mismatched
+        # (hashed) sonames. So we treat it as a runtime dependency provided by the
+        # ROCm install, like FFmpeg. decode_jpeg(device='cuda') therefore needs
+        # ROCm (with rocJPEG) present at runtime.
+        "librocjpeg*",
         # ROCm/HIP runtime and its system deps: provided by the torch-ROCm wheel
         # (torch/lib/) at runtime, exactly like the CUDA libs above. Never bundle
         # them — they'd duplicate torch's copies and bloat the wheel.

pyproject.toml

@@ -51,6 +51,7 @@ location = "source"
 
 [tool.scikit-build.cmake.define]
 ENABLE_CUDA = {env = "ENABLE_CUDA", default = ""}
+ENABLE_ROCM = {env = "ENABLE_ROCM", default = ""}
 TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR = {env = "TORCHCODEC_DISABLE_COMPILE_WARNING_AS_ERROR", default = "OFF"}
 TORCHCODEC_DISABLE_HOMEBREW_RPATH = {env = "TORCHCODEC_DISABLE_HOMEBREW_RPATH", default = "OFF"}
 
@@ -65,6 +66,7 @@ TORCHCODEC_BUILD_AVIF = {env = "TORCHCODEC_BUILD_AVIF", default = "AUTO"}
 TORCHCODEC_BUILD_GIF = {env = "TORCHCODEC_BUILD_GIF", default = "AUTO"}
 TORCHCODEC_BUILD_HEIC = {env = "TORCHCODEC_BUILD_HEIC", default = "AUTO"}
 TORCHCODEC_BUILD_NVJPEG = {env = "TORCHCODEC_BUILD_NVJPEG", default = "AUTO"}
+TORCHCODEC_BUILD_ROCJPEG = {env = "TORCHCODEC_BUILD_ROCJPEG", default = "AUTO"}
 
 [project.optional-dependencies]
 dev = [

src/torchcodec/_core/CMakeLists.txt

@@ -43,7 +43,7 @@ endif()
 if(NOT DEFINED TORCHCODEC_BUILD_IMAGE)
     set(TORCHCODEC_BUILD_IMAGE "ON")
 endif()
-foreach(_codec JPEG PNG WEBP AVIF GIF HEIC NVJPEG)
+foreach(_codec JPEG PNG WEBP AVIF GIF HEIC NVJPEG ROCJPEG)
     if(NOT DEFINED TORCHCODEC_BUILD_${_codec})
         set(TORCHCODEC_BUILD_${_codec} "AUTO")
     endif()
@@ -563,6 +563,41 @@ function(make_torchcodec_image_library)
         message(STATUS "Not building torchcodec with nvJPEG support (disabled via TORCHCODEC_BUILD_NVJPEG/TORCHCODEC_BUILD_IMAGE): decode_jpeg(device='cuda') will raise at runtime.")
     endif()
 
+    # ROCm counterpart of the nvJPEG block above: GPU JPEG decoding via rocJPEG.
+    # PyTorch exposes AMD GPUs under the "cuda" device string, so decode_jpeg(
+    # device="cuda") routes to the same decode_jpegs_cuda op, backed by rocJPEG.
+    resolve_image_codec("${TORCHCODEC_BUILD_ROCJPEG}" want_rocjpeg)
+    if(ENABLE_ROCM AND want_rocjpeg)
+        if(DEFINED ENV{ROCM_PATH})
+            set(ROCM_PATH "$ENV{ROCM_PATH}")
+        elseif(NOT DEFINED ROCM_PATH)
+            set(ROCM_PATH "/opt/rocm")
+        endif()
+        # So find_package(hip) locates ${ROCM_PATH}/lib/cmake/hip/hip-config.cmake.
+        list(APPEND CMAKE_PREFIX_PATH "${ROCM_PATH}")
+        find_package(hip REQUIRED)
+        find_path(ROCJPEG_INCLUDE_DIR
+            NAMES rocjpeg/rocjpeg.h
+            PATHS "${ROCM_PATH}/include")
+        find_library(ROCJPEG_LIBRARY
+            NAMES rocjpeg
+            PATHS "${ROCM_PATH}/lib")
+        if(NOT ROCJPEG_INCLUDE_DIR OR NOT ROCJPEG_LIBRARY)
+            message(FATAL_ERROR
+                "rocJPEG not found (looked under ${ROCM_PATH}), but ROCm GPU JPEG "
+                "decoding is enabled. Install the rocJPEG runtime and dev headers, "
+                "or set TORCHCODEC_BUILD_ROCJPEG=0 to build without it "
+                "(decode_jpeg(device='cuda') will raise at runtime).")
+        endif()
+        target_include_directories(${image_library_name} PRIVATE ${ROCJPEG_INCLUDE_DIR})
+        target_link_libraries(${image_library_name} PRIVATE ${ROCJPEG_LIBRARY} hip::host)
+        target_compile_definitions(${image_library_name} PRIVATE
+            TORCHCODEC_ENABLE_ROCJPEG=1 USE_ROCM __HIP_PLATFORM_AMD__)
+        message(STATUS "Building torchcodec with rocJPEG GPU JPEG decoding support.")
+    elseif(ENABLE_ROCM)
+        message(STATUS "Not building torchcodec with rocJPEG support (disabled via TORCHCODEC_BUILD_ROCJPEG/TORCHCODEC_BUILD_IMAGE): decode_jpeg(device='cuda') will raise at runtime.")
+    endif()
+
     if(want_jpeg)
         target_compile_definitions(${image_library_name} PRIVATE TORCHCODEC_ENABLE_JPEG=1)
         target_link_libraries(${image_library_name} PRIVATE JPEG::JPEG)

src/torchcodec/_core/DecodeJpegCommon.h

@@ -0,0 +1,59 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <cstring>
+
+#include "Exif.h"
+
+// Helpers shared by the GPU JPEG decoders (nvJPEG on CUDA, rocJPEG on ROCm).
+// The CPU JPEG decoder relies on libjpeg for EXIF handling; the GPU decoders
+// don't decode metadata, so we scan the bitstream for orientation ourselves.
+
+namespace facebook::torchcodec {
+
+// Scan a JPEG bitstream for the APP1/EXIF segment and return its orientation.
+inline exif_private::ExifOrientation fetch_exif_orientation_from_jpeg_bytes(
+    const unsigned char* jpeg,
+    size_t size) {
+  constexpr unsigned char MARKER_PREFIX = 0xFF;
+  constexpr unsigned char SOI = 0xD8;
+  constexpr unsigned char SOS = 0xDA; // start of scan: no more metadata markers
+  constexpr unsigned char EOI = 0xD9;
+  constexpr unsigned char APP1 = 0xE1;
+  constexpr size_t exif_header_size = 6; // "Exif\0\0"
+
+  if (size < 2 || jpeg[0] != MARKER_PREFIX || jpeg[1] != SOI) {
+    return exif_private::ExifOrientation::Unspecified;
+  }
+
+  size_t pos = 2;
+  while (pos + 4 <= size && jpeg[pos] == MARKER_PREFIX) {
+    unsigned char marker = jpeg[pos + 1];
+    if (marker == SOS || marker == EOI) {
+      break;
+    }
+    // Segment length is big-endian and includes the 2 length bytes themselves.
+    size_t segment_length =
+        (size_t(jpeg[pos + 2]) << 8) | size_t(jpeg[pos + 3]);
+    if (segment_length < 2 || pos + 2 + segment_length > size) {
+      break;
+    }
+
+    if (marker == APP1 && segment_length >= 2 + exif_header_size) {
+      const unsigned char* payload = jpeg + pos + 4;
+      if (std::memcmp(payload, "Exif\0\0", exif_header_size) == 0) {
+        return exif_private::fetch_exif_orientation(
+            payload + exif_header_size, segment_length - 2 - exif_header_size);
+      }
+    }
+    pos += 2 + segment_length;
+  }
+  return exif_private::ExifOrientation::Unspecified;
+}
+
+} // namespace facebook::torchcodec

src/torchcodec/_core/DecodeJpegCuda.cpp

@@ -11,7 +11,10 @@
 
 #include "StableABICompat.h"
 
-#if !TORCHCODEC_ENABLE_NVJPEG
+// The "no GPU JPEG support" stub is defined here, but only when neither GPU
+// backend is compiled in. When ROCm/rocJPEG is enabled, decode_jpegs_cuda is
+// defined in DecodeJpegRocm.cpp instead, so this file compiles to nothing.
+#if !TORCHCODEC_ENABLE_NVJPEG && !TORCHCODEC_ENABLE_ROCJPEG
 
 namespace facebook::torchcodec {
 
@@ -30,12 +33,10 @@ std::vector<torch::stable::Tensor> decode_jpegs_cuda(
 
 } // namespace facebook::torchcodec
 
-#else
-
-#include <cstring>
+#elif TORCHCODEC_ENABLE_NVJPEG
 
 #include "CUDACommon.h"
-#include "Exif.h"
+#include "DecodeJpegCommon.h"
 #include "ImageCommon.h"
 
 namespace facebook::torchcodec {
@@ -71,48 +72,6 @@ using namespace exif_private;
 
 namespace {
 
-// Scan a JPEG bitstream for the APP1/EXIF segment and return its orientation.
-// On the CPU jpeg decoder, we rely on libjpeg for that. Here, we have to parse
-// it ourselves.
-ExifOrientation fetch_exif_orientation_from_jpeg_bytes(
-    const unsigned char* jpeg,
-    size_t size) {
-  constexpr unsigned char MARKER_PREFIX = 0xFF;
-  constexpr unsigned char SOI = 0xD8;
-  constexpr unsigned char SOS = 0xDA; // start of scan: no more metadata markers
-  constexpr unsigned char EOI = 0xD9;
-  constexpr unsigned char APP1 = 0xE1;
-  constexpr size_t exif_header_size = 6; // "Exif\0\0"
-
-  if (size < 2 || jpeg[0] != MARKER_PREFIX || jpeg[1] != SOI) {
-    return ExifOrientation::Unspecified;
-  }
-
-  size_t pos = 2;
-  while (pos + 4 <= size && jpeg[pos] == MARKER_PREFIX) {
-    unsigned char marker = jpeg[pos + 1];
-    if (marker == SOS || marker == EOI) {
-      break;
-    }
-    // Segment length is big-endian and includes the 2 length bytes themselves.
-    size_t segment_length =
-        (size_t(jpeg[pos + 2]) << 8) | size_t(jpeg[pos + 3]);
-    if (segment_length < 2 || pos + 2 + segment_length > size) {
-      break;
-    }
-
-    if (marker == APP1 && segment_length >= 2 + exif_header_size) {
-      const unsigned char* payload = jpeg + pos + 4;
-      if (std::memcmp(payload, "Exif\0\0", exif_header_size) == 0) {
-        return fetch_exif_orientation(
-            payload + exif_header_size, segment_length - 2 - exif_header_size);
-      }
-    }
-    pos += 2 + segment_length;
-  }
-  return ExifOrientation::Unspecified;
-}
-
 // We cache decoder objects for the same reason we cache NVDEC decoders: they're
 // expensive to create and destroy. To determine the ideal cache size, I ran
 // benchmarks on a A100 where each thread calls its own decode_jpeg():
@@ -520,4 +479,4 @@ std::vector<torch::stable::Tensor> CUDAJpegDecoder::decode_images(
 
 } // namespace facebook::torchcodec
 
-#endif // !TORCHCODEC_ENABLE_NVJPEG
+#endif // GPU JPEG backend selection