Repository navigation
Fix UB in FFI State Get/Set Backwards Compatibility Case - #39155
Closed
mgoldfarb-nvidia wants to merge 3 commits into
Closed
mgoldfarb-nvidia wants to merge 3 commits into
mgoldfarb-nvidia wants to merge 3 commits into
Conversation
Also fixes the implementation to correctly set the state.
Contributor
Author
|
FYI @ezhulenev |
ezhulenev
reviewed
Mar 13, 2026
| auto* v02 = reinterpret_cast<XLA_FFI_State_Args_V02*>(args); | ||
| // memcpy to safely avoid U.B. | ||
| XLA_FFI_State_Args_V02 v02{}; | ||
| std::memcpy(&v02, args, XLA_FFI_State_Args_V02_STRUCT_SIZE); |
Member
There was a problem hiding this comment.
Why memcpy and what UB it avoids?
Contributor
Author
There was a problem hiding this comment.
Strict aliasing between two unrelated types causes UB. So reads and writes through args may not reflect . The memcpy is the pre-C++20 compliant way to circumvent this rule. In many cases compiler should optimize away an actual copy.
Member
There was a problem hiding this comment.
Can you use absl::bit_cast instead? It's a backport of std::bit_cast and is basically a wrapper around the memcpy pattern; plus some additional checks.
Contributor
Author
There was a problem hiding this comment.
One problem with absl::bit_cast (and std::bit_cast) is it (looks like it) requires the type sizes to be equal for writes. At least in the fallback path that uses std::memcpy
inline Dest bit_cast(const Source& source) {
Dest dest;
memcpy(static_cast<void*>(std::addressof(dest)),
static_cast<const void*>(std::addressof(source)), sizeof(dest));
return dest;
}
Its possible we could modify this function s.t. we pass void* from the FFI API then first cast to size_t to determine the size then cast to one concrete type. However, this could still be a U.B. violation and a bigger change than makes sense unless we start adding more cases where we are making ABI breaking changes.
Member
There was a problem hiding this comment.
Oh, I see. In that case the current version is fine.
ezhulenev
approved these changes
Mar 13, 2026
beckerhe
reviewed
Mar 13, 2026
| auto* v02 = reinterpret_cast<XLA_FFI_State_Args_V02*>(args); | ||
| // memcpy to safely avoid U.B. | ||
| XLA_FFI_State_Args_V02 v02{}; | ||
| std::memcpy(&v02, args, XLA_FFI_State_Args_V02_STRUCT_SIZE); |
Member
There was a problem hiding this comment.
Can you use absl::bit_cast instead? It's a backport of std::bit_cast and is basically a wrapper around the memcpy pattern; plus some additional checks.
beckerhe
approved these changes
Mar 13, 2026
copybara-service Bot
pushed a commit
that referenced
this pull request
Mar 13, 2026
Imported from GitHub PR #39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02 by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6 by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 Reverts d1bced5 FUTURE_COPYBARA_INTEGRATE_REVIEW=#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6 PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Mar 13, 2026
Imported from GitHub PR openxla/xla#39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02a0aa43ff92219ba2b1aa4ddf64306e1ca by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6bc35632fe2984021f43d9b1414d4c7c3b by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 Reverts 95aacc5 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6bc35632fe2984021f43d9b1414d4c7c3b PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Mar 16, 2026
Imported from GitHub PR openxla/xla#39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02a0aa43ff92219ba2b1aa4ddf64306e1ca by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6bc35632fe2984021f43d9b1414d4c7c3b by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6bc35632fe2984021f43d9b1414d4c7c3b PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
that referenced
this pull request
Mar 16, 2026
Imported from GitHub PR #39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02 by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6 by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 FUTURE_COPYBARA_INTEGRATE_REVIEW=#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6 PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
that referenced
this pull request
Mar 16, 2026
Imported from GitHub PR #39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02 by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6 by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 FUTURE_COPYBARA_INTEGRATE_REVIEW=#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6 PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Mar 16, 2026
Imported from GitHub PR openxla/xla#39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02a0aa43ff92219ba2b1aa4ddf64306e1ca by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6bc35632fe2984021f43d9b1414d4c7c3b by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#39155 from mgoldfarb-nvidia:mgoldfarb/fix_ffi_backwards_compat_bug_state d18d5b6bc35632fe2984021f43d9b1414d4c7c3b PiperOrigin-RevId: 883152448
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Mar 16, 2026
Imported from GitHub PR openxla/xla#39155 📝 Summary of Changes Replaces `reinterpret_cast` with `std::memcpy` to avoid U.B. Implementation of Get_State was not setting state pointer when using the backward compatibility path 🎯 Justification Support state with V0.2 FFI in latest JAX. 🚀 Kind of Contribution 🐛 Bug Fix Copybara import of the project: -- b0bed02a0aa43ff92219ba2b1aa4ddf64306e1ca by Michael Goldfarb <[email protected]>: Fix UB in state get/set when using backwards compat mode. Also fixes the implementation to correctly set the state. -- d18d5b6bc35632fe2984021f43d9b1414d4c7c3b by Michael Goldfarb <[email protected]>: Fix write to state. Merging this change closes #39155 PiperOrigin-RevId: 884428841
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📝 Summary of Changes
Replaces
reinterpret_castwithstd::memcpyto avoid U.B.Implementation of Get_State was not setting state pointer when using the backward compatibility path
🎯 Justification
Support state with V0.2 FFI in latest JAX.
🚀 Kind of Contribution
🐛 Bug Fix