namespace TFDevice {

// For architectures that support accelerated embedding lookups, this pass will
// rewrite the graph to use pipelining for better device utilization.
std::unique_ptr<OperationPass<ModuleOp>> CreateEmbeddingSequencingPass();

// This is a strictly sequential and formally correct fallback option for the
// embedding pipelining pass intended for debugging during pipelining
// development.

  }];
}

def EmbeddingSequencingPass : Pass<"tf-embedding-sequencing", "mlir::ModuleOp"> {
  let summary = "Rewrite graph for sequential execution of embeddings";
  let constructor = "TFDevice::CreateEmbeddingSequencingPass()";
    let description = [{
    This is a strictly sequential and formally correct fallback option for the
    embedding pipelining pass intended for debugging during pipelining
    development.
  }];
}

  // 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

  metadata_op->erase();
  compilation_op->erase();

  LOG(INFO) << "EmbeddingSequencingPass::runOnOperation done.";
}

}  // namespace

std::unique_ptr<OperationPass<mlir::ModuleOp>> CreateEmbeddingSequencingPass() {
  return std::make_unique<EmbeddingSequencingPass>();
}

}  // namespace TFDevice