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/setup.py b/setup.py index 63fd2312e59..cff0d28d93c 100644 --- a/setup.py +++ b/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), diff --git a/torchvision/_internally_replaced_utils.py b/torchvision/_internally_replaced_utils.py index 253920add71..7a564f75c87 100644 --- a/torchvision/_internally_replaced_utils.py +++ b/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) 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..bedbec0c6b8 100644 --- a/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp +++ b/torchvision/csrc/io/image/cuda/encode_jpegs_cuda.cpp @@ -1,47 +1,71 @@ #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 { -// 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; -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, @@ -49,27 +73,28 @@ std::vector encode_jpegs_cuda( // 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(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(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 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"