// nullptr will be returned.
TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValue(
    const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
    TF_Status* status);

TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValueWithTimeout(
    const char* key, int64_t key_size, int64_t timeout_seconds,
    TF_CoordinationServiceAgent* agent, TF_Status* status);

    AnalyzeArguments(func_op);
    AnalyzeBlock(&func_op.front());
  }

  int64_t GetCost(mlir::Operation* op) const;

 private:
  void AnalyzeArguments(mlir::func::FuncOp func_op);
  void AnalyzeBlock(mlir::Block* block);
  void EvaluateCost(mlir::Operation* op);

  int64_t max_arg_size_ = 1;
  llvm::DenseMap<mlir::Operation*, int64_t> cost_map_;
  const tfrt_stub::CostRecorder* cost_recorder_;
};

struct CostContext {

  monitoring::testing::CellReader<int64_t> mlir_function_pass_fallback_count_ =
      monitoring::testing::CellReader<int64_t>(
          /* metric name */
          "/tensorflow/core/mlir_function_pass_fallback_count");
  monitoring::testing::CellReader<int64_t>
      mlir_graph_optimization_pass_fallback_count_ =
          monitoring::testing::CellReader<int64_t>(
              /* metric name */

  static constexpr int64_t kDefaultMicros = 3000000;

  int64_t max_iters = 0;   // Maximum iterations to run, ignored if <= 0.
  int64_t max_micros = 0;  // Maximum microseconds to run, ignored if <= 0.
};

// Stats holds statistics collected during benchmarking.
struct Stats {
  std::vector<int64_t> per_iter_us;  // Per-iteration deltas in us.
  int64_t total_us;                  // Total time in us.

#include "tensorflow/core/platform/status.h"
#include "tsl/platform/mutex.h"

namespace tensorflow {
namespace {
bool ShouldBeMegamorphic(int64_t compile_count, int64_t execution_count) {
  const int64_t kCompileThreshold = 10;
  const int64_t kMinExecutionsPerCompile = 50;

  // This heuristic is trying to capture the following property: have we sunk a

  Value output_layer_norm_coefficients_;

  mlir::TFL::LSTMOp lstm_;

  Value none_;
  SmallVector<int64_t, 1> bias_slice_shape_;
  SmallVector<int64_t, 1> bias_size_values_;
  SmallVector<int64_t, 2> weight_slice_shape_;
  SmallVector<int64_t, 2> weight_slice_size_input_values_;
  SmallVector<int64_t, 2> weight_slice_size_recurrent_values_;
  OpBuilder builder_;
};

  int64_t x_row_dim = x_shape[x_shape.size() - 2];
  int64_t x_col_dim = x_shape[x_shape.size() - 1];
  int64_t y_row_dim = y_shape[y_shape.size() - 2];
  int64_t y_col_dim = y_shape[y_shape.size() - 1];
  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  int64_t expected_out_row_dim = op.getAdjX() ? x_col_dim : x_row_dim;

namespace {

std::vector<int64_t> TensorShapeAsVector(const tensorflow::TensorHandle& handle,
                                         tensorflow::Status* status) {
  std::vector<int64_t> shape;
  int rank = -1;
  *status = handle.NumDims(&rank);
  if (!status->ok()) {
    return shape;
  }
  shape.reserve(rank);
  for (int i = 0; i < rank; ++i) {
    int64_t dim;
    *status = handle.Dim(i, &dim);

    SmallVector<int64_t> start_indices(input_rank, 0);
    SmallVector<int64_t> slice_sizes = reshaped_permuted_shape;
    SmallVector<int64_t> strides(input_rank, 1);
    auto crop_values = llvm::to_vector<4>(crops.getValues<APInt>());
    for (int i = 0; i < block_rank; ++i) {
      int64_t crop_start = crop_values[i * 2].getSExtValue();
      int64_t crop_end = crop_values[i * 2 + 1].getSExtValue();

// 'shape' is computed from the original shape and the broadcast dimensions to
// match result shape.
int64_t GetElementIndex(llvm::SmallVectorImpl<int64_t> &shape,
                        llvm::SmallVectorImpl<int64_t> &current_index) {
  int64_t ind = 0;
  int64_t mul = 1;
  for (int i = shape.size() - 1; i >= 0; --i) {
    ind += (current_index[i] % shape[i]) * mul;
    mul *= shape[i];
  }