ONNXRuntimeのExecution Providerについて

CUDA版のカスタムオペレータを実装する場合、一応以下の2パターン方法が考えられる。

カスタムオペレータだけCUDAで実行し、他のオペレータはCPUで実行する
すべてCUDAで実行する

１のパターンはあまりうれしくないと思うので２のカスタムオペレータも他のオペレータもCUDAで動かすことを考える。

この時デフォルトだとONNXRuntimeはCPUで動くのでCUDA版のExecution Providerというのを指定して動かす必要がある。

Execution ProviderというのはONNXRuntime内のフレームワークらしくて、推論処理などのハードウェアアクセラレーションが受けられそうな箇所を切り替え可能な機能ブロックとして分離しているらしい。

詳細は以下の公式ドキュメントを参照のこと。 onnxruntime.ai

そして、CUDAExecutionProviderを使うと推論処理をCUDA版で実行できるらしい。

CUDA版のカスタムオペレータを実装する時にはCUDAExecutionProvider版のカスタムオペレータとして実装すれば全体がCUDAで実行できることになる。

CUDAExecutionProviderでカスタムオペレータを実装する際のポイント

少しだけ前回の復習をしておくとC++でカスタムオペレータを実装する時には主に以下の要素を実装した。

実装すべきもの	主な用途
kernelクラス	計算処理本体
オペレータクラス	カスタムオペレータの仕様（引数の個数など）に関する情報を返すクラス
Register関数など	ONNXRuntimeへの登録

CUDA版を実装する際には大雑把には以下のような実装内容の違いがある。

実装すべきもの	CPU版	CUDA版
kernelクラス	計算処理本体を実装する	CUDAホスト関数を呼び出す
オペレータクラス	CPUExecutionProviderとして実装する	CUDAExecutionProviderとして実装する
Register関数など	CPU/CUDAどちらも同じ	CPU/CUDAどちらも同じ
CUDAソースコード	実装不要	計算処理本体を実装する

具体的な変更点は以下の通り。

kernelクラスのCompute()メソッド
- GetTensorData()はCUDAのデバイスポインタが返ってくるのでそのままCUDA関数に渡してOK
- CUDAのStream IDはKernelContext_GetGPUComputeStream()で取得できる（該当コード）
オペレータクラスのGetExecutionProviderType()メソッド
- CPU版ではオーバーライド不要
- 文字列"CUDAExecutionProvider"を返すようにオーバーライドすることでCUDAExecutionProviderになる（該当コード）

Pythonの呼び出しコード

Pythonの推論処理コードはInferenceSessionのコンストラクタ引数にExecutionProviderのリストを追加で指定するようにすればOK。

import onnxruntime as ort


model = ...

# カスタムオペレータを実装したDLLをロードする
option = ort.SessionOptions()
option.register_custom_ops_library('./libmy_custom_multi_with_cuda.so')

# 使いたいExecutioinProviderを指定する
providers = ['CUDAExecutionProvider']

# カスタムオペレータDLLとExecutionProviderを指定してセッションを生成する
sess = ort.InferenceSession(model.SerializeToString(), option, providers)

# sess.run()で推論できる
...

お試しコードか上で紹介した公式ドキュメントも参考に。

maminusのメモか何か

CUDA実装のONNXRuntimeカスタムオペレータを実装してみた

ONNXRuntimeのExecution Providerについて

CUDAExecutionProviderでカスタムオペレータを実装する際のポイント

Pythonの呼び出しコード