Repository: tensorflow
Issue: SelfAdjointEigV2 GPU operation takes a lot of temporary memory.
Labels: stat:awaiting tensorflower,type:bug,comp:gpu,comp:core,TF 2.9

**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
- TensorFlow installed from (source or binary): source
- TensorFlow version (use command below): 2.9.0
- Python version: 3.8.10
- CUDA/cuDNN version: 11.5
- GPU model and memory: GTX 1660 Ti

**Describe the current behavior**
A single call to `tf.linalg.eigh()` takes a linear amount of memory in batch size, despite being processed matrix by matrix. I think the ScratchSpace is only freed in the end of the call after ALL matrices in the batch are processed, instead of on the fly matrix, per matrix. This I conclude from the enormous amount of allocations reported by the allocator during the OOM, and looking at the code.

**[Contributing](https://www.tensorflow.org/community/contribute)**

- Do you want to contribute a PR? (yes/no): maybe
- Briefly describe your candidate solution(if contributing):
  - **Best solution:** reuse the ScratchSpace for every matrix.  
  - **Next best solution:** free the ScratchSpace after every matrix in the batch.

**Standalone code to reproduce the issue**
```py
import tensorflow as tf
tensor = tf.random.uniform((1024 * 256, 4, 4)) # just make sure the batch size is big enough
tensor = tf.matmul(tensor, tensor, transpose_b=True)
t = tf.linalg.eigvalsh(tensor)
```

**Other info / logs**
See below the summary of the allocator:
```
2022-01-03 16:18:32.407928: W tensorflow/core/common_runtime/bfc_allocator.cc:462] Allocator (GPU_0_bfc) ran out of memory trying to allocate 215.4KiB (rounded to 220672)requested by op SelfAdjointEigV2

2022-01-03 16:18:32.518897: I tensorflow/core/common_runtime/bfc_allocator.cc:1071]      Summary of in-use Chunks by size: 
2022-01-03 16:18:32.518906: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 256 totalling 256B
2022-01-03 16:18:32.518916: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 1280 totalling 1.2KiB
2022-01-03 16:18:32.518924: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 65536 totalling 64.0KiB
2022-01-03 16:18:32.518930: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 11841 Chunks of size 220672 totalling 2.43GiB
2022-01-03 16:18:32.518937: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 251648 totalling 245.8KiB
2022-01-03 16:18:32.518944: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 410880 totalling 401.2KiB
2022-01-03 16:18:32.518951: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 4194304 totalling 4.00MiB
2022-01-03 16:18:32.518958: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 2 Chunks of size 268435456 totalling 512.00MiB
2022-01-03 16:18:32.518964: I tensorflow/core/common_runtime/bfc_allocator.cc:1078] Sum Total of in-use chunks: 2.94GiB
2022-01-03 16:18:32.518970: I tensorflow/core/common_runtime/bfc_allocator.cc:1080] total_region_allocated_bytes_: 3154771968 memory_limit_: 3154771968 available bytes: 0 curr_region_allocation_bytes_: 6309543936
2022-01-03 16:18:32.518982: I tensorflow/core/common_runtime/bfc_allocator.cc:1086] Stats: 
Limit:                      3154771968
InUse:                      3154771968
MaxInUse:                   3154771968
NumAllocs:                       11850
MaxAllocSize:                268435456
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0

2022-01-03 16:18:32.519244: W tensorflow/core/common_runtime/bfc_allocator.cc:474] ****************************************************************************************************
2022-01-03 16:18:32.519294: F ./tensorflow/core/util/gpu_solvers.h:533] Non-OK-status: context->allocate_temp(DataTypeToEnum<Scalar>::value, shape, &scratch_tensor_, alloc_attr) status: RESOURCE_EXHAUSTED: OOM when allocating tensor with shape[55137] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
```

Especially this line:  
```
...r.cc:1074] 11841 Chunks of size 220672 totalling 2.43GiB
```
So here, I conclude that it crashes when it's trying to process matrix number 11842 out of 16384.

Pointers:
 - ScratchSpace request in the `HeevdImpl` call (at line 635): https://github.com/tensorflow/tensorflow/blob/322cba072c9313689aeb2fc3174f18fce57194d7/tensorflow/core/util/cuda_solvers.cc#L620-L643
 - Batch processing matrix by matrix: https://github.com/tensorflow/tensorflow/blob/322cba072c9313689aeb2fc3174f18fce57194d7/tensorflow/core/kernels/linalg/self_adjoint_eig_v2_op_gpu.cc#L130-L140