| 79 | } |
| 80 | |
| 81 | void ExecStep::AddMemoryStats(const string& dev, |
| 82 | const NodeExecStats& step_stat) { |
| 83 | ExecMemory exec_mem; |
| 84 | if (step_stat.all_start_micros() > 0) { |
| 85 | exec_mem.set_memory_micros(step_stat.all_start_micros() + |
| 86 | step_stat.op_end_rel_micros()); |
| 87 | } else { |
| 88 | fprintf(stderr, "%s has no start time, skipping\n", |
| 89 | step_stat.node_name().c_str()); |
| 90 | return; |
| 91 | } |
| 92 | |
| 93 | int accelerator_allocator_cnt = 0; |
| 94 | for (const auto& mem : step_stat.memory()) { |
| 95 | // TODO(xpan): Fix this hack. Currently the allocator name seems quite |
| 96 | // ad-hoc. |
| 97 | if (mem.allocator_name().find("GPU") == mem.allocator_name().npos) { |
| 98 | continue; |
| 99 | } |
| 100 | ++accelerator_allocator_cnt; |
| 101 | exec_mem.set_allocator_bytes_in_use( |
| 102 | std::max(static_cast<int64>(exec_mem.allocator_bytes_in_use()), |
| 103 | static_cast<int64>(mem.allocator_bytes_in_use()))); |
| 104 | for (const auto& alloc : mem.allocation_records()) { |
| 105 | allocations_.push_back(alloc); |
| 106 | } |
| 107 | } |
| 108 | if (accelerator_allocator_cnt > 1) { |
| 109 | fprintf(stderr, "found %d gpu allocator for 1 node\n", |
| 110 | accelerator_allocator_cnt); |
| 111 | } |
| 112 | |
| 113 | int64 total_output_bytes = 0; |
| 114 | for (const auto& output : step_stat.output()) { |
| 115 | if (output.has_tensor_description() && |
| 116 | output.tensor_description().has_allocation_description()) { |
| 117 | // TODO(xpan): Maybe allocated_bytes. |
| 118 | int64 output_bytes = std::max(output.tensor_description() |
| 119 | .allocation_description() |
| 120 | .allocated_bytes(), |
| 121 | output.tensor_description() |
| 122 | .allocation_description() |
| 123 | .requested_bytes()); |
| 124 | uint64 output_ptr = |
| 125 | output.tensor_description().allocation_description().ptr(); |
| 126 | total_output_bytes += output_bytes; |
| 127 | |
| 128 | auto& mem = (*exec_mem.mutable_output_memory())[output.slot()]; |
| 129 | mem.set_ptr(output_ptr); |
| 130 | mem.set_bytes(output_bytes); |
| 131 | } |
| 132 | } |
| 133 | exec_mem.set_output_bytes(total_output_bytes); |
| 134 | |
| 135 | if (step_stat.has_memory_stats()) { |
| 136 | if (IsPlacedOnCPU(dev)) { |
| 137 | // Currently we assume ops placed on gpu only allocate memory on gpu. |
| 138 | exec_mem.set_host_temp_bytes(exec_mem.host_temp_bytes() + |
no test coverage detected