PR #10803: Expose CompiledMemoryStats on CPU and GPU
PR #10803: Expose CompiledMemoryStats on CPU and GPU

Imported from GitHub PR https://github.com/openxla/xla/pull/10803

[CompiledMemoryStats](https://github.com/openxla/xla/blob/01073cf10b57c64615483545bfb34fb9ae0d3c06/xla/pjrt/pjrt_executable.h#L269-L269) is supposed to summarize the amount of memory required to run a given executable. Buffer-related fields in CompiledMemoryStats are currently missing for both the GPU and the CPU backend. This PR computes the missing fields using the executable's underlying [BufferAssignment](https://github.com/openxla/xla/blob/50c6489cb058881cc65622605c9c55029abebc5b/xla/service/buffer_assignment.h#L491) / buffer allocations.

# Example
```python
import jax
import jax.numpy as jnp

def f(a, b):
    ms = jnp.sum(a, axis=1, keepdims=True) / a.shape[1]
    return a * jax.lax.rsqrt(ms) + b

a0 = jax.random.uniform(jax.random.PRNGKey(0), (128, 128))
b0 = jax.random.uniform(jax.random.PRNGKey(0), (128, 128))
f_jit = jax.jit(f, donate_argnums=(0,))

print(f_jit.lower(a0, b0).compile().memory_analysis())
# CompiledMemoryStats(
#     generated_code_size_in_bytes=5792,
#     argument_size_in_bytes=131072,
#     output_size_in_bytes=65536,
#     alias_size_in_bytes=65536,
#     temp_size_in_bytes=512,
#     host_generated_code_size_in_bytes=0,
#     host_argument_size_in_bytes=0,
#     host_output_size_in_bytes=0,
#     host_alias_size_in_bytes=0,
#     host_temp_size_in_bytes=0)
```
Note the use of `donate_argnums=(0,)` which allows XLA to modify the first input buffer in-place and reuse it as the output. This is reflected in `alias_size_in_bytes` -- without buffer donation that number is zero.

# Open questions
- Does this match the semantics for the TPU backend?
- `dynamic_cast`-ing on `gpu::Executable` ain't great. Do we want `Executable`s to expose their buffer allocations publicly (i.e. GpuExecutable's [GetAllocations](https://github.com/openxla/xla/blob/f6ce6c11ceac904c9c4840fcec5f1ee715291425/xla/service/gpu/gpu_executable.h#L164) but for all Executables)?
Copybara import of the project:

--
f9ae2c4e367aed813af74ae50edfd27a6f9304eb by Georg Stefan Schmid <gschmid@nvidia.com>:

Expose CompiledMemoryStats on CPU and GPU

Merging this change closes #10803

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/10803 from gspschmid:georg/compiledmemstats-gpu f9ae2c4e367aed813af74ae50edfd27a6f9304eb
