wgsl_types/
mem.rs

1//! Memory layout utilities.
2
3use itertools::Itertools;
4
5use crate::{
6    f16,
7    inst::{
8        ArrayInstance, AtomicInstance, Instance, LiteralInstance, MatInstance, StructInstance,
9        VecInstance,
10    },
11    ty::{Ty, Type},
12};
13
14impl Instance {
15    /// Memory representation of host-shareable instances.
16    ///
17    /// Returns `None` if the type is not host-shareable.
18    pub fn to_buffer(&self) -> Option<Vec<u8>> {
19        match self {
20            Instance::Literal(l) => l.to_buffer(),
21            Instance::Struct(s) => s.to_buffer(),
22            Instance::Array(a) => a.to_buffer(),
23            Instance::Vec(v) => v.to_buffer(),
24            Instance::Mat(m) => m.to_buffer(),
25            Instance::Ptr(_) => None,
26            Instance::Ref(_) => None,
27            Instance::Atomic(a) => a.inner().to_buffer(),
28            Instance::Opaque(_) => None,
29        }
30    }
31
32    /// Load an instance from a byte buffer.
33    ///
34    /// Returns `None` if the type is not host-shareable, or if the buffer is too small.
35    /// The buffer can be larger than the type; extra bytes will be ignored.
36    pub fn from_buffer(buf: &[u8], ty: &Type) -> Option<Self> {
37        match ty {
38            Type::Bool => None,
39            Type::AbstractInt => None,
40            Type::AbstractFloat => None,
41            Type::I32 => buf
42                .get(..4)?
43                .try_into()
44                .ok()
45                .map(|buf| LiteralInstance::I32(i32::from_le_bytes(buf)).into()),
46            Type::U32 => buf
47                .get(..4)?
48                .try_into()
49                .ok()
50                .map(|buf| LiteralInstance::U32(u32::from_le_bytes(buf)).into()),
51            Type::F32 => buf
52                .get(..4)?
53                .try_into()
54                .ok()
55                .map(|buf| LiteralInstance::F32(f32::from_le_bytes(buf)).into()),
56            Type::F16 => buf
57                .get(..2)?
58                .try_into()
59                .ok()
60                .map(|buf| LiteralInstance::F16(f16::from_le_bytes(buf)).into()),
61            #[cfg(feature = "naga-ext")]
62            Type::I64 => buf
63                .get(..8)?
64                .try_into()
65                .ok()
66                .map(|buf| LiteralInstance::I64(i64::from_le_bytes(buf)).into()),
67            #[cfg(feature = "naga-ext")]
68            Type::U64 => buf
69                .get(..8)?
70                .try_into()
71                .ok()
72                .map(|buf| LiteralInstance::U64(u64::from_le_bytes(buf)).into()),
73            #[cfg(feature = "naga-ext")]
74            Type::F64 => buf
75                .get(..8)?
76                .try_into()
77                .ok()
78                .map(|buf| LiteralInstance::F64(f64::from_le_bytes(buf)).into()),
79            Type::Struct(s) => {
80                let mut offset = 0;
81                let members = s
82                    .members
83                    .iter()
84                    .map(|m| {
85                        // handle the specific case of runtime-sized arrays.
86                        // they can only be the last member of a struct.
87                        let inst = if let Type::Array(_, None) = &m.ty {
88                            let buf = buf.get(offset as usize..)?;
89                            Instance::from_buffer(buf, &m.ty)?
90                        } else {
91                            // TODO: handle errors, check valid size...
92                            let size = m.size.or_else(|| m.ty.size_of())?;
93                            let align = m.align.or_else(|| m.ty.align_of())?;
94                            offset = round_up(align, offset);
95                            let buf = buf.get(offset as usize..(offset + size) as usize)?;
96                            offset += size;
97                            Instance::from_buffer(buf, &m.ty)?
98                        };
99                        Some(inst)
100                    })
101                    .collect::<Option<Vec<_>>>()?;
102                Some(StructInstance::new((**s).clone(), members).into())
103            }
104            Type::Array(ty, Some(n)) => {
105                let mut offset = 0;
106                let size = ty.size_of()?;
107                let stride = round_up(ty.align_of()?, size);
108                let mut comps = Vec::new();
109                while comps.len() != *n {
110                    let buf = buf.get(offset as usize..(offset + size) as usize)?;
111                    offset += stride;
112                    let inst = Instance::from_buffer(buf, ty)?;
113                    comps.push(inst);
114                }
115                Some(ArrayInstance::new(comps, false).into())
116            }
117            Type::Array(ty, None) => {
118                let mut offset = 0;
119                let size = ty.size_of()?;
120                let stride = round_up(ty.align_of()?, size);
121                let n = buf.len() as u32 / stride;
122                if n == 0 {
123                    // arrays must not be empty
124                    return None;
125                }
126                let comps = (0..n)
127                    .map(|_| {
128                        let buf = buf.get(offset as usize..(offset + size) as usize)?;
129                        offset += stride;
130                        Instance::from_buffer(buf, ty)
131                    })
132                    .collect::<Option<_>>()?;
133                Some(ArrayInstance::new(comps, true).into())
134            }
135            #[cfg(feature = "naga-ext")]
136            Type::BindingArray(_, _) => None,
137            Type::Vec(n, ty) => {
138                let mut offset = 0;
139                let size = ty.size_of()?;
140                let comps = (0..*n)
141                    .map(|_| {
142                        let buf = buf.get(offset as usize..(offset + size) as usize)?;
143                        offset += size;
144                        Instance::from_buffer(buf, ty)
145                    })
146                    .collect::<Option<Vec<_>>>()?;
147                Some(VecInstance::new(comps).into())
148            }
149            Type::Mat(c, r, ty) => {
150                let mut offset = 0;
151                let col_ty = Type::Vec(*r, ty.clone());
152                let col_size = col_ty.size_of()?;
153                let col_off = round_up(col_ty.align_of()?, col_size);
154                let cols = (0..*c)
155                    .map(|_| {
156                        let buf = buf.get(offset as usize..(offset + col_size) as usize)?;
157                        offset += col_off;
158                        Instance::from_buffer(buf, &col_ty)
159                    })
160                    .collect::<Option<Vec<_>>>()?;
161                Some(MatInstance::from_cols(cols).into())
162            }
163            Type::Atomic(ty) => {
164                let buf = buf.get(..4)?.try_into().ok()?;
165                let inst = match &**ty {
166                    Type::I32 => LiteralInstance::I32(i32::from_le_bytes(buf)).into(),
167                    Type::U32 => LiteralInstance::U32(u32::from_le_bytes(buf)).into(),
168                    _ => unreachable!("atomic type must be u32 or i32"),
169                };
170                Some(AtomicInstance::new(inst).into())
171            }
172            Type::Ptr(_, _, _)
173            | Type::Ref(_, _, _)
174            | Type::Texture(_)
175            | Type::Sampler(_)
176            | Type::Unknown => None,
177            #[cfg(feature = "naga-ext")]
178            Type::RayQuery(_) | Type::AccelerationStructure(_) => None,
179        }
180    }
181}
182
183impl LiteralInstance {
184    /// Memory representation of host-shareable instances.
185    ///
186    /// Returns `None` if the type is not host-shareable
187    fn to_buffer(self) -> Option<Vec<u8>> {
188        match self {
189            LiteralInstance::Bool(_) => None,
190            LiteralInstance::AbstractInt(_) => None,
191            LiteralInstance::AbstractFloat(_) => None,
192            LiteralInstance::I32(n) => Some(n.to_le_bytes().to_vec()),
193            LiteralInstance::U32(n) => Some(n.to_le_bytes().to_vec()),
194            LiteralInstance::F32(n) => Some(n.to_le_bytes().to_vec()),
195            LiteralInstance::F16(n) => Some(n.to_le_bytes().to_vec()),
196            #[cfg(feature = "naga-ext")]
197            LiteralInstance::I64(n) => Some(n.to_le_bytes().to_vec()),
198            #[cfg(feature = "naga-ext")]
199            LiteralInstance::U64(n) => Some(n.to_le_bytes().to_vec()),
200            #[cfg(feature = "naga-ext")]
201            LiteralInstance::F64(n) => Some(n.to_le_bytes().to_vec()),
202        }
203    }
204}
205
206// TODO: layout
207impl StructInstance {
208    /// Memory representation of host-shareable instances.
209    ///
210    /// Returns `None` if the type is not host-shareable.
211    fn to_buffer(&self) -> Option<Vec<u8>> {
212        let mut buf = Vec::new();
213        for (i, (inst, m)) in self.members.iter().zip(&self.ty.members).enumerate() {
214            let len = buf.len() as u32;
215            let size = m.size.or_else(|| m.ty.min_size_of())?;
216
217            // handle runtime-size arrays as last struct member
218            let size = match inst {
219                Instance::Array(a) if a.runtime_sized => {
220                    (i == self.members.len() - 1).then(|| a.n() as u32 * size)
221                }
222                _ => Some(size),
223            }?;
224
225            let align = m.align.or_else(|| m.ty.align_of())?;
226            let off = round_up(align, len);
227            if off > len {
228                buf.extend((len..off).map(|_| 0));
229            }
230            let mut bytes = inst.to_buffer()?;
231            let bytes_len = bytes.len() as u32;
232            if size > bytes_len {
233                bytes.extend((bytes_len..size).map(|_| 0));
234            }
235            buf.extend(bytes);
236        }
237        Some(buf)
238    }
239}
240
241impl ArrayInstance {
242    /// Memory representation of host-shareable instances.
243    ///
244    /// Returns `None` if the type is not host-shareable.
245    fn to_buffer(&self) -> Option<Vec<u8>> {
246        let mut buf = Vec::new();
247        let ty = self.inner_ty();
248        let size = ty.size_of()?;
249        let stride = round_up(ty.align_of()?, size);
250        for c in self.iter() {
251            buf.extend(c.to_buffer()?);
252            if stride > size {
253                buf.extend((size..stride).map(|_| 0))
254            }
255        }
256        Some(buf)
257    }
258}
259
260impl VecInstance {
261    /// Memory representation of host-shareable instances.
262    ///
263    /// Returns `None` if the type is not host-shareable.
264    fn to_buffer(&self) -> Option<Vec<u8>> {
265        Some(
266            self.iter()
267                .flat_map(|v| v.to_buffer().unwrap(/* SAFETY: vector elements must be host-shareable */).into_iter())
268                .collect_vec(),
269        )
270    }
271}
272
273impl MatInstance {
274    /// Memory representation of host-shareable instances.
275    ///
276    /// Returns `None` if the type is not host-shareable.
277    fn to_buffer(&self) -> Option<Vec<u8>> {
278        Some(
279            self.iter_cols()
280                .flat_map(|v| {
281                    // SAFETY: vector elements must be host-shareable
282                    let mut v_buf = v.to_buffer().unwrap();
283                    let len = v_buf.len() as u32;
284                    let align = v.ty().align_of().unwrap();
285                    if len < align {
286                        v_buf.extend((len..align).map(|_| 0));
287                    }
288                    v_buf.into_iter()
289                })
290                .collect_vec(),
291        )
292    }
293}
294
295fn round_up(align: u32, size: u32) -> u32 {
296    size.div_ceil(align) * align
297}
298
299impl Type {
300    /// Compute the size of the type.
301    ///
302    /// Return `None` if the type is not host-shareable, or if it contains a
303    /// runtime-sized array. See [`Type::min_size_of`] for runtime-sized arrays.
304    ///
305    /// Reference: <https://www.w3.org/TR/WGSL/#alignment-and-size>
306    pub fn size_of(&self) -> Option<u32> {
307        match self {
308            Type::Bool => Some(4),
309            Type::AbstractInt => None,
310            Type::AbstractFloat => None,
311            Type::I32 => Some(4),
312            Type::U32 => Some(4),
313            Type::F32 => Some(4),
314            Type::F16 => Some(2),
315            #[cfg(feature = "naga-ext")]
316            Type::I64 => Some(8),
317            #[cfg(feature = "naga-ext")]
318            Type::U64 => Some(8),
319            #[cfg(feature = "naga-ext")]
320            Type::F64 => Some(8),
321            Type::Struct(s) => {
322                let past_last_mem = s
323                    .members
324                    .iter()
325                    .map(|m| {
326                        // TODO: handle errors, check valid size...
327                        let size = m.size.or_else(|| m.ty.size_of())?;
328                        let align = m.align.or_else(|| m.ty.align_of())?;
329                        Some((size, align))
330                    })
331                    .try_fold(0, |offset, mem| {
332                        let (size, align) = mem?;
333                        Some(round_up(align, offset) + size)
334                    })?;
335                Some(round_up(self.align_of()?, past_last_mem))
336            }
337            Type::Array(ty, Some(n)) => {
338                let (size, align) = (ty.size_of()?, ty.align_of()?);
339                Some(*n as u32 * round_up(align, size))
340            }
341            Type::Array(_, None) => None,
342            #[cfg(feature = "naga-ext")]
343            Type::BindingArray(_, _) => None,
344            Type::Vec(n, ty) => {
345                let size = ty.size_of()?;
346                Some(*n as u32 * size)
347            }
348            Type::Mat(c, r, ty) => {
349                let align = Type::Vec(*r, ty.clone()).align_of()?;
350                Some(*c as u32 * align)
351            }
352            Type::Atomic(_) => Some(4),
353            Type::Ptr(_, _, _)
354            | Type::Ref(_, _, _)
355            | Type::Texture(_)
356            | Type::Sampler(_)
357            | Type::Unknown => None,
358            #[cfg(feature = "naga-ext")]
359            Type::RayQuery(_) | Type::AccelerationStructure(_) => None,
360        }
361    }
362
363    /// Variant of [`Type::size_of`], but for runtime-sized arrays, it returns the minimum
364    /// size of the array, i.e. the size of an array with one element.
365    pub fn min_size_of(&self) -> Option<u32> {
366        match self {
367            Type::Array(ty, None) => Some(round_up(ty.align_of()?, ty.size_of()?)),
368            // TODO: should we also compute for structs containing a runtime-sized array?
369            // This function is only used once, anyway.
370            _ => self.size_of(),
371        }
372    }
373
374    /// Compute the alignment of the type.
375    ///
376    /// Return `None` if the type is not host-shareable.
377    ///
378    /// Reference: <https://www.w3.org/TR/WGSL/#alignment-and-size>
379    pub fn align_of(&self) -> Option<u32> {
380        match self {
381            Type::Bool => Some(4),
382            Type::AbstractInt => None,
383            Type::AbstractFloat => None,
384            Type::I32 => Some(4),
385            Type::U32 => Some(4),
386            Type::F32 => Some(4),
387            Type::F16 => Some(2),
388            #[cfg(feature = "naga-ext")]
389            Type::I64 => Some(8),
390            #[cfg(feature = "naga-ext")]
391            Type::U64 => Some(8),
392            #[cfg(feature = "naga-ext")]
393            Type::F64 => Some(8),
394            Type::Struct(s) => s
395                .members
396                .iter()
397                // TODO: check valid align attr
398                .map(|m| m.align.or_else(|| m.ty.align_of()))
399                .try_fold(0, |a, b| Some(a.max(b?))),
400            Type::Array(ty, _) => ty.align_of(),
401            #[cfg(feature = "naga-ext")]
402            Type::BindingArray(_, _) => None,
403            Type::Vec(n, ty) => {
404                if *n == 3 {
405                    match **ty {
406                        Type::I32 | Type::U32 | Type::F32 => Some(16),
407                        Type::F16 => Some(8),
408                        _ => None,
409                    }
410                } else {
411                    self.size_of()
412                }
413            }
414            Type::Mat(_, r, ty) => Type::Vec(*r, ty.clone()).align_of(),
415            Type::Atomic(_) => Some(4),
416            Type::Ptr(_, _, _)
417            | Type::Ref(_, _, _)
418            | Type::Texture(_)
419            | Type::Sampler(_)
420            | Type::Unknown => None,
421            #[cfg(feature = "naga-ext")]
422            Type::RayQuery(_) | Type::AccelerationStructure(_) => None,
423        }
424    }
425}
wgsl_types/mem.rs

wgsl_types/
mem.rs
