diff --git a/CMakeLists.txt b/CMakeLists.txt index 658c265ce2c..fb5825435e0 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/packaging/wheel/relocate.py b/packaging/wheel/relocate.py index 4587f3798da..cb0ca762767 100644 --- a/packaging/wheel/relocate.py +++ b/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) diff --git a/setup.py b/setup.py index 63fd2312e59..2d8281816bc 100644 --- a/setup.py +++ b/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), diff --git a/torchvision/csrc/io/image/cuda/encode_decode_jpegs_cuda.h b/torchvision/csrc/io/image/cuda/encode_decode_jpegs_cuda.h deleted file mode 100644 index 3af51dd4e5f..00000000000 --- a/torchvision/csrc/io/image/cuda/encode_decode_jpegs_cuda.h +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include -#include "../common.h" -#include "encode_jpegs_cuda.h" - -namespace vision { -namespace image { - -/* -Fast jpeg encoding with CUDA. - -Args: - - decoded_images (const std::vector&): a vector of contiguous -CUDA tensors of dtype torch.uint8 to be encoded. - - quality (int64_t): 0-100, 75 is the default - -Returns: - - encoded_images (std::vector): a vector of CUDA -torch::Tensors of dtype torch.uint8 containing the encoded images - -Notes: - - If a single image fails, the whole batch fails. - - This function is thread-safe -*/ -C10_EXPORT std::vector encode_jpegs_cuda( - const std::vector& decoded_images, - const int64_t quality); - -} // namespace image -} // namespace vision diff --git a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp index c5cf4b63d75..7ce52f1b0d1 100644 --- a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp +++ b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp @@ -1,26 +1,29 @@ #include "encode_jpegs_cuda.h" +#include #include + #if !NVJPEG_FOUND namespace vision { namespace image { -std::vector encode_jpegs_cuda( - const std::vector& decoded_images, +std::vector encode_jpegs_cuda( + const std::vector& 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 -#include -#include +#else +#include +#include +#include +#include +#include #include -#include -#include "c10/core/ScalarType.h" - +#include +#include namespace vision { namespace image { @@ -29,19 +32,40 @@ namespace image { std::mutex encoderMutex; std::unique_ptr cudaJpegEncoder; -std::vector encode_jpegs_cuda( - const std::vector& 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 encode_jpegs_cuda( + const std::vector& 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 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 encode_jpegs_cuda( std::atexit([]() { delete cudaJpegEncoder.release(); }); } - std::vector contig_images; + std::vector 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 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 encoded_images; + std::vector 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 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(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(), ¤t_stream_ptr)); + current_stream = static_cast(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(); + target_image.channel[c] = + torch::stable::select(src_image, 0, c).mutable_data_ptr(); // 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(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(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(), + encoded_image.mutable_data_ptr(), &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 diff --git a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.h b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.h index 6ee0ad91df4..5937a4566e8 100644 --- a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.h +++ b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.h @@ -1,10 +1,12 @@ #pragma once -#include + +#include +#include #include + #if NVJPEG_FOUND -#include -#include +#include #include namespace vision { @@ -12,23 +14,28 @@ 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 encode_jpegs_cuda( + const std::vector& decoded_images, + const int64_t quality); + } // namespace image } // namespace vision #endif diff --git a/torchvision/csrc/io/image/image.cpp b/torchvision/csrc/io/image/image.cpp index ef4b2717d42..6b597ad2154 100644 --- a/torchvision/csrc/io/image/image.cpp +++ b/torchvision/csrc/io/image/image.cpp @@ -20,7 +20,6 @@ static auto registry = .op("image::write_file", &write_file) .op("image::decode_image(Tensor data, int mode, bool apply_exif_orientation=False) -> Tensor", &decode_image) - .op("image::encode_jpegs_cuda", &encode_jpegs_cuda) .op("image::_jpeg_version", &_jpeg_version) .op("image::_is_compiled_against_turbo", &_is_compiled_against_turbo); diff --git a/torchvision/csrc/io/image/image.h b/torchvision/csrc/io/image/image.h index 3f47fdec65c..991fa7b4611 100644 --- a/torchvision/csrc/io/image/image.h +++ b/torchvision/csrc/io/image/image.h @@ -8,4 +8,3 @@ #include "cpu/encode_jpeg.h" #include "cpu/encode_png.h" #include "cpu/read_write_file.h" -#include "cuda/encode_decode_jpegs_cuda.h"