Make include_spirv actually work statically (#8250)

Author: clarfonthey

CHANGELOG.md

@@ -241,6 +241,7 @@ By @wumpf in [#8282](https://github.com/gfx-rs/wgpu/pull/8282), [#8285](https://
 - Require new `F16_IN_F32` downlevel flag for `quantizeToF16`, `pack2x16float`, and `unpack2x16float` in WGSL input. By @aleiserson in [#8130](https://github.com/gfx-rs/wgpu/pull/8130).
 - The error message for non-copyable depth/stencil formats no longer mentions the aspect when it is not relevant. By @reima in [#8156](https://github.com/gfx-rs/wgpu/pull/8156).
 - Track the initialization status of buffer memory correctly when `copy_texture_to_buffer` skips over padding space between rows or layers, or when the start/end of a texture-buffer transfer is not 4B aligned. By @andyleiserson in [#8099](https://github.com/gfx-rs/wgpu/pull/8099).
+- Allow `include_spirv!` and `include_spirv_raw!` macros to be used in constants and statics. By @clarfonthey in [#8250](https://github.com/gfx-rs/wgpu/pull/8250).
 
 #### naga
 

wgpu-types/src/lib.rs

@@ -7835,7 +7835,7 @@ pub struct ShaderRuntimeChecks {
 impl ShaderRuntimeChecks {
     /// Creates a new configuration where the shader is fully checked.
     #[must_use]
-    pub fn checked() -> Self {
+    pub const fn checked() -> Self {
         unsafe { Self::all(true) }
     }
 
@@ -7846,7 +7846,7 @@ impl ShaderRuntimeChecks {
     /// See the documentation for the `set_*` methods for the safety requirements
     /// of each sub-configuration.
     #[must_use]
-    pub fn unchecked() -> Self {
+    pub const fn unchecked() -> Self {
         unsafe { Self::all(false) }
     }
 
@@ -7858,7 +7858,7 @@ impl ShaderRuntimeChecks {
     /// See the documentation for the `set_*` methods for the safety requirements
     /// of each sub-configuration.
     #[must_use]
-    pub unsafe fn all(all_checks: bool) -> Self {
+    pub const unsafe fn all(all_checks: bool) -> Self {
         Self {
             bounds_checks: all_checks,
             force_loop_bounding: all_checks,

wgpu/src/macros/be-aligned.spv

@@ -0,0 +1 @@
+#"3D

wgpu/src/macros/le-aligned.spv

@@ -0,0 +1 @@
+#D3"

wgpu/src/macros.rs renamed to wgpu/src/macros/mod.rs

@@ -77,6 +77,35 @@ fn test_vertex_attr_array() {
     assert_eq!(attrs[1].shader_location, 3);
 }
 
+#[macro_export]
+#[doc(hidden)]
+macro_rules! include_spirv_source {
+    ($($token:tt)*) => {
+        {
+            // FIXME(MSRV): when bumping to 1.89, use [u8; _] here
+            const SPIRV_SOURCE: [
+                u8;
+                $crate::__macro_helpers::include_bytes!($($token)*).len()
+            ] = *$crate::__macro_helpers::include_bytes!($($token)*);
+            const SPIRV_LEN: usize = SPIRV_SOURCE.len() / 4;
+            const SPIRV_WORDS: [u32; SPIRV_LEN] = $crate::util::make_spirv_const(SPIRV_SOURCE);
+            &SPIRV_WORDS
+        }
+    }
+}
+
+#[test]
+fn make_spirv_le_pass() {
+    static SPIRV: &[u32] = include_spirv_source!("le-aligned.spv");
+    assert_eq!(SPIRV, &[0x07230203, 0x11223344]);
+}
+
+#[test]
+fn make_spirv_be_pass() {
+    static SPIRV: &[u32] = include_spirv_source!("be-aligned.spv");
+    assert_eq!(SPIRV, &[0x07230203, 0x11223344]);
+}
+
 /// Macro to load a SPIR-V module statically.
 ///
 /// It ensures the word alignment as well as the magic number.
@@ -90,12 +119,18 @@ macro_rules! include_spirv {
             //log::info!("including '{}'", $($token)*);
             $crate::ShaderModuleDescriptor {
                 label: Some($($token)*),
-                source: $crate::util::make_spirv(include_bytes!($($token)*)),
+                source: $crate::ShaderSource::SpirV(
+                    $crate::__macro_helpers::Cow::Borrowed($crate::include_spirv_source!($($token)*))
+                ),
             }
         }
     };
 }
 
+#[cfg(all(feature = "spirv", test))]
+#[expect(dead_code)]
+static SPIRV: crate::ShaderModuleDescriptor<'_> = include_spirv!("le-aligned.spv");
+
 /// Macro to load raw SPIR-V data statically, for use with [`Features::EXPERIMENTAL_PASSTHROUGH_SHADERS`].
 ///
 /// It ensures the word alignment as well as the magic number.
@@ -108,13 +143,11 @@ macro_rules! include_spirv_raw {
             //log::info!("including '{}'", $($token)*);
             $crate::ShaderModuleDescriptorPassthrough {
                 label: $crate::__macro_helpers::Some($($token)*),
-                spirv: Some($crate::util::make_spirv_raw($crate::__macro_helpers::include_bytes!($($token)*))),
-
-                entry_point: "".to_owned(),
+                spirv: Some($crate::__macro_helpers::Cow::Borrowed($crate::include_spirv_source!($($token)*))),
+                entry_point: $crate::__macro_helpers::String::new(),
                 // This is unused for SPIR-V
                 num_workgroups: (0, 0, 0),
-                reflection: None,
-                shader_runtime_checks: $crate::ShaderRuntimeChecks::unchecked(),
+                runtime_checks: $crate::ShaderRuntimeChecks::unchecked(),
                 dxil: None,
                 msl: None,
                 hlsl: None,
@@ -125,6 +158,11 @@ macro_rules! include_spirv_raw {
     };
 }
 
+#[cfg(test)]
+#[expect(dead_code)]
+static SPIRV_RAW: crate::ShaderModuleDescriptorPassthrough<'_> =
+    include_spirv_raw!("le-aligned.spv");
+
 /// Load WGSL source code from a file at compile time.
 ///
 /// The loaded path is relative to the path of the file containing the macro call, in the same way
@@ -232,7 +270,7 @@ macro_rules! hal_type_gles {
 
 #[doc(hidden)]
 pub mod helpers {
-    pub use alloc::borrow::Cow;
+    pub use alloc::{borrow::Cow, string::String};
     pub use core::{include_bytes, include_str};
     pub use Some;
 }

wgpu/src/util/be-unaligned.spv

@@ -0,0 +1 @@
+#"

wgpu/src/util/empty.spv

@@ -0,0 +1 @@
+#"

wgpu/src/util/le-unaligned.spv

@@ -14,7 +14,7 @@ mod mutex;
 mod texture_blitter;
 
 use alloc::{borrow::Cow, format, string::String, vec};
-use core::ptr::copy_nonoverlapping;
+use core::{mem, ptr::copy_nonoverlapping};
 
 #[cfg(std)]
 pub use belt::StagingBelt;
@@ -46,18 +46,29 @@ pub fn make_spirv(data: &[u8]) -> super::ShaderSource<'_> {
     super::ShaderSource::SpirV(make_spirv_raw(data))
 }
 
+const SPIRV_MAGIC_NUMBER: u32 = 0x0723_0203;
+
+const fn check_spirv_len(data: &[u8]) {
+    assert!(
+        data.len() % size_of::<u32>() == 0,
+        "SPIRV data size must be a multiple of 4."
+    );
+    assert!(!data.is_empty(), "SPIRV data must not be empty.");
+}
+
+const fn verify_spirv_magic(words: &[u32]) {
+    assert!(
+        words[0] == SPIRV_MAGIC_NUMBER,
+        "Wrong magic word in data. Make sure you are using a binary SPIRV file.",
+    );
+}
+
 /// Version of `make_spirv` intended for use with [`Device::create_shader_module_passthrough`].
 /// Returns a raw slice instead of [`ShaderSource`](super::ShaderSource).
 ///
 /// [`Device::create_shader_module_passthrough`]: crate::Device::create_shader_module_passthrough
 pub fn make_spirv_raw(data: &[u8]) -> Cow<'_, [u32]> {
-    const MAGIC_NUMBER: u32 = 0x0723_0203;
-    assert_eq!(
-        data.len() % size_of::<u32>(),
-        0,
-        "data size is not a multiple of 4"
-    );
-    assert_ne!(data.len(), 0, "data size must be larger than zero");
+    check_spirv_len(data);
 
     // If the data happens to be aligned, directly use the byte array,
     // otherwise copy the byte array in an owned vector and use that instead.
@@ -76,21 +87,68 @@ pub fn make_spirv_raw(data: &[u8]) -> Cow<'_, [u32]> {
 
     // Before checking if the data starts with the magic, check if it starts
     // with the magic in non-native endianness, own & swap the data if so.
-    if words[0] == MAGIC_NUMBER.swap_bytes() {
+    if words[0] == SPIRV_MAGIC_NUMBER.swap_bytes() {
         for word in Cow::to_mut(&mut words) {
             *word = word.swap_bytes();
         }
     }
 
-    assert_eq!(
-        words[0], MAGIC_NUMBER,
-        "wrong magic word {:x}. Make sure you are using a binary SPIRV file.",
-        words[0]
-    );
+    verify_spirv_magic(&words);
+
+    words
+}
+
+/// Version of `make_spirv_raw` used for implementing [`include_spirv!`] and [`include_spirv_raw!`] macros.
+///
+/// Not public API. Also, don't even try calling at runtime; you'll get a stack overflow.
+///
+/// [`include_spirv!`]: crate::include_spirv
+#[doc(hidden)]
+pub const fn make_spirv_const<const IN: usize, const OUT: usize>(data: [u8; IN]) -> [u32; OUT] {
+    #[repr(align(4))]
+    struct Aligned<T: ?Sized>(T);
+
+    check_spirv_len(&data);
+
+    // NOTE: to get around lack of generic const expressions
+    assert!(IN / 4 == OUT);
+
+    let aligned = Aligned(data);
+    let mut words: [u32; OUT] = unsafe { mem::transmute_copy(&aligned) };
+
+    // Before checking if the data starts with the magic, check if it starts
+    // with the magic in non-native endianness, own & swap the data if so.
+    if words[0] == SPIRV_MAGIC_NUMBER.swap_bytes() {
+        let mut idx = 0;
+        while idx < words.len() {
+            words[idx] = words[idx].swap_bytes();
+            idx += 1;
+        }
+    }
+
+    verify_spirv_magic(&words);
 
     words
 }
 
+#[should_panic = "multiple of 4"]
+#[test]
+fn make_spirv_le_fail() {
+    let _: [u32; 1] = make_spirv_const([0x03, 0x02, 0x23, 0x07, 0x44, 0x33]);
+}
+
+#[should_panic = "multiple of 4"]
+#[test]
+fn make_spirv_be_fail() {
+    let _: [u32; 1] = make_spirv_const([0x07, 0x23, 0x02, 0x03, 0x11, 0x22]);
+}
+
+#[should_panic = "empty"]
+#[test]
+fn make_spirv_empty() {
+    let _: [u32; 0] = make_spirv_const([]);
+}
+
 /// CPU accessible buffer used to download data back from the GPU.
 pub struct DownloadBuffer {
     _gpu_buffer: super::Buffer,