在 ONNX Runtime 中使用设备张量

使用设备张量可能是构建高效 AI 管道的关键部分，尤其是在异构内存系统上。此类系统的典型示例是任何带有专用 GPU 的 PC。虽然最新的 GPU 本身具有约 1TB/s 的内存带宽，但到 CPU 的互连 PCI 4.0 x16 通常会成为限制因素，仅有约 32GB/s。因此，通常最好将数据保留在 GPU 本地，或者将缓慢的内存流量隐藏在计算之后，因为 GPU 能够同时执行计算和 PCI 内存流量。

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

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

CUDA

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

使用 Ort::Sessions 的分配器分配张量非常简单，使用 C++ 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 绑定，以消除网络上的复制操作，并将责任转移给用户。通过这种 I/O 绑定，可以进行更多性能调整

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

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

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

Python API

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

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)