| 1145 | } |
| 1146 | |
| 1147 | Status Service::GetComputationGraphStats( |
| 1148 | const ComputationGraphStatsRequest* arg, ComputationStatsResponse* result) { |
| 1149 | if (!arg->has_computation()) { |
| 1150 | return InvalidArgument("Computations may not be empty."); |
| 1151 | } |
| 1152 | if (!arg->computation().has_host_program_shape()) { |
| 1153 | return InvalidArgument("Program shape may not be empty."); |
| 1154 | } |
| 1155 | |
| 1156 | HloModuleConfig config(ProgramShape{arg->computation().host_program_shape()}); |
| 1157 | config.set_debug_options(arg->debug_options()); |
| 1158 | TF_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module, |
| 1159 | CreateModuleFromProto(arg->computation(), config)); |
| 1160 | DumpHloModuleIfEnabled(*module, kBeforeOptimizationsDumpName); |
| 1161 | |
| 1162 | // Run HLO analysis to get the computation statistics. |
| 1163 | HloCostAnalysis analysis( |
| 1164 | execute_backend_->compiler()->ShapeSizeBytesFunction()); |
| 1165 | |
| 1166 | TF_RETURN_IF_ERROR(module->entry_computation()->Accept(&analysis)); |
| 1167 | |
| 1168 | ComputationStats stats; |
| 1169 | stats.set_flop_count(analysis.flop_count()); |
| 1170 | stats.set_transcendental_count(analysis.transcendental_count()); |
| 1171 | *result->mutable_stats() = stats; |
| 1172 | return Status::OK(); |
| 1173 | } |
| 1174 | |
| 1175 | DeviceHandle Service::SingleComputationDeviceHandle() const { |
| 1176 | DeviceHandle device_handle; |
nothing calls this directly
no test coverage detected