Repository: tensorflow
Issue: tf.linalg.diag_part() is slow on TPU
Labels: stat:awaiting tensorflower,comp:ops,type:performance,TF 2.11

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Performance

### Have you reproduced the bug with TF nightly?

No

### Source

source

### Tensorflow Version

2.11.0

### Custom Code

Yes

### OS Platform and Distribution

Google colab

### Mobile device

_No response_

### Python version

3.9.16

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?


I am working on implementing efficient attention mechanisms using tensorflow. (such as longformer).

The local interaction in the longformer only requires tokens to interact with their neighboring ones, which could reduce the computation overhead calculating the full dot product between two tensors.

The problem I am currently experiencing is that tf.linalg.diag_part(), a function used for extracting the band diagonal from the key tensor, works extremely slow on TPU. 

In my test, the efficient attention takes 52sec to compute, whereas it takes 2sec to compute the regular dot product attention.

I tried warping it on @tf.function decorator jit_compile=True, but it is still slower than the full dot product calculation.

Using CPU, the efficient attention gets x4-5 performance boost depending on the settings.

There may be some limitations in tf.linalg.diag_part() in parallel use. The source code is written in C++, which I lack expertise in. 

Is there any way to enhance the performance of tf.linalg.diag_part() in parallel processing?


### Standalone code to reproduce the issue

```shell
import tensorflow as tf
from functools import partial

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))
strategy = tf.distribute.TPUStrategy(resolver)


def split_head(input, B, N, D, num_heads):
    assert D/num_heads == D//num_heads, "D must be divisible by num_heads"
   
    x = tf.reshape(input, 
               (B, N, num_heads, D//num_heads))
    return tf.transpose(x, perm = (0, 2, 1, 3))
       
B = 2048
N = 1024
D = 1024
H = 32
b_k = 15
input = tf.ones((B, N, D))

with strategy.scope():
      
 #tf.reshape(tf.range(B*N*D), (B,N,D))
  q = input
  k = input
  split_head = partial(split_head, B= B, N=N, D= D)
  q = split_head(q, num_heads= H)
  k = split_head(k, num_heads= H)
  print(B,N,D)
  k_T = tf.transpose(k, perm =(0, 1, 3, 2))
  print(k_T.shape)
  band_k = extract_band(input = k, b_k = b_k)
  print(band_k.shape)
  if N >= D//H:
      attention = q@tf.transpose(band_k, perm = (0, 1, 3, 2))
      print(attention.shape)
      attention = tf.reduce_mean(attention, axis = 1) # combine head
      attention_map = tf.linalg.diag(tf.transpose(attention, perm = (0, 2, 1) ), k = (-b_k//2, b_k//2))
```


### Relevant log output

_No response_</details>