chore: import upstream snapshot with attribution

tensorflow/python/data/benchmarks/BUILD

@@ -0,0 +1,132 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load(
+    "//tensorflow/tools/test:performance.bzl",
+    "tf_py_benchmark_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "benchmark_base",
+    srcs = ["benchmark_base.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/data/util:nest",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "meta_benchmark",
+    srcs = ["meta_benchmark.py"],
+    deps = [
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "batch_benchmark",
+    srcs = ["batch_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:random_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "filter_benchmark",
+    srcs = ["filter_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/ops:array_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "from_tensor_slices_benchmark",
+    srcs = ["from_tensor_slices_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/experimental/ops:get_single_element",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:structured_function",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:dataset_ops_gen",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "interleave_benchmark",
+    srcs = ["interleave_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/experimental/ops:interleave_ops",
+        "//tensorflow/python/data/experimental/ops:testing",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "list_files_benchmark",
+    srcs = ["list_files_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "map_benchmark",
+    srcs = ["map_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:map_fn",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops:while_loop",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "prefetch_benchmark",
+    srcs = ["prefetch_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+    ],
+)
+
+tf_py_benchmark_test(
+    name = "range_benchmark",
+    srcs = ["range_benchmark.py"],
+    deps = [
+        ":benchmark_base",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/data/ops:options",
+    ],
+)

tensorflow/python/data/benchmarks/batch_benchmark.py

@@ -0,0 +1,109 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.batch()`."""
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import random_ops
+
+
+class BatchBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.batch()`."""
+
+  def benchmark_batch_sparse(self):
+    non_zeros_per_row_values = [0, 1, 5, 10, 100]
+    batch_size_values = [1, 32, 64, 128, 1024]
+
+    for non_zeros_per_row in non_zeros_per_row_values:
+
+      tensor = sparse_tensor.SparseTensor(
+          indices=np.arange(non_zeros_per_row, dtype=np.int64)[:, np.newaxis],
+          values=np.arange(non_zeros_per_row, dtype=np.int64),
+          dense_shape=[1000])
+
+      for batch_size in batch_size_values:
+        dataset = dataset_ops.Dataset.from_tensors(tensor).repeat().batch(
+            batch_size)
+        self.run_and_report_benchmark(
+            dataset,
+            num_elements=100000 // batch_size,
+            iters=1,
+            extras={
+                "model_name": "batch.benchmark.1",
+                "parameters": "%d.%d" % (batch_size, non_zeros_per_row),
+            },
+            name="sparse_num_elements_%d_batch_size_%d" %
+            (non_zeros_per_row, batch_size))
+
+  def _benchmark_batch_dense(self, parallel_copy, benchmark_id):
+    for element_exp in [10, 12, 14, 16, 18, 20, 22]:
+      for batch_exp in [3, 6, 9]:
+        element_size = 1 << element_exp
+        batch_size = 1 << batch_exp
+        dataset = dataset_ops.Dataset.from_tensors(
+            np.random.rand(element_size)).repeat().batch(batch_size)
+        options = options_lib.Options()
+        options.experimental_optimization.parallel_batch = parallel_copy
+        dataset = dataset.with_options(options)
+        tag = "_parallel_copy" if parallel_copy else ""
+        self.run_and_report_benchmark(
+            dataset,
+            num_elements=(1 << (22 - ((batch_exp + element_exp) // 2))),
+            iters=1,
+            extras={
+                "model_name": "batch.benchmark.%d" % benchmark_id,
+                "parameters": "%d.%d" % (batch_size, element_size),
+            },
+            name="batch_element_size_%d_batch_size_%d%s" %
+            (element_size, batch_size, tag))
+
+  def benchmark_batch_dense(self):
+    self._benchmark_batch_dense(parallel_copy=False, benchmark_id=2)
+    #self._benchmark_batch_dense(parallel_copy=True, benchmark_id=3)
+
+  def benchmark_parallel_batch(self):
+    batch_size = 128
+    nums_parallel_calls = [None, 1, 4, 16, dataset_ops.AUTOTUNE]
+    num_range = 100000
+
+    def f(_):
+      return random_ops.random_uniform([224, 224, 3])
+
+    for num_parallel_calls in nums_parallel_calls:
+      num_parallel_calls_str = ("autotune"
+                                if num_parallel_calls == dataset_ops.AUTOTUNE
+                                else str(num_parallel_calls))
+      op_str = ("batch" if num_parallel_calls is None else
+                ("parallel_batch_num_parallel_calls_%s" %
+                 num_parallel_calls_str))
+
+      dataset = dataset_ops.Dataset.range(num_range).map(f).batch(
+          batch_size, num_parallel_calls=num_parallel_calls)
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=num_range // batch_size,
+          iters=1,
+          extras={
+              "model_name": "batch.benchmark.4",
+              "parameters": "%d.%s" % (batch_size, num_parallel_calls_str),
+          },
+          name="batch_size_%d_%s" % (batch_size, op_str))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/benchmark_base.py

@@ -0,0 +1,250 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test utilities for tf.data benchmarking functionality."""
+import time
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.data.util import nest
+from tensorflow.python.eager import context
+from tensorflow.python.platform import test
+
+
+class DatasetBenchmarkBase(test.Benchmark):
+  """Base class for dataset benchmarks."""
+
+  def _run_eager_benchmark(self, iterable, iters, warmup):
+    """Benchmark the iterable in eager mode.
+
+    Runs the iterable `iters` times. In each iteration, the benchmark measures
+    the time it takes to go execute the iterable.
+
+    Args:
+      iterable: The tf op or tf.data Dataset to benchmark.
+      iters: Number of times to repeat the timing.
+      warmup: If true, warms up the session caches by running an untimed run.
+
+    Returns:
+      A float, representing the median time (with respect to `iters`)
+      it takes for the iterable to be executed `iters` num of times.
+
+    Raises:
+      RuntimeError: When executed in graph mode.
+    """
+
+    deltas = []
+    if not context.executing_eagerly():
+      raise RuntimeError(
+          "Eager mode benchmarking is not supported in graph mode.")
+
+    for _ in range(iters):
+      if warmup:
+        iterator = iter(iterable)
+        next(iterator)
+
+      iterator = iter(iterable)
+      start = time.time()
+      next(iterator)
+      end = time.time()
+      deltas.append(end - start)
+    return np.median(deltas)
+
+  def _run_graph_benchmark(self,
+                           iterable,
+                           iters,
+                           warmup,
+                           session_config,
+                           initializer=None):
+    """Benchmarks the iterable in graph mode.
+
+    Runs the iterable `iters` times. In each iteration, the benchmark measures
+    the time it takes to go execute the iterable.
+
+    Args:
+      iterable: The tf op or tf.data Dataset to benchmark.
+      iters: Number of times to repeat the timing.
+      warmup: If true, warms up the session caches by running an untimed run.
+      session_config: A ConfigProto protocol buffer with configuration options
+        for the session. Applicable only for benchmarking in graph mode.
+      initializer: The initializer op required to initialize the iterable.
+
+    Returns:
+      A float, representing the median time (with respect to `iters`)
+      it takes for the iterable to be executed `iters` num of times.
+
+    Raises:
+      RuntimeError: When executed in eager mode.
+    """
+
+    deltas = []
+    if context.executing_eagerly():
+      raise RuntimeError(
+          "Graph mode benchmarking is not supported in eager mode.")
+
+    for _ in range(iters):
+      with session.Session(config=session_config) as sess:
+        if warmup:
+          # Run once to warm up the session caches.
+          if initializer:
+            sess.run(initializer)
+          sess.run(iterable)
+
+        if initializer:
+          sess.run(initializer)
+        start = time.time()
+        sess.run(iterable)
+        end = time.time()
+      deltas.append(end - start)
+    return np.median(deltas)
+
+  def run_op_benchmark(self, op, iters=1, warmup=True, session_config=None):
+    """Benchmarks the op.
+
+    Runs the op `iters` times. In each iteration, the benchmark measures
+    the time it takes to go execute the op.
+
+    Args:
+      op: The tf op to benchmark.
+      iters: Number of times to repeat the timing.
+      warmup: If true, warms up the session caches by running an untimed run.
+      session_config: A ConfigProto protocol buffer with configuration options
+        for the session. Applicable only for benchmarking in graph mode.
+
+    Returns:
+      A float, representing the per-execution wall time of the op in seconds.
+      This is the median time (with respect to `iters`) it takes for the op
+      to be executed `iters` num of times.
+    """
+
+    if context.executing_eagerly():
+      return self._run_eager_benchmark(iterable=op, iters=iters, warmup=warmup)
+
+    return self._run_graph_benchmark(
+        iterable=op, iters=iters, warmup=warmup, session_config=session_config)
+
+  def run_benchmark(self,
+                    dataset,
+                    num_elements,
+                    iters=1,
+                    warmup=True,
+                    apply_default_optimizations=False,
+                    session_config=None):
+    """Benchmarks the dataset.
+
+    Runs the dataset `iters` times. In each iteration, the benchmark measures
+    the time it takes to go through `num_elements` elements of the dataset.
+
+    Args:
+      dataset: Dataset to benchmark.
+      num_elements: Number of dataset elements to iterate through each benchmark
+        iteration.
+      iters: Number of times to repeat the timing.
+      warmup: If true, warms up the session caches by running an untimed run.
+      apply_default_optimizations: Determines whether default optimizations
+        should be applied.
+      session_config: A ConfigProto protocol buffer with configuration options
+        for the session. Applicable only for benchmarking in graph mode.
+
+    Returns:
+      A float, representing the per-element wall time of the dataset in seconds.
+      This is the median time (with respect to `iters`) it takes for the dataset
+      to go through `num_elements` elements, divided by `num_elements.`
+    """
+
+    # The options that have been applied to the dataset are preserved so that
+    # they are not overwritten while benchmarking.
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = (
+        apply_default_optimizations)
+    dataset = dataset.with_options(options)
+
+    # NOTE: We use `dataset.skip()` to perform the iterations in C++, avoiding
+    # the overhead of having to execute a TensorFlow op for each step of the
+    # input pipeline. Note that this relies on the underlying implementation of
+    # `skip` to execute upstream computation. If it is optimized in the future,
+    # we will have to change this code.
+    dataset = dataset.skip(num_elements - 1)
+
+    if context.executing_eagerly():
+      median_duration = self._run_eager_benchmark(
+          iterable=dataset, iters=iters, warmup=warmup)
+      return median_duration / float(num_elements)
+
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    next_element = iterator.get_next()
+    op = nest.flatten(next_element)[0].op
+    median_duration = self._run_graph_benchmark(
+        iterable=op,
+        iters=iters,
+        warmup=warmup,
+        session_config=session_config,
+        initializer=iterator.initializer)
+    return median_duration / float(num_elements)
+
+  def run_and_report_benchmark(self,
+                               dataset,
+                               num_elements,
+                               name,
+                               iters=5,
+                               extras=None,
+                               warmup=True,
+                               apply_default_optimizations=False,
+                               session_config=None):
+    """Benchmarks the dataset and reports the stats.
+
+    Runs the dataset `iters` times. In each iteration, the benchmark measures
+    the time it takes to go through `num_elements` elements of the dataset.
+    This is followed by logging/printing the benchmark stats.
+
+    Args:
+      dataset: Dataset to benchmark.
+      num_elements: Number of dataset elements to iterate through each benchmark
+        iteration.
+      name: Name of the benchmark.
+      iters: Number of times to repeat the timing.
+      extras: A dict which maps string keys to additional benchmark info.
+      warmup: If true, warms up the session caches by running an untimed run.
+      apply_default_optimizations: Determines whether default optimizations
+        should be applied.
+      session_config: A ConfigProto protocol buffer with configuration options
+        for the session. Applicable only for benchmarking in graph mode.
+
+    Returns:
+      A float, representing the per-element wall time of the dataset in seconds.
+      This is the median time (with respect to `iters`) it takes for the dataset
+      to go through `num_elements` elements, divided by `num_elements.`
+    """
+    wall_time = self.run_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        iters=iters,
+        warmup=warmup,
+        apply_default_optimizations=apply_default_optimizations,
+        session_config=session_config)
+    if extras is None:
+      extras = {}
+    if context.executing_eagerly():
+      name = "{}.eager".format(name)
+      extras["implementation"] = "eager"
+    else:
+      name = "{}.graph".format(name)
+      extras["implementation"] = "graph"
+    extras["num_elements"] = num_elements
+    self.report_benchmark(
+        wall_time=wall_time, iters=iters, name=name, extras=extras)
+    return wall_time

tensorflow/python/data/benchmarks/filter_benchmark.py

@@ -0,0 +1,45 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.filter()`."""
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.ops import array_ops
+
+
+class FilterBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.filter()`."""
+
+  def _benchmark(self, predicate, name, benchmark_id):
+    num_elements = 100000
+    dataset = dataset_ops.Dataset.from_tensors(True)
+    dataset = dataset.repeat().filter(predicate)
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "filter.benchmark.%d" % benchmark_id,
+            "parameters": "%d" % num_elements,
+        },
+        name=name)
+
+  def benchmark_simple_function(self):
+    self._benchmark(array_ops.identity, "simple_function", benchmark_id=1)
+
+  def benchmark_return_component_optimization(self):
+    self._benchmark(lambda x: x, "return_component", benchmark_id=2)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/from_tensor_slices_benchmark.py

@@ -0,0 +1,161 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.from_tensor_slices()`."""
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import get_single_element
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import structured_function
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import gen_dataset_ops
+
+
+class SingleThreadedFlatMapDataset(dataset_ops.UnaryDataset):
+  """A `Dataset` that maps a function over its input and flattens the result."""
+
+  def __init__(self, input_dataset, map_func):
+    """See `Dataset.flat_map()` for details."""
+    self._input_dataset = input_dataset
+    self._map_func = structured_function.StructuredFunctionWrapper(
+        map_func,
+        self._transformation_name(),
+        dataset=input_dataset,
+        defun_kwargs={"_executor": "SINGLE_THREADED_EXECUTOR"})
+    self._structure = self._map_func.output_structure._element_spec  # pylint: disable=protected-access
+    variant_tensor = gen_dataset_ops.flat_map_dataset(
+        input_dataset._variant_tensor,  # pylint: disable=protected-access
+        self._map_func.function.captured_inputs,
+        f=self._map_func.function,
+        **self._flat_structure)
+    super(SingleThreadedFlatMapDataset, self).__init__(input_dataset,
+                                                       variant_tensor)
+
+  def _functions(self):
+    return [self._map_func]
+
+  @property
+  def element_spec(self):
+    return self._structure
+
+  def _transformation_name(self):
+    return "SingleThreadedFlatMapDataset"
+
+
+class FromTensorSlicesBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.from_tensor_slices()`."""
+
+  def benchmark_slice_repeat_batch(self):
+    input_size = 10000
+    batch_size = 100
+    num_epochs = 100
+    num_elements = input_size * num_epochs // batch_size
+
+    input_data = np.random.randn(input_size)
+
+    dataset = dataset_ops.Dataset.from_tensor_slices(input_data)
+    dataset = dataset.repeat(num_epochs).batch(batch_size)
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "from_tensor_slices.benchmark.1",
+            "parameters": "%d.%d" % (input_size, batch_size),
+        },
+        name="slice_repeat_batch_input_%d_batch_%d" % (input_size, batch_size))
+
+  def benchmark_reshape_slice_repeat(self):
+    input_size = 10000
+    reshape_dim = [100, 100]
+    num_epochs = 100
+
+    num_elements = num_epochs * reshape_dim[0]
+
+    data = np.random.randn(input_size).reshape(*reshape_dim)
+    dataset = dataset_ops.Dataset.from_tensor_slices(data).repeat(num_epochs)
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "from_tensor_slices.benchmark.2",
+            "parameters": "%d" % input_size,
+        },
+        name="reshape_slice_repeat_input_%d" % input_size,
+    )
+
+  def benchmark_slice_repeat_sparse(self):
+    non_zeros_per_row_values = [0, 1, 5, 10, 100]
+    num_rows_values = [32, 64, 128, 1024]
+
+    for non_zeros_per_row in non_zeros_per_row_values:
+      tensor = sparse_tensor.SparseTensor(
+          indices=np.arange(non_zeros_per_row, dtype=np.int64)[:, np.newaxis],
+          values=np.arange(non_zeros_per_row, dtype=np.int64),
+          dense_shape=[1000])
+
+      for num_rows in num_rows_values:
+
+        # TODO(b/147153744): Function-valued attributes with their own
+        # attributes are currently only supported in graph mode.
+        @def_function.function
+        def make_dataset():
+          # pylint: disable=cell-var-from-loop
+          dataset = dataset_ops.Dataset.from_tensors(tensor)
+          dataset = dataset.repeat(num_rows).batch(num_rows)
+          batched_tensor = get_single_element.get_single_element(dataset)
+
+          dataset = dataset_ops.Dataset.from_tensors(batched_tensor).repeat()
+          return SingleThreadedFlatMapDataset(
+              dataset, dataset_ops.Dataset.from_tensor_slices)
+
+        self.run_and_report_benchmark(
+            make_dataset(),
+            num_elements=100000,
+            iters=5,
+            extras={
+                "model_name": "from_tensor_slices.benchmark.3",
+                "parameters": "%d.%d" % (non_zeros_per_row, num_rows),
+            },
+            name="slice_repeat_sparse_elements_per_row_%d_num_rows_%d" %
+            (non_zeros_per_row, num_rows))
+
+  def benchmark_slice_batch_cache_repeat(self):
+    input_size = 10000
+    batch_size = 100
+    num_epochs = 100
+    num_elements = input_size * num_epochs // batch_size
+
+    input_data = np.random.randn(input_size)
+
+    dataset = (
+        dataset_ops.Dataset.from_tensor_slices(input_data).batch(
+            batch_size).cache().repeat(num_epochs))
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "from_tensor_slices.benchmark.4",
+            "parameters": "%d.%d" % (input_size, batch_size),
+        },
+        name="slice_batch_cache_repeat_input_%d_batch_%d" %
+        (input_size, batch_size))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/interleave_benchmark.py

@@ -0,0 +1,183 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.interleave()`."""
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.ops import dataset_ops
+
+NON_PARALLEL = "non_parallel"
+EXPERIMENTAL_PARALLEL = "experimental_parallel"
+CORE_PARALLEL = "core_parallel"
+
+
+def _make_fake_dataset_fn(initial_delay_us, remainder_delay_us):
+  """Returns a dataset that emulates a remote storage data source.
+
+  Returns a dataset factory which creates a dataset with 100 elements that
+  emulates the performance characteristic of a file-based dataset stored in a
+  remote storage. In particular, the first element will take an order of
+  magnitude longer to produce than the remaining elements (100ms vs. 1ms).
+
+  Args:
+    initial_delay_us: How long to wait before producing the first element.
+    remainder_delay_us: How long to wait before producing subsequent elements.
+  """
+
+  def fake_dataset_fn(unused):
+    """Returns a function that creates a dataset with the specified delays."""
+    del unused
+
+    def make_dataset(time_us, num_elements):
+      dataset = dataset_ops.Dataset.range(num_elements)
+      if time_us > 0:
+        dataset = dataset.apply(testing.sleep(time_us))
+      return dataset
+
+    if not initial_delay_us:
+      return make_dataset(remainder_delay_us, 100)
+
+    return make_dataset(initial_delay_us,
+                        0).concatenate(make_dataset(remainder_delay_us, 100))
+
+  return fake_dataset_fn
+
+
+class ParallelInterleaveBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.experimental.parallel_interleave()`."""
+
+  def apply_interleave(self, interleave_version, dataset, interleave_fn,
+                       cycle_length, num_parallel_calls):
+    if interleave_version == NON_PARALLEL:
+      return dataset.interleave(interleave_fn, cycle_length=cycle_length)
+    elif interleave_version == EXPERIMENTAL_PARALLEL:
+      return dataset.apply(
+          interleave_ops.parallel_interleave(
+              interleave_fn, cycle_length=cycle_length))
+    elif interleave_version == CORE_PARALLEL:
+      if not num_parallel_calls:
+        num_parallel_calls = cycle_length
+      return dataset.interleave(
+          interleave_fn,
+          cycle_length=cycle_length,
+          num_parallel_calls=num_parallel_calls)
+    else:
+      raise ValueError("Unknown version: " + interleave_version)
+
+  def make_dataset(self,
+                   interleave_version,
+                   initial_delay,
+                   remainder_delay,
+                   cycle_length,
+                   num_parallel_calls=None):
+    dataset = dataset_ops.Dataset.range(1).repeat()
+    interleave_fn = _make_fake_dataset_fn(initial_delay, remainder_delay)
+    return self.apply_interleave(
+        interleave_version=interleave_version,
+        dataset=dataset,
+        interleave_fn=interleave_fn,
+        cycle_length=cycle_length,
+        num_parallel_calls=num_parallel_calls)
+
+  def _benchmark(self,
+                 interleave_version,
+                 num_elements,
+                 benchmark_id,
+                 benchmark_label,
+                 initial_delay_us=0,
+                 remainder_delay_us=0,
+                 cycle_length=10,
+                 iters=100,
+                 num_parallel_calls=None,
+                 name=None):
+    dataset = self.make_dataset(
+        interleave_version=interleave_version,
+        initial_delay=initial_delay_us,
+        remainder_delay=remainder_delay_us,
+        cycle_length=cycle_length,
+        num_parallel_calls=num_parallel_calls)
+
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        num_elements=num_elements,
+        iters=iters,
+        warmup=True,
+        extras={
+            "model_name":
+                "interleave.benchmark.%s.%d" % (benchmark_label, benchmark_id),
+            "parameters":
+                "%d.%d.%d.%s" %
+                (num_elements, cycle_length, iters, str(num_parallel_calls)),
+        },
+        name=name)
+
+  def benchmark_remote_file_simulation(self):
+    for i, version in enumerate([EXPERIMENTAL_PARALLEL, CORE_PARALLEL]):
+      self._benchmark(
+          interleave_version=version,
+          initial_delay_us=100 * 1000,
+          remainder_delay_us=1000,
+          num_elements=5000,
+          name="remote_file_simulation_" + version,
+          benchmark_id=i,
+          benchmark_label="remote_file")
+
+  def benchmark_fast_input(self):
+    for i, version in enumerate([EXPERIMENTAL_PARALLEL, CORE_PARALLEL]):
+      self._benchmark(
+          interleave_version=version,
+          num_elements=200000,
+          name="fast_input_" + version,
+          benchmark_id=i,
+          benchmark_label="fast_input")
+
+  # Measure the overhead of parallel interleaves compared to non-parallel
+  # interleave.
+  def benchmark_single_cycle(self):
+    for i, version in enumerate(
+        [NON_PARALLEL, EXPERIMENTAL_PARALLEL, CORE_PARALLEL]):
+      self._benchmark(
+          interleave_version=version,
+          cycle_length=1,
+          num_elements=200000,
+          name="single_cycle_" + version,
+          benchmark_id=i,
+          benchmark_label="single_cycle")
+
+  # Compare with a more reasonable cycle length. Experimental interleave
+  # cannot be compared here because it sets num_parallel_calls = cycle_length.
+  def benchmark_single_parallel_call(self):
+    self._benchmark(
+        interleave_version=CORE_PARALLEL,
+        num_elements=200000,
+        num_parallel_calls=1,
+        name="single_parallel_call_" + CORE_PARALLEL,
+        benchmark_id=1,
+        benchmark_label="single_parallel_call")
+
+  def benchmark_long_cycle(self):
+    for i, version in enumerate([EXPERIMENTAL_PARALLEL, CORE_PARALLEL]):
+      self._benchmark(
+          interleave_version=version,
+          cycle_length=1000,
+          num_elements=100000,
+          name="long_cycle_" + version,
+          benchmark_id=i,
+          benchmark_label="long_cycle")
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/list_files_benchmark.py

@@ -0,0 +1,65 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.list_files()`."""
+import os
+import shutil
+import tempfile
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+
+
+class ListFilesBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.list_files()`."""
+
+  def benchmark_nested_directories(self):
+    tmp_dir = tempfile.mkdtemp()
+    width = 1024
+    depth = 16
+    for i in range(width):
+      for j in range(depth):
+        new_base = os.path.join(tmp_dir, str(i),
+                                *[str(dir_name) for dir_name in range(j)])
+        os.makedirs(new_base)
+        child_files = ['a.py', 'b.pyc'] if j < depth - 1 else ['c.txt', 'd.log']
+        for f in child_files:
+          filename = os.path.join(new_base, f)
+          open(filename, 'w').close()
+    patterns = [
+        os.path.join(tmp_dir, os.path.join(*['**'
+                                             for _ in range(depth)]), suffix)
+        for suffix in ['*.txt', '*.log']
+    ]
+    # the num_elements depends on the pattern that has been defined above.
+    # In the current scenario, the num of files are selected based on the
+    # ['*.txt', '*.log'] patterns. Since the files which match either of these
+    # patterns are created once per `width`. The num_elements would be:
+    num_elements = width * 2
+
+    dataset = dataset_ops.Dataset.list_files(patterns)
+    self.run_and_report_benchmark(
+        dataset=dataset,
+        iters=3,
+        num_elements=num_elements,
+        extras={
+            'model_name': 'list_files.benchmark.1',
+            'parameters': '%d.%d' % (width, depth),
+        },
+        name='nested_directory(%d*%d)' % (width, depth))
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+
+if __name__ == '__main__':
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/map_benchmark.py

@@ -0,0 +1,194 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.map()`."""
+import numpy as np
+
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import map_op
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import while_loop
+
+
+class MapBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.map()`."""
+
+  def benchmark_chain_of_maps(self):
+
+    def benchmark_helper(chain_length, fn, use_inter_op_parallelism, label,
+                         benchmark_id):
+      dataset = dataset_ops.Dataset.range(10000)
+      for _ in range(chain_length):
+        dataset = map_op._MapDataset(  # pylint: disable=protected-access
+            dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=10000,
+          extras={
+              "model_name": "map.benchmark.%d" % benchmark_id,
+              "parameters": "%d" % chain_length,
+          },
+          name="chain_length_%d%s" % (chain_length, label))
+
+    chain_lengths = [0, 1, 2, 5, 10, 20, 50]
+    for chain_length in chain_lengths:
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x + 1,
+          use_inter_op_parallelism=True,
+          label="",
+          benchmark_id=1)
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x + 1,
+          use_inter_op_parallelism=False,
+          label="_single_threaded",
+          benchmark_id=2)
+      benchmark_helper(
+          chain_length=chain_length,
+          fn=lambda x: x,
+          use_inter_op_parallelism=True,
+          label="_short_circuit",
+          benchmark_id=3)
+
+  def benchmark_map_fan_out(self):
+    fan_outs = [1, 2, 5, 10, 20, 50, 100]
+
+    def benchmark_helper(fan_out, fn, use_inter_op_parallelism, label,
+                         benchmark_id):
+      dataset = dataset_ops.Dataset.from_tensors(
+          tuple(0 for _ in range(fan_out))).repeat(None)
+      dataset = map_op._MapDataset(  # pylint: disable=protected-access
+          dataset, fn, use_inter_op_parallelism=use_inter_op_parallelism)
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=10000,
+          extras={
+              "model_name": "map.benchmark.%d" % benchmark_id,
+              "parameters": "%d" % fan_out,
+          },
+          name="fan_out_%d%s" % (fan_out, label))
+
+    for fan_out in fan_outs:
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: [x + 1 for x in xs],
+          use_inter_op_parallelism=True,
+          label="",
+          benchmark_id=4)
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: [x + 1 for x in xs],
+          use_inter_op_parallelism=False,
+          label="_single_threaded",
+          benchmark_id=5)
+      benchmark_helper(
+          fan_out=fan_out,
+          fn=lambda *xs: xs,
+          use_inter_op_parallelism=True,
+          label="_short_circuit",
+          benchmark_id=6)
+
+  def benchmark_sequential_control_flow(self):
+    dataset = dataset_ops.Dataset.from_tensors(100000)
+
+    def fn(x):
+      i = constant_op.constant(0)
+
+      def body(i, x):
+        return math_ops.add(i, 1), x
+
+      return while_loop.while_loop(math_ops.less, body, [i, x])
+
+    num_elements = 1
+    dataset = dataset.map(fn)
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "map.benchmark.8",
+            "parameters": "%d" % num_elements,
+        },
+        name="sequential_control_flow",
+        apply_default_optimizations=True)
+
+  def benchmark_parallel_control_flow(self):
+    dataset = dataset_ops.Dataset.from_tensors(
+        random_ops.random_uniform([100, 10000000]))
+
+    def fn(x):
+      return map_fn.map_fn(
+          lambda y: y * array_ops.transpose(y), x, parallel_iterations=10)
+
+    num_elements = 1
+    dataset = dataset.map(fn)
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=1,
+        extras={
+            "model_name": "map.benchmark.9",
+            "parameters": "%d" % num_elements,
+        },
+        name="parallel_control_flow",
+        apply_default_optimizations=True)
+
+  def _benchmark_nested_parallel_map(self, cycle_length, num_parallel_calls):
+    k = 1024 * 1024
+    num_map_elements = 10
+    num_range_elements = 2000
+
+    def g(_):
+      return np.random.rand(50 * k).sum()
+
+    def f(_):
+      return dataset_ops.Dataset.range(num_map_elements).map(
+          g, num_parallel_calls=num_parallel_calls)
+
+    dataset = dataset_ops.Dataset.range(num_range_elements)
+    dataset = dataset.interleave(
+        f, cycle_length=cycle_length, num_parallel_calls=dataset_ops.AUTOTUNE)
+
+    cycle_length_str = ("default"
+                        if cycle_length is None else str(cycle_length))
+    num_parallel_calls_str = ("autotune"
+                              if num_parallel_calls == dataset_ops.AUTOTUNE else
+                              str(num_parallel_calls))
+    map_dataset_str = ("map" if num_parallel_calls is None else
+                       "parallel_map_num_parallel_calls_%s" %
+                       num_parallel_calls_str)
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_map_elements * num_range_elements,
+        extras={
+            "model_name": "map.benchmark.10",
+            "parameters": "%s_%s" % (cycle_length_str, num_parallel_calls_str),
+        },
+        name=("%s_cycle_length_%s" % (map_dataset_str, cycle_length_str)))
+
+  def benchmark_nested_parallel_map(self):
+    cycle_lengths = [None, 100]
+    nums_parallel_calls = [None, 1, 10, 100, dataset_ops.AUTOTUNE]
+    for cycle_length in cycle_lengths:
+      for num_parallel_calls in nums_parallel_calls:
+        self._benchmark_nested_parallel_map(cycle_length, num_parallel_calls)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/meta_benchmark.py

@@ -0,0 +1,149 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test utilities for tf.data benchmarking functionality."""
+import timeit
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.data.experimental.ops import testing
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+from tensorflow.python.eager import context
+from tensorflow.python.platform import test
+
+
+class MetaBenchmark(test.Benchmark):
+  """Benchmark that compares various ways of running tf.data benchmarks."""
+
+  # Note that each of these benchmarks is a separate method so that we can
+  # run them independently and collect a performance profile.
+
+  def setup_fast_dataset(self):
+    self.num_reps = 15
+    self.iters = 100000
+    options = options_lib.Options()
+    options.experimental_optimization.apply_default_optimizations = False
+    return dataset_ops.Dataset.range(10000**2).with_options(options)
+
+  def benchmark_fast_dataset_with_only_cpp_iterations(self):
+    dataset = self.setup_fast_dataset()
+    self.run_benchmark_with_only_cpp_iterations(dataset)
+
+  def benchmark_fast_dataset_with_session_run(self):
+    dataset = self.setup_fast_dataset()
+    self.run_benchmark_with_session_run(dataset)
+
+  def benchmark_fast_dataset_with_session_callable(self):
+    dataset = self.setup_fast_dataset()
+    self.run_benchmark_with_session_run(dataset, make_callable=True)
+
+  def benchmark_fast_dataset_in_eager(self):
+    with context.eager_mode():
+      dataset = self.setup_fast_dataset()
+      self.run_benchmark_in_eager(dataset)
+
+  def setup_slow_dataset(self):
+    dataset = self.setup_fast_dataset()
+    self.iters = 1000
+    # sleep for 1e-3s per iteration
+    return dataset.apply(testing.sleep(1000))
+
+  def benchmark_slow_dataset_with_only_cpp_iterations(self):
+    dataset = self.setup_slow_dataset()
+    self.run_benchmark_with_only_cpp_iterations(dataset)
+
+  def benchmark_slow_dataset_with_session_run(self):
+    dataset = self.setup_slow_dataset()
+    self.run_benchmark_with_session_run(dataset)
+
+  def benchmark_slow_dataset_with_session_callable(self):
+    dataset = self.setup_slow_dataset()
+    self.run_benchmark_with_session_run(dataset, make_callable=True)
+
+  def benchmark_slow_dataset_in_eager(self):
+    with context.eager_mode():
+      dataset = self.setup_slow_dataset()
+      self.run_benchmark_in_eager(dataset)
+
+  def report(self, deltas):
+    # Each `delta` is the time taken for `self.iters` iterations. Divide by the
+    # number of iterations here to get per-element iteration time.
+    deltas = np.array(deltas) / self.iters
+    # Discard the first 5 results from "warming up" the session.
+    deltas = deltas[5:]
+
+    median = np.median(deltas)
+    mean = np.mean(deltas)
+    min_val = np.min(deltas)
+    max_val = np.max(deltas)
+    extras = {
+        "iters_per_second": 1 / median,
+        "median": median,
+        "mean": mean,
+        "min": min_val,
+        "max": max_val,
+        "num_reps": self.num_reps - 5,
+    }
+    self.report_benchmark(wall_time=median, iters=self.iters, extras=extras)
+
+  def run_benchmark_in_eager(self, dataset):
+    deltas = []
+    for _ in range(self.num_reps):
+      iterator = iter(dataset)
+      deltas.append(timeit.timeit(lambda: next(iterator), number=self.iters))  # pylint: disable=cell-var-from-loop
+
+    self.report(deltas)
+
+  def run_benchmark_with_session_run(self, dataset, make_callable=False):
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    next_element = iterator.get_next()
+
+    with session.Session() as sess:
+      deltas = []
+      for _ in range(self.num_reps):
+        if make_callable:
+          get_next_element = sess.make_callable(next_element)
+        else:
+          # Note: session.run(next_element.op) is more performant than
+          # session.run(next_element) because we avoid the cost of copying the
+          # tensor from C++ to python.
+          get_next_element = lambda: sess.run(next_element.op)
+
+        sess.run(iterator.initializer)
+        deltas.append(timeit.timeit(get_next_element, number=self.iters))
+    self.report(deltas)
+
+  def run_benchmark_with_only_cpp_iterations(self, dataset):
+    """Benchmarks the dataset with the iterations performed in C++."""
+    # NOTE: We use `dataset.skip()` to perform the iterations in C++, avoiding
+    # the overhead of multiple `session.run()` calls. Note that this relies on
+    # the underlying implementation of `skip`: if it is optimized in the future,
+    # we will have to change this code.
+    dataset = dataset.skip(self.iters - 1)
+    iterator = dataset_ops.make_initializable_iterator(dataset)
+    next_element = iterator.get_next()
+
+    with session.Session() as sess:
+      deltas = []
+      for _ in range(self.num_reps):
+        sess.run(iterator.initializer)
+        deltas.append(
+            timeit.timeit(lambda: sess.run(next_element.op), number=1))
+    self.report(deltas)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/data/benchmarks/prefetch_benchmark.py

@@ -0,0 +1,40 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.prefetch()`."""
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+
+
+class PrefetchBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.prefetch()`."""
+
+  def benchmark_prefetch(self):
+    num_elements = 1000000
+    for prefetch_buffer in [1, 5, 10, 20, 100]:
+      dataset = dataset_ops.Dataset.range(num_elements)
+      dataset = dataset.prefetch(prefetch_buffer)
+
+      self.run_and_report_benchmark(
+          dataset,
+          num_elements=num_elements,
+          extras={
+              "model_name": "prefetch.benchmark.1",
+              "parameters": "%d" % prefetch_buffer,
+          },
+          name="prefetch_{}".format(prefetch_buffer))
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()

tensorflow/python/data/benchmarks/range_benchmark.py

@@ -0,0 +1,47 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Benchmarks for `tf.data.Dataset.range()`."""
+from tensorflow.python.data.benchmarks import benchmark_base
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import options as options_lib
+
+
+class RangeBenchmark(benchmark_base.DatasetBenchmarkBase):
+  """Benchmarks for `tf.data.Dataset.range()`."""
+
+  def _benchmark_range(self, num_elements, autotune, benchmark_id):
+    options = options_lib.Options()
+    options.autotune.enabled = autotune
+    dataset = dataset_ops.Dataset.range(num_elements)
+    dataset = dataset.with_options(options)
+
+    self.run_and_report_benchmark(
+        dataset,
+        num_elements=num_elements,
+        extras={
+            "model_name": "range.benchmark.%d" % benchmark_id,
+            "parameters": "%d.%s" % (num_elements, autotune),
+        },
+        name="modeling_%s" % ("on" if autotune else "off"))
+
+  def benchmark_range_with_modeling(self):
+    self._benchmark_range(num_elements=10000000, autotune=True, benchmark_id=1)
+
+  def benchmark_range_without_modeling(self):
+    self._benchmark_range(num_elements=50000000, autotune=False, benchmark_id=2)
+
+
+if __name__ == "__main__":
+  benchmark_base.test.main()