PiperOrigin-RevId: 340071266
Change-Id: Ic21209a25a1f8efa1122c9cee4a8ab3b8043c308

PiperOrigin-RevId: 333597626
Change-Id: I63be2d4b736d7d1423c93999a0ac2184257eadb4

This CL introduces the concept of a SplitProvider. A SplitProvider produces a sequence of "split" tensors which are interpreted by source datasets to produce dataset elements.

When we initialize an iterator, a SplitProvider can be passed through the IteratorContext to indicate that the iterator should only iterate through the splits provided by the SplitProvider.

This CL adds an optional DatasetBase::MakeSplitIterator method which creates a SplitIterator to create splits for the dataset. For non-source datasets, the proper implementation is generally just to call MakeSplitIterator on their input. To support this reasonable default, we add a `DatasetBase::InputDatasets` method, which produces the input datasets for a dataset. If a dataset implements InputDatasets and has a single input dataset, MakeSplitIterator will call delegate to the input by default.

This CL only implements splitting for range_dataset_op; other splitting implementation will come in later CLs. This CL also implements a `ShardingSplitProvider`, to better test the range_dataset_op splitting implementation. `ShardingSplitProvider` will be useful in its own right for implementing an alternative to AutoShard which leverages splitting.

PiperOrigin-RevId: 332056019
Change-Id: I73b9b03cb91ae689c57a72fa6ba0acd092cf4cbe

[tf.data] Fix the problem for the current `AsGraphDefInternal` function in `ModelDatasetOp` is not capturing all attrs.

PiperOrigin-RevId: 331643560
Change-Id: I221b7a5510ee78ffb7ccfbed051d398ca056f4eb

PiperOrigin-RevId: 329784678
Change-Id: I3ad64c99aae42906f1ed96c46dc0885b518978fb

[tf.data] Reverting portion of cl/322686110, changing back model input time to be zero (i.e. infinitely fast consumer).

PiperOrigin-RevId: 323628170
Change-Id: Ia6dd4fe48985c431b84c6189e5622e29fde4f8f9

[tf.data] Calculate the average input time for the root node of the data input pipeline. Also add a unit test for the `SelfProcessingTime` function.

PiperOrigin-RevId: 322686110
Change-Id: I72c0bddb2ed32570325c7d8023aaaffbeab87378

[tf.data] Adding a metric for bytes produced and consumed by individual transformations, refactoring infrastructure for recording tf.data metrics, and moving the metrics API and implementation from `common_runtime` to `framework`.

PiperOrigin-RevId: 305062865
Change-Id: I63911f00154baf36aa225f66dbef0843239b7392

PiperOrigin-RevId: 301443563
Change-Id: I852269b86039a71466ddeadfe3ce03d75dc45fda

PiperOrigin-RevId: 299110761
Change-Id: I66ecaa9d01dc441f091888bef3f24d220e9180c5

[tf.data] Adding information necessary for reconstructing the input pipeline graph from TraceMe metadata.

PiperOrigin-RevId: 293728779
Change-Id: Ib06750c4c360db603c00eb2133ee936c243fdf88

PiperOrigin-RevId: 282386008
Change-Id: I14c1ac544da14e536855e85a455b6cb4ed885467

[tf.data] Adds an upper bound for the total buffer limit of the model in `Model::Optimize` as % of available RAM.

PiperOrigin-RevId: 263789432

This CL:
- removes unused `DatasetBase::Save()` and related tests
- replaces `SerilizationContext::optimization_only` with multiple functionality specific flags (`check_external_state`, `fail_if_unimplemented`, and `serialize_data_tensors`)
- introduces `DatasetBase::CheckExternalState` as an error-raising replacement for `DatasetBase::IsStateful` to make it possible to communicate the reason for why serialization failed through the error status
- adds `IteratorBase::SaveInternal` and `IteratorBase::RestoreInternal` in preparation of making these methods pure virtual

PiperOrigin-RevId: 262235093

After this change, restoring an iterator from a checkpoint will require that the iterator is initialized using a dataset that matches the dataset used for initializing the iterator used to create the checkpoint. In other words, if the Python definition of the input pipeline changes, the restoration of the iterator will fail.

The motivation for this change is to make it possible to save (and restore) datasets whose graph cannot be serialized (e.g. because it contains ops with resource inputs). This will in turn allow tf.data to implement "reshuffle each iteration" or in-memory caching between different Python iterator for the same dataset.

PiperOrigin-RevId: 260144783

PiperOrigin-RevId: 254812919

[tf.data] Adjusting auto-tuning period to 1 minute (from previous incorrect value of 1000 minutes) and improving auto-tuning logging.

PiperOrigin-RevId: 241826050

[tf.data] Exposing an option for specifying the CPU budget for autotuning parallelism and nesting the autotuning-related options under experimental_optimization.

PiperOrigin-RevId: 238369191

The previous implementation of many core `tf.data` transformations
(e.g. `Dataset.prefetch()`) would create one or more threads each time
an iterator over those datasets is created
(e.g. `ds.prefetch(N).repeat(100)` would create and destroy 100
threads). In addition to the overhead of thread creation, this
interacts poorly with some malloc implementations, and can contribute
to memory fragmentation.

The new implementation maintains an unbounded pool of physical threads
in each iterator (or `MultiDeviceIterator`) resource, and returns logical
"threads" to that pool when their work is complete instead of exiting
from them.

PiperOrigin-RevId: 236413014

This CL:
- adds counters for tf.data elements, autotuning and optimizations
- sets the number of iterations of the `tf_data_meta_optimizer` to one -- the iteration of tf.data optimizations is handled by the tf.data meta optimizer itself
- adds the `alwayslink` attribute to all tf.data optimization BUILD targets to make sure they are always registered (without this, they would not be registered for the Tensorflow server binary I was using for local testing) and further cleans up visibility and dependencies of //third_party/tensorflow/core/grappler/optimizers/data/BUILD
- introduces TFDataOptimizerBase as a base class for tf.data optimizations
- moves TensorFlow metrics into tensorflow::metrics namespace

PiperOrigin-RevId: 229302097

PiperOrigin-RevId: 226402626

[tf.data] Adding `tf.data.experimental.cardinality()` which provides information about dataset cardinality.

PiperOrigin-RevId: 224030418

The wrapped iterator's GetNext method is thread-safe. Calling it with mu_ exclusively locked serializes calls to impl_->GetNext().

PiperOrigin-RevId: 220799359

PiperOrigin-RevId: 220125067

[tf.data] Refactoring of performance modeling implementation and adding performance modeling for all core and experimental tf.data kernels.

PiperOrigin-RevId: 219406929

PiperOrigin-RevId: 219390881

PiperOrigin-RevId: 218765742

PiperOrigin-RevId: 216260575

PiperOrigin-RevId: 216247929

PiperOrigin-RevId: 215607038

PiperOrigin-RevId: 215579950

[tf.data] Moving auto-tuning optimizations into a background thread, refactoring the API for exposing tunable parameters, and removing `model::Node` from the public API.

PiperOrigin-RevId: 213907565

[tf.data] Adding support for `tf.data.AUTOTUNE` as a special value for the `num_parallel_calls` argument of `tf.data.Dataset.map()`, `tf.data.Dataset.interleave()`, and `tf.contrib.data.map_and_batch()`.

When `tf.data.AUTOTUNE` is specified, the level of parallelism is determined at runtime. The underlying mechanism instruments the input pipeline to build a performance model and then uses the model to find the optimal values for the parallelism knobs.

PiperOrigin-RevId: 213283297

PiperOrigin-RevId: 212557406

PiperOrigin-RevId: 211733735

- replacing `OpKernelContext` with newly introduced `DatasetContext` in `DatasetBase` constructor to make it possible to instantiate `DatasetBase` in places where an instance of `OpKernelContext` is not available

- replacing `dataset::MakeIteratorContext(OpKernelContext* ctx)` factory with `IteratorContext(OpKernelContext *ctx)` constructor.

- folding `GraphDatasetBase` into `DataseBase` and removing the default implementation of `AsGraphDefInternal`, making it the responsibility of the derived class to implement it to encourage/hint developers to provide serialization logic

PiperOrigin-RevId: 208560010

This CL:
- changes the `OptimizeDataset` checkpointing logic to checkpoint the optimized dataset (as opposed to the original dataset + the optimizations, re-running optimization every time a checkpoint is restored)
- replaces `OpKernelContext` with newly introduced `SerializationContext` in the signature of `AsGraphDefInternal` to reduce the scope of the context and also simplify the logic for overriding the `FunctionLibraryDefinition` when optimizations take place

PiperOrigin-RevId: 208282562

Renaming `AddParentDataset`, `SaveParent`, and `RestoreParent` to `AddInputDataset`, `SaveInput`, and `RestoreInput`.

PiperOrigin-RevId: 208272695

By marking DebugString() as const we can make some error messages more descriptive. Because DatasetIterator marks the return value of the dataset() function const, DebugString() cannot be called.

PiperOrigin-RevId: 198796894