Memory leak when jit compiling
@andremfreitas The warnings you're seeing indicate that your two jit-compiled functions, likely containing some floating-point division (`__cuda_sm20_div_rn_f64_full)`, are running out of available registers on your GPU. Could you please provide a standalone code to replicate the issue reported here?
Thank you!
 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/66434">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/66434">No</a>
 I am having the same or a similar issue. I am also using a single GPU, and also not turning on jit compilation in my `train_step`, and with `model.compile(jit_compile=False)`

`
I0000 00:00:1717163798.560575  211396 asm_compiler.cc:369] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_112', 8 bytes spill stores, 8 bytes spill loads

I0000 00:00:1717163798.631148  211415 asm_compiler.cc:369] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_88', 4 bytes spill stores, 4 bytes spill loads

I0000 00:00:1717163798.845685  211466 asm_compiler.cc:369] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_73', 104 bytes spill stores, 96 bytes spill loads

I0000 00:00:1717163798.845731  211485 asm_compiler.cc:369] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_5', 72 bytes spill stores, 72 bytes spill loads

I0000 00:00:1717163798.944683  211412 asm_compiler.cc:369] ptxas warning : Registers are spilled to local memory in function 'triton_gemm_dot_3', 16 bytes spill stores, 16 bytes spill loads
`

I am confused because reducing batch size to 1 doesn't change anything and neither does reducing the number of my largest model blocks. Do you have any idea what this could be due to ? 


### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

2.16

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04.3 LTS

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I have a function that two functions that are jit compiled. I call them in my training loop function setp (which is not jit compiled, simply in graph mode). With these two functions jit compiled I get a registers spilled warning as follows:

`2024-04-24 12:58:52.550355: I tensorflow/stream_executor/gpu/asm_compiler.cc:323] ptxas warning : Registers are spilled to local memory in function '__cuda_sm20_div_rn_f64_full', 8 bytes spill stores, 8 bytes spill loads
ptxas warning : Registers are spilled to local memory in function '__cuda_sm20_div_rn_f64_full', 8 bytes spill stores, 8 bytes spill loads`

This makes the training much slower.
Without jit compiling those two functions I don't get this warning.

My problem is similar to the one in [this forum thread](https://discuss.tensorflow.org/t/xla-jit-compile-causing-memory-leak/22657).  In this thread some user redirects to this open issue [`#56423`](https://github.com/tensorflow/tensorflow/issues/56423). However this one relates this problem to the distribution strategy. In my case I am just using one GPU. In this issue someone, mentioned that if they performed the optimizer step outside of the jit function it worked. In my case, as I said above, my training step function is not jit compiled and so the optimizer step is already not inside a jit compiled function.



### Standalone code to reproduce the issue

```shell
Unfortunately, I cannot build a MWE at the moment.
```


### Relevant log output

_No response_