wgpuでcomputeシェーダの結果をホストメモリにコピーする方法

こちらのブログポストをRustで実装した際に、Computeシェーダの結果をホストメモリにコピーする方法が分からなかったので、その方法をメモ。

環境

  • rustc: 1.69.0
  • wgpu: 0.16.1

不明点

以下のTypescript版のコードでmapAsyncしてから結果を表示するところまでをRustのwgpuではどう書けばいいのか分からなかった。

...
await stagingBuffer.mapAsync(
  GPUMapMode.READ,
  0, // Offset
  BUFFER_SIZE // Length
 );
const copyArrayBuffer =
  stagingBuffer.getMappedRange(0, BUFFER_SIZE);
const data = copyArrayBuffer.slice();
stagingBuffer.unmap();
console.log(new Float32Array(data));

Rustでの実装

以下のように書くことで、Computeシェーダの結果をホストメモリにコピーできた。

...
    let staging_slice = staging_buffer.slice(..);
    staging_slice.map_async(wgpu::MapMode::Read, |_| {
        println!("Mapped!");
    });
    device.poll(wgpu::Maintain::Wait);

    let data = staging_slice
        .get_mapped_range()  // BufferViewを取得
        .chunks_exact(4)  // 4byteずつに分割
        .map(|b| f32::from_ne_bytes(b.try_into().unwrap()))  // Byte -> Floatに変換
        .collect::<Vec<_>>();  // Vecに変換
    staging_buffer.unmap();

    println!("Data: {:?}", data);
...

ポイントは以下の部分。

  • staging_slice.map_asyncでcallbackを指定する (callbackがないサンプルもあるが、wgpu:0.13からcallbackが追加された様子)
  • device.pollでmap_asyncのcallbackが呼ばれるまで待つ
  • get_mapped_rangeでBufferViewを取得したら、自力でByteに分割してFloatに変換する

参考