Add support for slicing dynamically-shaped memrefs + DMAs between them

This was a little difficult because our current dialect conversion setup assumes 1-1 type conversions.
I think everything works out fine for as long as we never pass memrefs between basic blocks (i.e.
for as long as we never have memrefs as loop carry or return them from conditionals).

TODO: I still need to make sure that the changes to the TPU dialect are backwards-compatible.
I am afraid that the signature change in MemRefSliceOp might not be.
PiperOrigin-RevId: 617755035

Author: apaszke

jax/_src/pallas/mosaic/lowering.py

@@ -833,7 +833,7 @@ def _slice_memref(ref: ir.Value, ref_aval: state.AbstractRef,
       memory_space=ref.type.memory_space)
   inner_aval = ref_aval.inner_aval
   out_aval = ref_aval.update(inner_aval=inner_aval.update(shape=target_shape))
-  out = tpu.MemRefSliceOp(target_ref_ty, ref, starts).result
+  out = tpu.MemRefSliceOp(target_ref_ty, ref, starts, []).result
   if any(squeeze_dims):
     # We need to squeeze out some dimensions
     squeezed_ref_ty = ir.MemRefType.get(

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -397,14 +397,16 @@ def TPU_AssumeMultipleOp : TPU_Op<"assume_multiple", [Pure, SameOperandsAndResul
   let results = (outs AnyTypeOf<[Index, AnyInteger]>:$result);
 }
 
-def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure]> {
+def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure, AttrSizedOperandSegments]> {
   let arguments = (ins
     AnyMemRef:$mem_ref,
-    Variadic<I32>:$base_idx
+    Variadic<I32>:$base_idx,
+    Variadic<I32>:$dynamic_sizes
   );
   let results = (outs AnyMemRef:$result);
   let assemblyFormat = [{
-    $mem_ref `[` $base_idx `]` attr-dict `:` type($mem_ref) `->` type($result)
+    $mem_ref `[` $base_idx `]` (`<` $dynamic_sizes^ `>`)?
+    attr-dict `:` type($mem_ref) `->` type($result)
   }];
   let hasVerifier = 1;
   let hasCanonicalizeMethod = 1;

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -66,7 +66,8 @@ LogicalResult MemRefSliceOp::verify() {
       (target_memory_space == nullptr ||
        target_memory_space == source_type.getMemorySpace()) &&
       ((isa<AffineMapAttr>(target_layout) && target_layout.isIdentity()) ||
-       target_type.getLayout() == source_type.getLayout()));
+       target_type.getLayout() == source_type.getLayout()) &&
+      getDynamicSizes().size() == target_type.getNumDynamicDims());
 }
 
 LogicalResult MemRefSliceOp::canonicalize(MemRefSliceOp op,
@@ -82,8 +83,9 @@ LogicalResult MemRefSliceOp::canonicalize(MemRefSliceOp op,
   auto new_result_type = MemRefType::get(
       op.getResult().getType().getShape(), layout_ty.getElementType(),
       layout_ty.getLayout(), layout_ty.getMemorySpace());
-  auto slice = rewriter.create<MemRefSliceOp>(op.getLoc(), new_result_type,
-                                              layout_ref, op.getBaseIdx());
+  auto slice =
+      rewriter.create<MemRefSliceOp>(op.getLoc(), new_result_type, layout_ref,
+                                     op.getBaseIdx(), op.getDynamicSizes());
   rewriter.replaceOpWithNewOp<EraseLayoutOp>(op, op.getType(), slice);
   return success();
 }

jaxlib/mosaic/dialect/tpu/transforms/apply_vector_layout.cc

@@ -367,7 +367,7 @@ FailureOr<std::pair<Value, SmallVector<int64_t>>> sliceRef(
   Value sliced_ref = builder.create<tpu::MemRefSliceOp>(
       MemRefType::get(pad_slice_shape, ref_ty.getElementType(),
                       ref_ty.getLayout(), ref_ty.getMemorySpace()),
-      base_ref, slice_base_indices);
+      base_ref, slice_base_indices, /*dynamic_sizes=*/ValueRange());
 
   return std::make_pair(sliced_ref, indices_within_slice);
 }