Add buffer pointer support for WGSL (SPV_EXT_physical_storage_buffer)

src/back/glsl/mod.rs

@@ -108,6 +108,7 @@ impl crate::AddressSpace {
             crate::AddressSpace::WorkGroup
             | crate::AddressSpace::Uniform
             | crate::AddressSpace::Storage { .. }
+            | crate::AddressSpace::PhysicalStorage { .. }
             | crate::AddressSpace::Handle
             | crate::AddressSpace::PushConstant => false,
         }
@@ -1072,6 +1073,8 @@ impl<'a, W: Write> Writer<'a, W> {
             crate::AddressSpace::Storage { .. } => {
                 self.write_interface_block(handle, global)?;
             }
+            // The PhysicalStorage address space is only for pointers.
+            crate::AddressSpace::PhysicalStorage { .. } => unreachable!(),
             // A global variable in the `Function` address space is a
             // contradiction in terms.
             crate::AddressSpace::Function => unreachable!(),
@@ -4180,6 +4183,7 @@ const fn glsl_storage_qualifier(space: crate::AddressSpace) -> Option<&'static s
         As::Function => None,
         As::Private => None,
         As::Storage { .. } => Some("buffer"),
+        As::PhysicalStorage { .. } => Some("buffer_reference"),
         As::Uniform => Some("uniform"),
         As::Handle => Some("uniform"),
         As::WorkGroup => Some("shared"),

src/back/hlsl/writer.rs

@@ -688,6 +688,11 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 write!(self.out, "{prefix}ByteAddressBuffer")?;
                 register
             }
+            crate::AddressSpace::PhysicalStorage { .. } => {
+                return Err(Error::Unimplemented(
+                    "physical storage address space is not implemented for HLSL".to_string(),
+                ));
+            }
             crate::AddressSpace::Handle => {
                 let handle_ty = match *inner {
                     TypeInner::BindingArray { ref base, .. } => &module.types[*base].inner,

src/back/msl/writer.rs

@@ -399,6 +399,7 @@ impl crate::AddressSpace {
         match *self {
             Self::Uniform
             | Self::Storage { .. }
+            | Self::PhysicalStorage { .. }
             | Self::Private
             | Self::WorkGroup
             | Self::PushConstant
@@ -414,7 +415,7 @@ impl crate::AddressSpace {
             // rely on the actual use of a global by functions. This means we
             // may end up with "const" even if the binding is read-write,
             // and that should be OK.
-            Self::Storage { .. } => true,
+            Self::Storage { .. } | Self::PhysicalStorage { .. } => true,
             // These should always be read-write.
             Self::Private | Self::WorkGroup => false,
             // These translate to `constant` address space, no need for qualifiers.
@@ -424,11 +425,14 @@ impl crate::AddressSpace {
         }
     }
 
-    const fn to_msl_name(self) -> Option<&'static str> {
+    fn to_msl_name(self) -> Option<&'static str> {
         match self {
             Self::Handle => None,
             Self::Uniform | Self::PushConstant => Some("constant"),
             Self::Storage { .. } => Some("device"),
+            Self::PhysicalStorage { .. } => {
+                unreachable!("physical storage address space is not implemented for MSL")
+            }
             Self::Private | Self::Function => Some("thread"),
             Self::WorkGroup => Some("threadgroup"),
         }
@@ -3677,6 +3681,7 @@ impl<W: Write> Writer<W> {
                         }
                         crate::AddressSpace::Function
                         | crate::AddressSpace::Private
+                        | crate::AddressSpace::PhysicalStorage { .. }
                         | crate::AddressSpace::WorkGroup => {}
                     }
                 }

src/back/spv/block.rs

@@ -213,6 +213,22 @@ impl<'w> BlockContext<'w> {
                 let arg = &self.ir_function.arguments[index as usize];
                 self.ir_module.types[arg.ty].inner.pointer_space().is_some()
             }
+            crate::Expression::Load { pointer } => {
+                match *self.fun_info[pointer].ty.inner_with(&self.ir_module.types) {
+                    crate::TypeInner::Pointer { base, .. } => {
+                        if let crate::TypeInner::Pointer {
+                            space: crate::AddressSpace::PhysicalStorage { .. },
+                            ..
+                        } = self.ir_module.types[base].inner
+                        {
+                            true
+                        } else {
+                            false
+                        }
+                    }
+                    _ => false,
+                }
+            }
 
             // The chain rule: if this `Access...`'s `base` operand was
             // previously omitted, then omit this one, too.
@@ -1082,15 +1098,22 @@ impl<'w> BlockContext<'w> {
                 match self.write_expression_pointer(pointer, block, None)? {
                     ExpressionPointer::Ready { pointer_id } => {
                         let id = self.gen_id();
-                        let atomic_space =
+                        let (atomic_space, physical_alignment) =
                             match *self.fun_info[pointer].ty.inner_with(&self.ir_module.types) {
                                 crate::TypeInner::Pointer { base, space } => {
-                                    match self.ir_module.types[base].inner {
+                                    let physical_alignment = match space {
+                                        crate::AddressSpace::PhysicalStorage { alignment } => {
+                                            Some(alignment)
+                                        }
+                                        _ => None,
+                                    };
+                                    let atomic_space = match self.ir_module.types[base].inner {
                                         crate::TypeInner::Atomic { .. } => Some(space),
                                         _ => None,
-                                    }
+                                    };
+                                    (atomic_space, physical_alignment)
                                 }
-                                _ => None,
+                                _ => (None, None),
                             };
                         let instruction = if let Some(space) = atomic_space {
                             let (semantics, scope) = space.to_spirv_semantics_and_scope();
@@ -1103,6 +1126,14 @@ impl<'w> BlockContext<'w> {
                                 scope_constant_id,
                                 semantics_id,
                             )
+                        } else if let Some(alignment) = physical_alignment {
+                            Instruction::load_aligned(
+                                result_type_id,
+                                id,
+                                pointer_id,
+                                alignment,
+                                None,
+                            )
                         } else {
                             Instruction::load(result_type_id, id, pointer_id, None)
                         };
@@ -1549,6 +1580,9 @@ impl<'w> BlockContext<'w> {
                 crate::Expression::FunctionArgument(index) => {
                     break self.function.parameter_id(index);
                 }
+                crate::Expression::Load { .. } => {
+                    break self.cached[expr_handle];
+                }
                 ref other => unimplemented!("Unexpected pointer expression {:?}", other),
             }
         };

src/back/spv/helpers.rs

@@ -26,6 +26,7 @@ pub(super) const fn map_storage_class(space: crate::AddressSpace) -> spirv::Stor
         crate::AddressSpace::Function => spirv::StorageClass::Function,
         crate::AddressSpace::Private => spirv::StorageClass::Private,
         crate::AddressSpace::Storage { .. } => spirv::StorageClass::StorageBuffer,
+        crate::AddressSpace::PhysicalStorage { .. } => spirv::StorageClass::PhysicalStorageBuffer,
         crate::AddressSpace::Uniform => spirv::StorageClass::Uniform,
         crate::AddressSpace::WorkGroup => spirv::StorageClass::Workgroup,
         crate::AddressSpace::PushConstant => spirv::StorageClass::PushConstant,

src/back/spv/instructions.rs

@@ -446,6 +446,25 @@ impl super::Instruction {
         instruction
     }
 
+    pub(super) fn load_aligned(
+        result_type_id: Word,
+        id: Word,
+        pointer_id: Word,
+        alignment: u32,
+        memory_access: Option<spirv::MemoryAccess>,
+    ) -> Self {
+        let mut instruction = Self::new(Op::Load);
+        instruction.set_type(result_type_id);
+        instruction.set_result(id);
+        instruction.add_operand(pointer_id);
+
+        let memory_access = memory_access.unwrap_or(spirv::MemoryAccess::ALIGNED);
+        instruction.add_operand(memory_access.bits());
+        instruction.add_operand(alignment);
+
+        instruction
+    }
+
     pub(super) fn atomic_load(
         result_type_id: Word,
         id: Word,

src/back/spv/writer.rs

@@ -879,6 +879,15 @@ impl Writer {
                     _ => {}
                 }
             }
+            crate::TypeInner::Pointer {
+                space: crate::AddressSpace::PhysicalStorage { .. },
+                ..
+            } => {
+                self.require_any(
+                    "Buffer Pointer",
+                    &[spirv::Capability::PhysicalStorageBufferAddresses],
+                )?;
+            }
             crate::TypeInner::AccelerationStructure => {
                 self.require_any("Acceleration Structure", &[spirv::Capability::RayQueryKHR])?;
             }
@@ -1905,6 +1914,17 @@ impl Writer {
             self.write_type_declaration_arena(&ir_module.types, handle)?;
         }
 
+        let has_physical_storage_pointer = self
+            .capabilities_used
+            .contains(&spirv::Capability::PhysicalStorageBufferAddresses);
+        if self.physical_layout.version < 0x10500 && has_physical_storage_pointer {
+            // enable the physical storage buffer class on < SPV-1.5
+            Instruction::extension("SPV_KHR_physical_storage_buffer")
+                .to_words(&mut self.logical_layout.extensions);
+            Instruction::extension("GL_EXT_buffer_reference")
+                .to_words(&mut self.logical_layout.extensions);
+        }
+
         // write all const-expressions as constants
         self.constant_ids
             .resize(ir_module.const_expressions.len(), 0);
@@ -1980,7 +2000,11 @@ impl Writer {
                 .to_words(&mut self.logical_layout.capabilities);
         }
 
-        let addressing_model = spirv::AddressingModel::Logical;
+        let addressing_model = if has_physical_storage_pointer {
+            spirv::AddressingModel::PhysicalStorageBuffer64
+        } else {
+            spirv::AddressingModel::Logical
+        };
         let memory_model = spirv::MemoryModel::GLSL450;
         //self.check(addressing_model.required_capabilities())?;
         //self.check(memory_model.required_capabilities())?;

src/back/wgsl/writer.rs

@@ -1870,9 +1870,7 @@ const fn sampling_str(sampling: crate::Sampling) -> &'static str {
     }
 }
 
-const fn address_space_str(
-    space: crate::AddressSpace,
-) -> (Option<&'static str>, Option<&'static str>) {
+fn address_space_str(space: crate::AddressSpace) -> (Option<&'static str>, Option<String>) {
     use crate::AddressSpace as As;
 
     (
@@ -1881,11 +1879,14 @@ const fn address_space_str(
             As::Uniform => "uniform",
             As::Storage { access } => {
                 if access.contains(crate::StorageAccess::STORE) {
-                    return (Some("storage"), Some("read_write"));
+                    return (Some("storage"), Some("read_write".to_string()));
                 } else {
                     "storage"
                 }
             }
+            As::PhysicalStorage { alignment } => {
+                return (Some("buffer"), Some(alignment.to_string()))
+            }
             As::PushConstant => "push_constant",
             As::WorkGroup => "workgroup",
             As::Handle => return (None, None),

src/front/wgsl/parse/conv.rs

@@ -9,6 +9,11 @@ pub fn map_address_space(word: &str, span: Span) -> Result<crate::AddressSpace,
         "storage" => Ok(crate::AddressSpace::Storage {
             access: crate::StorageAccess::default(),
         }),
+        "buffer" => Ok(crate::AddressSpace::PhysicalStorage {
+            // Follow the GLSL spec and use 16 byte alignment as the default
+            // https://github.com/KhronosGroup/GLSL/blob/ad383324612d8493c78d52d0c903969fe1dddd0a/extensions/ext/GLSL_EXT_buffer_reference.txt#L139
+            alignment: 16,
+        }),
         "push_constant" => Ok(crate::AddressSpace::PushConstant),
         "function" => Ok(crate::AddressSpace::Function),
         _ => Err(Error::UnknownAddressSpace(span)),

src/front/wgsl/parse/mod.rs

@@ -1220,12 +1220,32 @@ impl Parser {
                 let mut space = conv::map_address_space(ident, span)?;
                 lexer.expect(Token::Separator(','))?;
                 let base = self.type_decl(lexer, ctx)?;
-                if let crate::AddressSpace::Storage { ref mut access } = space {
-                    *access = if lexer.skip(Token::Separator(',')) {
-                        lexer.next_storage_access()?
-                    } else {
-                        crate::StorageAccess::LOAD
-                    };
+                match space {
+                    crate::AddressSpace::Storage { ref mut access } => {
+                        *access = if lexer.skip(Token::Separator(',')) {
+                            lexer.next_storage_access()?
+                        } else {
+                            crate::StorageAccess::LOAD
+                        };
+                    }
+                    crate::AddressSpace::PhysicalStorage { ref mut alignment } => {
+                        *alignment = if lexer.skip(Token::Separator(',')) {
+                            // TODO: we should probably use the constant expression evaluator here
+                            (match lexer.next() {
+                                (Token::Number(Ok(Number::I32(num))), _span) => {
+                                    u32::try_from(num).map_err(|_| Error::ExpectedNonNegative(span))
+                                }
+                                (Token::Number(Err(e)), span) => Err(Error::BadNumber(span, e)),
+                                other => Err(Error::Unexpected(
+                                    other.1,
+                                    ExpectedToken::PrimaryExpression,
+                                )),
+                            })?
+                        } else {
+                            *alignment
+                        };
+                    }
+                    _ => {}
                 }
                 lexer.expect_generic_paren('>')?;
                 ast::Type::Pointer { base, space }

src/lib.rs

@@ -403,6 +403,12 @@ pub enum AddressSpace {
     Handle,
     /// Push constants.
     PushConstant,
+    /// Buffer device pointers, mutable. Only used for pointers, not allowed on GlobalVariables.
+    /// Pointers in this address space reference buffer memory directly.
+    /// Requires capability [`crate::valid::Capabilities::PHYSICAL_STORAGE_BUFFER_ADDRESSES`].
+    /// Required for Vulkan extension `VK_KHR_buffer_device_address`:
+    /// <https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VK_KHR_buffer_device_address.html>
+    PhysicalStorage { alignment: u32 },
 }
 
 /// Built-in inputs and outputs.

src/proc/layouter.rs

@@ -14,6 +14,7 @@ impl Alignment {
     pub const EIGHT: Self = Self(unsafe { NonZeroU32::new_unchecked(8) });
     pub const SIXTEEN: Self = Self(unsafe { NonZeroU32::new_unchecked(16) });
 
+    pub const BUFFER_POINTER: Self = Self::EIGHT;
     pub const MIN_UNIFORM: Self = Self::SIXTEEN;
 
     pub const fn new(n: u32) -> Option<Self> {
@@ -200,6 +201,17 @@ impl Layouter {
                         alignment: Alignment::from(rows) * alignment,
                     }
                 }
+                Ti::Pointer {
+                    space: crate::AddressSpace::PhysicalStorage { .. },
+                    ..
+                }
+                | Ti::ValuePointer {
+                    space: crate::AddressSpace::PhysicalStorage { .. },
+                    ..
+                } => TypeLayout {
+                    size,
+                    alignment: Alignment::BUFFER_POINTER,
+                },
                 Ti::Pointer { .. } | Ti::ValuePointer { .. } => TypeLayout {
                     size,
                     alignment: Alignment::ONE,

src/proc/mod.rs

@@ -162,6 +162,7 @@ impl crate::Literal {
 }
 
 pub const POINTER_SPAN: u32 = 4;
+pub const BUFFER_POINTER_SPAN: u32 = 8;
 
 impl super::TypeInner {
     pub const fn scalar_kind(&self) -> Option<super::ScalarKind> {
@@ -218,6 +219,14 @@ impl super::TypeInner {
                 rows,
                 width,
             } => Alignment::from(rows) * width as u32 * columns as u32,
+            Self::Pointer {
+                space: crate::AddressSpace::PhysicalStorage { .. },
+                ..
+            }
+            | Self::ValuePointer {
+                space: crate::AddressSpace::PhysicalStorage { .. },
+                ..
+            } => BUFFER_POINTER_SPAN,
             Self::Pointer { .. } | Self::ValuePointer { .. } => POINTER_SPAN,
             Self::Array {
                 base: _,
@@ -346,6 +355,7 @@ impl super::AddressSpace {
             | crate::AddressSpace::WorkGroup => Sa::LOAD | Sa::STORE,
             crate::AddressSpace::Uniform => Sa::LOAD,
             crate::AddressSpace::Storage { access } => access,
+            crate::AddressSpace::PhysicalStorage { .. } => Sa::LOAD,
             crate::AddressSpace::Handle => Sa::LOAD,
             crate::AddressSpace::PushConstant => Sa::LOAD,
         }

src/valid/analyzer.rs

@@ -576,6 +576,7 @@ impl FunctionInfo {
                     As::Uniform | As::PushConstant => true,
                     // storage data is only uniform when read-only
                     As::Storage { access } => !access.contains(crate::StorageAccess::STORE),
+                    As::PhysicalStorage { .. } => false,
                     As::Handle => false,
                 };
                 Uniformity {