[Epic] Integrate NVIDIA libraries (cuBLAS, etc.) with CommandBuffer API

openxla / xla

A machine learning compiler for GPUs, CPUs, and ML accelerators

Apache License 2.0

2.75k stars 441 forks source link

[Epic] Integrate NVIDIA libraries (cuBLAS, etc.) with CommandBuffer API #5874

Open ezhulenev opened 1 year ago

ezhulenev commented 1 year ago

We will use cuGraphAddChildGraphNode API to insert CUDA graphs captured from library calls into "main" command buffers.

We need a not too CUDA graph specific APIs in StreamExecutor/CommandBuffer to make it potentially portable for other command-buffer-like APIs.

### Tasks
- [ ] https://github.com/openxla/xla/issues/6432

nouiz commented 1 year ago

What is the status here? Can you confirm it is done for cublas, and cudnn isn't done yet?

ezhulenev commented 1 year ago

It's work in progress. cuBLAS was done for the current runtime integration based on whole XLA program capture, this is for capturing library calls into child graphs