Asynchronous Dispatch in JAX

[GoktugGuvercin]

Hello JAX Community;

While testing asynchronous dispatch in JAX, I noticed that total duration for dot product of two matrices with and without "block_until_ready()" is almost same. It is like asynchronous dispatch does not work. I cannot see remarkable difference between two measured time slots. I executed the following code on Kaggle platform, and I started to ask myself whether asynchronous dispatch is usable for only GPUs, because I have only CPU.

x = random.uniform(random.PRNGKey(0), (1000, 1000))
%time jnp.dot(x, x)
%time jnp.dot(x, x).block_until_ready()

Time Results:

1. ** total: 59.2 ms **
2. ** total: 64.3 ms **

[jakevdp]

I think the asynchronous dispatch may only be important on GPU / TPU; for example, this is the output I get for your code on a standard Colab GPU runtime:

CPU times: user 101 µs, sys: 48 µs, total: 149 µs
Wall time: 161 µs
CPU times: user 1.03 ms, sys: 3 µs, total: 1.03 ms
Wall time: 1.03 ms

[zhangqiaorjc]

Twist...we have a bug in estimating flop counts, i missed a line and we have been getting 0 for flop counts. Will fix now, and it will be available in the next jaxlib release.

auto hlo_cost_analysis =
     std::make_unique<HloCostAnalysis>(cpu::CpuExecutable::ShapeSizeBytes);
+  TF_RETURN_IF_ERROR(cpu_executable->module().entry_computation()->Accept(
+      hlo_cost_analysis.get()));

[zhangqiaorjc]

Here's a simple way to find out the flop estimates from XLA for any JAX computation

x = np.random.uniform(size=(1000, 1000)).astype(np.float32)

def f(x):
  return jnp.dot(x, x)

m = jax.xla_computation(f)(x).as_hlo_module()
client = jax.lib.xla_bridge.get_backend()
analysis = jax.lib.xla_client._xla.hlo_module_cost_analysis(client, m)
flops = analysis['flops']
print(flops)

Notably, the flops value here is 2000000000.0.

[YouJiacheng]

Twist...we have a bug in estimating flop counts, i missed a line and we have been getting 0 for flop counts. Will fix now, and it will be available in the next jaxlib release.

auto hlo_cost_analysis =
     std::make_unique<HloCostAnalysis>(cpu::CpuExecutable::ShapeSizeBytes);
+  TF_RETURN_IF_ERROR(cpu_executable->module().entry_computation()->Accept(
+      hlo_cost_analysis.get()));

Does the bug only affect CPU backend?

[GoktugGuvercin]

I guess I have no question in my mind, thank you @mattjj and @zhangqiaorjc so much for your elegant explanation.

[zhangqiaorjc]

@YouJiacheng yes, only for CPU backend; for GPU and TPU, we always dispatch async from host to device with a configurable max inflight computations (32 is the default). CPU is special because sync dispatch for cheap computations is a performance optimization.