message: [MLIR] Add GetIndexTypeForKernelForMlir().
  This patch belongs to succeeding MLIR migration patches, that's why it doesn't have its own unit tests (nor is GetIndexTypeForKernel, the original version). It's not used anywhere in this patch either.
diff: @@ -405,6 +405,62 @@ llvm::Type* GetIndexTypeForKernel(const HloInstruction* hlo, int64 launch_size,
  return b->getInt32Ty();
  }
  +// The same as GetIndexTypeForKernel, but works with MLIR ops.
  +llvm::Type* GetIndexTypeForKernelFromMlir(mlir::Operation* op,
  + int64 launch_size,
  + llvm::IRBuilder<>* b) {
  + auto shape_in_range = [&](const Shape& s) {
  + bool in_range = true;
  + ShapeUtil::ForEachSubshape(s, [&](const Shape& sub_shape,
  + const ShapeIndex& /*index*/) {
  + if (sub_shape.IsArray() && !IsInt32(ShapeUtil::ElementsIn(sub_shape))) {
  + in_range = false;
  + }
  + });
  +
  + return in_range;
  + };
  +
  + llvm::Type* i64_ty = b->getInt64Ty();
  + // Check launch dimension
  + if (!IsInt32(launch_size)) {
  + return i64_ty;
  + }
  +
  + // Check the size of result tensors
  + for (auto result : op->getResults()) {
  + if (!shape_in_range(TypeToShape(result.getType()))) {
  + return i64_ty;
  + }
  + }
  +
  + auto hlo_shape_in_range = [&](mlir::Value operand) -> bool {
  + return shape_in_range(TypeToShape(operand.getType()));
  + };
  +
  + // Check the size of input tensors
  + if (!absl::c_all_of(op->getOperands(), hlo_shape_in_range)) {
  + return i64_ty;
  + }
  +
  + // Check the size of the internal result tensors
  + if (auto fusion = mlir::cast<mlir::lmhlo::FusionOp>(op)) {
  + auto result = fusion.region().walk([&](mlir::Operation* op) {
  + for (mlir::Value result : op->getResults()) {
  + if (!hlo_shape_in_range(result)) {
  + return mlir::WalkResult::interrupt();
  + }
  + }
  + return mlir::WalkResult::advance();
  + });
  + if (result.wasInterrupted()) {
  + return i64_ty;
  + }
  + }
  +
  + return b->getInt32Ty();
  +}
  +
  // Gets the input shape of the ROOT slices, which will be used as the kernel
  // launch dims. The slice input fusion requires the input shapes of the ROOT
  // slices to be the same although the (slice) output shapes can be different.
  