jnp.zeros((9500000,3), dtype=jnp.float32) lowers into slow broadcast.

[chihuahua]

Description

This call here ...

jax.numpy.zeros((9500000,3), dtype=jnp.float32)

lowers into this broadcast HLO call, which broadcasts a float constant of 0 to a giant ND array of (9500000,3).

@@f32[9500000,3]{1,0} fusion(), kind=kLoop, calls=
 {
   tmp_0 = f32[] constant(0)
   ROOT tmp_1 = f32[9500000,3]{1,0} broadcast(f32[] tmp_0), dimensions={}
 }

The broadcast call takes a while to execute on GPU (say a V100): 130 microseconds. Is there a faster way to generate the ND array of all zeros? This is slowing down some key models in Alphabet.

System info (python version, jaxlib version, accelerator, etc.)

• python3
• jaxlib version: jax at repository HEAD
• GPU: V100. I can reproduce this on other GPUs though.

[justinjfu]

Do you have a minimal code example that can reproduce?

Just profiling jax.block_until_ready(jax.numpy.zeros((9500000,3), dtype=jnp.float32)) in isolation takes <500us on my end on a V100.

[chihuahua]

Oh, XLA JIT compilation should be used to reproduce. Thank you! I can repro a compute time latency of 131 microseconds with

@jax.jit
def run_zeros() -> jax.Array:
  """Run the main logic to benchmark."""
  with jax.named_scope("benchmark_broadcast"):
    return jnp.zeros((9500000, 3), dtype=jnp.float32)

values = run_zeros()
  jax.block_until_ready(values)
  print(values)

The resulting HLO indicates the slow broadcast:

@@f32[9500000,3]{1,0} fusion(), kind=kLoop, calls=
 {
   tmp_0 = f32[] constant(0)
   ROOT tmp_1 = f32[9500000,3]{1,0} broadcast(f32[] tmp_0), dimensions={}
 }

[chihuahua]

Might it be possible that the CUDA implementation for broadcast here only uses 1 SM? And is thus not parallelized?

[hawkinsp]

Ah, it's 130 microseconds, not milliseconds?

I don't think it's physically possible to do better?

A V100 has 900GB/s memory bandwidth, per the specs. 9500000 * 3 * 4 bytes / 900e9 bytes/sec = 127us.

(I'm a bit surprised here, actually, I thought NVIDIA's published bandwidth numbers are bidirectional, which would mean twice as long, but your measurement seems to indicate otherwise.)

So... unless I'm mistaken, you're almost getting roofline memory bandwidth and you should not expect to do better?

[chihuahua]

Ohhh... Apologies. It is 131 microseconds! I have corrected the original post.

[chihuahua]

Thank you for the analysis! I suppose allocating a new memory block (versus reading from existing memory) also encounters this memory bandwidth. That would present a bottleneck indeed that means we can't do better.