Rust で wgpu を使った画像処理

Rust で wgpu を使った画像処理を試してみました。

wgpu は WebGPU をベースに Rust で実装された API で DirectX 12, Vulkan, Metal 等に対応しています。

テクスチャを画像ファイルへ出力1

wgpu では以下のような要素を使って処理を組み立てるようです。

要素	内容
Adapter（アダプター）	GPU ハードウェアを表現したもの
Device（デバイス）	GPU を操作する API を提供
Queue（キュー）	GPU へコマンドを送信

これらを使って、背景色で塗りつぶしたテクスチャを画像ファイルへ出力する処理を実装すると次のようになりました。

1. アダプター取得
2. アダプターからデバイスとキューを取得
3. デバイスを使ってテクスチャとその内容をコピーするバッファを作成
4. デバイスを使ってコマンドエンコーダーを作成し、GPU へ送信するコマンドを組み立て
   • レンダリングパスを開始して描画コマンドを組み立て
     • テクスチャを青色でクリア（塗りつぶす）
   • テクスチャの内容をバッファへコピー
5. キューを使ってコマンドエンコーダーの内容を GPU へ送信
6. バッファの内容を取り出して画像ファイルへ出力

sample1/src/main.rs

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

#[tokio::main]
async fn main() -> Result<()> {
    // ログの有効化
    env_logger::init();

    let w = 320;
    let h = 240;

    let instance = wgpu::Instance::default();
    // アダプターの取得
    let adapter = instance
        .request_adapter(&wgpu::RequestAdapterOptions::default())
        .await
        .ok_or("notfound adapter")?;

    println!("adapter: {:?}", adapter.get_info());

    // デバイスとキューの取得
    let (device, queue) = adapter
        .request_device(&wgpu::DeviceDescriptor::default(), None)
        .await?;

    // テクスチャの作成
    let texture = device.create_texture(&wgpu::TextureDescriptor {
        label: None,
        size: wgpu::Extent3d {
            width: w,
            height: h,
            depth_or_array_layers: 1,
        },
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D2,
        format: wgpu::TextureFormat::Rgba8UnormSrgb,
        usage: wgpu::TextureUsages::RENDER_ATTACHMENT | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[wgpu::TextureFormat::Rgba8UnormSrgb],
    });
    // バッファの作成（テクスチャのコピー用）
    let output_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: (w * h * 4) as u64, // 幅 * 高さ * 4バイト（RGBA）
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });

    let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());

    // コマンドエンコーダー作成
    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());

    {
        // レンダリングパスの開始
        encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &texture_view,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLUE), // 開始時にテクスチャを青色（背景色）でクリア
                    store: wgpu::StoreOp::Store, // 描画結果をテクスチャへ保存
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        });
    }

    // テクスチャの内容をバッファへコピー
    encoder.copy_texture_to_buffer(
        wgpu::ImageCopyTexture {
            texture: &texture,
            mip_level: 0,
            origin: wgpu::Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        },
        wgpu::ImageCopyBuffer {
            buffer: &output_buf,
            layout: wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: Some(w * 4), // 1行のバイトサイズ
                rows_per_image: None, // Some(h) でも可
            },
        },
        wgpu::Extent3d {
            width: w,
            height: h,
            depth_or_array_layers: 1,
        },
    );

    // GPU へコマンド送信
    queue.submit(Some(encoder.finish()));

    // バッファの内容取得
    let slice = output_buf.slice(..);
    slice.map_async(wgpu::MapMode::Read, |_| {});

    device.poll(wgpu::MaintainBase::Wait);

    // 画像ファイルへ出力
    image::save_buffer(
        "output.png",
        &slice.get_mapped_range().to_vec(),
        w,
        h,
        image::ColorType::Rgba8,
    )?;

    Ok(())
}

sample1/Cargo.toml

[dependencies]
env_logger = "0.10"
wgpu = "0.18"
image = "0.24"
tokio = { version = "1", features = ["full"] }

MacBook Air で実行した結果は次のようになりました。

実行結果

$ cd sample1
$ cargo run
・・・ 省略
adapter: AdapterInfo { name: "Apple M2", vendor: 0, device: 0, device_type: IntegratedGpu, driver: "", driver_info: "", backend: Metal }

画像ファイル output.png はこのようになりました。

テクスチャを画像ファイルへ出力2

次に、3つの三角形をテクスチャへ描画するようにしてみます。

まずは、シェーダーを WGSL で記述します。

@vertex を付けた vs_main が頂点シェーダーで頂点毎に呼び出されるようです。
頂点の座標変換を実施するもので、頂点をそのまま使った値を返すようにしています。

@fragment を付けた fs_main がフラグメントシェーダーで描画するピクセル毎に呼び出されるようです。
ピクセルの色を返すもので、常に黄色を返すようにしています。

sample2/src/shader.wgsl

@vertex
fn vs_main(@location(0) pos: vec2f) -> @builtin(position) vec4f {
    return vec4f(pos, 0.0, 1.0);
}

@fragment
fn fs_main() -> @location(0) vec4f {
    return vec4f(1.0, 1.0, 0.0, 1.0);
}

このシェーダーを使ってテクスチャへ描画する処理は次のようになります。

• シェーダーモジュールを作成し、それを使ってレンダリングパイプラインを作成
• 頂点の情報を書き込んだバッファを作成
• レンダリングパスでパイプラインとバッファ（頂点情報）を設定し、draw で描画する頂点を指定

sample2/src/main.rs

use std::borrow::Cow;
use wgpu::util::DeviceExt;

type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

#[tokio::main]
async fn main() -> Result<()> {
    env_logger::init();

    let w = 320;
    let h = 240;

    // 3つの三角形の頂点
    let verticies: Vec<f32> = vec![
        // 1つ目
        0.0, 0.0, 
        -0.6, -0.6, 
        0.5, -0.4,
        // 2つ目
        0.0, 0.0, 
        0.8, 0.4, 
        0.5, 0.9, 
        // 3つ目
        0.0, 0.0, 
        -0.7, 0.3, 
        -0.3, 0.7,
    ];

    let instance = wgpu::Instance::default();

    let adapter = instance
        .request_adapter(&wgpu::RequestAdapterOptions::default())
        .await
        .ok_or("notfound adapter")?;

    let (device, queue) = adapter
        .request_device(&wgpu::DeviceDescriptor::default(), None)
        .await?;

    // シェーダーモジュールの作成
    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: None,
        source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("shader.wgsl"))),
    });

    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
        label: None,
        bind_group_layouts: &[],
        push_constant_ranges: &[],
    });
    // レンダリングパイプラインの作成
    let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: None,
        layout: Some(&pipeline_layout),
        vertex: wgpu::VertexState {
            module: &shader,
            entry_point: "vs_main",
            buffers: &[wgpu::VertexBufferLayout {
                array_stride: 8,
                attributes: &[wgpu::VertexAttribute {
                    format: wgpu::VertexFormat::Float32x2,
                    offset: 0,
                    shader_location: 0,
                }],
                step_mode: wgpu::VertexStepMode::Vertex,
            }],
        },
        fragment: Some(wgpu::FragmentState {
            module: &shader,
            entry_point: "fs_main",
            targets: &[Some(wgpu::TextureFormat::Rgba8UnormSrgb.into())],
        }),
        primitive: wgpu::PrimitiveState::default(),
        depth_stencil: None,
        multisample: wgpu::MultisampleState::default(),
        multiview: None,
    });

    // 頂点を書き込んだバッファ作成
    let verticies_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None,
        contents: bytemuck::cast_slice(&verticies),
        usage: wgpu::BufferUsages::VERTEX,
    });

    let texture = device.create_texture(&wgpu::TextureDescriptor {
        ... 省略
    });

    let output_buf = device.create_buffer(&wgpu::BufferDescriptor {
        ... 省略
    });

    let texture_view = texture.create_view(&wgpu::TextureViewDescriptor::default());

    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());

    {
        let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: &texture_view,
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color::BLUE),
                    store: wgpu::StoreOp::Store,
                },
            })],
            depth_stencil_attachment: None,
            timestamp_writes: None,
            occlusion_query_set: None,
        });

        rpass.set_pipeline(&pipeline);
        rpass.set_vertex_buffer(0, verticies_buf.slice(..));
        // レンダリングする頂点を指定
        rpass.draw(0..9, 0..1);
    }

    ... 省略
}

sample2/Cargo.toml

[dependencies]
env_logger = "0.10"
wgpu = "0.18"
image = "0.24"
bytemuck = "1"
tokio = { version = "1", features = ["full"] }

実行した結果、画像ファイル output.png はこのようになりました。