Port CUDA jpeg decoder to stable ABI - #9533
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9533
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit bd9f4cd with merge base bc721a5 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
NicolasHug
left a comment
There was a problem hiding this comment.
Thanks a lot @adabeyta ! LGTM just left two small comments below
| // TODO(stable-abi): use a stable warn-once macro once one exists. | ||
| static bool warned = false; | ||
| if (!warned) { | ||
| warned = true; | ||
| fprintf( | ||
| stderr, | ||
| "Warning: There is a memory leak issue in the nvjpeg library for CUDA versions < 11.6. " | ||
| "Make sure to rely on CUDA 11.6 or above before using decode_jpeg(..., device='cuda').\n"); | ||
| } |
There was a problem hiding this comment.
let's actually remove this warning entirely, we only support CUDA >= 12.6 now, so we don't need it anymore
| just leak the libnvjpeg & cuda variables for the time being and hope | ||
| that the CUDA runtime handles cleanup for us. | ||
| Please send a PR if you have a solution for this problem. | ||
| */ |
There was a problem hiding this comment.
Let's keep the commented-out code below. No need to port it, we can leave it commented-out, but I think it's useful to keep to communicate that well, if we could, that's what we'd do.
3b950a8 to
a7d2768
Compare
|
Hey @NicolasHug! You merged this PR, but no labels were added. |
Ports the CUDA jpeg decoder to stable ABI. The CUDA jpeg decoder lives in the image extension, so this also stands up image_stable (same dual-extension pattern to
_C/_C_stable).Structure: Image_stable Extension
Mirrors the _C / _C_stable split.
torchvision.image: the existing extension, now built from all not-yet-ported image sources.torchvision.image_stable: the new extension, built only from migrated image sources (currently the jpeg decoder + its extension glue).Layout Changes
setup.py: splits the image sources between the two extensions. Migrated files added to STABLE_SOURCES, the existing _stable()/_not_stable() helpers route them into image_stable and keep everything else in image.make_image_stable_extension(): stable sibling of make_image_extension()).CMakeLists.txt: migrated files added toTORCHVISION_STABLE_SOURCESso the cmake build pins them identically.io/image.py: loads image_stable after image.Schema via
STABLE_TORCH_LIBRARY_FRAGMENT(image), this extends the image namespace co-owned with the legacy .so.common_stable.h: torch-free ImageReadMode constants, so the stable TU doesn't pull full libtorch in via common.h.common_stable.cpp: extension-level glue for image_stable, the stable counterpart of common.cpp. Holds the Windows PyInit_image_stable stub extension glue, not tied to any op, it lives there instead of in a kernel file.CUDA Stream/Event Model
cudaEventCreate → Record → cudaStreamWaitEvent → Destroysequence.syncStreams(extracted into asyncStreamshelper). Current stream viaaoti_torch_get_current_cuda_stream.Stable-ABI audit (
torch-abi-audit)Ran
torch-abi-auditagainst the builttorchvisionpackage to verify the migrated extension stays on the stable surface and never reaches intoat::/c10::/torch::jit::internals.image_stable.soaudits asSTABLE— 65 stable-shim symbols, 0 unstable. The legacyimage.so/_C.sostayUNSTABLE(expected: only nms and the jpeg decoder havemigrated), so the package-level verdict stays
UNSTABLEuntil the rest move — same mid-migration shape as torchcodec.