在 ONNX Runtime 中使用设备张量

在构建高效的 AI 流水线时，使用设备张量可能至关重要，尤其是在异构内存系统中。此类系统的典型例子是任何配备独立 GPU 的 PC。虽然最新的 GPU 本身具有约 1TB/s 的内存带宽，但连接 CPU 的 PCI 4.0 x16 互连通常是瓶颈，带宽仅为约 32GB/s。因此，通常最好尽可能将数据保存在 GPU 本地，或者在 GPU 执行计算和 PCI 内存流量的同时，通过计算来隐藏缓慢的内存流量。

在内存已位于推理设备本地的情况下，一个典型的用例是对编码视频流进行 GPU 加速视频处理，该视频流可以使用 GPU 解码器进行解码。另一个常见情况是迭代网络，例如扩散网络或大型语言模型，它们的中间张量不需要复制回 CPU。针对高分辨率图像的基于分块的推理是另一个用例，其中自定义内存管理对于减少 PCI 复制期间的 GPU 空闲时间至关重要。与顺序处理每个分块不同，可以在 GPU 上重叠进行 PCI 复制和处理，并通过这种方式流水线化工作。

Image of sequential PCI->Processing->PCI and another image of it being interleaved.

CUDA

ONNX Runtime 中的 CUDA 有两种自定义内存类型。"CudaPinned" 和 "Cuda" 内存，其中 CUDA pinned 实际上是 CPU 内存，GPU 可以直接访问它，从而允许使用 cudaMemcpyAsync 完全异步地上载和下载内存。普通 CPU 张量只允许从 GPU 到 CPU 的同步下载，而从 CPU 到 GPU 的复制总是可以异步执行的。

使用 Ort::Sessions 的分配器分配张量非常简单，可以使用 C++ API，该 API 直接映射到 C API。

Ort::Session session(ort_env, model_path_cstr, session_options);
Ort::MemoryInfo memory_info_cuda("Cuda", OrtArenaAllocator, /*device_id*/0,
                                 OrtMemTypeDefault);
Ort::Allocator gpu_allocator(session, memory_info_cuda);
auto ort_value = Ort::Value::CreateTensor(
        gpu_allocator, shape.data(), shape.size(),
        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);

外部分配的数据也可以包装到 Ort::Value 中而无需复制

Ort::MemoryInfo memory_info_cuda("Cuda", OrtArenaAllocator, device_id,
                                 OrtMemTypeDefault);
std::array<int64_t, 4> shape{1, 4, 64, 64};
size_t cuda_buffer_size = 4 * 64 * 64 * sizeof(float);
void *cuda_resource;
CUDA_CHECK(cudaMalloc(&cuda_resource, cuda_buffer_size));
auto ort_value = Ort::Value::CreateTensor(
    memory_info_cuda, cuda_resource, cuda_buffer_size,
    shape.data(), shape.size(),
    ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);

然后，这些分配的张量可以用作 I/O 绑定，以消除网络上的复制操作，并将责任转移给用户。通过这种 IO 绑定，可以进行更多性能调优

由于张量地址固定，可以捕获 CUDA 图以减少 CPU 上的 CUDA 启动延迟
由于对固定内存进行完全异步下载或通过使用设备本地张量消除内存复制，CUDA 可以在其给定的流上通过运行选项完全异步运行

要为 CUDA 设置自定义计算流，请参考 V2 选项 API，该 API 暴露了 Ort[CUDA|TensorRT]ProviderOptionsV2* 不透明结构体指针和函数 Update[CUDA|TensorRT]ProviderOptionsWithValue(options, "user_compute_stream", cuda_stream); 来设置其流成员。更多详细信息可在每个执行提供程序的文档中找到。

如果您想验证您的优化，Nsight System 有助于关联 CPU API 和 CUDA 操作的 GPU 执行。这也可以验证是否进行了所需的同步，并且没有异步操作回退到同步执行。它也用于本次 GTC 演讲，解释了设备张量的最佳用法。

Python API

Python API 支持与上述 C++ API 相同的性能机会。如此处所示，可以分配设备张量。除此之外，可以通过此 API 设置 user_compute_stream

sess = onnxruntime.InferenceSession("model.onnx", providers=["TensorrtExecutionProvider"])
option = {}
s = torch.cuda.Stream()
option["user_compute_stream"] = str(s.cuda_stream)                    
sess.set_providers(["TensorrtExecutionProvider"], [option])

通过与 C++ API 相同的运行选项，可以在 Python 中启用异步执行。

DirectML

通过 DirectX 资源可以实现相同的行为。要运行异步处理，进行与 CUDA 相同的执行流管理至关重要。对于 DirectX，这意味着管理设备及其命令队列，这可以通过 C API 实现。如何设置计算命令队列的详细信息在使用 SessionOptionsAppendExecutionProvider_DML1 的文档中有说明。

如果使用单独的命令队列进行复制和计算，则可以重叠 PCI 复制和执行，并使执行异步化。

#include <onnxruntime/dml_provider_factory.h>
Ort::MemoryInfo memory_info_dml("DML", OrtDeviceAllocator, device_id,
                                OrtMemTypeDefault);

std::array<int64_t, 4> shape{1, 4, 64, 64};
void *dml_resource;
size_t d3d_buffer_size = 4 * 64 * 64 * sizeof(float);
const OrtDmlApi *ort_dml_api;
Ort::ThrowOnError(Ort::GetApi().GetExecutionProviderApi(
                  "DML", ORT_API_VERSION, reinterpret_cast<const void **>(&ort_dml_api)));

// Create d3d_buffer using D3D12 APIs
Microsoft::WRL::ComPtr<ID3D12Resource> d3d_buffer = ...;

// Create the dml resource from the D3D resource.
ort_dml_api->CreateGPUAllocationFromD3DResource(d3d_buffer.Get(), &dml_resource);


Ort::Value ort_value(Ort::Value::CreateTensor(memory_info_dml, dml_resource,
                     d3d_buffer_size, shape.data(), shape.size(),
                     ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT));

在 GitHub 上可以找到一个单文件示例，展示了如何管理和创建复制和执行命令队列。

Python API

尽管从 Python 分配 DirectX 输入可能不是主要用例，但 API 是可用的。这可能会非常有益，特别是对于中间网络缓存，例如大型语言模型 (LLM) 中的键值缓存。

import onnxruntime as ort
import numpy as np

session = ort.InferenceSession("model.onnx",
                               providers=["DmlExecutionProvider"])

cpu_array = np.zeros((1, 4, 512, 512), dtype=np.float32)
dml_array = ort.OrtValue.ortvalue_from_numpy(cpu_array, "dml")

binding = session.io_binding()
binding.bind_ortvalue_input("data", dml_array)
binding.bind_output("out", "dml")
# if the output dims are known we can also bind a preallocated value
# binding.bind_ortvalue_output("out", dml_array_out)

session.run_with_iobinding(binding)