struct MoveTpuCompileToFrontPass
    : public impl::MoveTpuCompileToFrontPassBase<MoveTpuCompileToFrontPass> {
  void runOnOperation() override;
};

void MarkCompilationOps(Operation* func) {
  func->walk([&](Operation* op) {
    if (llvm::isa<TF::_TPUCompileMlirOp>(op)) {
      op->setAttr("_is_compilation", UnitAttr::get(func->getContext()));
      op = op->getParentOp();