microsoft / microsoft/onnxruntime
[Feature Request] Support defining memory type in custom op's I/O
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the feature request
Custom ops may have inputs/outputs on different memory locations. Internally we can specify inputs/outputs memory type by calling KernelDefBuilder__InputMemoryType() or KernelDefBuilder__OutputMemoryType 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
Contributor guide
Research direction
Start by tracing custom-op registration and the existing KernelDefBuilder__InputMemoryType() and KernelDefBuilder__OutputMemoryType APIs named in the issue. No file or test is identified, so first locate the custom-op kernel definition path and related memory-copy tests; done means custom ops can declare input and output memory types and the beam-search scenario avoids unnecessary host-device copies.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100