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

[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.