max_arg_size_ = std::max(max_arg_size_, GetRankedTensorSize(type));
    }
  }
}

void CostAnalysis::AnalyzeBlock(mlir::Block* block) {
  for (auto& op : *block) {
    EvaluateCost(&op);
  }
}

void CostAnalysis::EvaluateCost(mlir::Operation* op) {
  if (auto cost_function =
          mlir::dyn_cast<tfrt::compiler::CostFunctionInterface>(op)) {
    cost_map_[op] = cost_function.cost();
    return;
  }