Christian Sigg <******@****.***> 1714060318 -0700

Pat Notz <******@****.***> 1690925009 -0700

        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Support",
        "@llvm-project//mlir:TransformUtils",
    ],
)

cc_library(
    name = "embedding_sequencing",
    srcs = ["embedding_sequencing.cc"],
    hdrs = [
        "sparsecore_passes.h",
    ],
    deps = [
        ":sparsecore_passes_inc_gen",
        "//tensorflow/compiler/jit:flags_headers",

  // TensorCore sequencing will be used.
  bool tf_xla_disable_full_embedding_pipelining;

  // Force the WhileOps in embedding_pipelining and embedding_sequencing to use
  // this many parallel_iterations
  int tf_xla_embedding_parallel_iterations;
};

// Flags for common MLIR configurations.
struct MlirCommonFlags {

  pm.addNestedPass<FuncOp>(
      mlir::TFTPU::CreateTPUReorderReplicateAndPartitionedInputsPass());
  pm.addNestedPass<FuncOp>(mlir::TF::CreateDecomposeReduceDatasetPass());
  // Only one of EmbeddingSequencing and EmbeddingPipelining will actually
  // run and the logic is in EmbeddingPipeliningPass. If the pipelining pass
  // runs, embedding attributes are stripped and the sequencing pass will have

}

void EmbeddingPipeliningPass::runOnOperation() {
  LOG(INFO) << "EmbeddingPipeliningPass::runOnOperation()";
  ModuleOp module = getOperation();

  // We only use one of the EmbeddingPipelining and EmbeddingSequencing passes.
  if (!UseEmbeddingPipelining(module)) return;

  SymbolTable symbol_table(module);

  llvm::SetVector<Operation*> forward_pass_ops;
  llvm::SetVector<Operation*> backward_pass_ops;