| 188 | } |
| 189 | |
| 190 | static std::shared_ptr<HloSnapshot> DumpArguments( |
| 191 | const Backend* backend, const Executable* executable, |
| 192 | const absl::Span<const ShapedBuffer* const> arguments, se::Stream* stream) { |
| 193 | auto snapshot = std::make_shared<HloSnapshot>(); |
| 194 | snapshot->set_execution_platform(backend->platform()->Name()); |
| 195 | *snapshot->mutable_hlo() = *executable->hlo_proto(); |
| 196 | for (const ShapedBuffer* arg : arguments) { |
| 197 | auto literal = std::make_shared<Literal>(arg->on_host_shape()); |
| 198 | backend->transfer_manager()->TransferLiteralFromDevice( |
| 199 | stream, *arg, literal.get(), [snapshot, literal](Status status) { |
| 200 | if (!status.ok()) { |
| 201 | LOG(ERROR) << "TransferLiteralFromDevice for HLO snapshot inputs " |
| 202 | "failed: " |
| 203 | << status; |
| 204 | return; |
| 205 | } |
| 206 | *snapshot->add_arguments() = literal->ToProto(); |
| 207 | }); |
| 208 | } |
| 209 | return snapshot; |
| 210 | } |
| 211 | |
| 212 | static void DumpOutputsAndSaveSnapshot(const Backend* backend, |
| 213 | const ShapedBuffer& outputs, |
no test coverage detected