Add support for FFI calls with side effects via ffi_call #23982
Merged
+44
−5
Conversation
This file contains 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
hawkinsp
approved these changes
Sep 27, 2024
tests/extend_test.py
Outdated
| def testEffects(self): | ||
| def fun(): | ||
| jex.ffi.ffi_call("test_ffi", (), has_side_effect=True) | ||
| self.assertIn("test_ffi", jax.jit(fun).lower().as_text()) |
There was a problem hiding this comment.
Probably you should verify that it doesn't get DCEd from the compiled IR, not the input to the compiler?
Or at least check the input IR has the side effect annotation?
There was a problem hiding this comment.
Good suggestion! Done.
dfm
force-pushed
the
ffi-call-effects
branch
from
68e1544
to
d80a89d
Compare
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.
As discussed in #23963, the
jax.extend.ffi.ffi_callfunction doesn't currently support FFI calls with side effects. Since we already have support for lowering effectful custom calls, plumbing this through toffi_callis straightforward. I think it's reasonable to stick to the the same terminology here and call the boolean parameterhas_side_effect.If needed it would also be possible to add support for ordering via tokens since support for tokens was added in #23394, but I didn't want to do too much in this PR.