experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Look for TPU device planes if GPU planes are not found.

Browse Source 
PiperOrigin-RevId: 358097387
Change-Id: Icd774cd8e96566b177fc6ba2e7ad78b36c676311
This commit is contained in:
Yi Situ 2021-02-17 21:11:51 -08:00 committed by TensorFlower Gardener
parent b6fab0741d
commit ff6efc701e
3 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/core/profiler/convert/xplane_to_trace_events.cc
View File

@ -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();

1
tensorflow/core/profiler/utils/xplane_schema.cc
View File

@ -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";

2
tensorflow/core/profiler/utils/xplane_schema.h
View File

@ -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.