| 128 | } // namespace |
| 129 | |
| 130 | void StatSummarizer::ProcessStepStats(const StepStats& step_stats) { |
| 131 | int64 curr_total_us = 0; |
| 132 | int64 mem_total = 0; |
| 133 | |
| 134 | int64 first_node_start_us = |
| 135 | step_stats.dev_stats(0).node_stats(0).all_start_micros(); |
| 136 | |
| 137 | int node_num = 0; |
| 138 | for (const auto& ds : step_stats.dev_stats()) { |
| 139 | for (const auto& ns : ds.node_stats()) { |
| 140 | // NOTE(blackhc): To better support GPUs: |
| 141 | // GPU kernels are duplicated both in /stream:all and their |
| 142 | // /stream:$index. GPU memcpys are duplicated both in /memcpy and their |
| 143 | // /stream:$index. So only keep /stream:all and /memcpy and ignore all |
| 144 | // /stream:$index to only count GPU executions once. |
| 145 | if (ds.device().find("/stream") != std::string::npos && |
| 146 | ds.device().find("/stream:all") == std::string::npos) { |
| 147 | continue; |
| 148 | } |
| 149 | // NOTE(fishx): We will record ops execution time twice: one as CPU |
| 150 | // activity with device name "/host:CPU" and the other as TF runtime |
| 151 | // activity with device name started with "/job:*". It is safe to ignore |
| 152 | // CPU activties here. |
| 153 | // TODO(b/138729463): Read ops execution time from CPU activities instead |
| 154 | // of runtime acitivities. |
| 155 | if (ds.device().find("/host:CPU") != std::string::npos) { |
| 156 | continue; |
| 157 | } |
| 158 | |
| 159 | std::string name = ns.node_name(); |
| 160 | std::string op_type = "<>"; |
| 161 | // NOTE(blackhc): we have to ensure that all keys into the detail map |
| 162 | // are unique, so we add [Kernel] or [MemCpy] as a suffix to the name. |
| 163 | // To make the node type summary work better, we prefix "gpu:" to |
| 164 | // the op type when the info is from a /gpu/stream or /memcpy channel. |
| 165 | if (ds.device().find("/stream") != std::string::npos) { |
| 166 | // node_name: name ":" opType |
| 167 | auto parts = str_util::Split(ns.node_name(), ':'); |
| 168 | if (parts.size() == 2) { |
| 169 | name = parts[0] + " [Kernel]"; |
| 170 | op_type = "gpu:" + parts[1]; |
| 171 | } |
| 172 | } else if (ds.device().find("/memcpy") != std::string::npos) { |
| 173 | // node_name: name (":" opType)? ":" memCpyType |
| 174 | auto parts = str_util::Split(ns.node_name(), ':'); |
| 175 | if (parts.size() == 2 || parts.size() == 3) { |
| 176 | name = parts.front() + " [MemCpy]"; |
| 177 | // We don't care about the actual op type (it might not be available |
| 178 | // for edge_ memcpys). We only care that it's a memcpy for now. |
| 179 | op_type = "gpu:" + parts.back(); |
| 180 | } |
| 181 | } else { |
| 182 | op_type = OpType(ds, ns); |
| 183 | } |
| 184 | |
| 185 | ++node_num; |
| 186 | const int64 curr_time = ns.all_end_rel_micros(); |
| 187 | curr_total_us += curr_time; |