#include "tensorflow/compiler/mlir/lite/stablehlo/odml_converter/passes.h.inc"

constexpr float kOne = 1.0;
const float kOneOverRoot2 = kOne / std::sqrt(2);
constexpr float kHalf = kOne / 2.0;
constexpr float kTolerance = kOne / 1000.0;

// Gets the operation that uses the sole result of given operation
// if there is only one.
Operation* GetUserIfOnlyOne(Operation* op) {
  if (op->getNumResults() != 1) return nullptr;

namespace mlir {
namespace TFTPU {

namespace {

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

constexpr char kDeviceAttr[] = "device";
// Attribute of colocation classes.
constexpr char kClassAttr[] = "_class";

bool HasDevice(Operation* op) {
  auto attr = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!attr) return false;

 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlrt.tf_tensor";
};

// The MLIR type represents a tensorflow::Device*
class TFDeviceType
    : public mlir::Type::TypeBase<TFDeviceType, mlir::Type, mlir::TypeStorage> {
 public:
  using Base::Base;
  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlirt.tf_device";
};

class ControlType : public Type::TypeBase<ControlType, Type, TypeStorage> {
 public:
  using Base::Base;
  static constexpr ::mlir::StringLiteral name = "tf_executor.control";
};

class TokenType : public Type::TypeBase<TokenType, Type, TypeStorage> {
 public:
  using Base::Base;
  static constexpr ::mlir::StringLiteral name = "tf_executor.token";
};

}  // namespace tf_executor
}  // namespace mlir

    "/tensorflow/core/tf2xla/api/v2/dynamism_function_counter",
    "Counts how many functions are dynamic", "has_dynamism");

constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";

class InputMetricsLoweringPass
    : public impl::InputLoweringMetricsPassBase<InputMetricsLoweringPass> {
 public:

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

namespace mlir {
namespace TF {
namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kTFDeviceAttr[] = "tf.device";

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

class SimpleTFDeviceAssignmentPass

#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) {

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