Jitting nested functions

[AimbotParce]

I am trying to run a simulation with JAX for my Final Degree Project. This simulation uses two large functions to compute two sides of an equation, but these functions are interchangeable. For instance, I might have a computeLeft() function and a computeLeftDimensionless(), and I want to be able to switch them at will. Finally, a mainLoop() function solves the equation for a given ammount of iterations.

So, once I've selected which functions to use for the left and right of the equation, I want to jit them, and pass them to the mainLoop function, which should also be jitted. I think this is how it should be done:

jitCompLeft = jax.jit(computeLeft, backend="gpu")
jitCompRight = jax.jit(computeRight, backend="gpu")
jitMainLoop = jax.jit(mainLoop, backend="gpu", static_argnums=(0, 1, 2))

psi = jitMainLoop(computeLeft, computeRight, iterCount, psi, ... )

Where mainLoop is something like the following:

def mainLoop(computeLeft, computeRight, iterCount, psi, ... ):
    for iteration in range(0, iterCount):
        A = computeLeft( ... )
        B = computeRight( ... )

        right = B @ psi[iteration]
        psi = psi.at[iteration + 1].set(jnp.linalg.solve(A, right))
    return psi

I am pretty sure I might be missing something here, but my program just gets stuck after tracing the jitMainLoop function, and never seems to do anything.

[mattjj]

Thanks for the question!

I'm guessing your code is spending all its time compiling because there are large unrolled loops (i.e. Python loop construct which gets unrolled when JAX traces and compiles it). I can't tell about computeLeft and computeRight, but it sure looks like mainLoop has a potentially-large unrolled loop. Try removing the jit from mainLoop?

If that's the issue, and you still want mainLoop to be jit-compiled, try rewriting that loop to be a jax.lax.scan. That's how you stage out and compile rolled (rather than unrolled) loops.

You can also set jax.config.update('jax_log_compiles', True) to see if it's getting stuck in a compilation step.

[AimbotParce]

I am completely new here, but I am guessing that having the mainLoop unjitted means that psi, A and B will be allocated on my RAM, and they'll have to be transfered to the GPU's RAM every time A and B are computed, isn't it? The thing is that computeLeft and computeRight aren't that slow, and I wanted to speed up the B @ psi and solve part of the simulation, which looked to me like the slower parts of it.

Am I on the wrong here?

PS: Thanks for the quick reply :)

[mattjj]

psi, A and B will be allocated on my RAM, and they'll have to be transfered to the GPU's RAM every time A and B are computed, isn't it?

Actually, no! the outputs of any jax.jit computation will stay on the device. They're only transferred back to CPU memory lazily, i.e. only when something on the CPU really needs their value, e.g. for control flow, or plotting, or saving to disk.

and I wanted to speed up the B @ psi and solve part of the simulation, which looked to me like the slower parts of it.

Those will also happen on-device actually: any operation on a jax.Array, like those produced as output from a jax.jit-decorated function application, will be performed on device by default. If B is on-device, then the @ operator on it in the expression B @ anything will be an on-device jax.numpymatmul.

[AimbotParce]

Then I don't really understand why is the simulation running faster on my CPU rather than my GPU. I really thought it was due to memory being moved up and down. Do you have any idea?

Just for the record, here is a quick benchmark I did:

[mattjj]

Try profiling? https://jax.readthedocs.io/en/latest/profiling.html