Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix member access to packed CUDA struct #26691

Merged
merged 1 commit into from
Mar 5, 2025

Conversation

h-vetinari
Copy link
Contributor

The mosaic code uses the CUDA CUpti_ActivityKernel9 struct, which (at least as of CUDA 12.6) has been declared with __attribute__((packed)). This makes it UB to access a reference to a potentially unaligned field, causing GCC 13 to raise an error:

jaxlib/mosaic/gpu/mosaic_gpu_ext.cc: In function 'void jax::cuda::{anonymous}::callback_complete(CUcontext, uint32_t, uint8_t*, size_t, size_t)':
jaxlib/mosaic/gpu/mosaic_gpu_ext.cc:185:37: error: cannot bind packed field 'kernel->CUpti_ActivityKernel9::name' to 'const char*&'
  185 |             std::make_tuple(kernel->name, duration_ms));
      |                             ~~~~~~~~^~~~

Work around this by separating out the member access, which avoids the packed struct reference issue because then it's just a copy of the pointer, not a reference. We needed this to unblock the distribution of jax in conda-forge (where we're currently on GCC 13.3 & CUDA 12.6): conda-forge/jaxlib-feedstock#295

Copy link
Collaborator

@apaszke apaszke left a comment

Choose a reason for hiding this comment

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

Thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Mar 4, 2025
@copybara-service copybara-service bot merged commit a13b3ce into jax-ml:main Mar 5, 2025
17 of 22 checks passed
@h-vetinari h-vetinari deleted the packed branch March 6, 2025 02:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants