int new_idx = child_idx >= cluster_idx ? child_idx + num_cores_per_replica
                                           : child_idx;
    new_parallel_execute->getRegions()[new_idx].takeBody(
        old_parallel_execute.getRegions()[old_idx]);
  }

  return cluster_idx;
}

mlir::tf_device::LaunchOp WrapOpInLaunch(mlir::OpBuilder* builder,
                                         mlir::Location loc,

      }
    };

    // Assign a unique ID to each value seen in this function.
    func.walk([&](Operation* op) {
      // For all attached regions, assign ID to the region arguments.
      for (Region& region : op->getRegions()) {
        for (auto region_arg : TF::filter_resources(region.getArguments()))
          assign_id(region_arg);
      }

      // Assign ID for all results.

        all_internal_results.insert(result);
      }
      // We treat block arguments of inner blocks as "results", too, in
      // the sense that they're values produced inside this op.
      for (Region& region : op->getRegions()) {
        for (Block& block : region.getBlocks()) {
          for (BlockArgument& arg : block.getArguments()) {
            all_internal_results.insert(arg);
          }
        }
      }
    });
  }

// Returns true if the op contains a `CalibrationStatisticsSaverOp`.
bool ContainCalibrationStatisticsSaverOp(Operation* op) {
  // Check the region for CaseRegionOp, IfRegionOp and WhileRegionOp.
  for (Region& region : op->getRegions()) {
    if (!region.getOps<TF::CalibrationStatisticsSaverOp>().empty()) {
      return true;
    }
  }

  SymbolTable symbol_table(op->getParentOfType<ModuleOp>());

      op->getAttrs(), op->getPropertiesStorage(), op->getSuccessors(),
      op->getNumRegions());
  builder.insert(new_op);

  // Move region bodies to the new operation.
  for (auto it : llvm::zip(op->getRegions(), new_op->getRegions())) {
    Region &old_region = std::get<0>(it);
    Region &new_region = std::get<1>(it);
    new_region.takeBody(old_region);
  }

  // Replace used results and erase the old op.

  auto new_parallel_execute = builder.create<tf_device::ParallelExecuteOp>(
      parallel_execute.getLoc(), num_regions, new_result_types);

  for (auto region : llvm::zip(new_parallel_execute.getRegions(),
                               parallel_execute.getRegions()))
    std::get<0>(region).takeBody(std::get<1>(region));

  for (auto result :
       llvm::zip(results_to_remap, new_parallel_execute.getResults()))

  return parallel_execute;
}

LogicalResult RemoveSingletonParallelExecuteOp(
    tf_device::ParallelExecuteOp parallel_execute, OpBuilder* builder) {
  if (parallel_execute.getRegions().size() == 1) {
    builder->setInsertionPoint(parallel_execute);
    auto& block = parallel_execute.GetRegionBlockWithIndex(0);
    llvm::SmallVector<Operation*, 2> ops_move;
    for (Operation& op : block) {

  // Followed by everything else.
  return 0;
}

void MoveTpuCompileToFrontPass::runOnOperation() {
  MarkCompilationOps(getOperation());
  getOperation().walk([](Operation* op) {
    for (Region& region : op->getRegions()) {
      for (Block& block : region.getBlocks()) {
        if (block.empty()) continue;
        auto ops = SortBlockTopologically(block, OutsideCompilationOrdering);
        // Replace the block with the reordered block.

      info =
          &getAnalysis<mlir::TF::SideEffectAnalysis>().GetAnalysisForFunc(func);
    }
    func->walk([&](Operation* function) {
      for (Region& region : function->getRegions()) {
        for (Block& block : region.getBlocks()) {
          if (block.empty()) continue;
          auto ops = SortBlockTopologically(
              block, DialectOrdering,

    if (auto cast = dyn_cast<TF::CastOp>(op)) {
      this->join(results[0], *operands[0]);
    } else if (auto while_op = dyn_cast<TF::WhileRegionOp>(op)) {
      for (auto &region : while_op->getRegions()) {
        for (auto [arg, value] :
             llvm::zip(region.getArguments(), while_op->getOperands())) {
          this->join(getLatticeElement(arg), *getLatticeElement(value));
        }
      }