#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"

namespace mlir {
namespace TFDevice {
namespace {
constexpr char kReplicaIdAttr[] = "_xla_replica_id";
constexpr char kDeviceOrdinalAttr[] = "device_ordinal";

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

#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"

namespace mlir {
namespace TFL {
namespace tac {
namespace {

// These are just fake costs.
constexpr float kDequantCost = 2.0;
constexpr float kQuantCost = 2.0;
constexpr float kRequantCost = 2.0;

// TODO(renjieliu): Ideally this should consider different kinds of SOCs as
// well.

// Get total bytes transferred.

#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime

namespace tensorflow {
namespace {

constexpr char kCostAttrName[] = "_tfrt_cost";
constexpr char kOpKeyAttrName[] = "op_key";

absl::flat_hash_map<int64_t, uint64_t> GetOpCostMap(mlir::ModuleOp op) {
  absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {

  EXPECT_TRUE(uq_requant_op);

  return uq_requant_op;
}

TEST_F(TfToUniformAttributeUtilsTest, UniformQuantizedAddOpAttributes) {
  TfToUniformAttributeUtilsTestPeer test_peer(&ctx_);

  constexpr absl::string_view kAddOpExpr =
      R"mlir(
      %0 = "tf.Const"() {value = #tf_type<tensor_proto : "0x746674656"> : tensor<1x3x2x2x!tf_type.qint32>} : () -> tensor<1x3x2x2x!tf_type.qint32>

namespace mlir {
namespace TFTPU {
namespace {

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

constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";

struct TPUResourceReadsWritesPartitioningPass
    : public impl::TPUResourceReadsWritesPartitioningPassBase<

#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
#include "tensorflow/compiler/mlir/lite/quantization/numerical_utils.h"

namespace mlir {
namespace quant {

constexpr int k8Bits = 8;
constexpr int k32Bits = 32;
constexpr unsigned kSigned = QuantizationFlags::Signed;

DeviceTarget::DeviceTarget(MLIRContext* ctx) : ctx_(ctx) {
  f32_ = FloatType::getF32(ctx_);
  i8_ = IntegerType::get(ctx_, k8Bits);

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TF {
namespace {

// Location attribute.
constexpr StringRef kClassAttr = "_class";
constexpr StringRef kSharedNameAttr = "shared_name";
constexpr StringRef kLocationPrefix = "loc:@";

#define GEN_PASS_DEF_CONVERTREADONLYREFERENCEVARIABLESTORESOURCEVARIABLESPASS

    {tsl::monitoring::Buckets::Exponential(1, 1.5, 45)});

// There were no MLIR ops so the old bridge was called successfully.
constexpr char kOldBridgeNoMlirSuccess[] = "kOldBridgeNoMlirSuccess";
// There were no MLIR ops so the old bridge was called but it failed.
constexpr char kOldBridgeNoMlirFailure[] = "kOldBridgeNoMlirFailure";

namespace {

// Time the execution of kernels (in CPU cycles). Meant to be used as RAII.

#include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
#include "tensorflow/core/platform/test.h"

namespace mlir {
namespace {

class MlirTranslationTest : public ::testing::Test {
 private:
  static constexpr char kMlirToGraphFlag[] = "-mlir-to-graph";

 public:
  MlirTranslationTest() : translation_(RegisterTranslation()) {
    // Create fake command line args so that the parser gets chosen.

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

namespace mlir {
namespace odml {

static constexpr std::string_view kStablehloModuleDefaultEntryFuncName = "main";
static constexpr std::string_view kStablehloFuncNamePrefix = "XlaCallModule";
static constexpr char kShardingAttr[] = "mhlo.sharding";
static constexpr char kShardingName[] = "Sharding";

class RemoveCustomCallWithSharding