Skip to content

Port CUDA jpeg decoder to stable ABI - #9533

Merged
NicolasHug merged 2 commits into
pytorch:mainfrom
adabeyta:decoder-jpg-stable
Jun 26, 2026
Merged

NicolasHug merged 2 commits into
pytorch:mainfrom
adabeyta:decoder-jpg-stable

Conversation

@adabeyta

Copy link
Copy Markdown
Collaborator

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 to TORCHVISION_STABLE_SOURCES so 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

  • Decode runs on a dedicated stream, the caller's current stream then waits on it via the exact cudaEventCreate → Record → cudaStreamWaitEvent → Destroy sequence.
  • syncStreams (extracted into a syncStreams helper). Current stream via aoti_torch_get_current_cuda_stream.

Stable-ABI audit (torch-abi-audit)

Ran torch-abi-audit against the built torchvision package to verify the migrated extension stays on the stable surface and never reaches into at:: / c10:: / torch::jit:: internals.

image_stable.so audits as STABLE — 65 stable-shim symbols, 0 unstable. The legacy image.so / _C.so stay UNSTABLE (expected: only nms and the jpeg decoder have
migrated), so the package-level verdict stays UNSTABLE until the rest move — same mid-migration shape as torchcodec.

 Package: torchvision
   Torch ABI:   UNSTABLE
   CPython ABI: n/a
   Bundled libs: 4
   -- bundled libs --
     [UNSTABLE] [not-abi3        ] _C.so            (stable_shim=0,  unstable=110)
     [STABLE  ] [not-abi3        ] _C_stable.so     (stable_shim=67, unstable=0)
     [UNSTABLE] [uses-private-api] image.so         (stable_shim=0,  unstable=71)
     [STABLE  ] [not-abi3        ] image_stable.so  (stable_shim=65, unstable=0) 

@pytorch-bot

pytorch-bot Bot commented Jun 26, 2026

Copy link
Copy Markdown

🔗 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 Failures

As of commit bd9f4cd with merge base bc721a5 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the cla signed label Jun 26, 2026
@adabeyta
adabeyta requested a review from NicolasHug June 26, 2026 03:08

@NicolasHug NicolasHug left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks a lot @adabeyta ! LGTM just left two small comments below

Comment on lines +102 to +110
// 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");
}

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.
*/

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@adabeyta
adabeyta force-pushed the decoder-jpg-stable branch from 3b950a8 to a7d2768 Compare June 26, 2026 13:39
@NicolasHug
NicolasHug merged commit f098790 into pytorch:main Jun 26, 2026
56 checks passed
@github-actions

Copy link
Copy Markdown

Hey @NicolasHug!

You merged this PR, but no labels were added.
The list of valid labels is available at https://github.com/pytorch/vision/blob/main/.github/process_commit.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants