What if you can't use block_until_ready?

[ntlm1686]

I am trying to benchmark some code. However, it seems I can't use block_until_ready everywhere.

Is this a good way to benchmark this code?

@jit
def polyval_gpu():
    return [jnp.polyval(r, 0.1) for r in SOME_DATA]

%timeit polyval_gpu()

OR

@jit
def polyval_gpu():
    for r in SOME_DATA:
        ans = jnp.polyval(r, 0.1)
    return ans

%timeit polyval_gpu()

[jakevdp]

For non-array outputs, rather than calling the array.block_until_ready() method, you can call the jax.block_until_ready() function:

%timeit jax.block_until_ready(polyval_gpu())

[ntlm1686]

Are these non-array outputs moved from GPU to CPU after calculated?
For example, in the code below

@jit
def polymul_many_times_gpu():
    return [jnp.polymul(p, p) for p in P1]

[YouJiacheng]

Are these non-array outputs moved from GPU to CPU after calculated? For example, in the code below

@jit
def polymul_many_times_gpu():
    return [jnp.polymul(p, p) for p in P1]

No, unless you print them.

[jakevdp]

No, block_until_ready() does not materialize outputs to different devices, it simply blocks Python program execution until those arrays are finished computing.

[ntlm1686]

@jakevdp @YouJiacheng
Thank you very much. I think my question has been solved.

[YouJiacheng]

jax.block_until_ready will tree_map array.block_until_ready() for you.
But you still cannot(and should not) use block_until_ready inside jit, you should use it on the return value of jitted function.
BTW, useless computations are likely(almost always) optimized by JIT.

@jit
def polyval_gpu():
    for r in SOME_DATA:
        ans = jnp.polyval(r, 0.1)
    return ans

%timeit jax.block_until_ready(polyval_gpu())

polyval_gpu may only evaluate jnp.polyval(r, 0.1) at the last data in SOME_DATA, thus the performance measured may be unreliable.

[ntlm1686]

Thank you very much.