#include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"

#include <string>
#include <vector>

#include "absl/strings/str_cat.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project

// graph optimizers) into a single MLIR optimization pass.
class MlirGraphOptimizationPass : public ::tensorflow::MlirOptimizationPass {
 public:
  llvm::StringRef name() const override { return "graph_optimization"; }

  ::tensorflow::MlirOptimizationPassState GetPassState(
      const ::tensorflow::DeviceSet* device_set,
      const ::tensorflow::ConfigProto& config_proto,

#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;

namespace mlir {
namespace tfd {

// Dialect for TFRT delegate operations.
class RuntimeFallbackDialect : public Dialect {
 public:
  explicit RuntimeFallbackDialect(MLIRContext* context);
  static StringRef getDialectNamespace() { return "tfd"; }
};

}  // namespace tfd
}  // namespace mlir

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.h.inc"

namespace tensorflow {
namespace tf_mlrt_tpu {

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

}  // namespace tf_mlrt_tpu
}  // namespace tensorflow

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h.inc"

//   * has an attribute with the given `name`,
//   * and the attribute is an integer list of the form [1, X, Y, 1],
// and writes X, Y as 32-bit integer attribute to `x`, `y`.
bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
                     IntegerAttr *y);

// Returns true if the attribute is an integer list of the form [1, X, Y, 1].
bool TFIntListIs1XY1(Attribute attr);

#include <optional>
#include <string>
#include <unordered_set>

#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "llvm/Support/Debug.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/Location.h"  // from @llvm-project

#include "tensorflow/compiler/mlir/tf2xla/transforms/test_utils.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/strings/string_view.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/register_common_dialects.h"
#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"

namespace tfrt {
namespace fallback_sync {

// Dialect for fallback operations.
class FallbackSyncDialect : public Dialect {
 public:
  explicit FallbackSyncDialect(MLIRContext *context);
  static StringRef getDialectNamespace() { return "tfrt_fallback_sync"; }
};

}  // namespace fallback_sync
}  // namespace tfrt

#define GET_OP_CLASSES
#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h.inc"

}

struct GetOpCostPass
    : mlir::PassWrapper<GetOpCostPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetOpCostPass)

  llvm::StringRef getArgument() const final { return "tfl-get-op-cost"; }
  llvm::StringRef getDescription() const final {
    return "Get cost for every op";
  }
  void runOnOperation() override;
};

void GetOpCostPass::runOnOperation() {