explicit CostFunctionRegistration(
      std::function<int64_t(const CostContext&, OpType)> cost_function) {
    RegisterCostFunction<OpType>(std::move(cost_function));
  }
};

bool HasCostFunctionRegistered(absl::string_view op_name);

}  // namespace tfrt_compiler
}  // namespace tensorflow

                          CostFunction cost_function) {
  RegisterCostFunction(GetCostFunctionRegistry(), op_name,
                       std::move(cost_function));
}

bool HasCostFunctionRegistered(absl::string_view op_name) {
  return GetCostFunctionRegistry().contains(op_name);
}

int64_t CostAnalysis::GetCost(mlir::Operation* op) const {
  assert(cost_map_.count(op) > 0);