[Pallas] Fix fori_loop multi-dim index decomposition

[norx1991]

Summary

When _codegen_fori_loop handles a multi-dimensional inner loop (e.g. hl.tile([m, n])), the iteration was flattened into a single jax.lax.fori_loop with divmod to recover per-dimension indices. Previously, the raw flattened loop variable _j was used directly for all dimensions, producing incorrect indices for 2+ inner dims.

This PR fixes the indexing and uses nested fori_loop calls (one per dimension) instead of flat divmod decomposition. Benchmarking on TPU v7 showed nested fori_loop is ~5-10% faster than divmod.

Generated code for a 2D inner loop now looks like:

def _fori_body_0(_j0, _):
    def _fori_body_1(_j1, _):
        # DMA copies and compute using _j0, _j1 directly
    jax.lax.fori_loop(0, grid_n, _fori_body_1, None)
jax.lax.fori_loop(0, grid_m, _fori_body_0, None)

This issue was discovered while working toward enabling rms_norm_bwd on the Pallas/TPU backend, which uses a 2D inner tile.

[v0i0]

could you check if there's a perf difference between doing this w/ div&mod vs generating (a) nested fori (b) having running counters that get reset in a standalone jax-pallas sample?

[norx1991]

could you check if there's a perf difference between doing this w/ div&mod vs generating (a) nested fori (b) having running counters that get reset in a standalone jax-pallas sample?

Good catch! The result is here. Both are faster than div/mod. Which one would you suggest? Seems like nested fori_loop will be the more traditional path.

[v0i0]

could you check if there's a perf difference between doing this w/ div&mod vs generating (a) nested fori (b) having running counters that get reset in a standalone jax-pallas sample?

Good catch! The result is here. Both are faster than div/mod. Which one would you suggest? Seems like nested fori_loop will be the more traditional path.

nested fori seems the most natural