lax.scan very slow on GPU

[fbartolic]

I have an array of complex numbers z with shape (10, 2000). I need to permute each column of this array such that the element-wise distance the complex numbers in the i-th and the i+1th columns in minimized, starting with the first column sequantially. I came up with this solution involving lax.scan :

import jax.numpy as jnp
from jax import jit, vmap, lax, random
from jax.config import config
config.update('jax_platform_name', 'gpu')

key = random.PRNGKey(42)
key, subkey = random.split(key)
z = random.uniform(subkey, (10, 2000))

@jit
def linear_sum_assignment(a, b):
   idcs_d = jnp.argsort(jnp.abs(b - a[:, None]), axis=1)    
   idcs_final = jnp.repeat(999, len(a))
   
   def f(carry, idcs_d_row):
       i, idcs_final = carry
       cond1 = jnp.isin(idcs_d_row[0], jnp.array(idcs_final))
       cond2 = jnp.isin(idcs_d_row[1], jnp.array(idcs_final))
       
       idx_closest = jnp.where(
           cond1, jnp.where(cond2, idcs_d_row[2], idcs_d_row[1]), idcs_d_row[0]
       )
       idcs_final = idcs_final.at[i].set(idx_closest)
       return (i + 1, idcs_final), idx_closest
   
   _, res = lax.scan(f, (0, idcs_final), idcs_d)
   
   return res

@jit
def f(z):
    def apply_linear_sum_assignment(carry, z_slice):
        idcs = linear_sum_assignment(carry, z_slice)
        return z_slice[idcs], z_slice[idcs]

    _, z_permuted = lax.scan(
        apply_linear_sum_assignment, z[:, 0], z.T
    )
    return z_permuted
  
# Slow on GPU
res = f(z)

This code is more than an order of magnitude slower on the GPU than the CPU. Am I correct in thinking that this result is expected because the input to linear_sum_assignment depends on carry and GPUs just aren't suited for this kind of problem?

[YouJiacheng]

scan is slow on accelerator, especially when there is inplace-update to carry(strong data dependency).
And you use two nested scan, which cannot leverage parallelism capability of GPU at all.

[jakevdp]

To elaborate on this, the reason GPU is so fast for vectorized operations is not that individual floating point operations are particularly fast (they're actually often slower than similar operations on a CPU!), but rather that it can very efficiently run many operations in parallel. For an operation like scan in which each step depends on the output of the previous, the sequence of operations as a whole cannot be parallelized. So you end up not taking advantage of any of the GPU's inherent parallelism, and the result is slow execution.

Contrast this to CPU, where individual floating point operations are relatively fast, but there is no so much in-built parallelism available. Because of this, scan does not incur as much of a performance penalty.

[fbartolic]

Thank you, that makes sense!

After posting this, I realised that I can just specify the device for each JIT-ed function and execute this particular function on a CPU while using the GPU for everything else.

[YouJiacheng]

Do you need an exact minimization, or an approximation one is okay?
Maybe you can try simulated annealing algorithm!

[fbartolic]

It needs to be exact so I doubt there's a much for efficient solution to this problem than the brute-force search approach I've implemented. Each value in the array represents a point in the complex plane which needs to be assigned to a different curve consisting in the correct order.

[hawkinsp]

BTW, scan has an unroll parameter that may help in cases where the loop body is cheap. It's worth a try!