Print info about an instruction or computation.
| 249 | |
| 250 | // Print info about an instruction or computation. |
| 251 | void DoInfoCommand(const HloModule& module, const std::vector<string>& tokens) { |
| 252 | if (tokens.size() != 2) { |
| 253 | std::cerr << "Illegal info query syntax. Use " |
| 254 | << R"("info name".)"; |
| 255 | return; |
| 256 | } |
| 257 | string node_name = tokens[1]; |
| 258 | |
| 259 | const HloInstruction* instr = FindInstruction(module, node_name); |
| 260 | const HloComputation* comp = FindComputation(module, node_name); |
| 261 | if (!instr && !comp) { |
| 262 | std::cerr << "Couldn't find HloInstruction or HloComputation named " |
| 263 | << node_name << std::endl; |
| 264 | return; |
| 265 | } |
| 266 | |
| 267 | if (comp != nullptr) { |
| 268 | std::cout << "HloComputation " << comp->name() << std::endl; |
| 269 | if (comp->IsFusionComputation()) { |
| 270 | std::cout << " Fusion instruction: " << comp->FusionInstruction()->name() |
| 271 | << std::endl; |
| 272 | } |
| 273 | std::cout << " Parameters:" << std::endl; |
| 274 | for (const auto& param : comp->parameter_instructions()) { |
| 275 | std::cout << " " << param->name() << " (" |
| 276 | << ShapeUtil::HumanStringWithLayout(param->shape()) << ")" |
| 277 | << std::endl; |
| 278 | } |
| 279 | HloInstruction* root = comp->root_instruction(); |
| 280 | std::cout << " Root instruction: " << root->name() << " (" |
| 281 | << ShapeUtil::HumanStringWithLayout(root->shape()) << ")" |
| 282 | << std::endl; |
| 283 | |
| 284 | auto embedded_computations = comp->MakeEmbeddedComputationsList(); |
| 285 | std::cout << " " << embedded_computations.size() << " embedded computation" |
| 286 | << (embedded_computations.size() != 1 ? "s" : "") |
| 287 | << (!embedded_computations.empty() ? ":" : ".") << std::endl; |
| 288 | for (const HloComputation* c : embedded_computations) { |
| 289 | std::cout << " " << c->name() << std::endl; |
| 290 | } |
| 291 | |
| 292 | // Find which computations reference comp as an embedded computation. |
| 293 | std::vector<const HloComputation*> users; |
| 294 | for (const HloComputation* c : module.computations()) { |
| 295 | if (absl::c_linear_search(c->MakeEmbeddedComputationsList(), comp)) { |
| 296 | users.push_back(c); |
| 297 | } |
| 298 | } |
| 299 | std::cout << " Used by " << users.size() << " computation" |
| 300 | << (users.size() != 1 ? "s" : "") << (!users.empty() ? ":" : "."); |
| 301 | for (const HloComputation* c : users) { |
| 302 | std::cout << " " << c->name() << std::endl; |
| 303 | } |
| 304 | } else { |
| 305 | std::cout << "HloInstruction " << instr->name() << std::endl; |
| 306 | std::cout << " Parent computation: " << instr->parent()->name() |
| 307 | << std::endl; |
| 308 | std::cout << " Opcode: " << HloOpcodeString(instr->opcode()) << std::endl; |
no test coverage detected