[MLBuffer] Creation and representing MLBuffer on a XPU devices

[bbernhar]

Purpose/Motivation

Defines a device-based storage object that may be used by WebNN operations. This is a sub-issue of #482.

Proposed API

typedef unsigned long MLFlagsConstant;

[Exposed=(Window, DedicatedWorker)]
interface MLBuffer {
  readonly attribute MLFlagsConstant usage;
  readonly attribute MLOperandDescriptor descriptor;
  [CallWith=Isolate] void destroy();
}

[Exposed=(Window, DedicatedWorker), SecureContext]
namespace MLBufferUsage {
    // TBD
};

[Exposed=(Window, DedicatedWorker), SecureContext]
partial interface MLContext {
    Promise<MLBuffer> createBuffer(MLOperandDescriptor descriptor, MLBufferUsage usages);
};

Example JS

const ml_buffer = await mlContext.createBuffer(descriptor, usages);
ml_buffer.Destroy(); // invalid

• The buffer's allocation will be zeroed (as it is for WebGPU's createBuffer() method)
• Layout of MLBuffer is always known (and linear access is assumed).
• Destroy() gets called on the context timeline but doesn't actually release until the device signals completion.

Edits

• 5/14/24 - Removed "size" in favor or using MLOperandDescriptor
• 6/03/24 - Added usage flags and descriptor attributes
• 7/09/24 - createBuffer() now returns promise.

Alternative API proposals

N/A

Opens

1. Where will an MLBuffer's memory be allocated on systems where an MLContext may not be as closely tied to a given physical device as an IDMLDevice? See Need to understand how WebNN supports implementation that involves multiple devices and timelines #350 @a-sully
2. Must an MLBuffer only be used with an MLContext it was created from? @a-sully
3. Can an MLBuffer's size always be known at the time of buffer allocation? @a-sully
4. When will MLBuffer be deallocated if destroy() is not called? @a-sully
5. Does MLBuffer require explicit buffer usages (ex. input, output, or both)? @bbernhar
6. Does MLBuffer need to support being a staging buffer? @bbernhar
7. Is a zero sized MLBuffer allowed? @bbernhar

[a-sully]

My recent investigation into supporting MLBuffer on CoreML has lead me to the following two suggestions for createBuffer():

1. We need a WebGPU usage flag (at minimum)

The only zero-copy way to pass a buffer to both WebGPU (as an IOSurface) and CoreML (as an MLMultiArray) is to first allocate the buffer as an IOSurface containing "float16" data (IOSurface -> CVPixelBuffer -> MLMultiArray)

If the MLBuffer is to be used with WebGPU it must be allocated in this fashion (to be zero-copy, at least), whereas an MLBuffer which is only used within WebNN may be allocated as an MLMultiArray directly (more on that below)

2. MLBufferDescriptor should include an MLOperandDescriptor rather than an MLSize64

CoreML's inputs and outputs are given as MLMultiArrays, which require the data type and dimensions to be known. If we're to allocate a hardware buffer for createBuffer(), this information must be known.

Given that the dimensions + data type of input and output operands to an MLGraph are well-defined anyways, it seems reasonable to enforce that an MLBuffer must have matching constraints to be passed as an input or output to an MLGraph as #544 describes? Is there a reason why we should keep MLSize64?

[bbernhar]

Thanks @a-sully for delving into the CoreML side of things.

Regarding the need for a WebGPU usage flag:

Is it feasible for an MLBuffer to always be created as an MLMultiArray where, upon import to WebGPU, we could assign or request the usages? Assigning GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC appears to be sufficient.

As for the question about keeping MLSize64:

Without MLSize64, any ML framework that doesn't represent its tensor datatype like MLMultiArray would require re-architecting to avoid creating (especially output) tensors from raw allocations (or malloc). Alternatively, the web developer would need to defer calling createBuffer() until dispatch(), impacting the first inference-time. Could MLOperandDescriptor be made optional instead? The size could then be ignored where irrelevant.

[a-sully]

Is it feasible for an MLBuffer to always be created as an MLMultiArray where, upon import to WebGPU, we could assign or request the usages?

AFAICT an MLMultiArray cannot be handed off to WebGPU. It's a data type specific to CoreML. Importing to a type WebGPU can understand would always require a copy - even on UMA systems, which would be unfortunate!

Without MLSize64, any ML framework that doesn't represent its tensor datatype like MLMultiArray would require re-architecting to avoid creating (especially output) tensors from raw allocations (or malloc). Alternatively, the web developer would need to defer calling createBuffer() until dispatch(), impacting the first inference-time.

Hmm I'm not sure if I understand your concern.... Implementations are still welcome to allocate an MLBuffer as one contiguous block of memory. It's just that the WebNN "front end" would assert (when you call dispatch()) that the dtype and dimensions of the passed-in MLBuffer match what the graph expects, rather than just that the sizes are the same. Concretely, that maps to these checks in your prototype CL.

Is the use case you're referring to one where a single MLBuffer is assumed to be able to be contort into different dtype and dimensions? For example:

const mlBuffer = new MLBuffer({size:3*4*4});

// `graph1` expects a float32 output with shape [3, 4]
context.dispatch(graph1, inputs, {'out': mlBuffer});

// `graph2` expects a float16 input with shape [4, 3, 2]
context.dispatch(graph2, {'in': mlBuffer}, outputs);

This hides an internal reinterpretation of the data type and dimensions of what's assumed to be an opaque bag of bytes. I think there's a reasonable argument that this implementation detail should not make its way into the WebNN spec, which shouldn't prescribe a particular implementation.

WebNN has reshape and cast operators. In the example above, graph2 may use these operators to convert an input into whatever dtype and dimensions it needs, if it still wants to be able to use mlBuffer. An advantage of this approach is that the otherwise opaque reinterpretation of the buffer can be expressed in terms of other well-defined operators.

Could you elaborate on the use case(s) you have in mind?

Could MLOperandDescriptor be made optional instead? The size could then be ignored where irrelevant.

What would be the expected behavior on platforms which require a data type and dimensions when the buffer is allocated? An MLOperandDescriptor implies a size - but not the other way around.

[bbernhar]

AFAICT an MLMultiArray cannot be handed off to WebGPU.

I was expecting we start the allocation in CoreML via MLBuffer then import it as a MTLBuffer into WebGPU, using the GPUBuffer usages I mentioned.

Implementations are still welcome to allocate an MLBuffer as one contiguous block of memory

Consider a native C++ framework which implements a Tensor dtype as a bag of bytes. If you want to deploy this ML framework using WebNN JS API as an execution provider (or EP), it expects buffers will be allocated using a size. If we force createBuffer() to accept only a MLOperandDescriptor then this EP couldn't simply map Tensor allocation to createBuffer(). They would need to come up with a solution just for MLBuffer, preserving MLOperandDescriptor, or defer createBuffer(), which seems either burdensome or ineffective.

[a-sully]

AFAICT an MLMultiArray cannot be handed off to WebGPU.

I was expecting we start the allocation in CoreML via MLBuffer then import it as a MTLBuffer into WebGPU, using the GPUBuffer usages I mentioned.

Ah, I think my wording of "We need a WebGPU usage flag" above was misleading. I'm not suggesting that we need WebGPU usage flags here, but rather a usage flag saying "I want this MLBuffer to convertible to a GPUBuffer" (because the implementation may use that information to determine where/how the buffer should be allocated). Does that clear things up?

Could you also clarify what exactly you mean by "start the allocation in CoreML"? I assume you mean "as an MLMultiArray", but that would require the dtype and dimensions to be known, no?

Implementations are still welcome to allocate an MLBuffer as one contiguous block of memory

Consider a native C++ framework which implements a Tensor dtype as a bag of bytes. If you want to deploy this ML framework using WebNN JS API as an execution provider (or EP), it expects buffers will be allocated using a size. If we force createBuffer() to accept only a MLOperandDescriptor then this EP couldn't simply map Tensor allocation to createBuffer(). They would need to come up with a solution just for MLBuffer, preserving MLOperandDescriptor, or defer createBuffer(), which seems either burdensome or ineffective.

Thanks for the explanation. Could you provide a concrete example of where this concern is relevant? A quick glance at some common ML frameworks suggests that just size is often not sufficient to allocate a tensor. OnnxRuntime's JavaScript Tensor requires dtype and dimensions, for example. As does TFLite’s equivalent. Are there known examples where a size is available but not the dtype and dimensions? Presumably the MLBuffer is being allocated with use by some given MLGraph in mind, and the data types and dimensions of inputs and outputs must already be known? (input() and build() (for outputs) each require an MLOperandDescriptor)

Another consideration is that size may not be enough regardless of whether we want to replace size with an MLOperandDescriptor. As mentioned above, I expect we'll need usage flags, too. Does your concern still hold if arguments other than size become required?

[bbernhar]

Could you also clarify what exactly you mean by "start the allocation in CoreML"?

Could we pass a union to createBuffer() which specifies either the size or MLOperandDescriptor so MLBuffer could be always created as MLMultiArray? If not, another (possible) alt. solution is have createBuffer(size) defer creation of MLMultiArray until dispatch().

Are there known examples where a size is available but not the dtype and dimensions?

Yes, the ORT web tensor dtype can only be implemented behind a "malloc" like C inference. When WebNN is used as a EP, it exists within the ML runtime itself.

[reillyeon]

Yes, the ORT web tensor dtype can only be implemented behind a "malloc" like C inference.

I don't understand this comment because all of the Tensor constructors in that header take shape information. What am I missing?

[bbernhar]

I don't understand this comment because all of the Tensor constructors in that header take shape information. What am I missing?

Notice the Tensor constructor uses a IAllocator interface. That's the only way MLBuffer can be created from because it must own the buffer for the specified shape. Funny enough, the shape information is right there but the main point is ORT expects its possible to whip up a device buffer given only a size.

[a-sully]

Taking a step back, the Web Platform Design Principles implore us to "design based on user needs, not the underlying API or hardware":

This means newly proposed APIs should be designed with careful consideration on how they are intended to be used rather than how the underlying hardware, device, or native API available today.

The use cases for an MLBuffer - using some (hardware-optimized) buffer as an input or output to an ML graph - all require that the data type and dimensions of the buffer be known. We should not prescribe implementation details, such at that the buffer must be allocated contiguously, as this other design principle cautions:

Be particularly careful about exposing the exact lifecycle and data structures of the underlying native APIs. When possible, consider flexibility for new hardware.

The point about considering flexibility for new hardware is especially pertinent to WebNN :)

While I understand the desire to design a web platform API which (especially WASM) user-space frameworks can easily plug into, the web platform API should not bend over backwards to accommodate the implementation choices of any given framework. And the web platform API certainly should not bake in assumptions based on the current limitations of said frameworks! In this case, ORT does not support CoreML in cases where an MLMultiArray used as an output is not contiguously allocated. It seems likely that addressing that limitation would require changes to ORT which are ~the same as what would be needed to support MLBuffer if creating an MLBuffer required a dtype and dimensions?

[RafaelCintron]

[@a-sully wrote]

The only zero-copy way to pass a buffer to both WebGPU (as an IOSurface) and CoreML (as an MLMultiArray) is to first allocate the buffer as an IOSurface containing "float16" data (IOSurface -> CVPixelBuffer -> MLMultiArray)

For Apple platforms, my understanding is you can go from MLMultiArray -> MTLBuffer by calling getBytesWithHandler + newBufferWithBytesNoCopy. With an MTLBuffer you should be able to create a WebGPU buffer.

Why are IOSurfaces be required?

[huningxin]

another (possible) alt. solution is have createBuffer(size) defer creation of MLMultiArray until dispatch().

Seems doable. writeBuffer() may hold the BigBuffer with user data, then at dispatch(), create an MLMultiArray by initWithDataPointer?

[a-sully]

For Apple platforms, my understanding is you can go from MLMultiArray -> MTLBuffer by calling getBytesWithHandler + newBufferWithBytesNoCopy. With an MTLBuffer you should be able to create a WebGPU buffer.

Good question! I originally thought so too, but my current understanding is that this is not generically true (i.e. for all data types). If anyone can definitively confirm or dispute this understanding (@mwyrzykowski?) please speak up! Alright here goes...

The docs of newBufferWithBytesNoCopy say that it:

Creates a buffer that wraps an existing contiguous memory allocation

whereas the docs for getBytesWithHandler say of the buffer:

It may not store these scalar values contiguously

so I would assume that this would not be allowed (or at least not be zero-copy) unless the MLMultiArray was specifically allocated contiguously.

How can we ensure an MLMultiArray is allocated contiguously?

Of all the MLMultiArray constructors, the candidates for ensuring a contiguous memory allocation seem to be:

• init(dataPointer:shape:dataType:strides:deallocator:)
• init(pixelBuffer:shape:)

The first one looks promising! Unfortunately it seems - based on past offline discussions - that CoreML internally makes a copy of the bytes when using this constructor. That strides is a parameter seems to corroborate this.

So this would not be zero-copy:

another (possible) alt. solution is have createBuffer(size) defer creation of MLMultiArray until dispatch().

Seems doable. writeBuffer() may hold the BigBuffer with user data, then at dispatch(), create an MLMultiArray by initWithDataPointer?

The latter constructor takes a CVPixelBuffer, but this only works if the CVPixelBuffer is a "float16" IOSurface in disguise:

Use this initializer to create an IOSurface-backed MLMultiArray that reduces the inference latency by avoiding the buffer copy to and from some compute units.

The pixel buffer’s pixel format type must be kCVPixelFormatType_OneComponent16Half. The MLMultiArray data type is MLMultiArrayDataType.float16.

So eith regards to this question....

Why are IOSurfaces be required?

It seems that the only way to avoid copies of a backing memory which is to be shared as both an MLMultiArray and an MTLBuffer is to start with a float16 IOSurface. Unfortunately this suggests that zero-copy buffer sharing is only possible under certain dtype + "do we need to share with WebGPU" configurations. Of course, if we know the memory will stay within CoreML (i.e. it doesn't need to be shared with WebGPU) then we can allocate an MLMultiArray directly, though this would require dtype and shape to be known before writeBuffer()

Data Type	WebNN Use Only	WebGPU Interop
float16	✅ Zero copy (as MLMultiArray or IOSurface)	✅ Zero copy (as IOSurface)
float32	✅ Zero copy (as MLMultiArray)	🌕 Data copies (with initWithDataPointer)
float64	✅ Zero copy (as MLMultiArray)	🌕 Data copies (with initWithDataPointer)
int32	✅ Zero copy (as MLMultiArray)	🌕 Data copies (with initWithDataPointer)
other	❓ May be emulated as int32?	❓ Not sure

[mwyrzykowski]

For Apple platforms, my understanding is you can go from MLMultiArray -> MTLBuffer by calling getBytesWithHandler + newBufferWithBytesNoCopy. With an MTLBuffer you should be able to create a WebGPU buffer.

Good question! I originally thought so too, but my current understanding is that this is not generically true (i.e. for all data types). If anyone can definitively confirm or dispute this understanding (@mwyrzykowski?) please speak up! Alright here goes...

It is zero copy in CoreML but anything other than fp16 + CVPixelBuffer will result in a copy below CoreML