GPUDirect Storage を Rust で動かしてみた

この記事は TSG Advent Calendar 2022 の 8 日目の記事です。

はじめに

GPUDirect Storage とは？

従来、データをストレージから GPU に読み込んで処理するためには、データを CPU メモリに読み込んでから、GPU メモリに転送する必要がありました。

GPUDirect Storage を使うと、CPU メモリを介さずに直接 GPU にデータを読み込むことが可能になります。

（NVIDIA 技術ブログ より引用）

この記事の目的

GPUDirect Storage は、cuFile API という CUDA の API として実装されています。が、Rust で書きたい！ ということで、GPUDirect Storage を Rust で動かしていきます。この記事に CUDA のコードは出てきません。

この記事は、だいぶ自分のための忘備録的なところが大きいです。そもそも

• GPUDirect Storage を動かしてみた情報が少ない
• Rust と CUDA の組み合わせのための周辺環境が習熟しきっていない

という状況ですが、これらの組み合わせにより更なるニッチさを生み出していきます。

なお、本記事のコードは以下のリポジトリにあるので、そちらもご確認ください。

下準備

GPUDirect Storage を使えるようにする

GPUDirect Storage インストールガイド と CUDA インストールガイド に従ってインストール作業を進めていきます。

なお、今回用いた環境は以下の通りです。

• CPU: Intel Xeon W-2245
• GPU: NVIDIA RTX A4000
• SSD: WD_BLACK AN1500
• OS: Ubuntu 22.04

IOMMU を無効化する

まずは IOMMU を無効化する必要があります。以下を実行して、無効化されているかどうかを確認します。

dmesg | grep -i iommu

無効化されていなければ設定を変更します。

sudo vi /etc/default/grub

- GRUB_CMDLINE_LINUX_DEFAULT=""
+ GRUB_CMDLINE_LINUX_DEFAULT="intel_iommu=off"

Intel CPU の場合は intel_iommu、AMD CPU の場合は amd_iommu を off にします。

ブートローダーを更新し、再起動します。

sudo update-grub
sudo reboot

MLNX_OFED をインストールする

MLNX_OFED のダウンロードページ から、自分の環境に合ったものをダウンロードします。自分の場合は 5.8-1.1.2.1-LTS > Ubuntu > Ubuntu 22.04 > x86_64 を選択し、ISO イメージをダウンロードしました。

ダウンロードした ISO ファイルを インストール手順 に従ってインストールします。マウント先はどこでも良いです。インストールオプションは GPUDirect Storage インストールガイドに書かれています。

sudo mkdir -p /mnt/MLNX_OFED_LINUX-5.8-1.0.1.1-ubuntu22.04-x86_64
sudo mount -o ro,loop MLNX_OFED_LINUX-5.8-1.0.1.1-ubuntu22.04-x86_64.iso /mnt/MLNX_OFED_LINUX-5.8-1.0.1.1-ubuntu22.04-x86_64
sudo /mnt/MLNX_OFED_LINUX-5.8-1.0.1.1-ubuntu22.04-x86_64/mlnxofedinstall --with-nvmf --with-nfsrdma --enable-gds --add-kernel-support --dkms

設定を更新し、再起動します。

sudo update-initramfs -u -k `uname -r`
sudo reboot

CUDA Toolkit と GPUDirect Storage をインストールする

必要なバージョンのカーネルヘッダが入っていることを確かめます。Ubuntu の場合は以下の通りになります。

sudo apt-get install linux-headers-$(uname -r)

CUDA Toolkit ダウンロードページ から、自分の環境にあったものを選ぶと、インストールコマンドが表示されます。自分の場合は Linux > x86_64 > Ubuntu > 20.04 > deb (local) を選択しました。

コピペして実行します。

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin
sudo mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2204-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2204-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

GPUDirect Storage をインストールします（cuda と分けてインストールする必要があります）。

sudo apt-get -y install nvidia-gds

環境変数を設定する

CUDA の実行バイナリと共有ライブラリにパスを通します。~/.profile など好きなところに追記してください（宗教なので）。

~/.profile

+ export PATH="/usr/local/cuda-11.8/bin${PATH:+:${PATH}}"
+ export LD_LIBRARY_PATH="/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"

確認する

お疲れ様でした。以下のコマンドを叩いて、正常に動いているかどうか確認してください。

/usr/local/cuda/gds/tools/gdscheck.py -p

NVMe : Supported と表示されれば成功です。

Rust-CUDA をセットアップする

Rust で CUDA のカーネルや呼び出しを書くにあたって、Rust-GPU/Rust-CUDA を使用します。

これには何を使っても良いのですが、2022 年現在 deprecated でなく、最も進展がありそうなエコシステムとして Rust-CUDA を選びました。とはいえ Rust-CUDA も開発初期段階です。他のプロジェクトも含めて全体的に、Rust と CUDA を組み合わせるときの周辺環境は未だ習熟しきっていない感じがあります。

では、リポジトリ と ガイドブック に従って構成していきます。

LLVM 7.x.x を入れる

要件にあるように、LLVM 7.x.x が必要になります。Ubuntu 20.04 では apt で llvm-7 をインストールできるのですが、Ubuntu 22.04 では無くなっているようです。また、ビルド済みバイナリも Ubuntu 18.04 までにしか対応していません。Docker を使う手もありますが、ここではソースコードからビルドすることにしました。

LLVM インストールガイド と clang インストールガイド に従ってインストールしていきます。

バージョン 7 系で最も新しい 7.1.0 を指定してリポジトリをクローンします。

git clone --depth=1 https://github.com/llvm/llvm-project.git -b llvmorg-7.1.0

ビルドします。自分の環境では 30～40 分ほどかかったので、その間にご飯を食べました。

cd llvm-project
mkdir build
cd build
cmake -G "Unix Makefiles" -DCMAKE_BUILD_TYPE=Release ../llvm
make

インストールすると、/usr/local 以下に色々と入ります。

sudo make install

フォルダ構造

ガイドブックや examples を参考に、フォルダ構造を決めます。以下のように、CPU クレート（バイナリ）と GPU クレート（ライブラリ）が必要です。

.
├── cpu
│   ├── src
│   │   └── main.rs
│   ├── Cargo.toml
│   └── build.rs
├── gpu
│   ├── src
│   │   └── lib.rs
│   └── Cargo.toml
├── resources
│   └── gpu.ptx
├── Cargo.lock
├── Cargo.toml
└── rust-toolchain

ガイドブックから変えた点として、ワークスペースを使うようにしました。他にも GPU のコードを GPU 側でビルドするようにしたかったのですが、ビルドしたときに無限ループに陥ってしまったので、ガイドブックの通り CPU クレートから GPU クレートをビルドして使うようにします。

rust-toolchain

nightly でしか動かないと書いてあるので、ルートフォルダに rust-toolchain をおいて rustc のバージョンを指定します。新しくしすぎると Rust-CUDA がコンパイルエラーを吐いたので、examples にあるバージョンにします。

rust-toolchain

[toolchain]
channel = "nightly-2021-12-04"
components = ["rust-src", "rustc-dev", "llvm-tools-preview"]

GPU 側のコードを書く

ようやく Rust のコードを書いていきます。

まずはライブラリクレートを用意します。

cargo init gpu --lib

Cargo.toml

[workspace]
members = [
  "gpu",
]

gpu/Cargo.toml を編集します。rust-toolchain で指定している rustc のバージョンが古い場合、cargo add は使えません。

gpu/Cargo.toml

[package]
name = "gpu"
version = "0.1.0"
edition = "2021"

+ [lib]
+ crate-type = ["cdylib", "rlib"]

[dependencies]
+ cuda_std = "0.2.2"

カーネルのコードを書きます。

gpu/src/lib.rs

#![cfg_attr(
    target_os = "cuda",
    no_std,
    feature(register_attr),
    register_attr(nvvm_internal)
)]

use cuda_std::*;

#[kernel]
#[allow(improper_ctypes_definitions)]
pub unsafe fn add(a: &[f32], b: &[f32], c: *mut f32) {
    let idx = thread::index_1d() as usize;
    if idx < a.len() {
        let elem = &mut *c.add(idx);
        *elem = a[idx] + b[idx];
    }
}

examples にあるのは、与えられた 2 つの f32 の配列 a と b を足した結果を c に入れる関数です。

#![cfg_attr(...)] は lib.rs に必要で、#[kernel] はカーネル関数ごとに必要なアトリビュートです。

CPU 側のコードを書く

バイナリクレートを用意します。

cargo init cpu --bin

Cargo.toml

[workspace]
members = [
  "gpu",
+   "cpu",
]

cpu/Cargo.toml を編集します。

cpu/Cargo.toml

[package]
name = "cpu"
version = "0.1.0"
edition = "2021"

[dependencies]
cust = "0.3.2"
nanorand = "0.7.0"

+ [build-dependencies]
+ cuda_builder = "0.3.0"

GPU クレートをビルドするスクリプトを書きます。前述のように、CPU クレート側で生成するという点に注意してください。resources/gpu.ptx に PTX ファイルが生成されます。

cpu/build.rs

use cuda_builder::CudaBuilder;

fn main() {
    CudaBuilder::new("../gpu")
        .copy_to("../resources/gpu.ptx")
        .build()
        .unwrap();
}

ただし自分の環境では、生成された PTX のバージョンやターゲットを書き換えないと動きませんでした。エラーメッセージは Error: UnknownError としか出ませんが、内部的には CUDA_ERROR_UNSUPPORTED_PTX_VERSION (222) というエラーで、これを特定するのが大変でした。

以下の .version や .target は必要に応じて書き換えてください。

cpu/build.rs

use cuda_builder::CudaBuilder;
+ use std::process::Command;

fn main() {
    CudaBuilder::new("../gpu")
        .copy_to("../resources/gpu.ptx")
        .build()
        .unwrap();

+     Command::new("sed")
+         .arg("-i")
+         .args(&["-e", r#"s/^\.version .*/.version 7.6/"#])
+         .args(&["-e", r#"s/^\.target .*/.target sm_86/"#])
+         .arg("../resources/gpu.ptx")
+         .output()
+         .unwrap();
}

ホストのコードを書きます。長いので折りたたみます。Rust-CUDA の examples にあるコードから、PTX のパス部分を変更したものです。

cpu/src/main.rs を見る

cpu/src/main.rs

use cust::prelude::*;
use nanorand::{Rng, WyRand};
use std::error::Error;

/// How many numbers to generate and add together.
const NUMBERS_LEN: usize = 100_000;

static PTX: &str = include_str!("../../resources/gpu.ptx");

fn main() -> Result<(), Box<dyn Error>> {
    // generate our random vectors.
    let mut wyrand = WyRand::new();
    let mut lhs = vec![2.0f32; NUMBERS_LEN];
    wyrand.fill(&mut lhs);
    let mut rhs = vec![0.0f32; NUMBERS_LEN];
    wyrand.fill(&mut rhs);

    // initialize CUDA, this will pick the first available device and will
    // make a CUDA context from it.
    // We don't need the context for anything but it must be kept alive.
    let _ctx = cust::quick_init()?;

    // Make the CUDA module, modules just house the GPU code for the kernels we created.
    // they can be made from PTX code, cubins, or fatbins.
    let module = Module::from_ptx(PTX, &[])?;

    // make a CUDA stream to issue calls to. You can think of this as an OS thread but for dispatching
    // GPU calls.
    let stream = Stream::new(StreamFlags::NON_BLOCKING, None)?;

    // allocate the GPU memory needed to house our numbers and copy them over.
    let lhs_gpu = lhs.as_slice().as_dbuf()?;
    let rhs_gpu = rhs.as_slice().as_dbuf()?;

    // allocate our output buffer. You could also use DeviceBuffer::uninitialized() to avoid the
    // cost of the copy, but you need to be careful not to read from the buffer.
    let mut out = vec![0.0f32; NUMBERS_LEN];
    let out_buf = out.as_slice().as_dbuf()?;

    // retrieve the add kernel from the module so we can calculate the right launch config.
    let func = module.get_function("add")?;

    // use the CUDA occupancy API to find an optimal launch configuration for the grid and block size.
    // This will try to maximize how much of the GPU is used by finding the best launch configuration for the
    // current CUDA device/architecture.
    let (_, block_size) = func.suggested_launch_configuration(0, 0.into())?;

    let grid_size = (NUMBERS_LEN as u32 + block_size - 1) / block_size;

    println!(
        "using {} blocks and {} threads per block",
        grid_size, block_size
    );

    // Actually launch the GPU kernel. This will queue up the launch on the stream, it will
    // not block the thread until the kernel is finished.
    unsafe {
        launch!(
            // slices are passed as two parameters, the pointer and the length.
            func<<<grid_size, block_size, 0, stream>>>(
                lhs_gpu.as_device_ptr(),
                lhs_gpu.len(),
                rhs_gpu.as_device_ptr(),
                rhs_gpu.len(),
                out_buf.as_device_ptr(),
            )
        )?;
    }

    stream.synchronize()?;

    // copy back the data from the GPU.
    out_buf.copy_to(&mut out)?;

    println!("{} + {} = {}", lhs[0], rhs[0], out[0]);

    Ok(())
}

実行する

ここまでで、CPU クレートを実行するとしっかり動くはずです。

cargo run -p cpu

using 131 blocks and 768 threads per block
0.27493936 + 0.42648312 = 0.70142245

本題

cuFile API のラッパーを作る

C や C++ のヘッダファイルから Rust の FFI バインディングを生成してくれる便利ツール bindgen を使って、cufile.h のラッパーを作っていきます。bindgen は Rust-CUDA 内部の cust_raw クレートでも使われているので、オプションを参考にしていきます。

まずはライブラリクレートを作ります。cufile-sys とでも名付けます。

Cargo.toml

[workspace]
members = [
    "cpu",
+     "cufile-sys",
    "gpu",
]

cargo init cufile-sys --lib

ビルド時に bindgen を使うようにします。

cufile-sys/Cargo.toml

[package]
name = "cufile-sys"
version = "0.1.0"
edition = "2021"

[dependencies]

+ [build-dependencies]
+ bindgen = "0.63.0"

bindgen を用意します。cust_raw や以下を参考に、CLI から bindgen を使うことにします。

cargo install bindgen

cufile-sys/wrapper.h

#include "cufile.h"

cufile-sys/bindgen.sh

#!/bin/bash
set -exu

bindgen \
  --allowlist-type="^CU.*" \
  --allowlist-type="^cuuint(32|64)_t" \
  --allowlist-type="^cudaError_enum" \
  --allowlist-type="^cu.*Complex$" \
  --allowlist-type="^cuda.*" \
  --allowlist-type="^libraryPropertyType.*" \
  --allowlist-var="^CU.*" \
  --allowlist-function="^cu.*" \
  --default-enum-style=rust \
  --no-doc-comments \
  --with-derive-default \
  --with-derive-eq \
  --with-derive-hash \
  --with-derive-ord \
  --size_t-is-usize \
  wrapper.h -- -I/usr/local/cuda/include \
  > src/cufile.rs

./bindgen.sh を実行すると src/cufile.rs が生成されます。自分の環境では 10,184 行ほどありました。これを公開します。

cufile/src/lib.rs

#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

mod cufile;
pub use cufile::*;

ビルド時にリンクするようにします。

cufile-sys/build.rs

fn main() {
    println!("cargo:rustc-link-search=/usr/local/cuda/lib64");
    println!("cargo:rustc-link-lib=cufile");
}

これで、別のクレートから cufile-sys クレートを介して、cuFile API が使えるようになります。

CPU 側の呼び出しコードを書く

ようやく本題の本題です。

cuFile API のリファレンスガイド によると、cuFile API を使って GPU メモリからストレージ I/O を行うには、例として以下の手順が必要です。

• open("/path/to/file")
• cuFileDriverOpen()
• cuFileHandleRegister(handle, descr)
• cudaMalloc(buf_dev, buf_size)
• cuFileBufRegister(buf_dev, buf_size, 0)
• cuFileWrite(handle, buf_dev, io_size, offset, buf_offset)
• cuFileBufDeregister(buf_dev)
• cuFileFree(buf_dev)
• cuFileHandleDeregister()
• cuFileDriverClose()
• close(fd)

長いですね。C 言語風の API になっているので、open したら close、register したら deregister です。簡潔に言えば、ファイルディスクリプタを cuFile API 用のハンドラに変換する処理と、デバイス上のバッファを登録する処理が必要になります。

さらに言えば、open() する際に O_DIRECT というフラグが必要になります。

よって、リファレンスガイドにある gds_helloworld.cxx は、Rust では以下のようなコードで実現できます。

gds-helloworld/Cargo.toml

[package]
name = "gds-helloworld"
version = "0.1.0"
edition = "2021"

[dependencies]
cufile-sys = { path = "../cufile-sys", version = "0.1.0"}
cust = "0.3.2"
libc = "0.2.138"

gds-helloworld/src/main.rs

use cufile_sys as cufile;
use cust::prelude::*;
use std::env;
use std::error::Error;
use std::fs::OpenOptions;
use std::mem;
use std::os::unix::fs::OpenOptionsExt;
use std::os::unix::io::IntoRawFd;

fn main() -> Result<(), Box<dyn Error>> {
    let _ctx = cust::quick_init()?;

    let testfn = env::var("TESTFILE")?;

    println!("Opening File {}", testfn);
    let f = OpenOptions::new()
        .write(true)
        .create(true)
        .custom_flags(libc::O_DIRECT)
        .open(&testfn)?;

    let raw_fd = f.into_raw_fd();

    println!("Opening cuFileDriver.");
    unsafe {
        let status = cufile::cuFileDriverOpen();
        assert_eq!(status.err, cufile::CUfileOpError::CU_FILE_SUCCESS);
    }

    let mut cf_handle: cufile::CUfileHandle_t = ::std::ptr::null_mut();

    let mut descr = cufile::CUfileDescr_t::default();
    descr.handle.fd = raw_fd;
    descr.type_ = cufile::CUfileFileHandleType::CU_FILE_HANDLE_TYPE_OPAQUE_FD;

    println!("Registering cuFile handle to {}.", testfn);
    unsafe {
        let status = cufile::cuFileHandleRegister(&mut cf_handle, &mut descr);
        assert_eq!(status.err, cufile::CUfileOpError::CU_FILE_SUCCESS);
    }

    let io_size = 1 << 24;
    let buff_size = io_size + 0x1000;

    println!("Allocating CUDA buffer of {} bytes.", buff_size);
    let mut buf = unsafe { DeviceBuffer::<u8>::uninitialized(buff_size)? };

    println!("Registering Buffer of {} bytes.", buff_size);
    unsafe {
        let status = cufile::cuFileBufRegister(
            buf.as_device_ptr().as_raw() as *mut ::std::ffi::c_void,
            mem::size_of::<u8>() * buf.len(),
            0,
        );
        assert_eq!(status.err, cufile::CUfileOpError::CU_FILE_SUCCESS);
    }

    println!("Filling memory.");
    buf.set_8(0xab)?;

    let dev_ptr_offset = 0x1000;
    let file_offset = 0x2000;

    println!("Writing buffer to file.");
    unsafe {
        let ret = cufile::cuFileWrite(
            cf_handle,
            buf.as_device_ptr().as_raw() as *mut ::std::ffi::c_void,
            io_size,
            file_offset,
            dev_ptr_offset,
        );
        assert!(ret >= 0 && ret as usize == io_size);
    }

    println!("Releasing cuFile buffer.");
    unsafe {
        let status =
            cufile::cuFileBufDeregister(buf.as_device_ptr().as_raw() as *mut ::std::ffi::c_void);
        assert_eq!(status.err, cufile::CUfileOpError::CU_FILE_SUCCESS);
    }

    println!("Releasing file handle.");
    unsafe {
        cufile::cuFileHandleDeregister(cf_handle);
    }

    println!("Closing File Driver.");
    unsafe {
        let status = cufile::cuFileDriverClose();
        assert_eq!(status.err, cufile::CUfileOpError::CU_FILE_SUCCESS);
    }

    Ok(())
}

Rust-CUDA 特有の事情として、呼び出す前に、CUDA の初期化が必要です。コンテキストがデストラクトされないように、一度変数に保持しておく必要があります。

では実行してみます。

cargo build -p gds-helloworld
TESTFILE=./foo.txt ./target/debug/gds-helloworld

Opening File ./foo.txt
Opening cuFileDriver.
Registering cuFile handle to ./foo.txt.
Allocating CUDA buffer of 16781312 bytes.
Registering Buffer of 16781312 bytes.
Filling memory.
Writing buffer to file.
Releasing cuFile buffer.
Releasing file handle.
Closing File Driver.

成功です！

おわりに

雑感

ここまでで書いたコードはリポジトリに置いてあります。

cuFile API を Rust で書けることのメリットは、

• 強い静的型付け
• Result 型によるエラー処理
• ライフタイムによるリソースの自動的な確保・解放
• Cargo という利便性の高いエコシステム

など、色々あります。これらを達成した、より GPUDirect Storage が使いやすくなるようなラッパーライブラリを作ることも考えられます。これに関しては後述する kvikIO が参考になるでしょう。

Rust はいいぞ。

関連プロジェクト

kvikIO は、cuFile API を C++ 用と Python 用にバインディングしたライブラリです。生の cuFile API では open, register に対する close, deregister 処理や、エラー処理が煩雑だったのに対し、kvikIO では C++ の RAII や例外を活用して簡潔に書けるようになっています。また、並列な読み書きも独自に実装されています。ライブラリ名が可愛いですね。

謝辞

GPU をはじめ、今回使ったハードウェア環境は研究室のものです。ありがとう……。