FacetでWebGPUデータレイアウトをテストする

Matt Keeter // ブログ プロジェクト リサーチ ブログ リンク Testing WebGPU data layouts with Facet GPUコンピューティングシェーダーを扱う際、ホストとシェーダー間で設定データのスニペットを共有したいと考えるのは一般的です。以下に、fidget-wgpuから適応させた簡単な例を示します。 // WGSL struct Config { /// スクリーンからモデルへの変換行列 mat: mat3x3f, /// 画像サイズ（ピクセル単位） image_size: vec2u, /// 画像を描画するZ位置 z: f32, } // Rust #[derive( zerocopy::IntoBytes, zerocopy::Immutable, zerocopy::FromBytes, zerocopy::KnownLayout )] #[repr(C)] struct Config { mat: [[f32; 3]; 3], image_size: [u32; 2], z: f32, } zerocopyアノテーションのおかげで、as_bytes()を呼び出し、設定データを直接WebGPUバッファに書き込むことができ、非常に便利です！残念ながら、RustとWebGPUのレイアウト規則の間には、しばしば微妙な違いがあります。上記の例の何が間違っているか分かりますか？さらに残念なことに、私はコンピューティングシェーダーで実行されるバイトコードVMを書いており、その失敗モードはしばしば「おめでとうございます、あなたのGPUはコンピュータを再起動しないとキルできない、永続的にスレッドが回転しています」となります。最近の再起動（WGSLオフセットコンピューターのおかげで）をデバッグした後、この問題をより体系的な方法で修正することにしました。既存の選択肢があります：wgsl_to_wgpu、wgsl_bindgen、encaseはすべてこの問題に関連しています。しかし、私は外部依存関係やビルドスクリプトの煩雑さを避けるために、独自のものを開発することにしました。具体的には、通常のビルドにオーバーヘッドを追加しないように、単体テストを書くことにしました。（もしあなたが厳密になりたいなら、単体テストを使用することには、新しい設定オブジェクトをテストし忘れるという失敗モードがありますが、私はそれほど心配していません）nagaを使用してWGSL構造体のレイアウトを見つけることができます。これはすでにシェーダーコンパイルの依存関係ツリーに含まれています：let module = naga::front::wgsl::parse_str(code).expect("valid WGSL"); let members = module .types .iter() .find_map(|(_, ty)| { if ty.name.as_deref() == Some("Config") && let naga::TypeInner::Struct { members, .. } = &ty.inner { Some(members) } else { None } }) .expect("could not find struct");次に、構造体Configの各メンバーと照合できます：let expected_offsets = [ ("mat", std::mem::offset_of!(Config, mat)), ("image_size", std::mem::offset_of!(Config, image_size)), ("z", std::mem::offset_of!(Config, z)), ]; for (field_name, rust_offset) in expected_offsets { let wgsl_member = members .iter() .find(|m| m.name.as_deref() == Some(field_name)) .unwrap_or_else(|| { panic!("field `{field_name}` missing in WGSL struct") }); assert_eq!( wgsl_member.offset as usize, rust_offset, "offset mismatch for field `{field_name}`" ); }確かに、これは問題を発見しました！thread 'pixel::test::blog_test' (532515) panicked at fidget-wgpu/src/pixel/mod.rs:1762:13: assertion `left == right` failed: offset mismatch for field `image_size` left: 48 right: 36 WGSLでは、mat3x3の各行には4バイトのパディングがあるため、各行は合計16バイト、行列全体は48バイトになります。Rustでは、[[f32; 3]; 3]オブジェクトは密にパッキングされているため、36バイトしか占有しません。テストは機能しましたが、このアプローチには2つの問題があります：第一に、Configの各メンバーをハードコードする必要があり、これは厄介で間違いやすいです！複数の設定オブジェクトをテストしたい場合、それぞれを手書きで同期させておく必要があります。第二に、フィールドのサイズを直接比較できないため、エラーメッセージが良くありません。image_sizeのオフセットが間違っている理由は、matのサイズが一致しないためです。近因を報告すべきです。（これは、最後のメンバーのサイズが間違っている場合も報告しないことを意味します）どうすればよいでしょうか？facetはRustのランタイムリフレクションを提供するライブラリです。構造体に#[derive(facet::Facet)]をアノテートすることで、ランタイムで検査できるSHAPEという関連型が得られます。これを使用してConfigオブジェクトを自動的にチェックできます！ジェネリックチェッカー関数をステップバイステップで見ていきましょう。これはT: facet::Facetをパラメータ化しています。まずシェーダーを解析し、名前で設定構造体を抽出します。これは基本的に以前と同じです：pub(crate) fn compare_struct_layout<T: facet::Facet<'static>>( shader: &str, struct_name: &str, ) { // [1] WGSLを解析します。以前と同じです let module = naga::front::wgsl::parse_str(shader) .expect("valid WGSL"); let (members, span) = module .types .iter() .find_map(|(_, ty)| { if ty.name.as_deref() == Some(struct_name) && let naga::TypeInner::Struct { members, span } = &ty.inner { Some((members, *span)) } else { None } }) .expect("could not find struct");次に、オブジェクト全体のサイズをチェックします。ここには1つの微妙な点があります：私はしばしば、設定オブジェクトの最後のメンバーとしてランタイムサイズ配列を使用します。（これはC言語のフレキシブル配列メンバーやRustの動的サイズ型と考えてください）WGSLでは次のようになります：struct VoxelConfig { mat: mat4x4f, axes: vec3u, tape_data_offset: atomic<u32>, render_size: vec3u, tape_data_capacity: u32, image_size: vec3u, root_tape_len: atomic<u32>, tape_data: array<TapeWord>, // <- ランタイムサイズ配列 }Rust構造体で同等のものを作成する場合、最後のメンバーは省略されます。Nagaにランタイムサイズ配列がある場合、オブジェクトの合計サイズ（Nagaによって報告される）はチェックしません。代わりに、tape_dataのオフセットがRust構造体のサイズと一致することをチェックします。チェッカー関数のそのセクションは次のとおりです：// [2] 全体の`struct`サイズをチェックします // // 構造体の最後のメンバーがランタイムサイズ配列の場合、 // 配列の開始オフセットを構造体サイズとして扱います。 let dynamic_array_offset = members.last().and_then(|m| { let ty = &module.types[m.ty]; let naga::TypeInner::Array { base: _, size: naga::ir::ArraySize::Dynamic, stride: _, } = &ty.inner else { return None; }; Some(m.offset) }); if let Some(dynamic_array_offset) = dynamic_array_offset { assert_eq!( dynamic_array_offset as usize, std::mem::size_of::<T>() ); } else { assert_eq!( span as usize, std::mem::size_of::<T>() ); }次に、フィールドのサイズとオフセットをチェックします。これは以前のフィールドごとの手動チェックと同じ考え方ですが、Facetのイントロスペクションデータを使用してランタイムでチェックを実行します（メンバーサイズもチェックします）。// [3] RustとWGSL間のフィールドサイズとオフセットをチェックします let facet::Type::User(facet::UserType::Struct(shape)) = T::SHAPE.ty else { panic!("must build a struct"); }; let mut shape_field_names = HashSet::new(); // 後で使用するため for field in shape.fields { let field_name = field.name; shape_field_names.insert(field_name); let wgsl_member = members .iter() .find(|m| m.name.as_deref() == Some(field_name)) .unwrap_or_else(|| { panic!("field `{field_name}` missing in WGSL struct") }); assert_eq!( wgsl_member.offset as usize, field.offset, "offset mismatch for field `{field_name}`" ); assert_eq!( module.types[wgsl_member.ty] .inner .size(module.to_ctx()) as usize, field.shape().layout.sized_layout().unwrap().size(), "size mismatch for field `{field_name}`" ); }最後に、反対方向を確認し、すべてのWGSLメンバーがRust構造体に（末尾のランタイムサイズ配列を除く）存在することを確認します：// [4] すべてのWGSLメンバーがRustに存在することを確認します let slice_len = if dynamic_array_offset.is_some() { members.len() - 1 } else { members.len() }; for m in &members[..slice_len] { assert!( shape_field_names.contains( m.name .as_ref() .expect("cannot check unnamed WGSL fields") .as_str(), ), "field `{field_name}` missing in Rust struct" ); }このジェネリックチェッカー関数は少し手間がかかりましたが、一度書けば、すべての設定オブジェクトに簡単に適用できます。Facetは単体テストでのみ条件付きで派生させます。例：#[derive( zerocopy::IntoBytes, zerocopy::Immutable, zerocopy::FromBytes, zerocopy::KnownLayout )] #[cfg_attr(test, derive(facet::Facet))] // <- テストでのみ派生 #[repr(C)] struct Config { mat: [[f32; 3]; 3], image_size: [u32; 2], z: f32, }その後、ほぼ無料でテストを取得できます：#[test] fn color_config_layout() { crat