#include "mlir/Transforms/Passes.h"  // from @llvm-project
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h"
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

using mlir::OpPassManager;
using mlir::func::FuncOp;

constexpr char kMiniBatchSplitsAttr[] = "mini_batch_splits";
constexpr char kMiniBatchCsrAttr[] = "mini_batch_in_csr";

#define GEN_PASS_DEF_EMBEDDINGPROGRAMKEYPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

struct EmbeddingProgramKeyPass
    : public impl::EmbeddingProgramKeyPassBase<EmbeddingProgramKeyPass> {
  void runOnOperation() override;
};

// Checks if `op` is nested in `block`.

#include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"

#define GEN_PASS_DEF_EMBEDDINGSEQUENCINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/sparsecore/sparsecore_passes.h.inc"

static constexpr char kEmbeddingPipelining[] = "_embedding_pipelining";
static constexpr char kEmbeddingForward[] = "forward";
static constexpr char kEmbeddingBackward[] = "backward";