// and _xla_compile_device_type=TPU should be replaced with _tpu_replicate=X.
// This ensures the unified attributes not get exposed outside of the MLIR
// bridge with V1 pipeline in some cases.

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace tensorflow {
namespace tf_mlrt {

class TensorflowMlrtDialect : public mlir::Dialect {
 public:
  explicit TensorflowMlrtDialect(mlir::MLIRContext *context);
  static llvm::StringRef getDialectNamespace() { return "tf_mlrt"; }

  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
  void printType(mlir::Type type, mlir::DialectAsmPrinter &os) const override;
};

    if (IsArgConsumedByFallback(func, arg_index)) action(arg_index);
  }
}

void ForEachArgConsumedByFallback(
    mlir::ModuleOp module,
    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action) {
  for (auto func : module.getOps<mlir::func::FuncOp>()) {
    ForEachArgConsumedByFallback(
        func, [func_name = func.getName(), action](int arg_index) {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace tf_executor {

class TensorFlowExecutorDialect : public Dialect {
 public:
  static StringRef getDialectNamespace() { return "tf_executor"; }
  explicit TensorFlowExecutorDialect(MLIRContext *context);

  // Parses a type registered to this dialect.
  Type parseType(DialectAsmParser &parser) const override;

#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_EXPORT_GRAPHDEF_H_

#include "absl/base/attributes.h"
#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project

#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/protobuf.h"

namespace tensorflow {
namespace {
inline llvm::StringRef StringViewToRef(absl::string_view view) {
  return {view.data(), view.size()};
}
}  // namespace

Status LoadProtoFromBuffer(absl::string_view input, protobuf::Message* proto) {
  // Attempt to parse as text.

#include <algorithm>
#include <cstdint>
#include <memory>
#include <numeric>
#include <vector>

#include "absl/types/span.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project

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

namespace mlir {
namespace tf_executor {

namespace {
constexpr llvm::StringRef kNestedModule = "_tpu_v1_compat_outlined";
constexpr llvm::StringRef kOutlinedFuncPrefix = "_tpu_v1_compat_outlined_func";

#define GEN_PASS_DEF_TPUBRIDGEEXECUTORISLANDOUTLININGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_API_V2_TF_EXECUTOR_TO_GRAPH_H_

#include "absl/container/flat_hash_set.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project

  virtual void runOnFunction() = 0;

  void runOnOperation() final {
    if (!getFunction().isExternal()) runOnFunction();
  }

 protected:
  // Returns the derived pass name.
  StringRef getName() const override { return llvm::getTypeName<T>(); }

  // A clone method to create a copy of this pass.
  std::unique_ptr<Pass> clonePass() const override {
    return std::make_unique<T>(*static_cast<const T*>(this));