[TEST] Test stable cuda jpeg encoder

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

setup.py

@@ -170,6 +170,9 @@ 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")
 )
+# nvJPEG is CUDA-only (no ROCm encoder), so the GPU jpeg encoder builds only on CUDA.
+if not IS_ROCM:
+    STABLE_SOURCES.add(CSRS_DIR / "io/image/cuda/encode_jpegs_cuda.cpp")
 
 
 def _not_stable(paths):
@@ -393,8 +396,8 @@ 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")
+        # image.cpp has no CUDA now, so hipify no longer renames it to image_hip.cpp;
+        # keep the original so the legacy image op registrations are compiled in.
     else:
         sources += _not_stable(image_dir.glob("cuda/*.cpp"))
 
@@ -444,20 +447,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/_internally_replaced_utils.py

@@ -23,6 +23,41 @@ def _is_remote_location_available() -> bool:
     from torch.utils.model_zoo import load_url as load_state_dict_from_url  # noqa: 401
 
 
+def _preload_image_stable_cuda_libs():
+    # image_stable links nvjpeg, which torch does not bundle, so on a wheel install without a
+    # system CUDA it may not be on the loader path and image_stable fails to load. We try to
+    # locate nvjpeg in the CUDA toolkit (CUDA_PATH/CUDA_HOME, then nvcc) and preload it; on
+    # Windows a .pyd import does not search PATH, so we also scan PATH, where the CUDA
+    # installer's DLL dir lands. No-op if already loaded or not found.
+    import ctypes
+    import glob
+    import shutil
+
+    win = os.name == "nt"
+    subdirs = ("bin", os.path.join("bin", "x64")) if win else ("lib64", "lib")
+    pattern = "nvjpeg64_*.dll" if win else "libnvjpeg.so*"
+
+    def _load(path):
+        return ctypes.WinDLL(path) if win else ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+
+    cuda_homes = [os.environ.get("CUDA_PATH"), os.environ.get("CUDA_HOME")]
+    if nvcc := shutil.which("nvcc"):
+        cuda_homes.append(os.path.dirname(os.path.dirname(nvcc)))
+    lib_dirs = []
+    for cuda_home in filter(None, cuda_homes):
+        lib_dirs += [os.path.join(cuda_home, d) for d in subdirs]
+        lib_dirs += glob.glob(os.path.join(cuda_home, "targets", "*", "lib"))
+    if win:
+        lib_dirs += os.environ.get("PATH", "").split(os.pathsep)
+    for lib_dir in lib_dirs:
+        for path in sorted(glob.glob(os.path.join(lib_dir, pattern)), reverse=True):
+            try:
+                _load(path)
+                return
+            except OSError:
+                continue
+
+
 def _get_extension_path(lib_name):
 
     lib_dir = os.path.dirname(__file__)
@@ -40,6 +75,8 @@ def _get_extension_path(lib_name):
         os.add_dll_directory(lib_dir)
 
         kernel32.SetErrorMode(prev_error_mode)
+    if lib_name == "image_stable":
+        _preload_image_stable_cuda_libs()
 
     loader_details = (importlib.machinery.ExtensionFileLoader, importlib.machinery.EXTENSION_SUFFIXES)
 

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

@@ -1,75 +1,100 @@
 #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 {
 
-// We use global variables to cache the encoder and decoder instances and
-// reuse them across calls to the corresponding pytorch functions
+// We use a global to cache the encoder instance and reuse it across calls to
+// the corresponding pytorch function.
 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,
   // so we reuse it across calls. NB: the cached structures are device specific
   // and cannot be reused across devices
   if (cudaJpegEncoder == nullptr || device != cudaJpegEncoder->target_device) {
     if (cudaJpegEncoder != nullptr) {
-      delete cudaJpegEncoder.release();
+      cudaJpegEncoder.reset(new CUDAJpegEncoder(device));
+    } else {
+      cudaJpegEncoder = std::make_unique<CUDAJpegEncoder>(device);
+
+      // Unfortunately, we cannot rely on the smart pointer releasing the
+      // encoder object correctly upon program exit. This is because, when
+      // cudaJpegEncoder gets destroyed, the CUDA runtime may already be shut
+      // down, rendering all destroy* calls in the encoder destructor invalid.
+      // Instead, we use an atexit hook which executes after main() finishes,
+      // but hopefully before CUDA shuts down when the program exits. If CUDA is
+      // already shut down the destructor will detect this and will not attempt
+      // to destroy any encoder structures.
+      std::atexit([]() { cudaJpegEncoder.reset(); });
     }
-
-    cudaJpegEncoder = std::make_unique<CUDAJpegEncoder>(device);
-
-    // Unfortunately, we cannot rely on the smart pointer releasing the encoder
-    // object correctly upon program exit. This is because, when cudaJpegEncoder
-    // gets destroyed, the CUDA runtime may already be shut down, rendering all
-    // destroy* calls in the encoder destructor invalid. Instead, we use an
-    // atexit hook which executes after main() finishes, but hopefully before
-    // CUDA shuts down when the program exits. If CUDA is already shut down the
-    // destructor will detect this and will not attempt to destroy any encoder
-    // structures.
-    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