naga/back/msl/
ray.rs

1use alloc::{format, string::String};
2use core::fmt::Write;
3
4use crate::{
5    back::{
6        self,
7        msl::{
8            writer::{NameKeyExt, StatementContext, TypeContext, WrappedFunction},
9            BackendResult, Error, Writer, NAMESPACE,
10        },
11        INDENT,
12    },
13    Handle,
14};
15
16pub(super) const RT_NAMESPACE: &str = "metal::raytracing";
17
18/// The ray query type, needs to be a function so it can format the constants.
19pub(super) fn metal_intersector_ty() -> String {
20    format!("{RT_NAMESPACE}::intersection_query<{RT_NAMESPACE}::instancing, {RT_NAMESPACE}::triangle_data>")
21}
22
23pub(super) const INTERSECTION_FUNCTION_NAME: &str = "ray_query_get_intersection";
24pub(crate) const RAY_QUERY_TRACKER_VARIABLE_PREFIX: &str = "naga_query_init_tracker_for_";
25pub(crate) const RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX: &str = "naga_query_tmax_tracker_for_";
26
27impl<W: Write> Writer<W> {
28    fn write_not_finite(&mut self, expr: &str) -> BackendResult {
29        self.write_contains_flags(&format!("as_type<uint>({expr})"), 0x7f800000)
30    }
31
32    /// Checks whether `expr` does not have the bitpattern of IEEE f32 `NaN`.
33    ///
34    /// Note that this evaluates `expr` in the written code multiple times.
35    fn write_is_nan(&mut self, expr: &str) -> BackendResult {
36        write!(self.out, "(")?;
37        self.write_not_finite(expr)?;
38        write!(self.out, " && ((as_type<uint>({expr}) & 0x7fffff) != 0))")?;
39        Ok(())
40    }
41
42    fn write_contains_flags(&mut self, expr: &str, flags: u32) -> BackendResult {
43        write!(self.out, "(({expr} & {flags}) == {flags})")?;
44        Ok(())
45    }
46
47    /// Writes a function to get the current intersection from the ray query
48    ///
49    /// Like other backends, this is needed to have a single branch for constructing
50    /// the parts of the intersection that need to be checked whether they do or don't
51    /// hit.
52    pub(super) fn write_rq_get_intersection_function(
53        &mut self,
54        module: &crate::Module,
55        committed: bool,
56        options: &super::Options,
57    ) -> BackendResult {
58        let wrapped = WrappedFunction::RayQueryGetIntersection { committed };
59        if !self.wrapped_functions.insert(wrapped) {
60            return Ok(());
61        }
62
63        let ty = if committed { "committed" } else { "candidate" };
64        let intersection = TypeContext {
65            handle: module
66                .special_types
67                .ray_intersection
68                .expect("intersection ty should be there for intersection function"),
69            gctx: module.to_ctx(),
70            names: &self.names,
71            access: crate::StorageAccess::empty(),
72            first_time: false,
73        };
74        let mut base_level = back::Level(1);
75        writeln!(
76            self.out,
77            "{intersection} {INTERSECTION_FUNCTION_NAME}_{committed}({} intersector",
78            metal_intersector_ty()
79        )?;
80        if options.ray_query_initialization_tracking {
81            writeln!(self.out, ", uint intersector_tracker")?;
82        }
83        writeln!(self.out, ") {{")?;
84        // Initialize the intersection to its default values (which should be zero).
85        writeln!(
86            self.out,
87            "{base_level}{intersection} intersection = {intersection} {{}};"
88        )?;
89
90        if options.ray_query_initialization_tracking {
91            write!(self.out, "{base_level}if (")?;
92            if committed {
93                self.write_contains_flags(
94                    "intersector_tracker",
95                    back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
96                )?;
97            } else {
98                self.write_contains_flags(
99                    "intersector_tracker",
100                    back::RayQueryPoint::PROCEED.bits(),
101                )?;
102                write!(self.out, " && !")?;
103                self.write_contains_flags(
104                    "intersector_tracker",
105                    back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
106                )?;
107            }
108            writeln!(self.out, ") {{")?;
109            base_level = base_level.next();
110        }
111
112        writeln!(self.out, "{base_level}{RT_NAMESPACE}::intersection_type ty = intersector.get_{ty}_intersection_type();")?;
113        // If the ray hit a triangle, call all methods that require that and set the intersection type.
114        writeln!(
115            self.out,
116            "{base_level}if (ty == {RT_NAMESPACE}::intersection_type::triangle) {{"
117        )?;
118        writeln!(
119            self.out,
120            "{base_level}{INDENT}intersection.kind = {};",
121            crate::RayQueryIntersection::Triangle as u32
122        )?;
123        if !committed {
124            writeln!(
125                self.out,
126                "{base_level}{INDENT}intersection.t = intersector.get_candidate_triangle_distance();"
127            )?;
128        }
129        writeln!(self.out, "{base_level}{INDENT}intersection.barycentrics = intersector.get_{ty}_triangle_barycentric_coord();")?;
130        writeln!(
131            self.out,
132            "{base_level}{INDENT}intersection.front_face = intersector.is_{ty}_triangle_front_facing();"
133        )?;
134        // Otherwise, if the ray hit an AABB (called a bounding box in metal) set the intersection type
135        // (which depends on whether this is a committed or candidate intersection).
136        writeln!(
137            self.out,
138            "{base_level}}} else if (ty == {RT_NAMESPACE}::intersection_type::bounding_box) {{"
139        )?;
140        if committed {
141            writeln!(
142                self.out,
143                "{base_level}{INDENT}intersection.kind = {};",
144                crate::RayQueryIntersection::Generated as u32
145            )?;
146        } else {
147            writeln!(
148                self.out,
149                "{base_level}{INDENT}intersection.kind = {};",
150                crate::RayQueryIntersection::Aabb as u32
151            )?;
152        }
153        writeln!(self.out, "{base_level}}}")?;
154
155        // If the ray hit anything at all, call all methods that require that.
156        writeln!(
157            self.out,
158            "{base_level}if (ty != {RT_NAMESPACE}::intersection_type::none) {{"
159        )?;
160        if committed {
161            writeln!(
162                self.out,
163                "{base_level}{INDENT}intersection.t = intersector.get_committed_distance();"
164            )?;
165        }
166        writeln!(self.out, "{base_level}{INDENT}intersection.instance_custom_data = intersector.get_{ty}_user_instance_id();")?;
167        writeln!(
168            self.out,
169            "{base_level}{INDENT}intersection.instance_index = intersector.get_{ty}_instance_id();"
170        )?;
171        // Metal does not appear to support obtaining the intersection offset from a ray query.
172        //writeln!(self.out, "{level}{level}intersection.sbt_record_offset = intersector.get_{ty}_user_instance_id();")?;
173        writeln!(
174            self.out,
175            "{base_level}{INDENT}intersection.geometry_index = intersector.get_{ty}_geometry_id();"
176        )?;
177        writeln!(
178            self.out,
179            "{base_level}{INDENT}intersection.primitive_index = intersector.get_{ty}_primitive_id();"
180        )?;
181        writeln!(self.out, "{base_level}{INDENT}intersection.object_to_world = intersector.get_{ty}_object_to_world_transform();")?;
182        writeln!(self.out, "{base_level}{INDENT}intersection.world_to_object = intersector.get_{ty}_world_to_object_transform();")?;
183        writeln!(self.out, "{base_level}}}")?;
184
185        if options.ray_query_initialization_tracking {
186            writeln!(self.out, "{INDENT}}}")?;
187        }
188
189        writeln!(self.out, "{INDENT}return intersection;")?;
190        writeln!(self.out, "}}")?;
191
192        Ok(())
193    }
194
195    pub(super) fn write_ray_query_stmt(
196        &mut self,
197        level: back::Level,
198        context: &StatementContext,
199        query: Handle<crate::Expression>,
200        fun: &crate::RayQueryFunction,
201    ) -> BackendResult {
202        if context.expression.lang_version < (2, 4) {
203            return Err(Error::UnsupportedRayTracing);
204        }
205
206        // There are three possibilities for a ptr to be:
207        // 1. A variable
208        // 2. A function argument
209        // 3. part of a struct
210        //
211        // 2 and 3 are not possible, a ray query (in naga IR)
212        // is not allowed to be passed into a function, and
213        // all languages disallow it in a struct (you get fun results if
214        // you try it :) ).
215        //
216        // Therefore, the ray query expression must be a variable.
217        let crate::Expression::LocalVariable(query_var) =
218            context.expression.function.expressions[query]
219        else {
220            unreachable!()
221        };
222
223        let tracker_expr_name = format!(
224            "{RAY_QUERY_TRACKER_VARIABLE_PREFIX}{}",
225            self.names[&crate::proc::NameKey::local(context.expression.origin, query_var)]
226        );
227
228        let tmax_tracker_expr_name = format!(
229            "{RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX}{}",
230            self.names[&crate::proc::NameKey::local(context.expression.origin, query_var)]
231        );
232
233        // TODO: check for misuse.
234        match *fun {
235            crate::RayQueryFunction::Initialize {
236                acceleration_structure,
237                descriptor,
238            } => {
239                //TODO: how to deal with winding? Is it by default the same as the other APIs?
240
241                // Put everything in a block so that the variable names
242                // do not conflict with user variable names
243                writeln!(self.out, "{level}{{")?;
244
245                let inner_level = level.next();
246
247                let naga_ray_desc_ty = TypeContext {
248                    handle: context
249                        .expression
250                        .module
251                        .special_types
252                        .ray_desc
253                        .expect("ray desc is required as an argument so should be there"),
254                    gctx: context.expression.module.to_ctx(),
255                    names: &self.names,
256                    access: crate::StorageAccess::empty(),
257                    first_time: false,
258                };
259
260                write!(self.out, "{inner_level}{naga_ray_desc_ty} desc = ")?;
261                self.put_expression(descriptor, &context.expression, false)?;
262                writeln!(self.out, ";")?;
263
264                // Set up intersection parameters
265                writeln!(
266                    self.out,
267                    "{inner_level}{RT_NAMESPACE}::intersection_params params;"
268                )?;
269
270                {
271                    // Determine whether or not to cull opaque/non-opaques
272                    let f_opaque = back::RayFlag::CULL_OPAQUE.bits();
273                    let f_no_opaque = back::RayFlag::CULL_NO_OPAQUE.bits();
274                    writeln!(self.out, "{inner_level}{RT_NAMESPACE}::opacity_cull_mode cull_mode = 
275{inner_level}{INDENT}(desc.flags & {f_opaque}) != 0 ? {RT_NAMESPACE}::opacity_cull_mode::opaque : (
276{inner_level}{INDENT}{INDENT}(desc.flags & {f_no_opaque}) != 0 ? {RT_NAMESPACE}::opacity_cull_mode::non_opaque : {RT_NAMESPACE}::opacity_cull_mode::none
277{inner_level}{INDENT});")?;
278                    writeln!(
279                        self.out,
280                        "{inner_level}params.set_opacity_cull_mode(cull_mode);"
281                    )?;
282
283                    if context.expression.ray_query_initialization_tracking {
284                        writeln!(self.out, "{inner_level}bool force_opacity = cull_mode == {RT_NAMESPACE}::opacity_cull_mode::none;")?;
285                    }
286                }
287                {
288                    let mut current_level = inner_level;
289                    if context.expression.ray_query_initialization_tracking {
290                        writeln!(self.out, "{inner_level}if (force_opacity) {{")?;
291                        current_level = current_level.next();
292                    }
293                    // Determine whether to force a particular opacity
294                    let f_opaque = back::RayFlag::OPAQUE.bits();
295                    let f_no_opaque = back::RayFlag::NO_OPAQUE.bits();
296                    writeln!(self.out, "{current_level}params.force_opacity(
297{current_level}    (desc.flags & {f_opaque}) != 0 ? {RT_NAMESPACE}::forced_opacity::opaque : (
298{current_level}        (desc.flags & {f_no_opaque}) != 0 ? {RT_NAMESPACE}::forced_opacity::non_opaque : {RT_NAMESPACE}::forced_opacity::none
299{current_level}    )
300{current_level});")?;
301
302                    if context.expression.ray_query_initialization_tracking {
303                        writeln!(self.out, "{inner_level}}}")?;
304                    }
305                }
306                {
307                    let flag = back::RayFlag::TERMINATE_ON_FIRST_HIT.bits();
308                    writeln!(
309                        self.out,
310                        "{inner_level}params.accept_any_intersection((desc.flags & {flag}) != 0);"
311                    )?;
312                }
313
314                writeln!(
315                    self.out,
316                    "{inner_level}{RT_NAMESPACE}::ray ray = {RT_NAMESPACE}::ray(desc.origin, desc.dir, desc.tmin, desc.tmax);"
317                )?;
318
319                let mut init_level = inner_level;
320
321                // The `reset` function is virtually undocumented (many of the Metal ray tracing functions lack it), so to be safe,
322                // this assumes an invalid ray is UB (NOTE: invalid ray behaviour is defined for intersectors).
323                if context.expression.ray_query_initialization_tracking {
324                    write!(self.out, "{inner_level}bool invalid_nan_infs = ")?;
325                    // tmax needs special handling because it can be INF
326                    for (idx, &field_access) in [
327                        "origin.x", "origin.y", "origin.z", "dir.x", "dir.y", "dir.z", "tmin",
328                    ]
329                    .iter()
330                    .enumerate()
331                    {
332                        if idx != 0 {
333                            write!(self.out, " || ")?;
334                        }
335
336                        self.write_not_finite(&format!("desc.{field_access}"))?;
337                    }
338
339                    write!(self.out, " || ")?;
340                    self.write_is_nan("desc.tmax")?;
341                    writeln!(self.out, ";")?;
342
343                    // Metal also requires that tmax >= 0.0, but if tmax >= tmin and tmin >= 0.0, tmax must be >= 0.0
344                    writeln!(self.out, "{inner_level}bool invalid_t = (desc.tmin > desc.tmax) || (desc.tmin < 0.0);")?;
345                    // Metal requires that the length of the direction is not 0.0. This is the case only when all the
346                    // components are zero.
347                    //
348                    // Use absolute to cover signed zero.
349                    writeln!(self.out, "{inner_level}bool invalid_dir = {NAMESPACE}::all({NAMESPACE}::abs(desc.dir) == 0.0);")?;
350
351                    writeln!(
352                        self.out,
353                        "{inner_level}if (!(invalid_dir || invalid_t || invalid_nan_infs)) {{"
354                    )?;
355                    init_level = init_level.next();
356                }
357
358                write!(self.out, "{init_level}")?;
359                // A ray query can by initialized in metal by either using a "non-default constructor"
360                // or by calling reset. Ray queries cannot be assigned to in metal, so reset needs to
361                // be called.
362                self.put_expression(query, &context.expression, true)?;
363                write!(self.out, ".reset(ray,")?;
364                self.put_expression(acceleration_structure, &context.expression, true)?;
365                writeln!(self.out, ", desc.cull_mask, params);")?;
366                if context.expression.ray_query_initialization_tracking {
367                    // We don't set the initialization tracker to zero (uninitialized)
368                    // if the call fails. Resetting to uninitialized might be useful
369                    // for debugging, but for everything else it is just extra code.
370                    writeln!(
371                        self.out,
372                        "{init_level}{tracker_expr_name} = {};",
373                        back::RayQueryPoint::INITIALIZED.bits()
374                    )?;
375                    writeln!(
376                        self.out,
377                        "{init_level}{tmax_tracker_expr_name} = desc.tmax;"
378                    )?;
379                    writeln!(self.out, "{inner_level}}}")?;
380                }
381                writeln!(self.out, "{level}}}")?;
382            }
383            crate::RayQueryFunction::Proceed { result } => {
384                let mut current_level = level;
385                write!(self.out, "{current_level}")?;
386                let name = self.write_result_assignment(result, &context.expression)?;
387
388                writeln!(self.out, "false;")?;
389
390                if context.expression.ray_query_initialization_tracking {
391                    write!(self.out, "{level}if (")?;
392                    self.write_contains_flags(
393                        &tracker_expr_name,
394                        back::RayQueryPoint::INITIALIZED.bits(),
395                    )?;
396                    write!(self.out, " && !")?;
397                    self.write_contains_flags(
398                        &tracker_expr_name,
399                        back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
400                    )?;
401                    write!(self.out, ")")?;
402                    writeln!(self.out, " {{")?;
403                    current_level = current_level.next();
404                }
405                write!(self.out, "{current_level}{name} = ")?;
406                self.put_expression(query, &context.expression, true)?;
407                writeln!(self.out, ".next();")?;
408                if context.expression.ray_query_initialization_tracking {
409                    writeln!(self.out, "{current_level}{tracker_expr_name} = {tracker_expr_name} | ({name} ? {}: {});", back::RayQueryPoint::PROCEED.bits(), (back::RayQueryPoint::PROCEED | back::RayQueryPoint::FINISHED_TRAVERSAL).bits())?;
410                    writeln!(self.out, "{level}}}")?;
411                }
412            }
413            crate::RayQueryFunction::GenerateIntersection { hit_t } => {
414                let mut current_level = level;
415                if context.expression.ray_query_initialization_tracking {
416                    write!(self.out, "{level}if (")?;
417                    self.write_contains_flags(
418                        &tracker_expr_name,
419                        back::RayQueryPoint::PROCEED.bits(),
420                    )?;
421                    write!(self.out, " && !")?;
422                    self.write_contains_flags(
423                        &tracker_expr_name,
424                        back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
425                    )?;
426                    write!(self.out, ")")?;
427                } else {
428                    // For readability
429                    write!(self.out, "{level}")?;
430                }
431                writeln!(self.out, "{{")?;
432                current_level = current_level.next();
433                write!(self.out, "{current_level}float t = ")?;
434                self.put_expression(hit_t, &context.expression, true)?;
435                writeln!(self.out, ";")?;
436                if context.expression.ray_query_initialization_tracking {
437                    write!(
438                        self.out,
439                        "{current_level}float current_max_t = {tmax_tracker_expr_name};
440{current_level}if ("
441                    )?;
442                    self.put_expression(query, &context.expression, true)?;
443                    write!(self.out, ".get_committed_intersection_type() != {RT_NAMESPACE}::intersection_type::none) {{
444{current_level}{INDENT}current_max_t = ")?;
445                    self.put_expression(query, &context.expression, true)?;
446                    write!(
447                        self.out,
448                        ".get_committed_distance();
449{current_level}}}
450{current_level}if ("
451                    )?;
452                    self.put_expression(query, &context.expression, true)?;
453                    write!(self.out, ".get_candidate_intersection_type() == {RT_NAMESPACE}::intersection_type::bounding_box && (")?;
454                    self.put_expression(query, &context.expression, true)?;
455                    write!(self.out, ".get_ray_min_distance()")?;
456                    writeln!(self.out, " <= t) && (t <= current_max_t)) {{")?;
457                    current_level = current_level.next();
458                }
459                write!(self.out, "{current_level}")?;
460                self.put_expression(query, &context.expression, true)?;
461                writeln!(self.out, ".commit_bounding_box_intersection(t);")?;
462                if context.expression.ray_query_initialization_tracking {
463                    writeln!(self.out, "{level}{INDENT}}}")?;
464                }
465                writeln!(self.out, "{level}}}")?;
466            }
467            crate::RayQueryFunction::ConfirmIntersection => {
468                let mut current_level = level;
469                if context.expression.ray_query_initialization_tracking {
470                    write!(self.out, "{level}if (")?;
471                    self.write_contains_flags(
472                        &tracker_expr_name,
473                        back::RayQueryPoint::PROCEED.bits(),
474                    )?;
475                    write!(self.out, " && !")?;
476                    self.write_contains_flags(
477                        &tracker_expr_name,
478                        back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
479                    )?;
480                    writeln!(self.out, ") {{")?;
481                    current_level = current_level.next();
482                    write!(self.out, "{current_level}if (")?;
483                    self.put_expression(query, &context.expression, true)?;
484                    writeln!(self.out, ".get_candidate_intersection_type() == {RT_NAMESPACE}::intersection_type::triangle) {{")?;
485                }
486                write!(self.out, "{level}")?;
487                self.put_expression(query, &context.expression, true)?;
488                writeln!(self.out, ".commit_triangle_intersection();")?;
489                if context.expression.ray_query_initialization_tracking {
490                    writeln!(
491                        self.out,
492                        "{level}{INDENT}}}
493{level}}}"
494                    )?;
495                }
496            }
497            crate::RayQueryFunction::Terminate => {
498                let mut current_level = level;
499                if context.expression.ray_query_initialization_tracking {
500                    write!(self.out, "{level}if (")?;
501                    self.write_contains_flags(
502                        &tracker_expr_name,
503                        back::RayQueryPoint::PROCEED.bits(),
504                    )?;
505                    write!(self.out, " && !")?;
506                    self.write_contains_flags(
507                        &tracker_expr_name,
508                        back::RayQueryPoint::FINISHED_TRAVERSAL.bits(),
509                    )?;
510                    writeln!(self.out, ") {{")?;
511                    current_level = current_level.next();
512                }
513                write!(self.out, "{current_level}")?;
514                self.put_expression(query, &context.expression, true)?;
515                // Terminate appears to map to abort in spirv-cross, but metal only documents
516                // the existence of this method, not what it does.
517                writeln!(self.out, ".abort();")?;
518                // To get the committed intersection, an extra proceed must occur as specified in
519                // the API docs.
520                if context.expression.ray_query_initialization_tracking {
521                    writeln!(self.out, "{level}}}")?;
522                }
523            }
524            // Reset the initialization tracker to have no flags set which is equivalent to having
525            // a new variable as nothing else should be used if the ray query is marked as
526            // uninitialized.
527            crate::RayQueryFunction::Begin => {
528                if context.expression.ray_query_initialization_tracking {
529                    writeln!(self.out, "{level}{tracker_expr_name} = 0u;")?;
530                }
531            }
532        }
533
534        Ok(())
535    }
536}
naga/back/msl/ray.rs

naga/back/msl/
ray.rs
