Use llvm::cast/dyn_cast/isa since alternatives are deprecated in llvm/llvm-project#135556

PiperOrigin-RevId: 748551976

Author: Google-ML-Automation

jaxlib/mosaic/dialect/tpu/tpu_ops.cc

@@ -1096,7 +1096,7 @@ LogicalResult ConcatenateOp::verify() {
   if (getOperands().size() < 2) {
     return emitOpError("Expected at least 2 operands for concatenate op.");
   }
-  auto first_type = getOperand(0).getType().cast<VectorType>();
+  auto first_type = cast<VectorType>(getOperand(0).getType());
   auto first_shape = first_type.getShape();
   auto first_dtype = first_type.getElementType();
   for (auto operand : getOperands()) {

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

@@ -257,7 +257,7 @@ LogicalResult tpu_matmul_rule(const CanonicalizeContext &ctx,
       auto matmul_res =
           dot_dim_matmul(sliced_lhs.getResult(), sliced_rhs.getResult(),
                          sliced_acc.getResult());
-      auto res_ty = matmul_res.getType().cast<VectorType>();
+      auto res_ty = cast<VectorType>(matmul_res.getType());
       auto res_shape = res_ty.getShape();
       // reshape to 1x[prior_shape]
       auto reshape_shape = llvm::to_vector(res_shape);

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

@@ -136,10 +136,10 @@ class VectorLayoutInferer {
 
       bool has_vector_io = false;
       for (auto op : any_op.getOperands()) {
-        has_vector_io |= op.getType().isa<VectorType>();
+        has_vector_io |= isa<VectorType>(op.getType());
       }
       for (auto r : any_op.getResults()) {
-        has_vector_io |= r.getType().isa<VectorType>();
+        has_vector_io |= isa<VectorType>(r.getType());
       }
       if (!has_vector_io && any_op.getRegions().empty()) {
         SmallVector<Layout, 4> in_layout(any_op.getNumOperands(), kNoLayout);
@@ -1293,7 +1293,7 @@ class VectorLayoutInferer {
           (*(offsets.end() - 1) + *input_layout->offsets()[1]) % vreg_slice[1];
     }
     for (auto stride : strides_attr) {
-      TPU_CHECK_OP(stride.cast<IntegerAttr>().getInt() == 1,
+      TPU_CHECK_OP(cast<IntegerAttr>(stride).getInt() == 1,
                    "Only trivial strides supported.");
     }