Repository: tensorflow
Issue: tf.data.Dataset .map().batch() pattern is not matched to use fused implementation.
Labels: comp:tensorboard,type:bug,comp:data,TF 2.7

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
- Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
- TensorFlow installed from (source or binary): source
- TensorFlow version (use command below): v1.12.1-69264-g0cdf35562dc 2.9.0
- Python version: 3.8.10
- Bazel version (if compiling from source):
- GCC/Compiler version (if compiling from source):
- CUDA/cuDNN version: 11.5 / 8.3
- GPU model and memory: GTX1660 Ti

**Describe the current behavior**
combining `tf.data.Dataset.map()` with `.batch()` does not use the fused BatchAndMap implementation.

**Describe the expected behavior**
It does use the fused implementation. Currently, it's only possible to use the fused implementation when using the deprecated `experimental.map_and_batch()` transformation.

**[Contributing](https://www.tensorflow.org/community/contribute)**

- Do you want to contribute a PR? (yes/no): no
- Briefly describe your candidate solution(if contributing):

**Standalone code to reproduce the issue**

```py
import os
import datetime
from tqdm import tqdm
import numpy as np

import tensorflow as tf
print('TF version', tf.__version__)

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), 'Physical GPUs,', len(logical_gpus), 'Logical GPUs')
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)


@tf.function
def do_stuff(wmat, tf_var):
    with tf.device("/gpu:0"):
        S = tf.constant(0.0)
        for i in tf.range(4):
            fi = tf.cast(i, dtype=tf.float32)
            A = tf.math.lgamma(tf.tanh(tf.matmul(wmat + fi, tf.transpose(wmat - fi, [0, 2, 1]))))
            S += tf.reduce_sum(A)
        error = tf.reduce_mean(tf_var)
        return error, S

exp_uuid = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

n_batches = 512


def gen():
    for i in range(n_batches):
        with tf.device("/cpu:0"): # Make sure it comes from CPU
            r = tf.ones((400,800))
        yield r

option_names = ['map().batch()', 'map_and_batch()']
for option in range(2):

    with tf.device("/cpu:0"):
        dataset = tf.data.Dataset.from_generator(gen, output_types=tf.float32)

        def my_identity(x):
            with tf.device("/cpu:0"):
                print("my_identity input:", x, x.device)
                y = tf.identity(x)
                print("my_identity output:", y, y.device)
                return y

        if option == 0:
            ## Option 0: map().batch()
            dataset = dataset.map(my_identity).batch(16)

        elif option == 1:
            ## Option 1: deprecated map_and_batch()
            dataset = dataset.apply(tf.data.experimental.map_and_batch(my_identity, 16))

    gpu_transform = tf.data.experimental.prefetch_to_device('/gpu:0', buffer_size=4)
    dataset = dataset.apply(gpu_transform)


    tf_var = tf.Variable(tf.zeros(3))
    adam = tf.keras.optimizers.Adam(1e-4)
    logpath = os.path.join('data', 'logs', 'pa_' + exp_uuid + '_' + option_names[option])

    tf.profiler.experimental.start(logpath)
    start = datetime.datetime.now()
    for b, wmat in tqdm(enumerate(dataset)):
        with tf.GradientTape() as tape:

            if b == 0:
                print('\n\n dataset element device', wmat.device)
                print('\n')

            # Do some calculations
            result = do_stuff(wmat, tf_var)

        grads = tape.gradient(result[0], [tf_var])
        adam.apply_gradients(zip(grads, [tf_var]))
    stop = datetime.datetime.now()
    tf.profiler.experimental.stop()

    print(f'\n\nOption {option_names[option]}\n===========================\n')
    print(logpath)
    print('Time lapsed=', stop - start)
    print("\n\n")
```

**Other info / logs**

**Option 1**:
![image](https://user-images.githubusercontent.com/845012/147661299-a7f72017-00ff-47b6-bb71-8812bd5163d3.png)
Symptoms:
 - See the blocks `Iterator::FlapMap` and `Iterator::BatchV2` stacked on top of each other.
 - The MemcpyH2D (selected, see the details panel) is comping from pagable memory, instead of pinned memory (which is what MapAndBatch does). Because of the source being pagable memory, it can't overlap with kernel computations.
 
**Option 2**:
![image](https://user-images.githubusercontent.com/845012/147661558-9f316201-a4e5-4df1-a77b-032272537321.png)
Evidence:
 - The MapAndBatch block is used.
 - The MemcopyH2D comes from pinned memory (see details pane) and overlaps with kernel computations.

The whole deal about pinned memory is to allow parallel data upload and kernel computations. So the dataset needs to be produced into pinned host memory, which then can be uploaded asynchronously by the driver without an extra copy. See https://github.com/tensorflow/tensorflow/issues/43905#issuecomment-823675760 and https://github.com/tensorflow/tensorflow/issues/43905#issuecomment-824145184 and:
https://github.com/tensorflow/tensorflow/blob/40e9b534962989af7486bc6567ca472d71eb5049/tensorflow/core/kernels/data/experimental/map_and_batch_dataset_op.cc#L522

This is a follow up on https://github.com/tensorflow/tensorflow/issues/43905.