Port CUDA jpeg encoder to stable ABI.

CMakeLists.txt

@@ -103,6 +103,7 @@ set(TORCHVISION_STABLE_SOURCES
   ${TVCPP}/ops/cpu/nms_kernel.cpp
   ${TVCPP}/ops/cuda/nms_kernel.cu
   ${TVCPP}/io/image/cuda/decode_jpegs_cuda.cpp
+  ${TVCPP}/io/image/cuda/encode_jpegs_cuda.cpp
   ${TVCPP}/io/image/common_stable.cpp
   ${TVCPP}/vision_stable.cpp)
 # Pin them to torch 2.11. Per source, not library-wide: the pin makes ATen headers

packaging/wheel/relocate.py

@@ -216,7 +216,7 @@ def relocate_elf_library(patchelf, output_dir, output_library, binary):
             subprocess.check_output([patchelf, "--print-rpath", new_library_name], cwd=new_libraries_path)
 
     print("Update library dependencies")
-    library_dependencies = binary_dependencies[binary]
+    library_dependencies = binary_dependencies.get(binary, [])
     for dep in library_dependencies:
         new_dep = osp.basename(new_names[dep])
         print(f"{binary}: {dep} -> {new_dep}")
@@ -316,8 +316,9 @@ def patch_linux():
     output_dir = osp.join(PACKAGE_ROOT, "dist", ".wheel-process")
 
     image_binary = "image.so"
+    image_stable_binary = "image_stable.so"
     video_binary = "video_reader.so"
-    torchvision_binaries = [image_binary, video_binary]
+    torchvision_binaries = [image_binary, image_stable_binary, video_binary]
     for wheel in wheels:
         if osp.exists(output_dir):
             shutil.rmtree(output_dir)
@@ -352,8 +353,9 @@ def patch_win():
     output_dir = osp.join(PACKAGE_ROOT, "dist", ".wheel-process")
 
     image_binary = "image.pyd"
+    image_stable_binary = "image_stable.pyd"
     video_binary = "video_reader.pyd"
-    torchvision_binaries = [image_binary, video_binary]
+    torchvision_binaries = [image_binary, image_stable_binary, video_binary]
     for wheel in wheels:
         if osp.exists(output_dir):
             shutil.rmtree(output_dir)

setup.py

@@ -170,6 +170,7 @@ def get_macros_and_flags():
 STABLE_SOURCES.add(
     CSRS_DIR / ("io/image/hip/decode_jpegs_cuda.cpp" if IS_ROCM else "io/image/cuda/decode_jpegs_cuda.cpp")
 )
+STABLE_SOURCES.add(CSRS_DIR / "io/image/cuda/encode_jpegs_cuda.cpp")
 
 
 def _not_stable(paths):
@@ -393,8 +394,6 @@ def make_image_extension():
 
     if IS_ROCM:
         sources += _not_stable(image_dir.glob("hip/*.cpp"))
-        # we need to exclude this in favor of the hipified source
-        sources.remove(image_dir / "image.cpp")
     else:
         sources += _not_stable(image_dir.glob("cuda/*.cpp"))
 
@@ -444,20 +443,6 @@ def make_image_extension():
         else:
             warnings.warn("Building torchvision without WEBP support")
 
-    # NVJPEG is needed here for the GPU JPEG *encoder*, not supported by ROCM.
-    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:
-            print("Building torchvision with NVJPEG image support")
-            libraries.append("nvjpeg")
-            define_macros += [("NVJPEG_FOUND", 1)]
-            Extension = CUDAExtension
-        else:
-            warnings.warn("Building torchvision without NVJPEG support")
-    elif USE_NVJPEG and not IS_ROCM:
-        warnings.warn("Building torchvision without NVJPEG support")
-
     return Extension(
         name="torchvision.image",
         sources=sorted(str(s) for s in sources),

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

@@ -1,26 +1,29 @@
 #include "encode_jpegs_cuda.h"
 
+#include <torch/csrc/stable/library.h>
 #include <torch/headeronly/util/Exception.h>
+
 #if !NVJPEG_FOUND
 namespace vision {
 namespace image {
-std::vector<torch::Tensor> encode_jpegs_cuda(
-    const std::vector<torch::Tensor>& decoded_images,
+std::vector<torch::stable::Tensor> encode_jpegs_cuda(
+    const std::vector<torch::stable::Tensor>& decoded_images,
     const int64_t quality) {
   STD_TORCH_CHECK(
       false, "encode_jpegs_cuda: torchvision not compiled with nvJPEG support");
 }
 } // namespace image
 } // namespace vision
-#else
 
-#include <ATen/cuda/CUDAContext.h>
-#include <ATen/cuda/CUDAEvent.h>
-#include <iostream>
+#else
+#include <torch/csrc/inductor/aoti_torch/c/shim.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/c/shim.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/headeronly/core/ScalarType.h>
 #include <memory>
-#include <string>
-#include "c10/core/ScalarType.h"
-
+#include <mutex>
+#include <optional>
 namespace vision {
 namespace image {
 
@@ -29,19 +32,40 @@ namespace image {
 std::mutex encoderMutex;
 std::unique_ptr<CUDAJpegEncoder> cudaJpegEncoder;
 
-std::vector<torch::Tensor> encode_jpegs_cuda(
-    const std::vector<torch::Tensor>& decoded_images,
-    const int64_t quality) {
-  C10_LOG_API_USAGE_ONCE(
-      "torchvision.csrc.io.image.cuda.encode_jpegs_cuda.encode_jpegs_cuda");
+// Make waitingStream wait until all work currently enqueued on runningStream
+// has completed.
+// https://github.com/meta-pytorch/torchcodec/blob/1dc85b7a7900d91fee207ccdc02f211a051688fe/src/torchcodec/_core/CUDACommon.cpp#L30-L47
+static void syncStreams(
+    cudaStream_t runningStream,
+    cudaStream_t waitingStream) {
+  cudaEvent_t event;
+  cudaError_t err = cudaEventCreate(&event);
+  STD_TORCH_CHECK(
+      err == cudaSuccess, "cudaEventCreate failed: ", cudaGetErrorString(err));
+  err = cudaEventRecord(event, runningStream);
+  STD_TORCH_CHECK(
+      err == cudaSuccess, "cudaEventRecord failed: ", cudaGetErrorString(err));
+  err = cudaStreamWaitEvent(waitingStream, event, 0);
+  STD_TORCH_CHECK(
+      err == cudaSuccess,
+      "cudaStreamWaitEvent failed: ",
+      cudaGetErrorString(err));
+  cudaEventDestroy(event);
+}
 
+std::vector<torch::stable::Tensor> encode_jpegs_cuda(
+    const std::vector<torch::stable::Tensor>& decoded_images,
+    const int64_t quality) {
   // Some nvjpeg structures are not thread safe so we're keeping it single
   // threaded for now. In the future this may be an opportunity to unlock
   // further speedups
   std::lock_guard<std::mutex> lock(encoderMutex);
   STD_TORCH_CHECK(decoded_images.size() > 0, "Empty input tensor list");
-  torch::Device device = decoded_images[0].device();
-  at::cuda::CUDAGuard device_guard(device);
+  torch::stable::Device device = decoded_images[0].device();
+  STD_TORCH_CHECK(
+      device.is_cuda(),
+      "All input tensors must be on a CUDA device when encoding with nvjpeg");
+  torch::stable::accelerator::DeviceGuard device_guard(device.index());
 
   // lazy init of the encoder class
   // the encoder object holds on to a lot of state and is expensive to create,
@@ -65,11 +89,12 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
     std::atexit([]() { delete cudaJpegEncoder.release(); });
   }
 
-  std::vector<torch::Tensor> contig_images;
+  std::vector<torch::stable::Tensor> contig_images;
   contig_images.reserve(decoded_images.size());
   for (const auto& image : decoded_images) {
     STD_TORCH_CHECK(
-        image.dtype() == torch::kU8, "Input tensor dtype should be uint8");
+        image.scalar_type() == torch::headeronly::ScalarType::Byte,
+        "Input tensor dtype should be uint8");
 
     STD_TORCH_CHECK(
         image.device() == device,
@@ -85,21 +110,15 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
         image.size(0));
 
     // nvjpeg requires images to be contiguous
-    if (image.is_contiguous()) {
-      contig_images.push_back(image);
-    } else {
-      contig_images.push_back(image.contiguous());
-    }
+    contig_images.push_back(torch::stable::contiguous(image));
   }
 
   cudaJpegEncoder->set_quality(quality);
-  std::vector<torch::Tensor> encoded_images;
+  std::vector<torch::stable::Tensor> encoded_images;
   for (const auto& image : contig_images) {
     auto encoded_image = cudaJpegEncoder->encode_jpeg(image);
     encoded_images.push_back(encoded_image);
   }
-  at::cuda::CUDAEvent event;
-  event.record(cudaJpegEncoder->stream);
 
   // We use a dedicated stream to do the encoding and even though the results
   // may be ready on that stream we cannot assume that they are also available
@@ -108,21 +127,31 @@ std::vector<torch::Tensor> encode_jpegs_cuda(
   // do not want to block the host at this particular point
   // (which is what cudaStreamSynchronize would do.) Events allow us to
   // synchronize the streams without blocking the host.
-  event.block(cudaJpegEncoder->current_stream);
+  syncStreams(cudaJpegEncoder->stream, cudaJpegEncoder->current_stream);
   return encoded_images;
 }
 
-CUDAJpegEncoder::CUDAJpegEncoder(const torch::Device& target_device)
-    : original_device{torch::kCUDA, torch::cuda::current_device()},
-      target_device{target_device},
-      stream{
-          target_device.has_index()
-              ? at::cuda::getStreamFromPool(false, target_device.index())
-              : at::cuda::getStreamFromPool(false)},
-      current_stream{
-          original_device.has_index()
-              ? at::cuda::getCurrentCUDAStream(original_device.index())
-              : at::cuda::getCurrentCUDAStream()} {
+CUDAJpegEncoder::CUDAJpegEncoder(const torch::stable::Device& target_device)
+    : original_device(
+          torch::headeronly::DeviceType::CUDA,
+          torch::stable::accelerator::getCurrentDeviceIndex()),
+      target_device(target_device) {
+  torch::stable::accelerator::DeviceGuard device_guard(target_device.index());
+  // Pool-owned (not a raw leaked) stream; avoids a cross-DSO teardown hazard.
+  // https://github.com/pytorch/pytorch/blob/98e36864e640023a716e058d894ea2d20e76e5f7/torch/csrc/stable/c/shim.h#L127-L130
+  void* stream_ptr = nullptr;
+  TORCH_ERROR_CODE_CHECK(torch_get_cuda_stream_from_pool(
+      false, target_device.index(), &stream_ptr));
+  stream = static_cast<cudaStream_t>(stream_ptr);
+
+  // The caller's current stream, captured at construction. encode_jpegs_cuda
+  // makes it wait on our private `stream` (via syncStreams) before returning.
+  // https://github.com/pytorch/pytorch/blob/98e36864e640023a716e058d894ea2d20e76e5f7/torch/csrc/inductor/aoti_torch/c/shim.h#L573-L602
+  void* current_stream_ptr = nullptr;
+  TORCH_ERROR_CODE_CHECK(aoti_torch_get_current_cuda_stream(
+      original_device.index(), &current_stream_ptr));
+  current_stream = static_cast<cudaStream_t>(current_stream_ptr);
+
   nvjpegStatus_t status;
   status = nvjpegCreateSimple(&nvjpeg_handle);
   STD_TORCH_CHECK(
@@ -186,7 +215,8 @@ CUDAJpegEncoder::~CUDAJpegEncoder() {
   //     status == NVJPEG_STATUS_SUCCESS, "nvjpegDestroy failed: ", status);
 }
 
-torch::Tensor CUDAJpegEncoder::encode_jpeg(const torch::Tensor& src_image) {
+torch::stable::Tensor CUDAJpegEncoder::encode_jpeg(
+    const torch::stable::Tensor& src_image) {
   nvjpegStatus_t status;
   cudaError_t cudaStatus;
 
@@ -207,7 +237,8 @@ torch::Tensor CUDAJpegEncoder::encode_jpeg(const torch::Tensor& src_image) {
 
   nvjpegImage_t target_image;
   for (int c = 0; c < channels; c++) {
-    target_image.channel[c] = src_image[c].data_ptr<uint8_t>();
+    target_image.channel[c] =
+        torch::stable::select(src_image, 0, c).mutable_data_ptr<uint8_t>();
     // this is why we need contiguous tensors
     target_image.pitch[c] = width;
   }
@@ -242,20 +273,18 @@ torch::Tensor CUDAJpegEncoder::encode_jpeg(const torch::Tensor& src_image) {
   STD_TORCH_CHECK(cudaStatus == cudaSuccess, "CUDA ERROR: ", cudaStatus);
 
   // Reserve buffer for the encoded image
-  torch::Tensor encoded_image = torch::empty(
-      {static_cast<long>(length)},
-      torch::TensorOptions()
-          .dtype(torch::kByte)
-          .layout(torch::kStrided)
-          .device(target_device)
-          .requires_grad(false));
+  torch::stable::Tensor encoded_image = torch::stable::empty(
+      {static_cast<int64_t>(length)},
+      torch::headeronly::ScalarType::Byte,
+      std::nullopt,
+      target_device);
   cudaStatus = cudaStreamSynchronize(stream);
   STD_TORCH_CHECK(cudaStatus == cudaSuccess, "CUDA ERROR: ", cudaStatus);
   // Retrieve the encoded image
   status = nvjpegEncodeRetrieveBitstreamDevice(
       nvjpeg_handle,
       nv_enc_state,
-      encoded_image.data_ptr<uint8_t>(),
+      encoded_image.mutable_data_ptr<uint8_t>(),
       &length,
       stream);
   STD_TORCH_CHECK(
@@ -278,3 +307,17 @@ void CUDAJpegEncoder::set_quality(const int64_t quality) {
 } // namespace vision
 
 #endif // NVJPEG_FOUND
+
+namespace vision {
+namespace image {
+
+STABLE_TORCH_LIBRARY_FRAGMENT(image, m) {
+  m.def("encode_jpegs_cuda(Tensor[] decoded_images, int quality) -> Tensor[]");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) {
+  m.impl("encode_jpegs_cuda", TORCH_BOX(&encode_jpegs_cuda));
+}
+
+} // namespace image
+} // namespace vision

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

@@ -1,34 +1,41 @@
 #pragma once
-#include <torch/types.h>
+
+#include <torch/csrc/stable/device.h>
+#include <torch/csrc/stable/tensor.h>
 #include <vector>
+
 #if NVJPEG_FOUND
 
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime.h>
 #include <nvjpeg.h>
 
 namespace vision {
 namespace image {
 
 class CUDAJpegEncoder {
  public:
-  CUDAJpegEncoder(const torch::Device& device);
+  CUDAJpegEncoder(const torch::stable::Device& target_device);
   ~CUDAJpegEncoder();
 
-  torch::Tensor encode_jpeg(const torch::Tensor& src_image);
+  torch::stable::Tensor encode_jpeg(const torch::stable::Tensor& src_image);
 
   void set_quality(const int64_t quality);
 
-  const torch::Device original_device;
-  const torch::Device target_device;
-  const c10::cuda::CUDAStream stream;
-  const c10::cuda::CUDAStream current_stream;
+  const torch::stable::Device original_device;
+  const torch::stable::Device target_device;
+  cudaStream_t stream;
+  cudaStream_t current_stream;
 
  protected:
   nvjpegEncoderState_t nv_enc_state;
   nvjpegEncoderParams_t nv_enc_params;
   nvjpegHandle_t nvjpeg_handle;
 };
+
+std::vector<torch::stable::Tensor> encode_jpegs_cuda(
+    const std::vector<torch::stable::Tensor>& decoded_images,
+    const int64_t quality);
+
 } // namespace image
 } // namespace vision
 #endif