search

microsoft / onnxruntime

ONNX Runtime: cross-platform, high performance ML inferencing and training accelerator
https://onnxruntime.ai
MIT License
14.42k stars 2.89k forks source link

[Feature Request] Support defining memory type in custom op's I/O #13228

Open wangyems opened 2 years ago

wangyems commented 2 years ago

Describe the feature request

Custom ops may have inputs/outputs on different memory locations. Internally we can specify inputs/outputs memory type by calling KernelDefBuilderInputMemoryType() or KernelDefBuilderOutputMemoryType when registering kernels. It would be nice if this can be supported in custom op. The aim is to mitigate memory copy between device and host.

Describe scenario use case

custom beam search op

hariharans29 commented 2 years ago

Is this what you are looking for - https://github.com/microsoft/onnxruntime/pull/10879 ?

wangyems commented 2 years ago

Is this what you are looking for - #10879 ?

Exactly. I think we have a use case for it now.