Things compile twice... why? And how to track cache hits/misses?

[Jacob-Stevens-Haas]

TL;DR: I had assumed that jit(foo) would look up foo in some global memo of jit to give the same object (JitWrapped before execution, or PjitFunction after) each invocation, so that jit(foo)() would always access the same compilation cache. When I discovered that wasn't the case, I tried to check the cache hits, but how do I do that?

Situation

I was trying to understand which functions were being jitted, and how often, in a situation like:

from jax import jit

def do_foo(func: Callable, use_jit=True):
    if use_jit:
        jit(func)()

In my case, an iterative algorithm can either jit the iteration steps or not, depending on an argument (no jit is helpful for debugging, jit is faster in real problems). I solved a performance issue by looking at the debug logs: func was being defined locally (and redundantly) inside a function, so jit(func) thought was seeing a new object each time and was recompiling every iteration. I turned it into a module-level function so it wouldn't get redefined. Performance boosted dramatically.

Nevertheless, I was surprised to find that, if func is a module defined function, it appeared to get compiled exactly twice:

related logs from my project

Finished tracing + transforming _evaluate_objective for pjit in 0.15575790405273438 sec
Compiling _evaluate_objective with global shapes and types [ShapedArray(float64[3030])]. Argument mapping: [UnspecifiedValue].
Finished jaxpr to MLIR module conversion jit(_evaluate_objective) in 0.4484984874725342 sec
'jit__evaluate_objective' took at least 1.00 seconds to compile (1.18s)
Finished XLA compilation of jit(_evaluate_objective) in 1.1918911933898926 sec
Finished tracing + transforming _evaluate_objective for pjit in 0.2557542324066162 sec
Compiling _evaluate_objective with global shapes and types [ShapedArray(float64[3030])]. Argument mapping: [UnspecifiedValue].
Finished jaxpr to MLIR module conversion jit(_evaluate_objective) in 0.1098167896270752 sec
'jit__evaluate_objective' took at least 1.00 seconds to compile (1.33s)
Finished XLA compilation of jit(_evaluate_objective) in 1.3336119651794434 sec

Relatedly, multiple compilations of the same function were yielding exactly two distinct python objects. Here's a MWE:

MWE

from typing import Callable
from jax import jit

def foo():
    return 1

def do_foo(func: Callable, use_jit=True):
    if use_jit:
        func = jit(func)
    func()
    print(id(func))



do_foo(foo)
do_foo(foo)
do_foo(foo)
do_foo(foo)

print(jit(foo)._cache_size())

126667197689392
126667208965904
126667740725312
126667208965904
1

HOWEVER, this changed when I changed information around the call to jit():

def do_foo(func: Callable, use_jit=True):
    if use_jit:
        func = jit(func)
    func()
    print(id(func))
    return func

Running do_foo(foo) now gives four different object ids:

126667197691552
126667197690832
126667208965904
126667197689392

What's going on here?

Now there's four objects, one of which existed in the previous invocation, for five objects total (so it's not one per device, of which I have 4). I had previously thought jit(foo)() would avoid recompilation because either jit had a memo to lookup foo or there was a global memo, and then jit(foo) had a cache of compilations against different types/sizes of arguments to foo(). But obviously some global cacheing exists because the first MWE reused object ids. What's going on here?

[jakevdp]

The JIT cache is in play here, and as you've seen you end up with only one compilation (which can be tracked via func._cache_size()). You're correct that in both cases you've created four distinct Python function objects with different object ids, but each of these when called will execute the same cached XLA operation, and thus avoid recompilation. The JIT cache is tied to the id of the original, wrapped function, not the functions returned by jit. I hope that's clear!

[Jacob-Stevens-Haas]

That is clear, thank you! Out of curiosity, then, is there a way to track cache hits/misses?

[yashk2810]

There's no public API for it but you can do tricks like this:

tracing_count = 0

@jax.jit
def f():
  tracing_count += 1
  ...

f()
f()
f()

assert tracing_count == 1

[jakevdp]

There's also jax.log_compiles