voicevox_coreを改造してGPUを選択できるようにしたい

はじめに

voicevox_coreの中で使用するGPUは0番目に固定されています
Multi-GPU環境でVOICEVOXが使用するGPUを選べると便利だとおもいませんか？
voicevox_coreを改造して、使用するGPUを選べるようにします！
※そのうち公式でGPUを選択できるAPIが実装されるかもしれませんが暫定対応ということで

ビルド環境構築

Build Tools for Visual Studio 2022

https://visualstudio.microsoft.com/ja/downloads/#build-tools-for-visual-studio-2022
このあたりからダウンロードしてインストーラーを実行

「C++によるデスクトップ開発」にのみチェックを入れてインストール

rustup(Rust)

https://rustup.rs/
このあたりからダウンロードしてインストーラーを実行

デフォルトでおっけーなので、何も入力せずにEnter

CMake

このあたりからダウンロードしてインストーラーを実行
特につまるところはない

Git

このあたりからダウンロードしてインストーラーを実行
なんかいろいろオプションがあるが全部デフォでおっけー

voicevox_core

https://github.com/VOICEVOX/voicevox_core/releases
このあたりからダウンロードしてソースコードを展開
Gitを入れたのでgit cloneとかでとってきてもいい
※voicevox_coreのために入れたわけではなく、後々cargoで必要になる

展開場所はお好みで！

まずは普通にビルド

まずは、構築した環境の確認のために普通にビルドしてみる
展開したソースコードの場所をPowerShellとかで開いてビルドコマンド実行

ビルドコマンドはこんな感じで

cargo build -p voicevox_core_c_api --features load-onnxruntime --target x86_64-pc-windows-msvc --release

voicevox_core.dllができてればおっけー

次は改造してビルドしてみる

環境変数からデバイスを指定するように変更
修正をかけるファイルは↓

crates/voicevox_core/src/core/infer/runtimes/onnxruntime.rs

～中略～
//環境変数からdevice_idを取得する関数を追加
fn get_directml_device_id() -> Option<i32> {
    env::var("VOICEVOX_DIRECTML_DEVICE_ID")
        .ok()
        .and_then(|s| s.parse().ok())
}

impl InferenceRuntime for self::blocking::Onnxruntime {
    type Session = async_lock::Mutex<ort::session::Session>; // WASMでは`ort`を利用しないので、ここはasync-lockを用いてよいはず
    type RunContext = OnnxruntimeRunContext;

    const DISPLAY_NAME: &'static str = if cfg!(feature = "load-onnxruntime") {
        "現在ロードされているONNX Runtime"
    } else if cfg!(feature = "link-onnxruntime") {
        "現在リンクされているONNX Runtime"
    } else {
        panic!("either `load-onnxruntime` or `link-onnxruntime` must be enabled");
    };

    fn supported_devices(&self) -> crate::Result<SupportedDevices> {
        (|| {
            let cpu = CPUExecutionProvider::default().is_available()?;
            let cuda = CUDAExecutionProvider::default().is_available()?;

            //修正箇所１
            let dml = if let Some(device_id) = get_directml_device_id() {
                DirectMLExecutionProvider::default().with_device_id(device_id).is_available()?
            } else {
               DirectMLExecutionProvider::default().is_available()?
            };

            ensure!(cpu, "missing `CPUExecutionProvider`");

            Ok(SupportedDevices {
                cpu: true,
                cuda,
                dml,
            })
        })()
        .map_err(ErrorRepr::GetSupportedDevices)
        .map_err(Into::into)
    }

    fn test_gpu(&self, gpu: GpuSpec) -> anyhow::Result<()> {
        let sess_builder = &mut ort::session::builder::SessionBuilder::new()?;
        match gpu {
            GpuSpec::Cuda => CUDAExecutionProvider::default()
                .with_conv_algorithm_search(CuDNNConvAlgorithmSearch::Default)
                .register(sess_builder),

            //修正箇所２
            GpuSpec::Dml => {
                let dml = if let Some(device_id) = get_directml_device_id() {
                    DirectMLExecutionProvider::default().with_device_id(device_id)
                } else {
                    DirectMLExecutionProvider::default()
                };
                dml.register(sess_builder)
            }
        }
        .map_err(Into::into)
    }

    fn new_session(
        &self,
        model: &ModelBytes,
        options: InferenceSessionOptions,
    ) -> anyhow::Result<(
        Self::Session,
        Vec<ParamInfo<InputScalarKind>>,
        Vec<ParamInfo<OutputScalarKind>>,
    )> {
        let mut builder = ort::session::Session::builder()?
            .with_optimization_level(GraphOptimizationLevel::Level1)?
            .with_intra_threads(options.cpu_num_threads.into())?;

        match options.device {
            DeviceSpec::Cpu => {}
            DeviceSpec::Gpu(GpuSpec::Cuda) => {
                CUDAExecutionProvider::default()
                    .with_conv_algorithm_search(CuDNNConvAlgorithmSearch::Default)
                    .register(&mut builder)?;
            }

            //修正箇所３
            DeviceSpec::Gpu(GpuSpec::Dml) => {
                builder = builder
                    .with_parallel_execution(false)?
                    .with_memory_pattern(false)?;
                    
                if let Some(device_id) = get_directml_device_id() {
                    DirectMLExecutionProvider::default().with_device_id(device_id).register(&mut builder)?;
                } else {
                    DirectMLExecutionProvider::default().register(&mut builder)?;
                };
            }
        };

～中略～

修正したら、再ビルドする
ビルドコマンドは変わらず

cargo build -p voicevox_core_c_api --features load-onnxruntime --target x86_64-pc-windows-msvc --release

voicevox_core.dllが再作成できてればおっけー

実際に使ってみる

voicevox_core.dllを差し替える

環境変数を VOICEVOX_DIRECTML_DEVICE_ID=0 にセットして実行
GPU1が使われている（※DIRECTML的には0番目のデバイス）

環境変数を VOICEVOX_DIRECTML_DEVICE_ID=1 にセットして実行
GPU0が使われている（※DIRECTML的には1番目のデバイス）

想定通り！