[TPU][Mosaic][Easy] Add verification for AssumeMultipleOp.

A user must use AssumeMultipleOp to annotate integer constants that are divisible by the given multiple.

PiperOrigin-RevId: 726994750

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -608,6 +608,7 @@ def TPU_AssumeMultipleOp : TPU_Op<"assume_multiple", [Pure, SameOperandsAndResul
     I32Attr:$multiple
   );
   let results = (outs AnyTypeOf<[Index, AnyInteger]>:$result);
+  let hasVerifier = 1;
 }
 
 def TPU_MemRefSliceOp : TPU_Op<"memref_slice", [Pure, AttrSizedOperandSegments]> {

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -1225,6 +1225,27 @@ LogicalResult DynamicGatherOp::verify() {
   return success();
 }
 
+LogicalResult AssumeMultipleOp::verify() {
+  auto operand_value = getValue();
+  auto divisor = getMultiple();
+  if (auto cst_op = operand_value.getDefiningOp<arith::ConstantOp>()) {
+    auto int_attr = dyn_cast<IntegerAttr>(cst_op.getValue());
+    // Illegal usage of AssumeMultipleOp.
+    if (!int_attr) {
+      return emitOpError(
+                 "Illegal user annotation, expected an integer, but got ")
+             << cst_op.getValue();
+    }
+    if (int_attr.getInt() % divisor != 0) {
+      return emitOpError(
+                 "Illegal user annotation, expected an integer that is "
+                 "divisible by the multiple, but got ")
+             << int_attr.getInt() << " % " << divisor;
+    }
+  }
+  return success();
+}
+
 }  // namespace tpu
 }  // namespace mlir