[Mosaic] Several fixes/improvements for the new TPU interpret mode.

 - Checks bounds for reads and writes to shared memory.
 - Pads kernel arguments when necessary.
 - Fix support for input-output aliasing.
 - Fix handling of vmap'ed dimensions.
 - Supports un-masked `pl.load` and masked or un-masked `pl.swap`.
 - Switch to using single integer device IDs instead of tuples.
 - Better error messages for unsupported primitives: `for_p`, `atomic_rmw_p`, and `atomic_cas_p` .

PiperOrigin-RevId: 725784360