[Mosaic] Add extra memref_slice verification and a memory space check helper

PiperOrigin-RevId: 702883469

Author: naummo

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -93,6 +93,25 @@ LogicalResult MemRefSliceOp::verify() {
   auto target_type = getType();
   auto target_layout = target_type.getLayout();
   auto target_memory_space = target_type.getMemorySpace();
+  auto indices = getBaseIdx();
+  auto slice_shape = getResult().getType().getShape();
+  if (!source_type.hasStaticShape()) {
+    return emitOpError(
+        "Only slicing of memrefs with static shapes is supported.");
+  }
+  auto source_shape = source_type.getShape();
+  bool is_semaphore =
+      HasMemorySpace(source_type, tpu::MemorySpace::kSemaphoreMem);
+  if (is_semaphore &&
+      !isa<SemaphoreType, DMASemaphoreType>(source_type.getElementType())) {
+    return emitOpError(
+        "References to semaphore memory space must have a semaphore element "
+        "type.");
+  }
+  if (indices.size() != slice_shape.size() ||
+      indices.size() != source_shape.size()) {
+    return emitOpError("Indices and slice shapes must match.");
+  }
   // TODO(apaszke): Check that the result has a smaller shape.
   // TODO(apaszke): Check that strides are equivalent.
   // Source and target attributes may be different before propagation is done by

jaxlib/mosaic/dialect/tpu/util.cc

@@ -141,4 +141,10 @@ bool canReinterpretToUntiledMemref(TypedValue<MemRefType> tiled_memref,
   return *(tiled_layout.getTileStrides().end() - 1) == 1 &&
          *(tiled_layout.getTileStrides().end() - 2) == 1;
 }
+
+bool HasMemorySpace(MemRefType ty, tpu::MemorySpace space) {
+  auto memory_space =
+      dyn_cast_or_null<tpu::MemorySpaceAttr>(ty.getMemorySpace());
+  return memory_space && memory_space.getValue() == space;
+}
 }  // namespace mlir::tpu

jaxlib/mosaic/dialect/tpu/util.h

@@ -115,6 +115,8 @@ bool canReinterpretToUntiledMemref(TypedValue<MemRefType> tiled_memref,
                                    const std::array<int64_t, 2> &target_shape,
                                    bool allow_minormost_padding = false);
 
+// Determines whether the given MemRefType has the given memory space.
+bool HasMemorySpace(MemRefType ty, tpu::MemorySpace space);
 }  // namespace mlir::tpu
 
 #endif  // THIRD_PARTY_PY_JAX_JAXLIB_MOSAIC_DIALECT_TPU_UTIL_H_