[MLBuffer] Uploading/downloading tensor data

[bbernhar]

Purpose/Motivation

Provides a means to upload or download tensor data to/from a created MLBuffer via #542 . This is a sub-issue of #482.

Proposed API

Example JS

ml_buffer = ml_context.createBuffer(...);
ml_context.writeBuffer(ml_buffer, 0, srcData, ...);
/* use ml_context */
await data = ml_context.readBuffer(ml_buffer, dstData, ...);

• Execute on the device timeline in the same order they were enqueued on the context timeline.
• A copy of srcData is always made and returns control back to the web developer immediately.
• readBuffer can avoid an additional allocation and copy to WASM heap by specifying dstData, which is transferable.
• Destroying a buffer must reject pending promises returned by readBuffer().

Opens

• If MLBuffer are mappable, should they support mappedAtCreation? @a-sully
• What happens if the array buffer gets de-attached then used by writeBuffer or are "zero-writes" allowed? @bbernhar
• Should [write|read]Buffer arguments use a dictionary per https://w3ctag.github.io/design-principles/#prefer-dictionaries (note: WebGPU does not follow this).

Edits

• 7/23 - removed alt. proposal and added clarification on destruction during pending reads.

[bbernhar]

@anssiko FYI, this issue should be tagged as non-interop

[fdwr]

are "zero-writes" allowed

While pondering that question, please consider the following issue too #391, where accepting 0-byte size transfers as silent nops probably makes life easier for callers (they can share more of the same generic code paths with less special case handling).

[reillyeon]

The current note on dstData is that if it is transferable then we can avoid a copy to the WASM heap, however as I understand it WASM heap memory is not transferable so this zero-copy option is not available as proposed. However I believe we can actually avoid the transferability requirement by specifying that dstData is written to as a task on the ML task source. This works because it ensures that script can't observe a state in which dstData has been partially written to.

[bbernhar]

Good eye. I think that "no copy" note was meant for JavaScript, not WASM. To allow readBuffer() be functionally equivalent to how compute() works today (ex. Promise<ArrayBufferView> readBuffer(MLBuffer srcBuffer, ArrayBufferView dstData);).

Do you think "no copy" for JS is still useful (no WASM)?

However I believe we can actually avoid the transferability requirement by specifying that dstData is written to as a task on the ML task source.

How would a task source be used to grant readBuffer() access to safely write into the WASM heap?

[reillyeon]

How would a task source be used to grant readBuffer() access to safely write into the WASM heap?

It's important to consider why the design required the buffers to be transferred to begin with: We want to prevent concurrent access to the buffer from "the implementation" and JS/WASM. For output buffers, since compute() is asynchronous, we don't want both script and the underlying compute engine to be able to make changes to the ArrayBuffer at the same time, so we transfer it away so that only the latter has access. For readBuffer() there's a simpler solution: We can simply specify that writing into the buffer provided by script happens atomically and while script is not running (i.e. all at once in a single task).

Edited to add: The steps become basically "in parallel, gather the bytes that are contained in the MLBuffer and then queue a task to copy those bytes into the provided ArrayBuffer". While compute() could theoretically be zero-copy readBuffer() implies at least one copy and so we can specify when it happens in a way that benefits us by avoiding the transferability requirement.

[huningxin]

@reillyeon

The steps become basically "in parallel, gather the bytes that are contained in the MLBuffer and then queue a task to copy those bytes into the provided ArrayBuffer".

Should we consider whether the buffer is SharedArrayBuffer that a worker may see the partially written bytes? AFAIK, ORT Web multi-threaded build uses SharedArrayBuffer as Wasm heap.

[reillyeon]

Should we consider whether the buffer is SharedArrayBuffer that a worker may see the partially written bytes? AFAIK, ORT Web multi-threaded build uses SharedArrayBuffer as Wasm heap.

That is the behavior a developer opts into by using SharedArrayBuffer, though again since the write will happen on one of the worker threads with access to the buffer it's identical to script writing to the buffer itself.

[huningxin]

"in parallel, gather the bytes that are contained in the MLBuffer and then queue a task to copy those bytes into the provided ArrayBuffer".

Would the "copy those bytes into the provided ArrayBuffer" and "worker writes to ArrayBuffer" happen at the same time, if the dst is a SharedArrayBuffer?

[reillyeon]

Would the "copy those bytes into the provided ArrayBuffer" and "worker writes to ArrayBuffer" happen at the same time, if the dst is a SharedArrayBuffer?

For a SharedArrayBuffer the spec could allow implementations to write into the buffer in parallel with the worker that made the readBuffer() call accessing the buffer, however in Chromium it would be difficult to implement that behavior so we would do the copy on the worker thread instead. This would be perfectly spec compliant. Workers other than the one that made the request could access the buffer in parallel to the copy, as that is allowed for a SharedArrayBuffer.

[fdwr]

(naming) Given it's been renamed to MLTensor rather than MLBuffer (per a comment in Austin's explainer from Corentin), renaming the read/write functions accordingly would make sense.

await mlTensor = mlContext.createTensor(...);
mlContext.writeTensor(mlTensor, sourceData);
/* use mlContext */
await data = mlContext.readTensor(mlTensor, destinationData);

Otherwise it's a little confusing because writeBuffer doesn't actually write into the passed thing called a buffer (the ArrayBuffer), but rather it reads from it.

- void writeBuffer(MLTensor destinationTensor, ArrayBuffer sourceData);
+ void writeTensor(MLTensor destinationTensor, ArrayBuffer sourceData);

[inexorabletash]

I believe we can close this out, now that the spec is updated with MLTensor. Please re-open if i'm mistaken.

[fdwr]

I believe we can close this out, now that the spec is updated with MLTensor. Please re-open if i'm mistaken.

I was wondering that too yesterday, thinking everything was complete 🤔. @bbernhar to confirm? Should any open issues in your "Opens" section be migrated to a separate issue?

[bbernhar]

@inexorabletash @fdwr Yes, thanks for checking. No opens need follow-up atm.