Slow down using GPU

[jecampagne]

Hello,
Here is a snippet:

import jax
jax.config.update("jax_enable_x64", True)
import numpy as np
import jax.numpy as jnp

jax.clear_caches()
 
@jax.jit
def wrap_nonhermitian(im, xmin, ymin, nxwrap, nywrap):
    def _body_j(j, vals):
        i, im = vals

        ii = (i - ymin) % nywrap + ymin
        jj = (j - xmin) % nxwrap + xmin

        im = jax.lax.cond(
            # weird way to say if ii != i and jj != j
            # I tried other ways and got test failures
            jnp.abs(ii - i) + jnp.abs(jj - j) != 0,
            lambda im, i, j, ii, jj: im.at[ii, jj].add(im[i, j]),
            lambda im, i, j, ii, jj: im,
            im,
            i,
            j,
            ii,
            jj,
        )

        return [i, im]

    def _body_i(i, vals):
        im = vals
        _, im = jax.lax.fori_loop(0, im.shape[1], _body_j, [i, im],unroll=False)
        return im

    im = jax.lax.fori_loop(0, im.shape[0], _body_i, im, unroll=False)
    return im


@jax.jit
def expand_hermitian_x(im):
    return jnp.concatenate([im[:, 1:][::-1, ::-1].conjugate(), im], axis=1)


@jax.jit
def contract_hermitian_x(im):
    return im[:, im.shape[1] // 2 :]


@jax.jit
def wrap_hermitian_x(im, im_xmin, im_ymin, wrap_xmin, wrap_ymin, wrap_nx, wrap_ny):
    im_exp = expand_hermitian_x(im)
    im_exp = wrap_nonhermitian(
        im_exp, wrap_xmin - im_xmin, wrap_ymin - im_ymin, wrap_nx, wrap_ny
    )
    return contract_hermitian_x(im_exp)

Using V100 GPUs using jax v0.4.24 (or T2 in Colab with jax v0.4.26) with observe an expected slow down After warm-up

%timeit tmp_arr = wrap_hermitian_x(\
                jnp.zeros((501, 381),dtype=jnp.complex128),\
                -380,\
                -250,\
                -79,\
                -49,\
                160,\
               10,\
            ).block_until_ready()

leading to 22 sec per iteration while on CPU it is much faster.

We wander if there were new coding in lax that would explain such behaviour?

[jakevdp]

For your condition, did you try (ii != i) & (jj != j)?

Regarding the slowdown: in general scan (which is what fori_loop lowers to here in the case of static endpoints) carries a performance penalty on GPU, because each scan step requires a kernel launch and a host round-trip (this is discussed in #2491 (comment), which is somewhat old but I believe still more-or-less accurate).

You can improve on this in general by passing the unroll parameter to fori_loop, which essentially lets you balance the trade-off between compilation cost and runtime overhead. The default unroll=1 leads to fast compilation but more runtime overhead, while at the other end unroll=im.shape[0] leads to slower compilation (because the compiler optimizes across multiple unrolled steps) but faster runtime (because of fewer kernel launches).

[jecampagne]

Thanks @jakevdp BTW how one can use im.shape[0] at compilation time as im.shape can vary at execution time
?

[jakevdp]

Array shapes are static by definition. If you change them, it leads to recompilation.

[jecampagne]

Hi @jakevdp , I share a Colab nb that for me at least is quite pedagogic. One can see 3 methods of the same task (to sum patches of identical size from a 2D image) that have different timing on CPU & GPU. This is to show that "loop" based methods experience a huge compilation time (especially on T4) and that the nicknamed "roll" method is my favourite (up to now). Of course there might be other methods that one can elaborate.

[jakevdp]

Thanks! Yes this is an important consideration in the tradeoff between fori_loop and Python for loop. for loop is flattened, so leads to longer compile time. fori_loop leads to shorter compilation, but because the compiler cannot fuse operations across loop iterations, runtime may be longer. The unroll option of fori_loop lets you adjust between these extremes.

[jecampagne]

I have put also a method based on index broadcasting and vmaping which seems to get reasonable timing both on CPU and GPU w/o exploding when the image size get large... Anyway, it stresses that tuning the code wrt the use-case is of course an important aspect for choosing an algo (even in this simple task sizeable differences are shown).