mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Pallas/Fuser] DCE fusion jaxprs before pulling (to avoid unnecessary computations being staged out in block functions)

Browse Source 
PiperOrigin-RevId: 738218113
This commit is contained in:
Sharad Vikram 2025-03-18 18:59:50 -07:00 committed by jax authors
parent 4d715753c4
commit e949effcda
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/_src/pallas/fuser/block_spec.py
View File

@ -244,6 +244,13 @@ def pull_block_spec(
_unwrap_block_spec_scalar_prefetch, out_block_specs
)
flat_block_specs, out_tree = jax.tree.flatten(block_specs_)
jaxpr, used_consts, used_invars = pe.dce_jaxpr_consts(
jaxpr,
used_outputs=[True] * len(jaxpr.outvars),
instantiate=True,
)
assert all(used_invars)
assert all(used_consts)
in_block_specs, env, read_usage_env = _pull_block_spec(
jaxpr,
tuple(flat_block_specs),