[Mosaic TPU] Lock down local DMAs with an SMEM source buffer temporarily

PiperOrigin-RevId: 678616118
jax-ml · Sep 25, 2024 · 4e3de3a · 4e3de3a
1 parent a43c7f2
commit 4e3de3a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_ops.cc b/jaxlib/mosaic/dialect/tpu/tpu_ops.cc
@@ -562,6 +562,13 @@ LogicalResult EnqueueDMAOp::verify() {
           "device_id or core_id is specified");
     }
   }
+  // TODO(b/369517284): Add proper support for local SMEM DMAs.
+  auto src_memory_space = dyn_cast_if_present<tpu::MemorySpaceAttr>(
+      getSource().getType().getMemorySpace());
+  if (src_memory_space &&
+      src_memory_space.getValue() == tpu::MemorySpace::smem) {
+    return emitOpError("DMAs with an SMEM source are not fully supported");
+  }
   return success();
 }