// This is a strictly sequential and formally correct fallback option for the
// embedding pipelining pass intended for debugging during pipelining
// development.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingPipeliningPass();

// Passes in the program key to embedding ops, by moving the embedding ops
// after the _TPUCompileMlir op.
std::unique_ptr<OperationPass<func::FuncOp>> CreateEmbeddingProgramKeyPass();

  // preserved and the sequencing rewrite will trigger.
  pm.addPass(mlir::TFDevice::CreateEmbeddingPipeliningPass());
  pm.addPass(mlir::TFDevice::CreateEmbeddingSequencingPass());
  pm.addPass(tensorflow::tf2xla::internal::CreateTPUClusterFormationPass(
      strict_clusters));
  // CreateEmbeddingPipeliningPass may have created more functions, but
  // TPUClusterCleanup and OutsideCompiledToHostLaunch need every function to be

include "mlir/Pass/PassBase.td"

def EmbeddingPipeliningPass : Pass<"tf-embedding-pipelining", "mlir::ModuleOp"> {
  let summary = "Rewrite graph for embedding pipelining";
  let constructor = "TFDevice::CreateEmbeddingPipeliningPass()";
    let description = [{
    For architectures that support accelerated embedding lookups, this pass will
    rewrite the graph to use pipelining for better device utilization.
  }];
}

  orig_while_op.erase();

  LOG(INFO) << "EmbeddingPipeliningPass::runOnOperation done.";
}
}  // namespace

std::unique_ptr<OperationPass<mlir::ModuleOp>> CreateEmbeddingPipeliningPass() {
  return std::make_unique<EmbeddingPipeliningPass>();
}

}  // namespace TFDevice