Support ParallelBatch in tf.data bottleneck analysis.
PiperOrigin-RevId: 376220320
Change-Id: Ic56f880067462419adcc81d800f4cb9336fb408b
diff --git a/tensorflow/core/profiler/convert/xplane_to_tf_data_stats.cc b/tensorflow/core/profiler/convert/xplane_to_tf_data_stats.cc
index 91f48bf..218af0b 100644
--- a/tensorflow/core/profiler/convert/xplane_to_tf_data_stats.cc
+++ b/tensorflow/core/profiler/convert/xplane_to_tf_data_stats.cc
@@ -52,7 +52,8 @@
bool IsAsyncIterator(absl::string_view iterator_event_name) {
static auto* kAsyncIterators = new absl::flat_hash_set<absl::string_view>(
{"Prefetch", "ParallelInterleave", "ParallelMap", "ParseExample",
- "MapAndBatch", "DataService", "LegacyParallelInterleave"});
+ "MapAndBatch", "DataService", "LegacyParallelInterleave",
+ "ParallelBatch"});
return kAsyncIterators->contains(iterator_event_name);
}
diff --git a/tensorflow/core/profiler/utils/group_events.cc b/tensorflow/core/profiler/utils/group_events.cc
index 28df9bd..9d79ebc 100644
--- a/tensorflow/core/profiler/utils/group_events.cc
+++ b/tensorflow/core/profiler/utils/group_events.cc
@@ -774,7 +774,8 @@
{HostEventType::kPrefetchProduce,
HostEventType::kParallelInterleaveProduce,
HostEventType::kParallelMapProduce, HostEventType::kMapAndBatchProduce,
- HostEventType::kParseExampleProduce}) {
+ HostEventType::kParseExampleProduce,
+ HostEventType::kParallelBatchProduce}) {
auto produce_event_list = gtl::FindOrNull(event_node_map_, event_type);
if (!produce_event_list) continue;
VLOG(1) << produce_event_list->size() << " "
@@ -803,7 +804,8 @@
{HostEventType::kPrefetchConsume,
HostEventType::kParallelInterleaveConsume,
HostEventType::kParallelMapConsume, HostEventType::kMapAndBatchConsume,
- HostEventType::kParseExampleConsume}) {
+ HostEventType::kParseExampleConsume,
+ HostEventType::kParallelBatchConsume}) {
auto consume_event_list = gtl::FindOrNull(event_node_map_, event_type);
if (!consume_event_list) continue;
VLOG(1) << consume_event_list->size() << " "
diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc
index 9a36f82..aed84fa 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.cc
+++ b/tensorflow/core/profiler/utils/xplane_schema.cc
@@ -113,6 +113,8 @@
{"MapAndBatchConsume", kMapAndBatchConsume},
{"ParseExampleProduce", kParseExampleProduce},
{"ParseExampleConsume", kParseExampleConsume},
+ {"ParallelBatchProduce", kParallelBatchProduce},
+ {"ParallelBatchConsume", kParallelBatchConsume},
// Batching related.
{"BatchingSessionRun", kBatchingSessionRun},
{"ProcessBatch", kProcessBatch},
diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h
index 939dd86..1962a27 100644
--- a/tensorflow/core/profiler/utils/xplane_schema.h
+++ b/tensorflow/core/profiler/utils/xplane_schema.h
@@ -107,6 +107,8 @@
kMapAndBatchConsume,
kParseExampleProduce,
kParseExampleConsume,
+ kParallelBatchProduce,
+ kParallelBatchConsume,
// Batching related.
kBatchingSessionRun,
kProcessBatch,
