| 193 | } |
| 194 | |
| 195 | void TFGraphNode::AddStepStat(int64 step, const string& device, |
| 196 | const NodeExecStats& step_stat) { |
| 197 | string dev = absl::AsciiStrToLower(device); |
| 198 | |
| 199 | // TODO(xpan): Make this more robust? |
| 200 | // See run_metadata_test.py |
| 201 | // It can be /job:0/replica:0/xxxx/device:GPU:0, or simply /device:GPU:0. |
| 202 | // It can has some ad-hoc suffix, such as /stream:xx or /memcpy:xx. |
| 203 | if (IsCanonicalDevice(dev)) { |
| 204 | if (!node_.canonical_device().empty()) { |
| 205 | if (node_.canonical_device() != dev) { |
| 206 | // TODO(xpan): Some RunMetadata node appears at multiple devices. |
| 207 | // Need to address it. |
| 208 | return; |
| 209 | } |
| 210 | } else { |
| 211 | node_.set_canonical_device(dev); |
| 212 | // TODO(xpan): Support things other than gpu? |
| 213 | if (dev.find("sycl") != dev.npos) { |
| 214 | node_.set_host_device(StringReplace(dev, "device:sycl:\\d+", "cpu:0")); |
| 215 | } else { |
| 216 | node_.set_host_device(StringReplace(dev, "gpu:\\d+", "cpu:0")); |
| 217 | } |
| 218 | AddOpType(node_.canonical_device()); |
| 219 | } |
| 220 | } |
| 221 | |
| 222 | auto exec = execs_.find(step); |
| 223 | if (exec == execs_.end()) { |
| 224 | execs_.insert(std::pair<int64, ExecStep>(step, ExecStep())); |
| 225 | exec = execs_.find(step); |
| 226 | } |
| 227 | |
| 228 | exec->second.AddTimeStats(dev, step_stat); |
| 229 | |
| 230 | if (dev == node_.canonical_device()) { |
| 231 | exec->second.AddMemoryStats(dev, step_stat); |
| 232 | } |
| 233 | } |
| 234 | |
| 235 | int64 ExecStep::exec_micros() const { |
| 236 | return accelerator_exec_micros() + cpu_exec_micros(); |