[TPU[Mosaic] Fix missing sfences in smem DMAs

PiperOrigin-RevId: 724545200

jax/_src/pallas/mosaic/lowering.py

@@ -151,7 +151,6 @@ class LoweringRuleContext:
   avals_in: Sequence[jax_core.AbstractValue]
   avals_out: Sequence[jax_core.AbstractValue]
   block_shapes: Sequence[tuple[int | pallas_core.Mapped, ...] | None]
-
   replace = dataclasses.replace
 
   @property
@@ -3362,22 +3361,26 @@ def _dma_start_lowering_rule(ctx: LoweringRuleContext, *args, tree,
     device_id = _device_id_to_logical(ctx, device_id, device_id_type)
   tpu.enqueue_dma(src_ref, dst_ref, sem, source_semaphore=src_sem,
                   device_id=device_id)
+
   return []
 lowering_rules[tpu_primitives.dma_start_p] = _dma_start_lowering_rule
 
 
 def _dma_wait_lowering_rule(ctx: LoweringRuleContext, *args, tree,
                             device_id_type: tpu_primitives.DeviceIdType):
   del device_id_type
-  (_, _, ref, transforms, sem, sem_transforms, _, _, _) = tree_util.tree_unflatten(
-      tree, args)
-  (_, _, ref_aval, _, sem_aval, _, _, _, _) = tree_util.tree_unflatten(
-      tree, ctx.avals_in)
+  (src, src_transforms, dst, transforms, sem, sem_transforms, _, _, _) = (
+      tree_util.tree_unflatten(tree, args)
+  )
+  (src_aval, _, dst_aval, _, sem_aval, _, _, _, _) = tree_util.tree_unflatten(
+      tree, ctx.avals_in
+  )
   block_shapes = tree_util.tree_unflatten(tree, ctx.block_shapes)
   ref_block_shape = block_shapes[2]
-  ref, _ = _transform_ref(ref, ref_aval.dtype, ref_block_shape, transforms)
+  src, _ = _transform_ref(src, src_aval.dtype, src_aval.shape, src_transforms)
+  dst, _ = _transform_ref(dst, dst_aval.dtype, ref_block_shape, transforms)
   sem, _ = _transform_ref(sem, sem_aval.dtype, sem_aval.shape, sem_transforms)
-  tpu.wait_dma(sem, ref)
+  tpu.wait_dma(sem, src, dst)
   return []
 lowering_rules[tpu_primitives.dma_wait_p] = _dma_wait_lowering_rule
 

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -737,7 +737,8 @@ def TPU_EnqueueDMAOp : TPU_Op<"enqueue_dma", [AttrSizedOperandSegments]> {
 def TPU_WaitDMAOp : TPU_Op<"wait_dma"> {
   let arguments = (ins
     MemRefOf<[TPU_DMASemaphoreType]>:$semaphore,
-    AnyMemRef:$ref
+    AnyMemRef:$src,
+    AnyMemRef:$dst
   );
   let hasVerifier = 1;
 }