jax

d12cbffd - [Mosaic GPU] Refactor and generalize code in `optimization_barrier`.

[Mosaic GPU] Refactor and generalize code in `optimization_barrier`. The change in `utils.py` is to enable the use of `bitwidth` when the mlir dialect is not registered. PiperOrigin-RevId: 745060221