[Pallas:MGPU] Allow indexing to appear anywhere in the list of transforms

We only need to exchange the transforms preceding the indexer, while
the rest can remain unmodified.

PiperOrigin-RevId: 688112088

Author: apaszke

jax/_src/pallas/mosaic_gpu/core.py

@@ -226,11 +226,11 @@ def tree_unflatten(cls, metadata, arrays):
 
 
 def transpose_ref(
-    ref: pallas_core.TransformedRef | pallas_core.AbstractMemoryRef,
+    ref: pallas_core.TransformedRef | Any,
     permutation: tuple[int, ...],
 ) -> pallas_core.TransformedRef:
   if not isinstance(ref, pallas_core.TransformedRef):
-    if not isinstance(ref, pallas_core.AbstractMemoryRef):
+    if not isinstance(jax_core.get_aval(ref), pallas_core.AbstractMemoryRef):
       raise TypeError("ref must be a reference")
     ref = pallas_core.TransformedRef(ref, transforms=())
   return pallas_core.TransformedRef(

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -871,22 +871,24 @@ def _handle_indexing(
 ) -> tuple[ir.Value, Sequence[gpu_core.Transform]]:
   if not transforms:
     pass
-  if not any(isinstance(t, indexing.NDIndexer) for t in transforms):
+  indexer_idxs = [
+      i for i, t in enumerate(transforms) if isinstance(t, indexing.NDIndexer)
+  ]
+  if not indexer_idxs:
     return ref, transforms
-  if any(
-      isinstance(t, indexing.NDIndexer) for t in transforms[:-1]
-  ) or not isinstance(transforms[-1], indexing.NDIndexer):
+  if len(indexer_idxs) > 1:
     raise NotImplementedError("Only one level of indexing supported.")
-
-  indexer = cast(indexing.NDIndexer, transforms[-1])
+  [indexer_idx] = indexer_idxs
+  indexer = cast(indexing.NDIndexer, transforms[indexer_idx])
   if indexer.int_indexer_shape:
     raise NotImplementedError("int_indexer_shape non-empty")
   indices = _ndindexer_indices(indexer)
   new_transforms_rev = []
-  for t in reversed(transforms[:-1]):
+  for t in reversed(transforms[:indexer_idx]):
     indices, new_t = t.untransform_index(indices)
     new_transforms_rev.append(new_t)
-  return mgpu.memref_slice(ref, indices), new_transforms_rev[::-1]
+  new_transforms = [*reversed(new_transforms_rev), *transforms[indexer_idx + 1:]]
+  return mgpu.memref_slice(ref, indices), new_transforms
 
 
 def _ndindexer_indices(indexer: indexing.NDIndexer) -> tuple[gpu_core.Index, ...]:

jax/_src/pallas/mosaic_gpu/primitives.py

@@ -107,10 +107,10 @@ def copy_smem_to_gmem(
   if dst.memory_space is not gpu_core.GMEM:
     raise ValueError(f"dst must be a GMEM reference, got {dst.memory_space}")
   src, src_transforms = state_primitives.get_ref_and_transforms(
-      src, None, "copy_smem_to_gmem"
+      src, None, "copy_smem_to_gmem", force_trailing_indexer=False,
   )
   dst, dst_transforms = state_primitives.get_ref_and_transforms(
-      dst, None, "copy_smem_to_gmem"
+      dst, None, "copy_smem_to_gmem", force_trailing_indexer=False,
   )
   flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
       src_transforms
@@ -193,10 +193,10 @@ def copy_gmem_to_smem(
   if dst.memory_space is not gpu_core.SMEM:
     raise ValueError(f"dst must be a SMEM reference, got {dst.memory_space}")
   src, src_transforms = state_primitives.get_ref_and_transforms(
-      src, None, "copy_gmem_to_smem"
+      src, None, "copy_gmem_to_smem", force_trailing_indexer=False,
   )
   dst, dst_transforms = state_primitives.get_ref_and_transforms(
-      dst, None, "copy_gmem_to_smem"
+      dst, None, "copy_gmem_to_smem", force_trailing_indexer=False,
   )
   flat_src_transforms, src_transforms_treedef = tree_util.tree_flatten(
       src_transforms
@@ -205,7 +205,7 @@ def copy_gmem_to_smem(
       dst_transforms
   )
   barrier, barrier_transforms = state_primitives.get_ref_and_transforms(
-      barrier, None, "copy_gmem_to_smem"
+      barrier, None, "copy_gmem_to_smem", force_trailing_indexer=False,
   )
   flat_barrier_transforms, barrier_transforms_treedef = tree_util.tree_flatten(
       barrier_transforms
@@ -284,7 +284,7 @@ def _barrier_arrive_lowering(
 def barrier_arrive(barrier: pallas_core.AbstractMemoryRef) -> None:
   """Arrives at the given barrier."""
   barrier, transforms = state_primitives.get_ref_and_transforms(
-      barrier, None, "barrier_arrive"
+      barrier, None, "barrier_arrive", force_trailing_indexer=False,
   )
   flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
   barrier_arrive_p.bind(
@@ -321,7 +321,7 @@ def _barrier_wait_lowering(
 def barrier_wait(barrier: pallas_core.AbstractMemoryRef) -> None:
   """Waits on the given barrier."""
   barrier, transforms = state_primitives.get_ref_and_transforms(
-      barrier, None, "barrier_wait"
+      barrier, None, "barrier_wait", force_trailing_indexer=False,
   )
   flat_transforms, transforms_treedef = tree_util.tree_flatten(transforms)
   barrier_wait_p.bind(

jax/_src/state/primitives.py

@@ -73,6 +73,7 @@ def get_ref_and_transforms(
     ref_or_view: Any,
     idx: Indexer | tuple[Indexer, ...] | None,
     function_name: str,
+    force_trailing_indexer: bool = True,  # TODO(apaszke): Clean this up.
 ) -> tuple[Any, tuple[Transform, ...]]:
   if isinstance(ref_or_view, TransformedRef):
     ref, transforms = ref_or_view.ref, ref_or_view.transforms
@@ -89,6 +90,8 @@ def get_ref_and_transforms(
   elif not isinstance(idx, tuple):
     idx = (idx,)
 
+  if not idx and not force_trailing_indexer:
+    return ref, transforms
   if not idx and transforms and isinstance(transforms[-1], indexing.NDIndexer):
     return ref, transforms
   nd_indexer = indexing.NDIndexer.from_indices_shape(idx, ref_or_view.shape)

jax/experimental/mosaic/gpu/core.py

@@ -400,6 +400,12 @@ def async_copy(
           "Expected the SMEM reference to have the same shape as the"
           f" transformed slice: {tuple(smem_ref_ty.shape)} != {slice_shape}"
       )
+    smem_strides, _ = smem_ref_ty.get_strides_and_offset()
+    if smem_strides != utils.get_contiguous_strides(smem_ref_ty.shape):
+      raise ValueError(
+          "async_copy needs the SMEM reference to be contiguous, but got"
+          f" strides {smem_strides} for shape {smem_ref_ty.shape}"
+      )
 
     dyn_base_indices = list(dyn_base_indices)
     slice_shape = list(slice_shape)

tests/pallas/mosaic_gpu_test.py

@@ -314,6 +314,29 @@ def kernel(x_ref, o_ref, barrier_ref):
     x = jnp.arange(128 * 128, dtype=jnp.float32).reshape(128, 128)
     np.testing.assert_array_equal(f(x), np.stack([x, x], axis=0))
 
+  def test_indexing_before_transpose(self):
+    def kernel(x_ref, o_ref, barrier_ref):
+      for i in range(2):
+        plgpu.copy_gmem_to_smem(
+            x_ref, plgpu.transpose_ref(o_ref.at[i], (1, 0, 2)), barrier=barrier_ref
+        )
+        plgpu.barrier_wait(barrier_ref)
+
+    in_spec = pl.BlockSpec(memory_space=plgpu.GMEM)
+    out_spec = plgpu.GPUBlockSpec(
+        (2, 64, 2, 128), lambda: (0, 0, 0, 0), memory_space=plgpu.SMEM,
+    )
+    f = pl.pallas_call(
+        kernel,
+        out_shape=jax.ShapeDtypeStruct([2, 64, 2, 128], jnp.float32),
+        in_specs=(in_spec,),
+        out_specs=out_spec,
+        scratch_shapes=[plgpu.Barrier(num_arrivals=1)],
+    )
+    x = jnp.arange(2 * 64 * 128, dtype=jnp.float32).reshape(2, 64, 128)
+    xt = x.transpose((1, 0, 2))
+    np.testing.assert_array_equal(f(x), np.stack([xt, xt], axis=0))
+
   def test_copy_gmem_to_smem_in_run_scoped(self):
     @functools.partial(
         pl.pallas_call,