Add nd_loop and Enable block_n tiling for all_gather_lhs_matmul

[hanzlfs]

1. Enable block_n tiling
2. use plgpu.nd_loop
   [PAIR] justinfu@google.com

[justinjfu]

Thanks for the fixes!

[justinjfu]

nit: Add a space for formatting: rhs_ref.at[:,n_tile_slice] -> rhs_ref.at[:, n_tile_slice]

[apaszke]

I don't think that's the right way to do it. The gathers happen only on the M dimension and adding the N loop in the same place will perform the same gather n // block_n times. Instead, for every M chunk we gather, we should run an inner loop that steps over all the N blocks that need to be multiplied with it

[hanzlfs]

I don't think that's the right way to do it. The gathers happen only on the M dimension and adding the N loop in the same place will perform the same gather n // block_n times. Instead, for every M chunk we gather, we should run an inner loop that steps over all the N blocks that need to be multiplied with it

Thanks I will update this part tmr. I met another issue, if I want to use a 2d mesh x, y: (2, 4), or even x, y : (1, 8), to use axis_name = (x, y) it will give me following error

    out = core_map_p.bind(*consts, jaxpr=jaxpr, mesh=mesh,
jax._src.source_info_util.JaxStackTraceBeforeTransformation: ValueError: Failed to recompute the async_copy peer id on the host``` 
what's the corresponding change would be needed? cc @justinjfu 

[apaszke]

Could you please send a small PR with the change needed to reproduce the problem? I can take a look

[hanzlfs]

Could you please send a small PR with the change needed to reproduce the problem? I can take a look
This one should do
#29849