Does compilation_cache not work for gpus and cpus?

[OrenLeung]

Does complication_cache not for gpus and cpus? if not, can this be supported?

Code for initializing persistent cache

from jax.experimental.compilation_cache import compilation_cache as cc

cc.initialize_cache("./cache")

@jax.jit
def my_func(...

call my function

Command

JAX_LOG_COMPILES=1 python flax_sharding.py 

Output (every single run)

Initialized persistent compilation cache at ./cache
Finished tracing + transforming jit(create_state) in 0.041712045669555664 sec
Compiling create_state (140247611946464) for with global shapes and types (). Argument mapping: ().
Finished XLA compilation of jit(create_state) in 0.2413489818572998 sec

it never gets to the cache hit/miss log even with logging INFO enabled as suggested in #13325

even with XLA_FLAGS="--xla_gpu_enable_xla_runtime_executable=true" flag it does not work

[jakevdp]

Yes, the compilation cache is an experimental feature that is currently only implemented for TPU backends: https://github.com/google/jax/blob/357d04496579826aaa1d1bde009b458bd20b56d3/jax/_src/dispatch.py#L1040-L1042

@skye might be able to give more information on the status for other platforms.

[hawkinsp]

It's actually not far off from working on GPU also. Perhaps early next year? It depends on a XLA change that is not yet the default but should be soon. I'm not sure about CPU at this point.

[skye]

@ezhulenev do you know if there are any plans for serializable CPU executables in the near future? (This is the underlying feature needed for the compilation_cache to work on cpu)

[ezhulenev]

On GPU it should work today, it is just behind a check in dispatch.py, we will remove check soon, just waiting to see if xla runtime causes any unexpected problems.

[ezhulenev]

It will be available when we’ll switch to a new xla runtime + compiler for XLA:CPU, we still have few failing tests, but we are mostly there. I think that it might be enabled in Jan-Feb, unless we’ll find large blockers.

…

On Tue, Dec 20, 2022 at 1:34 PM Skye Wanderman-Milne < ***@***.***> wrote: @ezhulenev <https://github.com/ezhulenev> do you know if there are any plans for serializable CPU executables in the near future? (This is the underlying feature needed for the compilation_cache to work on cpu) — Reply to this email directly, view it on GitHub <#13736 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAI6W2VIFEE5MOTNTR2YXDDWOIQ6HANCNFSM6AAAAAATE5FLQ4> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[EelcoHoogendoorn]

Whats the status on this now? Cant seem to find any mention of it in the release notes.

[emchristiansen]

Any updates for CPU?

[lfaucheux]

It looks like we'll have to keep a close eye on this.

[nouiz]

Note, it is now working on GPU, but not enabled by default: https://github.com/google/jax/blob/main/jax/_src/compilation_cache.py#L56

[samhita-alla]

Hey @nouiz! Should setting cc.initialize_cache("./cache") work on GPUs?

[nouiz]

It should. I did a small tests and it worked.
If you see issue, raise it.
Note I heard that the handling of XLA_FLAGS should be validated. So if you use that, just make a different cache when you change that.
Otherwise, you should be good.