microsoft / microsoft/onnxruntime
[Performance] Access to shared GPU cublas/cublasLt/cudnn library handles from custom_op kernels
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
The GPU custom_op examples only shows direct CUDA programming examples, where the CUDA stream handle is accessible via the API. The provider and contrib_ops show access to cublas, cublasLt, and cudnn NVidia library handles. How might one gain access to these library handles from a custom_op?
While one could create the library handles within each custom_op instance, then associate them with a stream independent of the main compute_stream, that would both be inefficient and interfere with the chain of implicit synchronization of the node GPU operations before and after the custom_op node. Not to mention potentially invalidating the input and output the tensor memory pointers passed to the custom_op kernel if they belong to a different workspace.
### To reproduce
N/A
### Urgency
With the optimization effort put into libraries like cublasLt, cudnn, and Cutlass, it seems a shame to not provide the capability of tapping into those resources as part of a custom_op. Our deployment vehicle will be a Triton server, and having to come in as a provider just to be able to call workhorse GPU library routines would hurt the modularity of our application, as well as having to be unduly sensitive to all the versioning compatibilities that a provider must support. And a contrib_ops contributor route is not in line with our intellectual property objectives.
### Platform
Linux
### OS Version
Ubuntu
### ONNX Runtime Installation
Released Package
### ONNX Runtime Version or Commit ID
1.12.1, 1.14.1 if necessary
### ONNX Runtime API
C++
### Architecture
X64
### Execution Provider
CUDA
### Execution Provider Library Version
CUDA 11.6 or higher
### Model File
_No response_
### Is this a quantized model?
No
Contributor guide
Research direction
Start by comparing the GPU custom_op examples with the provider and contrib_ops implementations, which the issue identifies as examples of library-handle access. Determine the API and synchronization requirements for exposing shared cublas, cublasLt, and cudnn handles to custom_op kernels. Done means the supported access path and its stream and workspace behavior are specified and covered by relevant tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend-api-design, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100