builder.create<tf_device::ReturnOp>(op->getLoc(), op->getResults())
          .getOperation();
  MLIRContext* context = launch_op.getContext();
  op->removeAttr(StringAttr::get(context, kDeviceAttr));
  op->moveBefore(return_op);
}

void DeviceAttributeToLaunch::runOnOperation() {
  const Dialect* tf_dialect = getContext().getLoadedDialect("tf");

  getOperation().walk([&](Operation* op) {

  llvm::SmallVector<Operation*, 4> non_terminator_ops;
  for (Operation& op : block.without_terminator()) {
    non_terminator_ops.push_back(&op);
  }
  for (auto op : non_terminator_ops) {
    op->moveBefore(cluster_op);
  }

  auto& return_op = cluster_op.GetBody().front();
  // This is the last op, should be tf_device::ReturnOp.
  assert(mlir::isa<mlir::tf_device::ReturnOp>(return_op));

  Value old_val, new_val;

          return;
        }
        if (constant.hasOneUse()) {
          LLVM_DEBUG(llvm::dbgs() << "Moved constant " << constant << "\n");
          const_op.getOperation()->moveBefore(&body.begin()->front());
          return;
        }
        map_entry.first->getSecond() = const_op.clone();
        body.begin()->getOperations().insert(body.begin()->begin(),

  launch.getBody().push_back(new mlir::Block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<mlir::tf_device::ReturnOp>(loc, op->getResults());

  // Move op inside cluster.
  op->moveBefore(launch.GetBody().getTerminator());

  builder->restoreInsertionPoint(insert_point);

  return launch;
}

  op->replaceAllUsesWith(launch);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(loc, op->getResults());

  // Move op inside cluster.
  op->moveBefore(launch.GetBody().getTerminator());
}

llvm::SmallVector<Operation*, 4> GetResourceOpsUsingCompositeArgsInReplicate(
    tf_device::ReplicateOp replicate) {
  llvm::SmallVector<Operation*, 4> resource_users;

      if (read->isBeforeInBlock(peer)) continue;
      if (llvm::isa<ReadVariableOp>(peer)) continue;
      has_conflicting_write = true;
    }
    if (has_conflicting_write) continue;
    read->moveBefore(replicate);
  }
}

// Hoist `ReadVariableOp`s above the `tf_device.replicate`s.
void HoistBroadcastRead::runOnOperation() {
  FuncOp func = getOperation();

  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      call_op->getLoc(), call_op->getResultTypes());
  cluster.getBody().push_back(new Block);
  call_op->replaceAllUsesWith(cluster.getResults());
  call_op->moveBefore(&cluster.GetBody(), cluster.GetBody().end());
  builder.setInsertionPointToEnd(&cluster.GetBody());
  builder.create<mlir::tf_device::ReturnOp>(call_op->getLoc(),
                                            call_op->getResults());

      // Move AssignVariableOp and update the value to be written to the
      // resource variable to be the non forwarded value from within the
      // parallel_execute region.
      assign_var.getOperation()->moveBefore(terminator);
      assign_var.getValueMutable().assign(
          terminator->getOperand(result.index()));
      results_to_remove.push_back(result.index());
    }

    rewrite |= !results_to_remove.empty();

        num_partitions_after += op_partition && (scheduled_something == false);
        scheduled_something = true;
        // Schedule the operation by moving it to the start.
        unscheduled_op.moveBefore(block, next_unscheduled_op);
        if (&unscheduled_op == &*next_unscheduled_op) {
          ++next_unscheduled_op;
        }
      }
    }

    // defined by tf_device.replicate.
    if (DescendantUsesVirtualDevice(virtual_device_list, &inner_op)) continue;

    if (IsOpReplicateInvariant(replicate_region, &inner_op))
      inner_op.moveBefore(replicate_op);
  }
}

void ReplicateInvariantOpHoistingPass::runOnOperation() {
  getOperation().walk(
      [](tf_device::ReplicateOp op) { HoistReplicateInvariantOps(op); });
}