Port CPU jpeg encoder to stable ABI. (#9543)

Author: adabeyta

CMakeLists.txt

@@ -105,6 +105,7 @@ set(TORCHVISION_STABLE_SOURCES
   ${TVCPP}/io/image/cuda/decode_jpegs_cuda.cpp
   ${TVCPP}/io/image/cuda/encode_jpegs_cuda.cpp
   ${TVCPP}/io/image/cpu/encode_png.cpp
+  ${TVCPP}/io/image/cpu/encode_jpeg.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

@@ -166,6 +166,7 @@ def get_macros_and_flags():
     CSRS_DIR / "ops/quantized/cpu/qnms_kernel.cpp",
     CSRS_DIR / "io/image/common_stable.cpp",
     CSRS_DIR / "io/image/cpu/encode_png.cpp",
+    CSRS_DIR / "io/image/cpu/encode_jpeg.cpp",
 }
 STABLE_SOURCES.add(CSRS_DIR / ("ops/hip/nms_kernel.hip" if IS_ROCM else "ops/cuda/nms_kernel.cu"))
 STABLE_SOURCES.add(
@@ -470,6 +471,8 @@ def make_image_stable_extension():
         + _stable(image_dir.glob("cpu/*.cpp"))
         + _stable(image_dir.glob("hip/*.cpp" if IS_ROCM else "cuda/*.cpp"))
     )
+    # Shared libjpeg glue. Legacy decode_jpeg still needs it, so it builds into both extensions.
+    sources.append(image_dir / "cpu/common_jpeg.cpp")
 
     Extension = CppExtension
 
@@ -484,6 +487,18 @@ def make_image_stable_extension():
         else:
             warnings.warn("Building torchvision without PNG support")
 
+    if USE_JPEG:
+        jpeg_found, jpeg_include_dir, jpeg_library_dir = find_library(header="jpeglib.h")
+        if jpeg_found:
+            print("Building torchvision with JPEG image support")
+            if jpeg_include_dir is not None and jpeg_library_dir is not None:
+                include_dirs.append(jpeg_include_dir)
+                library_dirs.append(jpeg_library_dir)
+            libraries.append("jpeg")
+            define_macros += [("JPEG_FOUND", 1)]
+        else:
+            warnings.warn("Building torchvision without JPEG support")
+
     if USE_NVJPEG 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:

torchvision/csrc/io/image/cpu/encode_jpeg.cpp

@@ -1,17 +1,22 @@
 #include "encode_jpeg.h"
 
+#include <torch/csrc/stable/library.h>
+#include <torch/csrc/stable/ops.h>
 #include <torch/headeronly/util/Exception.h>
 
 #include <optional>
 
+#include "../common_stable.h"
 #include "common_jpeg.h"
 
 namespace vision {
 namespace image {
 
 #if !JPEG_FOUND
 
-torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
+torch::stable::Tensor encode_jpeg(
+    const torch::stable::Tensor& data,
+    int64_t quality) {
   STD_TORCH_CHECK(
       false, "encode_jpeg: torchvision not compiled with libjpeg support");
 }
@@ -28,9 +33,9 @@ using JpegSizeType = size_t;
 
 using namespace detail;
 
-torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
-  C10_LOG_API_USAGE_ONCE(
-      "torchvision.csrc.io.image.cpu.encode_jpeg.encode_jpeg");
+torch::stable::Tensor encode_jpeg(
+    const torch::stable::Tensor& data,
+    int64_t quality) {
   // Define compression structures and error handling
   struct jpeg_compress_struct cinfo {};
   struct torch_jpeg_error_mgr jerr {};
@@ -43,7 +48,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
   // unwind C++ stack frames, so destructors of objects created after setjmp
   // won't run. We use std::optional to declare tensors before setjmp while
   // deferring construction, and explicitly reset them on the error path.
-  std::optional<torch::Tensor> input;
+  std::optional<torch::stable::Tensor> input;
 
   cinfo.err = jpeg_std_error(&jerr.pub);
   jerr.pub.error_exit = torch_jpeg_error_exit;
@@ -63,12 +68,12 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
   }
 
   // Check that the input tensor is on CPU
-  STD_TORCH_CHECK(
-      data.device() == torch::kCPU, "Input tensor should be on CPU");
+  STD_TORCH_CHECK(data.is_cpu(), "Input tensor should be on CPU");
 
   // Check that the input tensor dtype is uint8
   STD_TORCH_CHECK(
-      data.dtype() == torch::kU8, "Input tensor dtype should be uint8");
+      data.scalar_type() == torch::headeronly::ScalarType::Byte,
+      "Input tensor dtype should be uint8");
 
   // Check that the input tensor is 3-dimensional
   STD_TORCH_CHECK(
@@ -78,7 +83,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
   int channels = data.size(0);
   int height = data.size(1);
   int width = data.size(2);
-  input = data.permute({1, 2, 0}).contiguous();
+  input = torch::stable::contiguous(stable_permute(data, {1, 2, 0}));
 
   STD_TORCH_CHECK(
       channels == 1 || channels == 3,
@@ -104,7 +109,7 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
   jpeg_start_compress(&cinfo, TRUE);
 
   auto stride = width * channels;
-  auto ptr = input->data_ptr<uint8_t>();
+  auto ptr = input->mutable_data_ptr<uint8_t>();
 
   // Encode JPEG file
   while (cinfo.next_scanline < cinfo.image_height) {
@@ -115,13 +120,25 @@ torch::Tensor encode_jpeg(const torch::Tensor& data, int64_t quality) {
   jpeg_finish_compress(&cinfo);
   jpeg_destroy_compress(&cinfo);
 
-  torch::TensorOptions options = torch::TensorOptions{torch::kU8};
-  auto out_tensor =
-      torch::from_blob(jpegBuf, {(long)jpegSize}, ::free, options);
+  auto out_tensor = torch::stable::from_blob(
+      jpegBuf,
+      {static_cast<int64_t>(jpegSize)},
+      {1},
+      torch::stable::Device(torch::headeronly::DeviceType::CPU),
+      torch::headeronly::ScalarType::Byte,
+      ::free);
   jpegBuf = nullptr;
   return out_tensor;
 }
 #endif
 
+STABLE_TORCH_LIBRARY_FRAGMENT(image, m) {
+  m.def("encode_jpeg(Tensor data, int quality) -> Tensor");
+}
+
+STABLE_TORCH_LIBRARY_IMPL(image, CompositeExplicitAutograd, m) {
+  m.impl("encode_jpeg", TORCH_BOX(&encode_jpeg));
+}
+
 } // namespace image
 } // namespace vision

torchvision/csrc/io/image/cpu/encode_jpeg.h

@@ -1,12 +1,12 @@
 #pragma once
 
-#include <torch/types.h>
+#include <torch/csrc/stable/tensor.h>
 
 namespace vision {
 namespace image {
 
-C10_EXPORT torch::Tensor encode_jpeg(
-    const torch::Tensor& data,
+torch::stable::Tensor encode_jpeg(
+    const torch::stable::Tensor& data,
     int64_t quality);
 
 } // namespace image

torchvision/csrc/io/image/image.cpp

@@ -14,7 +14,6 @@ static auto registry =
             &decode_jpeg)
         .op("image::decode_webp(Tensor encoded_data, int mode) -> Tensor",
             &decode_webp)
-        .op("image::encode_jpeg", &encode_jpeg)
         .op("image::read_file", &read_file)
         .op("image::write_file", &write_file)
         .op("image::decode_image(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor",

torchvision/csrc/io/image/image.h

@@ -5,5 +5,4 @@
 #include "cpu/decode_jpeg.h"
 #include "cpu/decode_png.h"
 #include "cpu/decode_webp.h"
-#include "cpu/encode_jpeg.h"
 #include "cpu/read_write_file.h"