diff options
| author | Yi Situ <yisitu@google.com> | 2021-02-17 21:11:51 -0800 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2021-02-17 21:16:51 -0800 |
| commit | ff6efc701e0a711c1c9642eaa0b84dec3bb7956c (patch) | |
| tree | a1bfe9ce429cac9eba610cf59ebc82d0aac3915a | |
| parent | b6fab0741dc847c3061119082598f3a2bea058a9 (diff) | |
| download | tensorflow-ff6efc701e0a711c1c9642eaa0b84dec3bb7956c.tar.gz | |
Look for TPU device planes if GPU planes are not found.
PiperOrigin-RevId: 358097387
Change-Id: Icd774cd8e96566b177fc6ba2e7ad78b36c676311
| -rw-r--r-- | tensorflow/core/profiler/convert/xplane_to_trace_events.cc | 7 | ||||
| -rw-r--r-- | tensorflow/core/profiler/utils/xplane_schema.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/profiler/utils/xplane_schema.h | 2 |
3 files changed, 8 insertions, 2 deletions
diff --git a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc index cf8e5c7c54a..2733cf28fec 100644 --- a/tensorflow/core/profiler/convert/xplane_to_trace_events.cc +++ b/tensorflow/core/profiler/convert/xplane_to_trace_events.cc @@ -126,9 +126,12 @@ void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace) { XPlaneVisitor xplane = CreateTfXPlaneVisitor(host_plane); ConvertXPlaneToTraceEvents(kHostThreadsDeviceId, xplane, trace); } - - const std::vector<const XPlane*> device_planes = + std::vector<const XPlane*> device_planes = FindPlanesWithPrefix(xspace, kGpuPlanePrefix); + // We don't expect GPU and TPU planes to be present in the same XSpace. + if (device_planes.empty()) { + device_planes = FindPlanesWithPrefix(xspace, kTpuPlanePrefix); + } for (const XPlane* device_plane : device_planes) { XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane); uint32 device_id = kFirstDeviceId + xplane.Id(); diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc index 18265ddeca3..c463cc94ae0 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.cc +++ b/tensorflow/core/profiler/utils/xplane_schema.cc @@ -28,6 +28,7 @@ namespace profiler { const absl::string_view kHostThreadsPlaneName = "/host:CPU"; const absl::string_view kGpuPlanePrefix = "/device:GPU:"; +const absl::string_view kTpuPlanePrefix = "/device:TPU:"; const absl::string_view kCuptiDriverApiPlaneName = "/host:CUPTI"; const absl::string_view kMetadataPlaneName = "/host:metadata"; const absl::string_view kTFStreamzPlaneName = "/host:tfstreamz"; diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h index 718797dff42..35b61ac51a6 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.h +++ b/tensorflow/core/profiler/utils/xplane_schema.h @@ -31,6 +31,8 @@ namespace profiler { TF_CONST_INIT extern const absl::string_view kHostThreadsPlaneName; // Name prefix of XPlane that contains GPU events. TF_CONST_INIT extern const absl::string_view kGpuPlanePrefix; +// Name prefix of XPlane that contains TPU events. +TF_CONST_INIT extern const absl::string_view kTpuPlanePrefix; // Name of XPlane that contains CUPTI driver API generated events. TF_CONST_INIT extern const absl::string_view kCuptiDriverApiPlaneName; // Name of XPlane that contains profile metadata such as XLA debug info. |