[tf.data] This CL enables a 5% experiment which uses max available parallelism for ParallelMapDataset. Empirical result indicated that parallelism can be more aggressive for ParallelMap.
PiperOrigin-RevId: 381292922
Change-Id: Iafe891de78aa84cb2d7d031aa11f190559b21183
diff --git a/tensorflow/core/data/dataset_utils.cc b/tensorflow/core/data/dataset_utils.cc
index 139cf8b..60c01de 100644
--- a/tensorflow/core/data/dataset_utils.cc
+++ b/tensorflow/core/data/dataset_utils.cc
@@ -893,6 +893,7 @@
REGISTER_DATASET_EXPERIMENT("enable_gradient_descent", 0);
REGISTER_DATASET_EXPERIMENT("parallelize_batch_copy", 50);
-}
+REGISTER_DATASET_EXPERIMENT("max_parallelism", 5);
+} // namespace
} // namespace data
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
index b432087..7980de9 100644
--- a/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_map_dataset_op.cc
@@ -724,7 +724,11 @@
&captured_func));
if (num_parallel_calls == model::kAutotune) {
- metrics::RecordTFDataAutotune(kDatasetType);
+ if (GetExperiments().contains("max_parallelism")) {
+ num_parallel_calls = port::MaxParallelism();
+ } else {
+ metrics::RecordTFDataAutotune(kDatasetType);
+ }
}
*output =
