re: Scatter

+1 to support scatter, in particular, scatterND

MLOperand scatterNd(MLOperand input, MLOperand indices, MLOperand updates);

When prototyping Whisper decoders' inference with MLBuffer, I found scatterND operator is useful to insert KV values by position into the pre-allocated static KV cache (WebNN requires static shape).

builder.scatterNd(past_key, position_ids, present_key);
builder.scatterNd(past_value, position_ids, present_value);

This avoids reading the KV cache tensor back to CPU and improves the performance. The initial prototype is available at: https://github.com/huningxin/onnxruntime-inference-examples/blob/whisper-mlbuffer/js/whisper-demo/whisper.js

Platforms' support:

DirectML DML_SCATTER_ND_OPERATOR_DESC: Copies the whole input tensor to the output, then overwrites selected indices with corresponding values from the updates tensor.
CoreML scatter_nd: Scatter updates to data at locations indices. Support "update" and other modes.
TFLite scatter_nd: Scatter sparse updates according to individual values at the specified indices. Only support "add"/"sum" mode. Need the emulation to support scattering to input tensor.

Where is scatter and gather op? #467

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions