Skip to content

Add iree_metal to platforms with buffer donation support #2

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.

Sign up for GitHub

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

Jump to bottom
Open
robtaylor wants to merge 380 commits into
base: main
Choose a base branch
from
Open

Add iree_metal to platforms with buffer donation support #2

robtaylor wants to merge 380 commits into from

Conversation

@robtaylor
Copy link
Owner

@robtaylor robtaylor commented Dec 24, 2025

Summary

  • Add iree_metal to _platforms_with_donation list in jax/_src/interpreters/mlir.py

The IREE Metal PJRT plugin now supports buffer donation, which allows JAX to reuse input buffers for outputs when the input is no longer needed. This optimization reduces memory allocation overhead and improves performance for workloads that can donate buffers.

Background

Buffer donation is an optimization where JAX can inform the runtime that an input buffer can be reused for an output, avoiding extra memory allocations. The IREE Metal PJRT plugin implements:

  • Duplicate buffer detection (same buffer donated multiple times)
  • Post-execution buffer invalidation
  • Runtime metadata extraction from compiled modules

Test plan

  • Run buffer donation tests with IREE Metal backend
  • Verify donated buffers are properly invalidated after execution
  • Confirm double-donation is detected and raises appropriate errors

🤖 Generated with Claude Code

mwhittaker and others added 30 commits December 8, 2025 15:46
@mwhittaker @Google-ML-Automation
Change PjRt to use new copy of coordination service.
0339ad1 
PiperOrigin-RevId: 841944431
@sharadmv @Google-ML-Automation
[Pallas SC] Allow semaphores to be returned by SCS kernels
2a6de35 
PiperOrigin-RevId: 841972235
@yashk2810 @Google-ML-Automation
Deprecate with mesh: context manager. Use `with jax.set_mesh(mesh):…
890ccd2 
…` instead

PiperOrigin-RevId: 841977281
@Google-ML-Automation
Add a jax network transfer benchmark script
c74b2a5 
PiperOrigin-RevId: 841983057
@pschuh @Google-ML-Automation
Add ResultHandler.wrap which allows us to avoid exploding PRNG keys.
249b7a7 
PiperOrigin-RevId: 842010889
@Google-ML-Automation
[Pallas] Device Id dict to mesh fastpath for power of twos
eff53a5 
PiperOrigin-RevId: 842041847
@Google-ML-Automation
Update XLA dependency to use revision http://github.com/openxla/xla/c…
9be13f9 
…ommit/91b3f740b75d1d932a12fb0886338f84f856a453

PiperOrigin-RevId: 842096427
@bchetioui @Google-ML-Automation
[Mosaic GPU] Remove {Least,Most}ReplicatedExpression constructors.
3bec82d 
After getting rid of hints, these constructors are no longer necessary.

PiperOrigin-RevId: 842112040
@allanrenucci @Google-ML-Automation
[Mosaic GPU] Support FragmentedArray.broadcast_in_dim for splat to …
e288c27 
…other layouts.

PiperOrigin-RevId: 842145211
@jakevdp
[maint] clean up jnp.arange implementation
1bfd8dc
@cantonios @Google-ML-Automation
Update generated stubs.
ea5aee9 
PiperOrigin-RevId: 842264374
@yashk2810 @Google-ML-Automation
Fix back compat test to ignore warnings
2f62fb1 
PiperOrigin-RevId: 842272730
@superbobry @Google-ML-Automation
[pallas:mosaic_gpu] Slightly tweaked the error messages in a few places
50bdc72 
PiperOrigin-RevId: 842272894
@gnecula
[export] Add backwards compatibility tests for serialized exports
6cd6cc7 
Recently we have found the need to evolve the serialization of
`jax.export.Exported`. E.g., in jax-ml#33942 we have added a 32-bit
representation for `nr_devices`. This introduced a compatiblity
bug that was found by usersm and fixed in in jax-ml#33685.

Here we add backwards compatibility tests. See the description
in the `export_serialization_back_compat_test.py` module docstring.

Note that this is separate from our previous set of backwards
compatibility tests for the lowering of custom calls
(in `export_back_compat_test.py`). However, we reuse some
of the same ideas, and we use the same directory for
storing saved old serializations.
@krishnaharidasan @Google-ML-Automation
Use ifrt::AttributeMap::Get instead of directly accessing map
14f072b 
Introduces a variant of Get in AttributeMap that returns the value variant as is.

PiperOrigin-RevId: 842283537
@Google-ML-Automation
Merge pull request jax-ml#33758 from gnecula:export_back_compat
64b16f0 
PiperOrigin-RevId: 842284947
@yashk2810 @mattjj @Google-ML-Automation
Mention pmap is in maintenance mode and point to shard_map and the mi…
f7e3bdb 
…gration guide

Co-authored-by: Matthew Johnson <mattjj@google.com>
PiperOrigin-RevId: 842298387
@superbobry @Google-ML-Automation
[mosaic] Added documentation and a few useful methods to tpu.tiled at…
6cb91ed 
…tribute

PiperOrigin-RevId: 842307512
@emilyfertig @Google-ML-Automation
Remove multiprocess tests from TPU presubmit due to latency.
863e4e7 
They still run in CI.

PiperOrigin-RevId: 842313199
@mwhittaker
Adds docs on multi-controller JAX fault tolerance.
3a36656
@jakevdp
[test] fix signatures test for NumPy nightly
f1bc1a5
@Google-ML-Automation
Merge pull request jax-ml#33819 from jakevdp:arange-followup
eeab9e4 
PiperOrigin-RevId: 842326801
@jakevdp
[test] remove pre-NumPy 2.0 API skips
b8cab94
@Google-ML-Automation
Merge pull request jax-ml#32218 from mwhittaker:fault_tolerance_docs
4e15a87 
PiperOrigin-RevId: 842328162
@Google-ML-Automation
Merge pull request jax-ml#33834 from jakevdp:numpy-nightly-sig
3e56355 
PiperOrigin-RevId: 842328715
@dougalm @mattjj
Remove dynamic shapes. Dead weight at this point.
75697ef
@jakevdp
jnp.arange: deprecate passing complex arguments
c006c30
@Google-ML-Automation
Merge pull request jax-ml#33836 from jakevdp:old-np-apis
b39d43a 
PiperOrigin-RevId: 842344094
@Google-ML-Automation
Merge pull request jax-ml#33565 from jax-ml:remove-dynamic-shapes
53ca7bb 
PiperOrigin-RevId: 842347956
@Google-ML-Automation
Merge pull request jax-ml#33837 from jakevdp:arange-complex-dep
f374387 
PiperOrigin-RevId: 842354532
jakevdp and others added 30 commits December 19, 2025 09:53
@jakevdp
[dep] remove several dozen finalized deprecations for v0.9.0
e27c8dd 
These names were already deprecated in v0.8.0 or prior; this change means that the custom
AttributeError will be replaced with a generic AttributeError.
@belitskiy @Google-ML-Automation
Remove an erroneously added license.
a26a544 
PiperOrigin-RevId: 846781195
@yashk2810 @Google-ML-Automation
Default make_mesh to Explicit axes for the upcoming 0.9.0 release
89401bd 
PiperOrigin-RevId: 846787187
@Google-ML-Automation
Merge pull request jax-ml#33828 from rao-ashish:asrao/cross_host_tran…
a82d51d 
…sfer_batching

PiperOrigin-RevId: 846791377
@Google-ML-Automation
Merge pull request jax-ml#33884 from melissawm:advanced-ad
0d0d491 
PiperOrigin-RevId: 846791567
@ybaturina @Google-ML-Automation
Enable using custom hermetic NCCL version.
cfc997c 
The NCCL version can be chosen via `HERMETIC_NCCL_VERSION` env var.

See docs [here](https://github.com/google-ml-infra/rules_ml_toolchain/blob/main/gpu/README.md#environment-variables-controlling-the-hermetic-cudacudnnnvshmem-versions).

PiperOrigin-RevId: 846797606
@Google-ML-Automation
Merge pull request jax-ml#34058 from jakevdp:finalize-deps
f29de06 
PiperOrigin-RevId: 846799030
@Google-ML-Automation
[dep] remove deprecated jax_safer_randint configuration for JAX v0.9.0
a9dac11 
PiperOrigin-RevId: 846803402
@Google-ML-Automation
[dep] finalize deprecation of lax.dot positional arguments for JAX v0…
f314722 
….9.0.

PiperOrigin-RevId: 846811504
@jakevdp
[dep] remove references to the already-deprecated interpolation argument
d6944d3
@belitskiy @Google-ML-Automation
Remove pre-0.8.2 version guards.
1085d3a 
PiperOrigin-RevId: 846830895
@Google-ML-Automation
[dep] Remove deprecated jax.lib submodules for JAX v0.9.0.
82f02eb 
The following modules are removed:

- `jax.lib.xla_bridge`
- `jax.lib.xla_client`
- `jax.lib.xla_extension`

All contents of these submodules were deprecated and removed as of JAX v0.8.0; the modules themselves have been raising warnings on import since this release.

PiperOrigin-RevId: 846831442
@Google-ML-Automation
Merge pull request jax-ml#34060 from jakevdp:interpolation-dep
3f8f860 
PiperOrigin-RevId: 846841818
@mattjj
[hijax] upgrade VJPHiPrimitive.def_fwd to support symbolic zeros
ec5015f
@Google-ML-Automation
Merge pull request jax-ml#34065 from mattjj:custom-vjp3-revival
69518c5 
PiperOrigin-RevId: 846849490
@Google-ML-Automation
Integrate LLVM at llvm/llvm-project@7d381f2a5634
4f0f93e 
Updates LLVM usage to match
[7d381f2a5634](llvm/llvm-project@7d381f2a5634)

PiperOrigin-RevId: 846858892
@Google-ML-Automation
[dep] make jax_default_dtype_bits config a no-op for JAX v0.9.0.
514b7c4 
This has been deprecated since JAX v0.8.0; after this change the flag still exists, but setting it raises a warning and otherwise has no effect. It will be removed in JAX v0.10.0.

PiperOrigin-RevId: 846877149
@pschuh @Google-ML-Automation
Support weakref_lru_cache.evict.
dd7ff37 
PiperOrigin-RevId: 846909748
@yashk2810 @Google-ML-Automation
Add replicated -> unreduced test coverage
d3ed2f4 
PiperOrigin-RevId: 846930629
@yashk2810 @Google-ML-Automation
Fix a broken sharded -> unreduced test
20bdca7 
PiperOrigin-RevId: 846986071
@Google-ML-Automation
Update XLA dependency to use revision http://github.com/openxla/xla/c…
6469f17 
…ommit/66dbbf501ffd74f83c6a5d8fc201c756b1198d64

PiperOrigin-RevId: 847021722
@Google-ML-Automation
Update XLA dependency to use revision http://github.com/openxla/xla/c…
0ad46d9 
…ommit/08872f587d442a05802cbdb052e8c9e6e87423f4

PiperOrigin-RevId: 847325574
@Google-ML-Automation
Update XLA dependency to use revision http://github.com/openxla/xla/c…
f29645b 
…ommit/7521349eccb22a780f50fa5f6f09dbaa1d09f470

PiperOrigin-RevId: 847636908
@allanrenucci @Google-ML-Automation
[Mosaic GPU] Add nvshmemx_cumodule_finalize support.
2da9713 
PiperOrigin-RevId: 847665237
@derdrdirk @Google-ML-Automation
[Autotuner] Prepare the CuDNN fusion test for the new autotuner.
6d41fa0 
PiperOrigin-RevId: 847713430
@subhamsoni-google @Google-ML-Automation
Add programmatic profiling test with session_id support in JAX.
c25a514 
PiperOrigin-RevId: 847801848
@WindQAQ @Google-ML-Automation
[Mosaic] Extend tpu.pack_elementwise to support non-32-bit integers.
9fe0430 
PiperOrigin-RevId: 847871002
@Google-ML-Automation
Adjust tile_n assert to accommodate 2 cta in tcgen05 blockscale mma.
5173d18 
PiperOrigin-RevId: 847908521
@Google-ML-Automation
Update XLA dependency to use revision http://github.com/openxla/xla/c…
82ae1b1 
…ommit/d1635d1c99de225c8029d82e56c4dd03f90b013f

PiperOrigin-RevId: 848049355
@robtaylor @claude
Add iree_metal to platforms with buffer donation support
5c2e646 
The IREE Metal PJRT plugin now supports buffer donation, which allows
JAX to reuse input buffers for outputs when the input is no longer
needed. This optimization reduces memory allocation overhead.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.