#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/core/platform/status.h"

namespace tensorflow {

inline constexpr absl::string_view kCrashReproducerStdErr = "-";
inline constexpr absl::string_view kCrashReproducerCrashAnalysis =
    "crash_analysis";

// Creates a file to use for dumping and returns success if a file could be

using mlir::TF::ReadVariableOp;
using mlir::tf_device::ReplicateOp;

#define GEN_PASS_DEF_HOISTBROADCASTREADPASS
#include "tensorflow/compiler/mlir/tf2xla/internal/passes/clustering_passes.h.inc"

constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kCpuDeviceType[] = "CPU";

struct HoistBroadcastRead
    : public impl::HoistBroadcastReadPassBase<HoistBroadcastRead> {
  void runOnOperation() override;
};

#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"

namespace mlir {
namespace TFL {
namespace tac {

class ExampleHardware : public SimpleHardware {
 public:
  static constexpr char kId[] = "ExampleHardware";

  mlir::RewritePatternSet GetTransformations(
      MLIRContext* context) const override;

  mlir::TypeID GetTypeId() const override {
    return mlir::TypeID::get<ExampleHardware>();

#include "tensorflow/compiler/mlir/lite/experimental/tac/hardwares/simple_hardware.h"

namespace mlir {
namespace TFL {
namespace tac {

class NNAPIHardware : public SimpleHardware {
 public:
  static constexpr char kId[] = "NNAPI";

  mlir::RewritePatternSet GetTransformations(
      MLIRContext* context) const override;

  mlir::TypeID GetTypeId() const override {
    return mlir::TypeID::get<NNAPIHardware>();
  }

#include "xla/xla_data.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {

inline constexpr llvm::StringRef kInputShardingAttr =
    "input_sharding_configuration";
inline constexpr llvm::StringRef kOutputShardingAttr =
    "output_sharding_configuration";

// Parses the sharding string. This sharding string can be binary (serialized)
// or human readable.

namespace {

constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";
constexpr StringRef kUsesShapePolymorphismAttr = "jax.uses_shape_polymorphism";

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Platforms for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
constexpr StringRef kPlatformTpu = "TPU";

#include "tensorflow/compiler/mlir/lite/experimental/tac/transforms/device_transform_patterns.h"
#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"

namespace mlir {
namespace TFL {
namespace tac {

constexpr char ExampleHardware::kId[];  // Define kId.

mlir::RewritePatternSet ExampleHardware::GetTransformations(
    MLIRContext* context) const {
  mlir::RewritePatternSet patterns(context);

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

#include <string>

#include <gtest/gtest.h>
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "xla/xla_data.pb.h"

inline constexpr llvm::StringRef kXlaShardingAttrName = "_XlaSharding";

namespace tensorflow {
namespace {

TEST(DecodeShardingAttributeTest, CheckInvalidString) {
  xla::OpSharding sharding;

// both implement "func_1".
//
// The attribute on the FuncOp means what it actually implements while the
// attribute on the CallOp means what it actually looks for.
constexpr char kInterfaceNameAttr[] = "tac.interface_name";

inline std::optional<std::string> GetInterFaceName(Operation* op) {
  auto name_attr = op->getAttrOfType<StringAttr>(kInterfaceNameAttr);
  if (!name_attr) return std::nullopt;

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/core/lib/io/path.h"

namespace mlir {
namespace TF {
namespace {

static constexpr int kTextFileIndex_WholeLine = -2;
static constexpr int kTextFileIndex_LineNumber = -1;

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