Repository navigation
Support for device-based tensor storage objects #482
Description
Activity
@wchao1115 Sure.
WebNN (as spec'd) and WebGPU lack a way of sharing tensor data on-device directly with each other:
GPUBufferis inaccessible to WebNN and WebGPU does not support NPU buffers. I believe a sharable NPU or GPU buffer type that WebGPU could use is the only path forward for us to support WebGPU interop for WebNN (ex. custom ops). The other problem is chained inferences: WebNN has no means to re-use existing GPU or NPU results betweencompute()calls without copying everything back to the CPU.- Give WebNN developer control of device-storage to avoid round-trips to/from CPU.
I think this feature would be critical for some language models' performance on device (GPU/NPU) where the outputs of the previous inference, e.g., hidden state or KV pairs, will be used as inputs of the next inference. For such use case, frameworks usually allow to allocate on-device tensors and use these tensors as inputs/outputs for model inference, for example ONNXRuntime I/O Binding.
Reacted by Dwayne RobinsonSince an
MLBufferis created from anMLContext, I infer from the proposal that the buffer is always bound to the context. Is that correct?If so, the read/write operations could be simplified by moving them onto the interface rather than needing to pass the buffer.
Nitpick on proposed API shape:
MLBuffer? createBuffer(...)impliesnullcan be returned; exceptions should be preferred for synchronous error cases.Reacted by Zoltan KisThanks @inexorabletash for the feedback.
Read/write ops must occur in the domain of the context because only the context determines device-execution order, not
MLBuffer. We could consider support for direct mappedMLBufferwhich supports that line of thinking.If so, the read/write operations could be simplified by moving them onto the interface rather than needing to pass the buffer.
My thoughts, too, but MLContext is indeed the main encapsulator here, and IIUC
MLBufferis a control object used for identification and indirect control of the underlying opaque data. So passing anMLBufferas argument does not necessarily involve any copying, just recording what to do with the given buffer. It is the explicit methods that control the content of the buffer.Ah, thank you for clarifying. The timeline mentions are subtle, I hadn't internalized that yet. One more thing to specify concretely in the spec. :)
FWIW, it's best to make API proposals by leading with examples of the JS code using the proposed API, and only worry about providing the IDL later.
WIW, it's best to make API proposals by leading with examples of the JS code using the proposed API, and only worry about providing the IDL later.
Unfortunately, no code example tells you which timeline gets used where, only the WebNN spec can describe this behavior: which state is available to which operations. The WebNN programming model should probably concretely define these "timelines" then describe the entire API using it.
Reacted by Joshua BellHi @bbernhar, thanks for this proposal (and the Chromium prototype)! I agree that specifying a clear WebNN <-> WebGPU interop is needed. I have a comment and a few of questions for ya
WebNN (as spec'd) and WebGPU lack a way of sharing tensor data on-device directly with each other:
GPUBufferis inaccessible to WebNN and WebGPU does not support NPU buffers. I believe a sharable NPU or GPU buffer type that WebGPU could use is the only path forward for us to support WebGPU interop for WebNN (ex. custom ops). The other problem is chained inferences: WebNN has no means to re-use existing GPU or NPU results betweencompute()calls without copying everything back to the CPU.This proposal includes a clear way to read back to JS/CPU using
readBuffer()andwriteBuffer(), but doesn't mention how data is expected to be passed to/from WebGPU. Could you elaborate on the plans to interopMLBufferwith WebGPU?In particular, I'm curious about:
- How do we plan to prevent WebNN and WebGPU from stomping over shared resources, as you had discussed here on an earlier proposal? CommandBuffer usage clarification: internal, external, both? #264 (comment)
- Are there plans to support other types of
MLBuffercreation? e.g. converting a WebGPU buffer to anMLBufferor importing aGPUExternalTexture(perhaps as a not-yet-createdGPUExternalBuffer) are the first things which come to mind. Meanwhile, we should consider how this might interact with e.g. the Wasm Memory Control proposal (@dtig FYI)
As @inexorabletash mentioned, code snippets showing WebGPU <-> WebNN interop would be extremely helpful here :)
a sharable NPU or GPU buffer type that WebGPU could use
Are we expecting to provide any guarantees about where this buffer resides? Would an
MLBufferlive on CPU if created from anMLContextwith a CPU backend, for example?We need to also take into consideration other platforms where ML execution is not so closely tied to a single "device" as DirectML is (e.g. Core ML on Mac - this is related to discussions in #322 around
MLContextcreation). I assume we're not expecting to promise zero-copy buffer transfer between WebNN and WebGPU in all scenarios, right?For synchronous compute, use the read-back functions for window and workers, async and sync, respectively.
Just a heads up that with JSPI coming soon, I would expect pushback on adding this sync interface, even in a worker :)
Some questions:
How to transfer MLBuffer between devices?
I think
read/writeBuffermakes sense for XPU <-> CPU transfer.Perhaps a bit future looking, how do we support GPU <-> NPU/GPU transfers? (e.g. GPU/NPU cooperation, iGPU <-> dGPU, multi-GPU)
From the current design,looks like developers need to:
- Read GPUBuffer to CPU, block until completion
- Write the buffer on CPU to NPUBuffer, block until completion
- Use NPUBuffer
Is there a faster path (or do we anticipate one) for inter-device transfer? Can we use Intel GPU <-> NPU transfer as an example?
Simplified types in compute()
Should we change
compute()anddispatch()to accept only MLBuffer (i.e. drop TypedArray and ArrayBuffers)?MLBuffer usage scope
Is MLBuffer only permitted for binding input/output buffers to a built-graph during compute?
Can MLBuffer be used where a MLOperand is accepted, like in
conv2d(inputNode, /*filters=*/ mlBuffer)?read/writeBuffer definition
Should these be defined on the MLBuffer themselves? Looks like read/write operations is dependent on the context the MLBuffer is associated with.
Defining read/write on MLBuffer removes the need to check
MLBuffer.context == thisContextin MLContext.read/writeBuffer.MLBuffer memory management
When is MLBuffer's memory allocated on device? Is it during writeBuffer?
Should MLBuffer.destroy() returns a Promise to tell the caller that the memory has been deallocated?
I also wonder if the continuous memory model is too simplified. What is different device use different channel ordering or endianness?
Are we expecting developers to perform said conversion on CPU manually?
Thanks, @a-sully for raising these questions.
How do we plan to prevent WebNN and WebGPU from stomping over shared resources
If we allow
MLBufferto be sharable or GPU transferable; unlikeGPUBuffer, it can synchronize itself when used in WebGPU operations.Then I think adding a new WebGPU API,
GPUDevice.importExternalBuffercould convert aMLBufferdirectly toGPUBuffer, leaving theMLBuffer"detached" or as if it was destroyed. To restore access [to WebNN], one could re-import it usingMLContext.importExternalBufferor disposeGPUBuffer(exact names are TBD).That code could look like this:
// Create sharable buffer in WebNN ml_context = ML.createContext(wgpuDevice); ml_buffer = ml_context.createBuffer({size:size, forExport:true}); // Import buffer to WebGPU gpu_buffer = wgpuDevice.importExternalBuffer(ml_buffer); pipeline = wgpuDevice.createComputePipeline(/* pipeline with compute shader that updates gpu_buffer */); bind_group = wgpuDevice.createBindGroup(/* create bind group for gpu_buffer */); command_encoder = wgpuDevice.createCommandEncoder(); pass = command_encoder.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(/* buffer index in shader */, bind_group); pass.dispatchWorkgroups(/* sizes */); pass.end(); wgpuDevice.queue.submit([command_encoder.finish()]); // Export buffer from WebGPU ml_buffer = ml_context.importExternalBuffer(gpu_buffer)
Are we expecting to provide any guarantees about where this buffer resides?
Yes, it will reside on the same device used to create the MLContext. If it is a CPU-only ML context, then WebNN should create a CPU backed MLBuffer.
I assume we're not expecting to promise zero-copy buffer transfer between WebNN and WebGPU in all scenarios, right?
Right. If you are on the same GPU, the result allows zero-copy. Otherwise, a GPU copy is usually required for NPU-to-GPU or iGPU-to-dGPU or video-to-tensor conversions.
Just a heads up that with JSPI coming soon, I would expect pushback on adding this sync interface, even in a worker :)
Thanks for the heads up.
Reacted by Zoltan Kis and mwyrzykowski@wacky6 Great questions, my thoughts below.
Can we use Intel GPU <-> NPU transfer as an example?
The only true "no copy" path I'm aware of is CPU/iGPU. I believe all other scenarios require GPU/NPU copy.
Should we change compute() and dispatch() to accept only MLBuffer (i.e. drop TypedArray and ArrayBuffers)?
I am also in favor of using
MLBuffereverywhere and only havingdispatch(). However, I'm told WebNN developers may prefer to stick with ArrayBuffers or TypedArray since usingMLBuffereverywhere creates a inconvenience.Is MLBuffer only permitted for binding input/output buffers to a built-graph during compute?
Currently, yes. In the future,
MLBuffercould also be permitted for interop API. See #482 (comment).Should these be defined on the
MLBufferthemselves?Perhaps the earlier response addresses this? #482 (comment).
When is MLBuffer's memory allocated on device? Is it during writeBuffer?
No, it would be on buffer creation. This avoids generating a fatal OOM where the WebNN developer wouldn't expect.
Should MLBuffer.destroy() returns a Promise to tell the caller that the memory has been deallocated?
MLBuffer.destroy()would not necessarily guarantee de-allocation since its preferable to re-use memory.continuous memory model is too simplified
I will follow up with our Intel NPU teams if they plan to introduce complex formats.
@wacky6 and @a-sully , thank you for your feedback.
wacky6 wrote:
How to transfer MLBuffer between devices?
I think read/writeBuffer makes sense for XPU <-> CPU transfer.
Perhaps a bit future looking, how do we support GPU <-> NPU/GPU transfers? (e.g. GPU/NPU cooperation, iGPU <-> dGPU, multi-GPU)From the current design,looks like developers need to:
- Read GPUBuffer to CPU, block until completion
- Write the buffer on CPU to NPUBuffer, block until completion
- Use NPUBuffer
Is there a faster path (or do we anticipate one) for inter-device transfer? Can we use Intel GPU <-> NPU transfer as an example?
In the current proposal, there is no "block until completion" on the JS side for steps 1 or 2. After developers call
writeBufferto transfer memory, they're free to use theMLBufferin one or moredispatchcalls before callingreadBuffer. For WebGPU interop,readBuffermay not be necessary depending on the scenario.The proposal does not talk about WebGPU/WebNN interop but I agree with Bryan about having an
importExternalBufferAPI in WebGPU which will turn anMLBufferinto aGPUBuffer. The implementation ofimportExternalBufferwould handle synchronizing the NPU and GPU device such that WebGPU reads see WebNN writes. While the buffer is imported as aGPUBuffer, developers would not be able to useMLBufferuntil it is relinquished byWebGPUusing a different API. Similar synchronization would need to happen such that WebNN reads see WebGPU writes.FYI, WebGPU already has a similar relationship with other web APIs such as video frames. See Importing External Textures.
wacky6 wrote:
MLBuffer usage scope
Is MLBuffer only permitted for binding input/output buffers to a built-graph during compute?
Can MLBuffer be used where a MLOperand is accepted, like in conv2d(inputNode, /filters=/ mlBuffer)?As Bryan says
MLBuffercan only be used as an input/output of a graph. If we allow anMLBufferto be used as anMLOperanddown the road, we need to make the spec clear (as is already the case for JS arrays) that a copy is made of the contents of theMLBufferat compilation time. Since graph compilation uses the contents of the buffer to make optimizations, anywriteBufferchanges made to theMLBufferafter compilation would be ignored.wacky6 wrote:
read/writeBuffer definition
Should these be defined on the MLBuffer themselves? Looks like read/write operations is dependent on the context the MLBuffer is associated with.Defining read/write on MLBuffer removes the need to check MLBuffer.context == thisContext in MLContext.read/writeBuffer.
I would prefer that we keep read/writeBuffer on the context so that it is more clear to web developers that those operations are queued relative to dispatch operations. WebGPU works in a similar manner. See GPUQueue
wacky6 wrote:
MLBuffer memory management
When is MLBuffer's memory allocated on device? Is it during writeBuffer?
Should MLBuffer.destroy() returns a Promise to tell the caller that the memory has been deallocated?I agree with Bryan the memory should be allocated when the buffer is created.
Both WebGPU and WebGL have similar
destroymethods. In neither case is a promise returned. When do you expect a WebNN developer would use this?@bbernhar and @RafaelCintron thanks for the explanations! The code example is very helpful
TLDR I'd like to raise some issues which I think are necessary to resolve before this proposal can move forward. I don't have any concrete proposals since I'm still learning this area, but I would appreciate confirmation that the raised issues do need to be tackled. I'm also very happy to help work these out together :)
The proposal does not talk about WebGPU/WebNN interop
I believe that if we're to go forward with
MLBuffer, WebGPU interop needs to be considered from the start rather than assuming we can patch it on later. If I'm reading between the lines correctly here, the WebNN "timelines" mentioned above will have to be closely integrated with WebGPU timelines. Also, given that we'll be hooking into the internals of WebGPU, we need to play by WebGPU's rules, e.g. around buffer usage. I think we need to at minimum:- Define usage of an
MLBufferat creation - Explicitly define WebNN's timelines and how they interact with WebGPU's timelines
With regards to (1) let's look at a snippet from the example above:
// ... wgpuDevice.queue.submit([command_encoder.finish()]); // Export buffer from WebGPU ml_buffer = ml_context.importExternalBuffer(gpu_buffer);
Presumably this code does not suggest that we are synchronously mapping/copying the
GPUBufferinto theMLBuffer(or else some synchronization would be required in JS between these statements), but rather that thegpu_bufferwill be mapped/copied toml_bufferoncegpu_buffer's contents are ready to be accessed. So a contrived example to read the contents of the GPU buffer via anMLBuffermight look like:// `gpuBuffer` is used in some WebGPU work submitted here wgpuDevice.queue.submit([commandEncoder.finish()]); // Inform WebGPU to map/copy `gpuBuffer` to `mlBuffer` once // `gpuBuffer`'s contents are ready to be accessed. const mlBuffer = mlContext.importExternalBuffer(gpuBuffer); // Queue this work behind the importExternalBuffer() call on a WebNN timeline. // This implicitly awaits all WebGPU work involving `gpuBuffer` const gpuBufferContentsCopiedToJsBuffer = await mlContext.readBuffer(mlBuffer);
Note that
readBufferSync()would be functionally equivalent to the much-discussedGPUBuffer.mapSync()if the import doesn't require a copy, which is why I expect it to receive pushback :)What's actually happening here? How can the user agent know whether it can map
gpuBuffertomlBufferor whether it will need to make a copy? This operation should only be valid if:- the
GPUBufferis read-only, - the
GPUBufferis invalidated afterwards, or - we explicitly want to copy the buffer contents
Since the usages of a
GPUBufferare known, we may be able to do this. That being said, mapping aGPUBufferto anMLBufferwill still require abiding by all of WebGPU's constraints - e.g. that these usage flags must be constant within a usage scope, and changing the state of aGPUBufferwill need to be scheduled on a queue timeline.Let's think about the reverse scenario of WebNN -> WebGPU mapping:
// Inform WebNN to map/copy `mlBuffer` to `gpuBuffer` once // `mlBuffer`'s contents are ready to be accessed const gpuBuffer = wgpuDevice.importExternalBuffer(mlBuffer);
How can the user agent know whether it can map
mlBuffertogpuBufferor whether it will need to make a copy?To import a
GPUExternalTexture, that texture is a snapshot which may not change. Presumably the importedMLBuffermust be guaranteed to be read-only to be mapped to a WebGPU buffer, as well?Buffer usage is always assumed on first access (ex. passed as
outputsassumes output usage).This does not seem feasible - especially if we expect the MLBuffer's memory to be allocated on buffer creation. For example, what's the implied usage here?
mlContext.dispatch( graph, /*inputs=*/{buffer: someMlBuffer}, /*outputs=*/{buffer: someMlBuffer}, );
Edge cases aside, let's look at an example of chained inference - the other use case for
MLBuffer:const inputMlBuffer = mlContext.createBuffer({inputSize}); const intermediateMlBuffer = mlContext.createBuffer({intermediateSize}); const outputMlBuffer = mlContext.createBuffer({outputSize}); mlContext.writeBuffer( inputMlBuffer, /*dstOffset=*/0, /*srcData=*/someJsArrayBuffer, ); mlContext.dispatch( graph, /*inputs=*/{buffer: inputMlBuffer}, /*outputs=*/{buffer: intermediateMlBuffer}, ); // Feed the output of one execution as the input to the next. Chained inference! mlContext.dispatch( graph, /*inputs=*/{buffer: intermediateMlBuffer}, /*outputs=*/{buffer: outputMlBuffer}, ); const resultBuffer = await mlContext.readBuffer(outputMlBuffer);
Seems great! Now, where exactly will these buffers be allocated?
This snippet from the WebGPU explainer gives us a hint that we can't both (1) allocate on creation and (2) not know the usage upfront - at least, not without sacrificing something (e.g. performance, extra copies):
The physical memory location for a GPUBuffer’s underlying buffer depends on whether it should be mappable and whether it is mappable for reading or writing
To make this concrete - the Chromium prototype's DML implementation allocates memory for an
MLBufferin the GPU process using theD3D12_RESOURCE_STATE_UNORDERED_ACCESSflag and withD3D12_HEAP_FLAG_NONE. Depending on the usage (and the device architecture), this may be suboptimal. For example, on discrete GPUs, generally resources that can be mapped to the CPU will not perform well when used asD3D12_RESOURCE_STATE_UNORDERED_ACCESSon the GPU. And if theMLBufferis to be used for mapping to CPU buffers, "upload" or "readback" heaps are more appropriate.This proposal doesn't use the words "mapping", but what's being proposing here is effectively mapping for
MLBuffers:Buffer type From JS to *Buffer From *Buffer to JS GPUBufferGPUQueue.writeBuffer()GPUBuffer.mapAsync()MLBufferMLContext.writeBuffer()MLContext.readBuffer()It seems clear to me that we need to define usage of
MLBufferat creation. How we define this mapping might be different if we're designing exclusively for WebNN <-> CPU (JS) interop vs. if we want to support mapping buffers between WebNN <-> WebGPU, which is why I think we should take WebGPU into account early on.
With regards to (2), let's take the real-time video processing use case as another example. Using
MLBuffer, this might look like:const applyEffectToFrame = () => { // Get the frame data as a GPU buffer // Some way to import directly into an MLBuffer directly would avoid this step const gpuExternalBuffer = device.importExternalBuffer({source: video}); // Get the frame data into WebNN. The imported buffer is read-only, so this should // hopefully not require a copy if `mlContext` tied to the same GPU as `gpuExternalBuffer` const inputMlBuffer = mlContext.importExternalBuffer(gpuExternalBuffer); const outputMlBuffer = mlContext.createBuffer({size: inputMlBuffer.size}); // Perform some effects described by `graph` on the frame (e.g. background blur) const inputs = {buffer: inputMlBuffer}; const outputs = {buffer: outputMlBuffer}; mlContext.dispatch(graph, inputs, outputs); // Inform WebNN to map/copy `outputMlBuffer` - which contains the resulting // frame after effects have been applied - to `gpuBufferToRender` once // `outputMlBuffer`'s contents are ready to be accessed // // To avoid a copy, `outputMlBuffer`'s contents must be guaranteed not to change const gpuBufferToRender = wgpuDevice.importExternalBuffer(outputMlBuffer); // create a bind group for `gpuBufferToRender`, create a command encoder, etc. // asking WebGPU to render `gpuBufferToRender` // ... // These queued commands must block on completion of the `dispatch()` call above wgpuDevice.queue.submit([commandEncoder.finish()]); // Call this method for each frame video.requestVideoFrameCallback(applyEffectToFrame); }
Without any additional synchronization, the commands submitted to the
GPUQueuemust block on completion ofMLContext.dispatch(). It seems that theGPUQueuemust either:- block, waiting for some signal from the
MLContextthatoutputMlBufferis available, or - ...be the same queue the
MLContextis running on?
My understanding is that #264 was attempting to specify the latter by describing WebNN execution in a
MLCommandEncoderto be submitted to aGPUQueue. That would naturally queue thecommandEncoder's commands behind the WebNN workload, which would effectively handle the synchronization. But how would this work if the WebNN workload cannot be expressed in terms of GPU commands?FYI, WebGPU already has a similar relationship with other web APIs such as video frames. See Importing External Textures.
My (limited) understanding of WebGPU's relationship with video frames is that the former behavior does not exist? Consider a basic rendering loop with WebGPU:
const render = () => { // Get the frame data as a GPU buffer const gpuExternalBuffer = device.importExternalBuffer({source: video}); // create a bind group for `gpuExternalBuffer`, create a command encoder, // beginRenderPass, etc // ... // Queue a bunch of commands to the GPUQueue, which will eventually render to // a WebGPU canvas wgpuDevice.queue.submit([commandEncoder.finish()]); // This method registers a callback which will be fired once a new frame is // sent to the compositor video.requestVideoFrameCallback(render); }
The encoded GPU commands will eventually "update the rendering of a WebGPU canvas", which in turn calls these steps in the HTML spec, which in turn (eventually) runs the animation frame or video request frame callbacks... which triggers the
render()function and so forth. There are hooks into other APIs, but (as far as I'm aware) theGPUQueuedoes not pause execution while waiting on other APIs.I think we need more details as to how WebGPU and WebNN synchronization will work :)
- Define usage of an
71 remaining items
- added 2 commits that reference this issue
on Apr 17, 2024 - added a commit that references this issue
on May 6, 2024 - added 2 commits that reference this issue
on May 6, 2024 - added a commit that references this issue
on May 14, 2024 Closing, this issue has been replaced by smaller sub-issues which I encourage we use for discussion instead.
https://github.com/webmachinelearning/webnn/issues?q=is%3Aissue+is%3Aopen+in%3Atitle+MLBuffer+
This issue proposes a new opaque device-specific storage type in WebNN,
MLBuffer.MLBufferis a backend-agnostic storage type (CPU, GPU, NPU, etc) which can be used in WebNN operations.MLBufferwould be the solution to:Construction/Destruction
MLBufferis always known (and linear access is assumed).Upload/Download tensor data
srcDatais always made and returns control back to the web developer immediately.Binding to graphs
dictionary MLBufferView { required MLBuffer buffer; MLSize64 offset = 0; MLSize64 size; }; typedef record<DOMString, MLBufferView> MLNamedMLBufferViews; undefined dispatch( MLGraph graph, MLNamedMLBufferViewsinputs, MLNamedMLBufferViews outputs);outputsassumes output usage).Edits:
dispatchinstead of overloadingcompute()per https://www.w3.org/2023/12/14-webmachinelearning-minutes.html