[Mosaic] Verify that the target IDs are provided in remote DMAs

Adds an extra verification check. Since the source semaphore is used only for remote DMAs, we should check that device or core IDs are also provided when source semaphore is provided.

PiperOrigin-RevId: 707675228

Author: naummo

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -935,6 +935,13 @@ LogicalResult EnqueueDMAOp::verify() {
           "device_id or core_id is specified");
     }
   }
+  if (getSourceSemaphore()) {
+    if (!getDeviceId() && !getCoreId()) {
+      return emitOpError(
+          "DMA destination device_id or core_id must be specified when source "
+          "semaphore is specified");
+    }
+  }
   return success();
 }