#include <vector>

namespace mlir {
namespace TF {

const llvm::SmallDenseSet<mlir::TypeID>&
TPUEmbeddingOpsRegistry::GetOpsTypeIds() {
  return ops_type_ids_;
}

// static
TPUEmbeddingOpsRegistry& TPUEmbeddingOpsRegistry::Global() {
  static TPUEmbeddingOpsRegistry* registry = new TPUEmbeddingOpsRegistry;
  return *registry;
}
}  // namespace TF

namespace mlir {
namespace TF {

// A global ops registry that is used to hold TPU embedding ops.
//
// Example:
//    TPUEmbeddingOpsRegistry::Global().Add<TF::FooOp>();
//    for (auto op_type_id : TPUEmbeddingOpsRegistry::Global().GetOpsTypeIds())
//    {
//      ...
//    }
class TPUEmbeddingOpsRegistry {
 public:
  // Add the op to the registry.
  //
  // Adding an op here will allow use old bridge legalization from the MLIR