const int64_t max_us = (options.max_micros <= 0 && options.max_iters <= 0)
                             ? Options::kDefaultMicros
                             : options.max_micros;
  // NOLINTNEXTLINE
  printf("Running benchmark for %lld us\n", static_cast<long long>(max_us));
  const int64_t start_us = NowMicros();
  int64_t iters = 0;
  while (true) {

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<string, 4>>, 4>
      arg_grads;
  llvm::SmallVector<std::pair<int64_t, int64_t>, 4> ret_forward_input;
};

// Updates a called function's input signature by adjusting resource types, and
// adding required gradient arguments.
void ChangeFunctionInputSignature(
    func::FuncOp func,
    const llvm::SmallDenseMap<int64_t, llvm::SmallVector<string, 4>>& grads,

      /*shape=*/Create1DConstValue<int64_t>(builder, loc, {operand_rank}),
      /*value=*/CreateScalarConstValue<int64_t>(builder, loc, 0));

  // Converts start_index_map proto to tensor.
  const int64_t index_map_size = dimension_numbers.start_index_map().size();
  SmallVector<int64_t> indices(index_map_size);
  for (int64_t i = 0; i < index_map_size; i++) {

  const ArrayRef<int64_t> lhs_contracting_dims =
      dot_dimension_numbers.getLhsContractingDimensions();
  const ArrayRef<int64_t> rhs_contracting_dims =
      dot_dimension_numbers.getRhsContractingDimensions();
  const int64_t input_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(0).getType())
          .getRank();
  const int64_t filter_rank =

  // Figure out output shapes.
  SmallVector<int64_t, 4> concat_out_shape;
  SmallVector<int64_t, 4> pack_out_shape;

  const int64_t rank = input_type.getRank();
  int64_t pack_axis = pack_op.getAxis();
  size_t count = pack_inputs.size();
  if (pack_axis < 0) {
    pack_axis += rank;
  }

  // Concat out shape.
  for (int i = 0; i < rank; ++i) {
    int64_t dim_size = input_type.getDimSize(i);

namespace tensorflow {
namespace tfrt_compiler {
namespace {

constexpr int64_t kDefaultCheapCost = 1;

int64_t GetRankedTensorSize(mlir::TensorType type) {
  auto shape = type.getShape();

  int64_t size = 1;
  for (int64_t dim : shape) {
    // For unknown dimensions, use 1 as the size because it is usually the batch
    // dimension.
    //

inline ShapedType GetTransposedType(Value input,
                                    llvm::ArrayRef<int64_t> permutation_array) {
  auto input_type = input.getType().cast<ShapedType>();
  if (permutation_array.size() != input_type.getRank()) {
    return nullptr;
  }
  llvm::SmallVector<int64_t> transposed_shape(permutation_array.size());
  for (int64_t i = 0; i < permutation_array.size(); ++i) {

// body and condition. The sets of keys of the two maps must be the same.
llvm::SmallDenseMap<int64_t, ResourceArgUseInfo> MergeArgResourceUseInfo(
    const llvm::SmallDenseMap<int64_t, ResourceArgUseInfo>& infos0,
    const llvm::SmallDenseMap<int64_t, ResourceArgUseInfo>& infos1) {
  llvm::SmallDenseMap<int64_t, ResourceArgUseInfo> result;
  for (const auto& entry : infos0) {
    auto info1_it = infos1.find(entry.getFirst());

  LogicalResult RewriteOp(
      Operation* op, PatternRewriter& rewriter,
      const std::function<bool(ArrayRef<int64_t>, ArrayRef<int64_t>,
                               SmallVectorImpl<int64_t>&)>&
          get_broadcasted_shape) const;

  LogicalResult RewriteBatchMatMulV2Op(Operation* op,
                                       PatternRewriter& rewriter) const;
};

    llvm::function_ref<std::optional<Type>(int64_t)> arg_to_buffer_type,
    llvm::function_ref<bool(int64_t)> arg_buffer_size_is_fixed) {
  auto new_input_types = llvm::to_vector<8>(func.getFunctionType().getInputs());
  int64_t original_arg_count = new_input_types.size();
  Location loc = func.getLoc();
  for (int64_t i = 0; i < original_arg_count; ++i) {
    auto buffer_type = arg_to_buffer_type(i);