コンテンツにスキップ
ONNX Runtime

デバイス テンソル

ONNX Runtime でのデバイス テンソルの使用

Section titled “ONNX Runtime でのデバイス テンソルの使用”

デバイス テンソルを使用することは、特に異種メモリ システム上で効率的な AI パイプラインを構築する上で重要な部分となります。 このようなシステムの典型的な例は、専用 GPU を搭載した PC です。 最近の GPU 自体は約 1TB/s のメモリ帯域幅を持っていますが、CPU への相互接続 PCI 4.0 x16 は、わずか 32GB/s 程度で制限要因となることがよくあります。 したがって、GPU は計算と PCI メモリ トラフィックを同時に実行できるため、データを可能な限り GPU にローカルに保つか、計算の背後に遅いメモリ トラフィックを隠すのが最善です。

メモリがすでに推論デバイスにローカルであるこれらのシナリオの典型的なユースケースは、GPU デコーダでデコードできるエンコードされたビデオ ストリームの GPU アクセラレーションによるビデオ処理です。 もう 1 つの一般的なケースは、中間テンソルを CPU にコピーする必要がない拡散ネットワークや大規模言語モデルなどの反復ネットワークです。 高解像度画像のタイルベースの推論も、PCI コピー中の GPU のアイドル時間を削減するためにカスタム メモリ管理が重要なもう 1 つのユースケースです。各タイルを順次処理するのではなく、GPU 上で PCI コピーと処理をオーバーラップさせ、その方法で作業をパイプライン化することが可能です。

順次 PCI->処理->PCI の画像と、それがインターリーブされている別の画像。

CUDA

Section titled “CUDA”

ONNX Runtime の CUDA には 2 つのカスタム メモリ タイプがあります。 "CudaPinned" と "Cuda" メモリで、CUDA ピン留め は実際には GPU が直接アクセスできる CPU メモリであり、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 バインディング として使用できます。 このような IO バインディングを使用すると、より多くのパフォーマンス チューニングが可能です。

  • 固定テンソル アドレスにより、CUDA グラフをキャプチャして CPU 上の CUDA 起動レイテンシを削減できます。
  • ピン留めメモリへの完全に非同期なダウンロード、またはデバイス ローカル テンソルを使用してメモリ コピーを排除することにより、CUDA は指定されたストリーム上で 実行オプションを介して完全に非同期 で実行できます。

CUDA のカスタム計算ストリームを設定するには、Ort[CUDA|TensorRT]ProviderOptionsV2* 不透明構造体ポインタを公開する V2 オプション API と、そのストリーム メンバーを設定する関数 Update[CUDA|TensorRT]ProviderOptionsWithValue(options, "user_compute_stream", cuda_stream); を参照してください。 詳細については、各実行プロバイダーのドキュメントを参照してください。

最適化を検証したい場合は、Nsight System が CPU API と CUDA 操作の GPU 実行を関連付けるのに役立ちます。 これにより、目的の同期が行われたかどうか、および非同期操作が同期実行にフォールバックしないかどうかを検証できます。 これは、デバイス テンソルの最適な使用法を説明する この GTC の講演でも使用されています。

Python API

Section titled “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])

python での非同期実行を有効にすることは、C++ API と同じ 実行オプション を介して可能です。

DirectML

Section titled “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)));


// D3D12 API を使用して d3d_buffer を作成
Microsoft::WRL::ComPtr<ID3D12Resource> d3d_buffer = ...;


// D3D リソースから dml リソースを作成
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

Section titled “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")
# 出力次元が既知の場合は、事前に割り当てられた値をバインドすることもできます
# binding.bind_ortvalue_output("out", dml_array_out)


session.run_with_iobinding(binding)