Committed arrays unable to be donated.

[aplund]

I'm very confused as to how array donation is working in jit compiled functions. It seems that, even though buffers can be donated, they can't be returned when execution actually occurs. I've tried to make a script that shows what I mean:

from functools import partial
import mmap

import psutil

import numpy as np
import jax

process = psutil.Process()

def dump_mappings(s):
    print("Large mappings:")
    for x in process.memory_maps(grouped=False):
        if x.size >= s:
            print(x)

# Jax jit compiled function
# Donate input buffer to output
@partial(jax.jit, donate_argnums=1)
def doit(arr, out):
    return out

print(f"Virtual memory size (start): {process.memory_info().vms:_d}")

# Set up data buffer
sz = 1024*16
# (I am unreliably informed that mmap will be page boundary aligned.)
buf_sz = np.dtype('complex64').itemsize*sz**2
print(f"numpy_array byte size {buf_sz:_d}")
data = mmap.mmap(-1, buf_sz)
numpy_array = np.frombuffer(data, dtype=np.complex64).reshape((sz,)*2)
numpy_array[0,0] = np.nan + 1j*np.nan
numpy_array[0,1] = np.inf
numpy_array[1,0] = complex(0.,np.inf)
orig_pointer = numpy_array.ctypes.data
print(f"numpy_array location {hex(orig_pointer)} ")

print(f"Virtual memory size (numpy array): {process.memory_info().vms:_d}")
dump_mappings(buf_sz)

cpu_device = jax.devices('cpu')[0]
print(f"Commiting array to {cpu_device}")
big_array = jax.dlpack.from_dlpack(numpy_array, device=cpu_device, copy=False)
put_pointer = big_array.unsafe_buffer_pointer()

assert put_pointer == orig_pointer

print(f"big_array2 location {hex(put_pointer)}")
print(big_array.at[:2,:2].get())

print(f"Virtual memory size (commited array): {process.memory_info().vms:_d}")
dump_mappings(buf_sz)

big_array_old = big_array
big_array = doit(10.-10j, big_array)

print(f"Is the input buffer deleted? {big_array_old.is_deleted()}")

new_pointer = big_array.unsafe_buffer_pointer()
print(f" out array location {hex(new_pointer)}")
print(big_array[:2,:2])

print(f"Virtual memory size (function run): {process.memory_info().vms:_d}")
dump_mappings(buf_sz)

assert new_pointer == orig_pointer

and a typical output I get is:

Virtual memory size (start): 835_125_248
numpy_array byte size 2_147_483_648
numpy_array location 0x7f80d98bd000
Virtual memory size (numpy array): 2_982_608_896
Large mappings:
pmmap_ext(addr='7f80d98bd000-7f81598bd000', perms='rw-s', path='/dev/zero', rss=8192, size=2147483648, pss=8192, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=8192, referenced=8192, anonymous=0, swap=0)
Commiting array to TFRT_CPU_0
big_array2 location 0x7f80d98bd000
[[nan+nanj inf +0.j]
 [ 0.+infj  0. +0.j]]
Virtual memory size (commited array): 5_818_396_672
Large mappings:
pmmap_ext(addr='7f80d98bd000-7f81598bd000', perms='rw-s', path='/dev/zero', rss=8192, size=2147483648, pss=8192, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=8192, referenced=8192, anonymous=0, swap=0)
Is the input buffer deleted? True
 out array location 0x7f7faffff040
[[nan+nanj inf +0.j]
 [ 0.+infj  0. +0.j]]
Virtual memory size (function run): 7_967_866_880
Large mappings:
pmmap_ext(addr='7f7faffff000-7f8030021000', perms='rw-p', path='[anon]', rss=2147504128, size=2147622912, pss=2147504128, shared_clean=0, shared_dirty=0, private_clean=0, private_dirty=2147504128, referenced=2147504128, anonymous=2147504128, swap=0)
pmmap_ext(addr='7f80d98bd000-7f81598bd000', perms='rw-s', path='/dev/zero', rss=2147483648, size=2147483648, pss=2147483648, shared_clean=0, shared_dirty=0, private_clean=2147475456, private_dirty=8192, referenced=2147483648, anonymous=0, swap=0)
Traceback (most recent call last):
  File "/home/lund/src/Upstream/xanadu/MrKite/jaxtest2.py", line 70, in <module>
    assert new_pointer == orig_pointer
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

The returned buffer from the trivial function is not the donated one from the input, but is copied into a new buffer. This is probably ok for small buffers, but how is one meant to deal with this situation for large buffers?