TF >= 2.7 slowdown tf.data API
Hi, @chnzhangrui!  I tried this provided code in TF v[2.8](https://colab.research.google.com/gist/sushreebarsa/12b8d058b114457a93847bdccc7ea87c/60822-2-12.ipynb#scrollTo=tuajrZc2nTS1),  TF v[2.12](https://colab.research.google.com/gist/sushreebarsa/c1ad2acaa810562a063e26bce51a2a49/60822-2-12.ipynb) and tf-[nightly](https://colab.research.google.com/gist/sushreebarsa/10f739e75d7dc06531aca5155e40fcef/untitled793.ipynb#scrollTo=wTUkrK8AnQ4D) using colab. The issue is not replicating as described here. Could you please have a look at these gists and confirm the same. Thank you! Hi @sushreebarsa, thanks for the quick response. I checked your gists and I can also confirm your (fast) result. However, it is not clear to me if it was running on GPU or CPU for you. The problem persist only if I run on GPU and I don't know if I can connect to a GPU for test on colab. Let's focus on v2.8. Hi @chnzhangrui I have tested the issue on colab using GPU runtime with TF v2.8 as well which is also not replicating the issue reported here. Please find the gist [here](https://colab.research.google.com/gist/sushreebarsa/661982f1752ab6fe6f1b4c78bd35fabf/60822-2-12.ipynb) for reference. Thank you! Hi @sushreebarsa, thanks for the confirmation. I am trying to confirm this with my system again and it currently encountered an issue. Will update once I have some idea. I am a little concerned that the fast result you see is really from GPU or just CPU, because this code runs very fast on CPU. I may need to come up with a code that doesn't work on CPU. @chnzhangrui Thank you for your response. Please let us know if there is any new update?
Thank you! Sorry for the silence. I am still waiting for response from our IT support. @chnzhangrui Is there any update on this issue?
 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60822">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60822">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Performance

### Have you reproduced the bug with TF nightly?

Yes

### Source

installed from pip

### Tensorflow Version

tf >= 2.7

### Custom Code

No

### OS Platform and Distribution

Linux, CentOS Linux 7

### Mobile device

No

### Python version

3.8, 3.9

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

Since TF >= 2.7, I observed a significant slowness when running `ds_iter.get_next()` where `ds_iter` is an iterator of tf.data.Dataset defined as `dataset.shuffle(...).cache().repeat(...).batch(..., drop_remainder=True).prefetch(tf.data.AUTOTUNE)`. The same code runs much faster on the same machine using TF <= 2.6. A reproducer is included below.

Looking at the [change log from TF2.8](https://www.exxactcorp.com/blog/Deep-Learning/tensorflow-2-8-0-released), I see the following related to Dataset API `(since v2.7) Stateful ops used in tf.data.Dataset`. However I am not entirely sure if there is a workaround in higher TF version to reach the same speed.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import time

X = tf.random.uniform(shape=[21000, 373])
Labels = tf.random.uniform(shape=[21000, 1])
batchsize = 8192
n_shuffles = 500000
n_samples = 21000
ds = tf.data.Dataset.from_tensor_slices((X, Labels))
ds = ds.shuffle(buffer_size=n_samples).cache().repeat(n_shuffles).batch(batchsize, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
ds_iter = iter(ds)

num_iter = 100
loop_start = time.time()
for i in range(num_iter):
    ds_iter.get_next()
loop_stop  = time.time()
print('loop time', f'{loop_stop-loop_start : .4f}')
```


### Relevant log output

```shell
output of the reproducer:
- with TF >= 2.7: loop time  15.8203
- with TF <= 2.6: loop time  1.0519

device to run the above two tests:
NVIDIA A100-PCIE-40GB NVIDIA-SMI 530.30.02 Driver Version: 530.30.02 CUDA Version: 12.1

$$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Jan_28_19:32:09_PST_2021
Cuda compilation tools, release 11.2, V11.2.142
Build cuda_11.2.r11.2/compiler.29558016_0
```
</details>