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); XPlaneVisitor xplane = CreateTfXPlaneVisitor(host_plane);
ConvertXPlaneToTraceEvents(kHostThreadsDeviceId, xplane, trace); ConvertXPlaneToTraceEvents(kHostThreadsDeviceId, xplane, trace);
} }
std::vector<const XPlane*> device_planes =
const std::vector<const XPlane*> device_planes =
FindPlanesWithPrefix(xspace, kGpuPlanePrefix); 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) { for (const XPlane* device_plane : device_planes) {
XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane); XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane);
uint32 device_id = kFirstDeviceId + xplane.Id(); 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 kHostThreadsPlaneName = "/host:CPU";
const absl::string_view kGpuPlanePrefix = "/device:GPU:"; 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 kCuptiDriverApiPlaneName = "/host:CUPTI";
const absl::string_view kMetadataPlaneName = "/host:metadata"; const absl::string_view kMetadataPlaneName = "/host:metadata";
const absl::string_view kTFStreamzPlaneName = "/host:tfstreamz"; 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; TF_CONST_INIT extern const absl::string_view kHostThreadsPlaneName;
// Name prefix of XPlane that contains GPU events. // Name prefix of XPlane that contains GPU events.
TF_CONST_INIT extern const absl::string_view kGpuPlanePrefix; 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. // Name of XPlane that contains CUPTI driver API generated events.
TF_CONST_INIT extern const absl::string_view kCuptiDriverApiPlaneName; TF_CONST_INIT extern const absl::string_view kCuptiDriverApiPlaneName;
// Name of XPlane that contains profile metadata such as XLA debug info. // Name of XPlane that contains profile metadata such as XLA debug info.