Question about the Time Alignment of Traces from CPU and TPU

[yutation]

Hi,

I'm running a workload that involves both CPU and TPU processing, where there's a true dependency from CPU to TPU—the CPU must complete its task before the TPU begins, as the TPU relies on data prepared by the CPU.

However, the actual trace shows the opposite, with the CPU appearing to finish after the TPU has started. I’d like to ask whether the profiling traces from the CPU and TPU are accurately time-aligned. Could this be a trace misalignment issue?

The example program is below, and the TPU kernel has a true dependency on the CPU kernel's result.
The figure below is the trace for the example program generated by jax.profiler.trace(). In the trace, we can see the TPU kernel starts before the CPU kernel finishes. Moreover, the block the red arrow points to, I think, is the CPU calling the TPU kernel, and it is also after the corresponding TPU kernel runs.
Thus, we suspect there is a time shift between the CPU and TPU traces.

Thanks,
Jingtian

def function_to_test(size: int, iterations: int, input_list,  state_matrix: jnp.ndarray) -> jnp.ndarray:
    for i in range(iterations):
        matrix_a_np = cpu_kernel(input_list[i], size)
        matrix_a = jnp.array(matrix_a_np, dtype=jnp.uint32)
        
        state_matrix = tpu_kernel_jit(state_matrix, matrix_a)
    return state_matrix