[Pallas:MGPU] Support not tiled transposed loads in swap LANE lowering rule.

PiperOrigin-RevId: 842740700

Author: allanrenucci

jax/_src/pallas/mosaic_gpu/lowering.py

@@ -1717,9 +1717,12 @@ def _swap_lowering_rule(
           layout=value.layout,
       )
       value.store_tiled(x_smem, swizzle=swizzle)
-    case ():
+    case () | (gpu_core.TransposeRef((1, 0)),):
+      transposed = bool(transforms)
       match value.layout:
         case mgpu.TiledLayout():
+          if transposed:
+            x_smem = mgpu.memref_transpose(x_smem, (1, 0))
           old_value = mgpu.FragmentedArray.load_untiled(
               x_smem,
               layout=value.layout,
@@ -1728,6 +1731,8 @@ def _swap_lowering_rule(
           )
           value.store_untiled(x_smem, optimized=False)
         case _:
+          if transposed:
+            raise NotImplementedError(f"Unsupported transforms: {transforms}")
           old_value = mgpu.FragmentedArray.load_strided(
               x_smem, is_signed=mgpu_utils.is_signed(v_aval.dtype)
           )

tests/pallas/mosaic_gpu_test.py

@@ -1056,12 +1056,6 @@ def test_transposed_load_store(self, src_layout, dst_layout):
     def is_transposed(layout):
       return layout == plgpu.Layout.WGMMA_TRANSPOSED
 
-    if (
-        self.LOWERING_SEMANTICS == mgpu.LoweringSemantics.Lane
-        and is_transposed(dst_layout)
-    ):
-      self.skipTest("Not implemented: transposed, not tiled")
-
     shape, dtype = (128, 128), jnp.float32
 
     @functools.partial(