naga/back/hlsl/
immediates.rs

1/*!
2Lowering of [`Immediate`] globals.
3
4WGSL lays out immediates like `storage` buffers, but HLSL can only read root constants
5through a `cbuffer`, which packs values differently: every nested struct, array and matrix
6starts on a new 16-byte register. Padding cannot fix this, since it only moves members later.
7
8Instead, we declare a flat array of `uint4` registers, and assemble the typed value in a
9`static` variable at the start of each entry point that uses it. All byte offsets are
10compile-time constants, much like [`storage`] loads from `ByteAddressBuffer`s:
11
12```ignore
13struct Inner { a: u32, b: u32 }
14struct Data { x: u32, inner: Inner, y: u32 }
15var<immediate> data: Data;
16```
17
18```ignore
19cbuffer data_block : register(b0) { uint4 data_raw[1]; };
20static Data data;
21
22void main() {
23    data = ConstructData(
24        asuint(data_raw[0].x),
25        ConstructInner(asuint(data_raw[0].y), asuint(data_raw[0].z)),
26        asuint(data_raw[0].w)
27    );
28    // ...
29}
30```
31
32[`Immediate`]: crate::AddressSpace::Immediate
33[`storage`]: super::storage
34*/
35
36use alloc::{format, string::String};
37use core::fmt::Write;
38
39use super::{help::WrappedConstructor, BackendResult, Error};
40use crate::{
41    back,
42    proc::{Alignment, NameKey},
43    Handle,
44};
45
46/// The number of 4-byte words in one register of the raw immediates `cbuffer`.
47const WORDS_PER_REGISTER: u32 = 4;
48
49/// Component names of a register, indexed by the word's position within it.
50const COMPONENTS: [char; WORDS_PER_REGISTER as usize] = ['x', 'y', 'z', 'w'];
51
52/// Returns an expression for the 4-byte word with index `word` in the raw immediates `raw`.
53fn immediate_word(raw: &str, word: u32) -> String {
54    let component = COMPONENTS[(word % WORDS_PER_REGISTER) as usize];
55    format!("{raw}[{}].{component}", word / WORDS_PER_REGISTER)
56}
57
58impl<W: Write> super::Writer<'_, W> {
59    /// Writes the declarations for an [`Immediate`] global.
60    ///
61    /// This declares the raw `cbuffer` that the root constants are bound to,
62    /// and the `static` variable that holds the typed value.
63    /// See the [module documentation](self) for background.
64    ///
65    /// [`Immediate`]: crate::AddressSpace::Immediate
66    pub(super) fn write_global_immediate(
67        &mut self,
68        module: &crate::Module,
69        handle: Handle<crate::GlobalVariable>,
70        global: &crate::GlobalVariable,
71    ) -> BackendResult {
72        let target = self
73            .options
74            .immediates_target
75            .as_ref()
76            .expect("No bind target was defined for the immediates block");
77        let (register, space) = (target.register, target.space);
78
79        let word_count = module.types[global.ty]
80            .inner
81            .size(module.to_ctx())
82            .div_ceil(4);
83        let register_count = word_count.div_ceil(WORDS_PER_REGISTER);
84
85        let name = self.names[&NameKey::GlobalVariable(handle)].clone();
86        let block_name = self.namer.call(&format!("{name}_block"));
87        let raw_name = self.namer.call(&format!("{name}_raw"));
88
89        write!(self.out, "cbuffer {block_name} : register(b{register}")?;
90        if space != 0 {
91            write!(self.out, ", space{space}")?;
92        }
93        writeln!(self.out, ") {{ uint4 {raw_name}[{register_count}]; }};")?;
94
95        write!(self.out, "static ")?;
96        self.write_type(module, global.ty)?;
97        writeln!(self.out, " {name};")?;
98
99        self.immediate_raw_names.insert(handle, raw_name);
100
101        Ok(())
102    }
103
104    /// Writes the wrapped constructor functions needed to assemble the
105    /// [`Immediate`] globals' values out of their raw words.
106    ///
107    /// [`Immediate`]: crate::AddressSpace::Immediate
108    pub(super) fn write_immediate_constructors(&mut self, module: &crate::Module) -> BackendResult {
109        for (_, global) in module.global_variables.iter() {
110            if global.space == crate::AddressSpace::Immediate {
111                self.write_wrapped_constructors_for_type(module, global.ty)?;
112            }
113        }
114        Ok(())
115    }
116
117    /// Writes the statements at the start of an entry point that assemble
118    /// the value of each [`Immediate`] global that the entry point uses.
119    ///
120    /// [`Immediate`]: crate::AddressSpace::Immediate
121    pub(super) fn write_immediates_initialization(
122        &mut self,
123        module: &crate::Module,
124        info: &crate::valid::FunctionInfo,
125    ) -> BackendResult {
126        for (handle, global) in module.global_variables.iter() {
127            if global.space != crate::AddressSpace::Immediate || info[handle].is_empty() {
128                continue;
129            }
130
131            let name = self.names[&NameKey::GlobalVariable(handle)].clone();
132            let raw = self.immediate_raw_names[&handle].clone();
133
134            write!(self.out, "{}{name} = ", back::INDENT)?;
135            self.write_immediate_load(module, &raw, global.ty, 0)?;
136            writeln!(self.out, ";")?;
137        }
138        Ok(())
139    }
140
141    /// Writes an expression that loads a value of type `ty` from the raw
142    /// immediate words in `raw`, starting `offset` bytes into the immediate data.
143    fn write_immediate_load(
144        &mut self,
145        module: &crate::Module,
146        raw: &str,
147        ty: Handle<crate::Type>,
148        offset: u32,
149    ) -> BackendResult {
150        match module.types[ty].inner {
151            crate::TypeInner::Scalar(scalar) => self.write_immediate_scalar(raw, scalar, offset),
152            crate::TypeInner::Vector { size, scalar } => {
153                self.write_immediate_vector(raw, size, scalar, offset)
154            }
155            crate::TypeInner::Matrix {
156                columns,
157                rows,
158                scalar,
159            } => {
160                write!(
161                    self.out,
162                    "{}{}x{}(",
163                    scalar.to_hlsl_str()?,
164                    columns as u8,
165                    rows as u8,
166                )?;
167
168                // Note: Matrices containing vec3s, due to padding, act like they contain vec4s.
169                let column_stride = Alignment::from(rows) * scalar.width as u32;
170                for column in 0..columns as u32 {
171                    if column != 0 {
172                        write!(self.out, ", ")?;
173                    }
174                    self.write_immediate_vector(
175                        raw,
176                        rows,
177                        scalar,
178                        offset + column * column_stride,
179                    )?;
180                }
181                write!(self.out, ")")?;
182                Ok(())
183            }
184            crate::TypeInner::Struct { ref members, .. } => {
185                let constructor = WrappedConstructor { ty };
186                self.write_wrapped_constructor_function_name(module, constructor)?;
187                write!(self.out, "(")?;
188                for (i, member) in members.iter().enumerate() {
189                    if i != 0 {
190                        write!(self.out, ", ")?;
191                    }
192                    self.write_immediate_load(module, raw, member.ty, offset + member.offset)?;
193                }
194                write!(self.out, ")")?;
195                Ok(())
196            }
197            ref other => Err(Error::Unimplemented(format!(
198                "immediate data of type {other:?}"
199            ))),
200        }
201    }
202
203    /// Writes an expression that loads a scalar from the raw immediate words in `raw`.
204    fn write_immediate_scalar(
205        &mut self,
206        raw: &str,
207        scalar: crate::Scalar,
208        offset: u32,
209    ) -> BackendResult {
210        use crate::ScalarKind as Kind;
211
212        let word = offset / 4;
213        match (scalar.kind, scalar.width) {
214            (Kind::Float | Kind::Sint | Kind::Uint, 4) => {
215                let cast = scalar.kind.to_hlsl_cast();
216                write!(self.out, "{cast}({})", immediate_word(raw, word))?;
217            }
218            // 16-bit values share a word with their neighbor.
219            (Kind::Float | Kind::Sint | Kind::Uint, 2) => {
220                let bits = format!("{} >> {}u", immediate_word(raw, word), (offset % 4) * 8);
221                if scalar.kind == Kind::Float {
222                    write!(self.out, "asfloat16(uint16_t({bits}))")?;
223                } else {
224                    write!(self.out, "{}({bits})", scalar.to_hlsl_str()?)?;
225                }
226            }
227            // 64-bit values occupy two words, with the low word first.
228            (Kind::Float | Kind::Sint | Kind::Uint, 8) => {
229                let low = immediate_word(raw, word);
230                let high = immediate_word(raw, word + 1);
231                let bits = format!("uint64_t({low}) | (uint64_t({high}) << 32u)");
232                match scalar.kind {
233                    Kind::Float => write!(self.out, "asdouble({low}, {high})")?,
234                    Kind::Sint => write!(self.out, "int64_t({bits})")?,
235                    _ => write!(self.out, "({bits})")?,
236                }
237            }
238            _ => return Err(Error::UnsupportedScalar(scalar)),
239        }
240        Ok(())
241    }
242
243    /// Writes an expression that loads a vector from the raw immediate words in `raw`.
244    ///
245    /// The components are loaded one at a time, which works wherever they happen to be
246    /// relative to register boundaries. The shader compiler merges them back into a
247    /// single constant buffer load when they share a register.
248    fn write_immediate_vector(
249        &mut self,
250        raw: &str,
251        size: crate::VectorSize,
252        scalar: crate::Scalar,
253        offset: u32,
254    ) -> BackendResult {
255        let count = size as u32;
256
257        write!(self.out, "{}{}(", scalar.to_hlsl_str()?, count)?;
258        for component in 0..count {
259            if component != 0 {
260                write!(self.out, ", ")?;
261            }
262            self.write_immediate_scalar(raw, scalar, offset + component * scalar.width as u32)?;
263        }
264        write!(self.out, ")")?;
265        Ok(())
266    }
267}
naga/back/hlsl/immediates.rs

naga/back/hlsl/
immediates.rs
