Correlate the events of XLA's PyLocalClient.
PiperOrigin-RevId: 300793558
Change-Id: I8902ca2dbe253378593451f88f25f0dece04d969
diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index 7f8e9cd..e9219f6 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -339,7 +339,10 @@
{StatType::kStepId, StatType::kIterNum}},
{HostEventType::kKernelLaunch,
HostEventType::kKernelExecute,
- {StatType::kCorrelationId}}});
+ {StatType::kCorrelationId}},
+ {HostEventType::kLocalExecutableExecuteOnLocalDevice,
+ HostEventType::kLocalExecutableExecute,
+ {StatType::kRunId}}});
const std::vector<int64 /*EventType*/> root_event_types(
{HostEventType::kTraceContext, HostEventType::kFunctionRun,
HostEventType::kSessionRun, HostEventType::kHostTrainingLoopIteration});
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index c8c7020..8eb7bd7 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -80,6 +80,10 @@
{"WhileOp-StartBody", kWhileOpStartBody},
{"ForOp", kForOp},
{"PartitionedCallOp", kPartitionedCallOp},
+ // XLA related.
+ {"LocalExecutable::ExecuteOnLocalDevices",
+ kLocalExecutableExecuteOnLocalDevice},
+ {"LocalExecutable::Execute", kLocalExecutableExecute},
// tf.data related.
{"IteratorGetNextOp::DoCompute", kIteratorGetNextOp},
// Virtual events for grouping.
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 8e9dea7..ef8ec92 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -74,6 +74,9 @@
kWhileOpStartBody,
kForOp,
kPartitionedCallOp,
+ // XLA related.
+ kLocalExecutableExecuteOnLocalDevice,
+ kLocalExecutableExecute,
// tf.data related.
kIteratorGetNextOp,
// Virtual events for grouping.
