jax

24884071 - [MosaicGPU] Remove the single_thread context from top-level dialect code.

[MosaicGPU] Remove the single_thread context from top-level dialect code. - Change the `async_load` lowering to manage the single thread context. - Use a predicate for the top-level arrive_expect. If we want to hide this further, we can have a warp-group level op that lowers to a single-threaded context. PiperOrigin-RevId: 716219730