(shortfin.array) What if we had a context manager for briefly getting a host copy of a device array?

[renxida]

It might be nice if we have something that lets us do

with GetMapOfViewOnHost(device_array, index, discard=True) as m:
   m.fill(0)

instead of

seq_lens_host = seq_lens.for_transfer()
with seq_lens_host.map(discard=True) as m:
    m.fill(0)
    m.items = [len(req.input_token_ids) for req in self.exec_requests]
seq_lens_host.copy_to(seq_lens)

or, alternatively, for a lighter-weight easier-to-implement more flexible approach...

I wonder if something like

class GetArrayOnHost:
    def __enter__(self, array: shortfin.array.device_array):
        self.array = array
        self.host_array = array.for_transfer()
        return self.host_array

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.host_array.copy_to(self.array)

That enables

with GetArrayOnHost(device_array) as host_array:
    ... read / write array

Instead of

host_array = device_array.for_transfer()
... read / write array
host_array.copy_to(device_array)

would be nice.

[stellaraccident]

Don't think too hard about making what is there nicer. We need a new API for managed DMA. it's not just ergonomics. The current approach creates global locks on the memory system that slow everything down. Was just the easiest thing to make and is always need as a fallback.