1use alloc::{
2 format,
3 string::{String, ToString},
4 vec,
5 vec::Vec,
6};
7use core::{
8 cmp::Ordering,
9 fmt::{Display, Error as FmtError, Formatter, Write},
10 iter,
11};
12use num_traits::real::Real as _;
13
14use half::f16;
15
16use super::{
17 ray::RT_NAMESPACE, sampler as sm, Error, LocationMode, Options, PipelineOptions,
18 TranslationInfo, NAMESPACE, WRAPPED_ARRAY_FIELD,
19};
20use crate::{
21 arena::{Handle, HandleSet},
22 back::{
23 self, get_entry_points,
24 msl::{mesh_shader::NestedFunctionInfo, BackendResult, EntryPointArgument},
25 Baked,
26 },
27 common,
28 proc::{
29 self, concrete_int_scalars,
30 index::{self, BoundsCheck},
31 ExternalTextureNameKey, NameKey, TypeResolution,
32 },
33 valid, FastHashMap, FastHashSet,
34};
35
36const ATOMIC_REFERENCE: &str = "&";
40
41pub(crate) const ATOMIC_COMP_EXCH_FUNCTION: &str = "naga_atomic_compare_exchange_weak_explicit";
42pub(crate) const MODF_FUNCTION: &str = "naga_modf";
43pub(crate) const FREXP_FUNCTION: &str = "naga_frexp";
44pub(crate) const ABS_FUNCTION: &str = "naga_abs";
45pub(crate) const DIV_FUNCTION: &str = "naga_div";
46pub(crate) const DOT_FUNCTION_PREFIX: &str = "naga_dot";
47pub(crate) const MOD_FUNCTION: &str = "naga_mod";
48pub(crate) const NEG_FUNCTION: &str = "naga_neg";
49pub(crate) const F2I32_FUNCTION: &str = "naga_f2i32";
50pub(crate) const F2U32_FUNCTION: &str = "naga_f2u32";
51pub(crate) const F2I64_FUNCTION: &str = "naga_f2i64";
52pub(crate) const F2U64_FUNCTION: &str = "naga_f2u64";
53pub(crate) const IMAGE_LOAD_EXTERNAL_FUNCTION: &str = "nagaTextureLoadExternal";
54pub(crate) const IMAGE_SIZE_EXTERNAL_FUNCTION: &str = "nagaTextureDimensionsExternal";
55pub(crate) const IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION: &str =
56 "nagaTextureSampleBaseClampToEdge";
57pub(crate) const ARGUMENT_BUFFER_WRAPPER_STRUCT: &str = "NagaArgumentBufferWrapper";
65pub(crate) const EXTERNAL_TEXTURE_WRAPPER_STRUCT: &str = "NagaExternalTextureWrapper";
70pub(crate) const COOPERATIVE_LOAD_FUNCTION: &str = "NagaCooperativeLoad";
71pub(crate) const COOPERATIVE_MULTIPLY_ADD_FUNCTION: &str = "NagaCooperativeMultiplyAdd";
72
73fn put_numeric_type(
82 out: &mut impl Write,
83 scalar: crate::Scalar,
84 sizes: &[crate::VectorSize],
85) -> Result<(), FmtError> {
86 match (scalar, sizes) {
87 (scalar, &[]) => {
88 write!(out, "{}", scalar.to_msl_name())
89 }
90 (scalar, &[rows]) => {
91 write!(
92 out,
93 "{}::{}{}",
94 NAMESPACE,
95 scalar.to_msl_name(),
96 common::vector_size_str(rows)
97 )
98 }
99 (scalar, &[rows, columns]) => {
100 write!(
101 out,
102 "{}::{}{}x{}",
103 NAMESPACE,
104 scalar.to_msl_name(),
105 common::vector_size_str(columns),
106 common::vector_size_str(rows)
107 )
108 }
109 (_, _) => Ok(()), }
111}
112
113const fn scalar_is_int(scalar: crate::Scalar) -> bool {
114 use crate::ScalarKind::*;
115 match scalar.kind {
116 Sint | Uint | AbstractInt | Bool => true,
117 Float | AbstractFloat => false,
118 }
119}
120
121const CLAMPED_LOD_LOAD_PREFIX: &str = "clamped_lod_e";
123
124const REINTERPRET_PREFIX: &str = "reinterpreted_";
126
127struct ClampedLod(Handle<crate::Expression>);
133
134impl Display for ClampedLod {
135 fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
136 self.0.write_prefixed(f, CLAMPED_LOD_LOAD_PREFIX)
137 }
138}
139
140struct ArraySizeMember(Handle<crate::GlobalVariable>);
155
156impl Display for ArraySizeMember {
157 fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
158 self.0.write_prefixed(f, "size")
159 }
160}
161
162#[derive(Clone, Copy)]
167struct Reinterpreted<'a> {
168 target_type: &'a str,
169 orig: Handle<crate::Expression>,
170}
171
172impl<'a> Reinterpreted<'a> {
173 const fn new(target_type: &'a str, orig: Handle<crate::Expression>) -> Self {
174 Self { target_type, orig }
175 }
176}
177
178impl Display for Reinterpreted<'_> {
179 fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
180 f.write_str(REINTERPRET_PREFIX)?;
181 f.write_str(self.target_type)?;
182 self.orig.write_prefixed(f, "_e")
183 }
184}
185
186pub(super) struct TypeContext<'a> {
187 pub handle: Handle<crate::Type>,
188 pub gctx: proc::GlobalCtx<'a>,
189 pub names: &'a FastHashMap<NameKey, String>,
190 pub access: crate::StorageAccess,
191 pub first_time: bool,
192}
193
194impl TypeContext<'_> {
195 fn scalar(&self) -> Option<crate::Scalar> {
196 let ty = &self.gctx.types[self.handle];
197 ty.inner.scalar()
198 }
199
200 fn vector_size(&self) -> Option<crate::VectorSize> {
201 let ty = &self.gctx.types[self.handle];
202 match ty.inner {
203 crate::TypeInner::Vector { size, .. } => Some(size),
204 _ => None,
205 }
206 }
207
208 fn unwrap_array(self) -> Self {
209 match self.gctx.types[self.handle].inner {
210 crate::TypeInner::Array { base, .. } => Self {
211 handle: base,
212 ..self
213 },
214 _ => self,
215 }
216 }
217}
218
219impl Display for TypeContext<'_> {
220 fn fmt(&self, out: &mut Formatter<'_>) -> Result<(), FmtError> {
221 let ty = &self.gctx.types[self.handle];
222 if ty.needs_alias() && !self.first_time {
223 let name = &self.names[&NameKey::Type(self.handle)];
224 return write!(out, "{name}");
225 }
226
227 match ty.inner {
228 crate::TypeInner::Scalar(scalar) => put_numeric_type(out, scalar, &[]),
229 crate::TypeInner::Atomic(scalar) => {
230 write!(out, "{}::atomic_{}", NAMESPACE, scalar.to_msl_name())
231 }
232 crate::TypeInner::Vector { size, scalar } => put_numeric_type(out, scalar, &[size]),
233 crate::TypeInner::Matrix {
234 columns,
235 rows,
236 scalar,
237 } => put_numeric_type(out, scalar, &[rows, columns]),
238 crate::TypeInner::CooperativeMatrix {
240 columns,
241 rows,
242 scalar,
243 role: _,
244 } => {
245 write!(
246 out,
247 "{NAMESPACE}::simdgroup_{}{}x{}",
248 scalar.to_msl_name(),
249 columns as u32,
250 rows as u32,
251 )
252 }
253 crate::TypeInner::Pointer { base, space } => {
254 let sub = Self {
255 handle: base,
256 first_time: false,
257 ..*self
258 };
259 let space_name = match space.to_msl_name() {
260 Some(name) => name,
261 None => return Ok(()),
262 };
263 write!(out, "{space_name} {sub}&")
264 }
265 crate::TypeInner::ValuePointer {
266 size,
267 scalar,
268 space,
269 } => {
270 match space.to_msl_name() {
271 Some(name) => write!(out, "{name} ")?,
272 None => return Ok(()),
273 };
274 match size {
275 Some(rows) => put_numeric_type(out, scalar, &[rows])?,
276 None => put_numeric_type(out, scalar, &[])?,
277 };
278
279 write!(out, "&")
280 }
281 crate::TypeInner::Array { base, .. } => {
282 let sub = Self {
283 handle: base,
284 first_time: false,
285 ..*self
286 };
287 write!(out, "{sub}")
290 }
291 crate::TypeInner::Struct { .. } => unreachable!(),
292 crate::TypeInner::Image {
293 dim,
294 arrayed,
295 class,
296 } => {
297 let dim_str = match dim {
298 crate::ImageDimension::D1 => "1d",
299 crate::ImageDimension::D2 => "2d",
300 crate::ImageDimension::D3 => "3d",
301 crate::ImageDimension::Cube => "cube",
302 };
303 let (texture_str, msaa_str, scalar, access) = match class {
304 crate::ImageClass::Sampled { kind, multi } => {
305 let (msaa_str, access) = if multi {
306 ("_ms", "read")
307 } else {
308 ("", "sample")
309 };
310 let scalar = crate::Scalar { kind, width: 4 };
311 ("texture", msaa_str, scalar, access)
312 }
313 crate::ImageClass::Depth { multi } => {
314 let (msaa_str, access) = if multi {
315 ("_ms", "read")
316 } else {
317 ("", "sample")
318 };
319 let scalar = crate::Scalar {
320 kind: crate::ScalarKind::Float,
321 width: 4,
322 };
323 ("depth", msaa_str, scalar, access)
324 }
325 crate::ImageClass::Storage { format, .. } => {
326 let access = if self
327 .access
328 .contains(crate::StorageAccess::LOAD | crate::StorageAccess::STORE)
329 {
330 "read_write"
331 } else if self.access.contains(crate::StorageAccess::STORE) {
332 "write"
333 } else if self.access.contains(crate::StorageAccess::LOAD) {
334 "read"
335 } else {
336 log::warn!(
337 "Storage access for {:?} (name '{}'): {:?}",
338 self.handle,
339 ty.name.as_deref().unwrap_or_default(),
340 self.access
341 );
342 unreachable!("module is not valid");
343 };
344 ("texture", "", format.into(), access)
345 }
346 crate::ImageClass::External => {
347 return write!(out, "{EXTERNAL_TEXTURE_WRAPPER_STRUCT}");
348 }
349 };
350 let base_name = scalar.to_msl_name();
351 let array_str = if arrayed { "_array" } else { "" };
352 write!(
353 out,
354 "{NAMESPACE}::{texture_str}{dim_str}{msaa_str}{array_str}<{base_name}, {NAMESPACE}::access::{access}>",
355 )
356 }
357 crate::TypeInner::Sampler { comparison: _ } => {
358 write!(out, "{NAMESPACE}::sampler")
359 }
360 crate::TypeInner::AccelerationStructure { vertex_return } => {
361 if vertex_return {
362 unimplemented!("metal does not support vertex ray hit return")
363 }
364 write!(out, "{RT_NAMESPACE}::instance_acceleration_structure")
365 }
366 crate::TypeInner::RayQuery { vertex_return } => {
367 if vertex_return {
368 unimplemented!("metal does not support vertex ray hit return")
369 }
370 write!(out, "{}", super::ray::metal_intersector_ty())
371 }
372 crate::TypeInner::BindingArray { base, .. } => {
373 let base_inner = &self.gctx.types[base].inner;
374 let base_tyname = Self {
375 handle: base,
376 first_time: false,
377 ..*self
378 };
379 match *base_inner {
380 crate::TypeInner::Struct { .. } => {
381 write!(
384 out,
385 "device {ARGUMENT_BUFFER_WRAPPER_STRUCT}<device {base_tyname}*>*"
386 )
387 }
388 _ => {
389 write!(
390 out,
391 "constant {ARGUMENT_BUFFER_WRAPPER_STRUCT}<{base_tyname}>*"
392 )
393 }
394 }
395 }
396 }
397 }
398}
399
400pub(super) struct TypedGlobalVariable<'a> {
401 pub module: &'a crate::Module,
402 pub names: &'a FastHashMap<NameKey, String>,
403 pub handle: Handle<crate::GlobalVariable>,
404 pub usage: valid::GlobalUse,
405 pub reference: bool,
406}
407
408struct TypedGlobalVariableParts {
409 ty_name: String,
410 var_name: String,
411}
412
413impl TypedGlobalVariable<'_> {
414 fn to_parts(&self) -> Result<TypedGlobalVariableParts, Error> {
415 let var = &self.module.global_variables[self.handle];
416 let name = &self.names[&NameKey::GlobalVariable(self.handle)];
417
418 let storage_access = match var.space {
419 crate::AddressSpace::Storage { access } => access,
420 _ => match self.module.types[var.ty].inner {
421 crate::TypeInner::Image {
422 class: crate::ImageClass::Storage { access, .. },
423 ..
424 } => access,
425 crate::TypeInner::BindingArray { base, .. } => {
426 match self.module.types[base].inner {
427 crate::TypeInner::Image {
428 class: crate::ImageClass::Storage { access, .. },
429 ..
430 } => access,
431 _ => crate::StorageAccess::default(),
432 }
433 }
434 _ => crate::StorageAccess::default(),
435 },
436 };
437 let ty_name = TypeContext {
438 handle: var.ty,
439 gctx: self.module.to_ctx(),
440 names: self.names,
441 access: storage_access,
442 first_time: false,
443 };
444
445 let (coherent, space, access, reference) = if matches!(
446 self.module.types[var.ty].inner,
447 crate::TypeInner::BindingArray { .. }
448 ) {
449 ("", "", "", "")
450 } else {
451 let access = if var.space.needs_access_qualifier()
452 && !self.usage.intersects(valid::GlobalUse::WRITE)
453 {
454 "const"
455 } else {
456 ""
457 };
458 match (var.space.to_msl_name(), var.space) {
459 (Some(space), crate::AddressSpace::WorkGroup) => {
460 ("", space, access, if self.reference { "&" } else { "" })
461 }
462 (Some(space), _) if self.reference => {
463 let coherent = if var
464 .memory_decorations
465 .contains(crate::MemoryDecorations::COHERENT)
466 {
467 "coherent "
468 } else {
469 ""
470 };
471 (coherent, space, access, "&")
472 }
473 _ => ("", "", "", ""),
474 }
475 };
476
477 let ty = format!(
478 "{coherent}{space}{}{ty_name}{}{access}{reference}",
479 if space.is_empty() { "" } else { " " },
480 if access.is_empty() { "" } else { " " },
481 );
482
483 Ok(TypedGlobalVariableParts {
484 ty_name: ty,
485 var_name: name.clone(),
486 })
487 }
488 pub(super) fn try_fmt<W: Write>(&self, out: &mut W) -> BackendResult {
489 let parts = self.to_parts()?;
490
491 Ok(write!(out, "{} {}", parts.ty_name, parts.var_name)?)
492 }
493}
494
495#[derive(Eq, PartialEq, Hash)]
496pub(super) enum WrappedFunction {
497 UnaryOp {
498 op: crate::UnaryOperator,
499 ty: (Option<crate::VectorSize>, crate::Scalar),
500 },
501 BinaryOp {
502 op: crate::BinaryOperator,
503 left_ty: (Option<crate::VectorSize>, crate::Scalar),
504 right_ty: (Option<crate::VectorSize>, crate::Scalar),
505 },
506 Math {
507 fun: crate::MathFunction,
508 arg_ty: (Option<crate::VectorSize>, crate::Scalar),
509 },
510 Cast {
511 src_scalar: crate::Scalar,
512 vector_size: Option<crate::VectorSize>,
513 dst_scalar: crate::Scalar,
514 },
515 ImageLoad {
516 class: crate::ImageClass,
517 },
518 ImageSample {
519 class: crate::ImageClass,
520 clamp_to_edge: bool,
521 },
522 ImageQuerySize {
523 class: crate::ImageClass,
524 },
525 CooperativeLoad {
526 space_name: &'static str,
527 columns: crate::CooperativeSize,
528 rows: crate::CooperativeSize,
529 scalar: crate::Scalar,
530 },
531 CooperativeMultiplyAdd {
532 space_name: &'static str,
533 columns: crate::CooperativeSize,
534 rows: crate::CooperativeSize,
535 intermediate: crate::CooperativeSize,
536 ab_scalar: crate::Scalar,
537 c_scalar: crate::Scalar,
538 },
539 RayQueryGetIntersection {
540 committed: bool,
541 },
542}
543
544#[expect(missing_debug_implementations, reason = "would be way too verbose?")]
545pub struct Writer<W> {
546 pub(super) out: W,
547 pub(super) names: FastHashMap<NameKey, String>,
548 pub(super) named_expressions: crate::NamedExpressions,
549 need_bake_expressions: back::NeedBakeExpressions,
551 pub(super) namer: proc::Namer,
552 pub(super) wrapped_functions: FastHashSet<WrappedFunction>,
553 emit_int_div_checks: bool,
554 struct_member_pads: FastHashSet<(Handle<crate::Type>, u32)>,
557 needs_object_memory_barriers: bool,
558}
559
560impl crate::Scalar {
561 pub(super) fn to_msl_name(self) -> &'static str {
562 use crate::ScalarKind as Sk;
563 match self {
564 Self {
565 kind: Sk::Float,
566 width: 4,
567 } => "float",
568 Self {
569 kind: Sk::Float,
570 width: 2,
571 } => "half",
572 Self {
573 kind: Sk::Sint,
574 width: 2,
575 } => "short",
576 Self {
577 kind: Sk::Uint,
578 width: 2,
579 } => "ushort",
580 Self {
581 kind: Sk::Sint,
582 width: 4,
583 } => "int",
584 Self {
585 kind: Sk::Uint,
586 width: 4,
587 } => "uint",
588 Self {
589 kind: Sk::Sint,
590 width: 8,
591 } => "long",
592 Self {
593 kind: Sk::Uint,
594 width: 8,
595 } => "ulong",
596 Self {
597 kind: Sk::Bool,
598 width: _,
599 } => "bool",
600 Self {
601 kind: Sk::AbstractInt | Sk::AbstractFloat,
602 width: _,
603 } => unreachable!("Found Abstract scalar kind"),
604 _ => unreachable!("Unsupported scalar kind: {:?}", self),
605 }
606 }
607}
608
609const fn separate(need_separator: bool) -> &'static str {
610 if need_separator {
611 ","
612 } else {
613 ""
614 }
615}
616
617fn should_pack_struct_member(
618 members: &[crate::StructMember],
619 span: u32,
620 index: usize,
621 module: &crate::Module,
622) -> Option<crate::Scalar> {
623 let member = &members[index];
624
625 let ty_inner = &module.types[member.ty].inner;
626 let last_offset = member.offset + ty_inner.size(module.to_ctx());
627 let next_offset = match members.get(index + 1) {
628 Some(next) => next.offset,
629 None => span,
630 };
631 let is_tight = next_offset == last_offset;
632
633 match *ty_inner {
634 crate::TypeInner::Vector {
635 size: crate::VectorSize::Tri,
636 scalar: scalar @ crate::Scalar { width: 4 | 2, .. },
637 } if is_tight => Some(scalar),
638 _ => None,
639 }
640}
641
642impl crate::AddressSpace {
643 const fn needs_pass_through(&self) -> bool {
647 match *self {
648 Self::Uniform
649 | Self::Storage { .. }
650 | Self::Private
651 | Self::WorkGroup
652 | Self::Immediate
653 | Self::Handle
654 | Self::TaskPayload => true,
655 Self::Function => false,
656 Self::RayPayload | Self::IncomingRayPayload => unreachable!(),
657 }
658 }
659
660 const fn needs_access_qualifier(&self) -> bool {
662 match *self {
663 Self::Storage { .. } => true,
668 Self::TaskPayload => true,
669 Self::RayPayload | Self::IncomingRayPayload => unimplemented!(),
670 Self::Private | Self::WorkGroup => false,
672 Self::Uniform | Self::Immediate => false,
674 Self::Handle | Self::Function => false,
676 }
677 }
678
679 const fn to_msl_name(self) -> Option<&'static str> {
680 match self {
681 Self::Handle => None,
682 Self::Uniform | Self::Immediate => Some("constant"),
683 Self::Storage { .. } => Some("device"),
684 Self::Private | Self::Function | Self::RayPayload => Some("thread"),
688 Self::WorkGroup => Some("threadgroup"),
689 Self::TaskPayload => Some("object_data"),
690 Self::IncomingRayPayload => Some("ray_data"),
691 }
692 }
693}
694
695impl crate::Type {
696 const fn needs_alias(&self) -> bool {
698 use crate::TypeInner as Ti;
699
700 match self.inner {
701 Ti::Scalar(_)
703 | Ti::Vector { .. }
704 | Ti::Matrix { .. }
705 | Ti::CooperativeMatrix { .. }
706 | Ti::Atomic(_)
707 | Ti::Pointer { .. }
708 | Ti::ValuePointer { .. } => self.name.is_some(),
709 Ti::Struct { .. } | Ti::Array { .. } => true,
711 Ti::Image { .. }
713 | Ti::Sampler { .. }
714 | Ti::AccelerationStructure { .. }
715 | Ti::RayQuery { .. }
716 | Ti::BindingArray { .. } => false,
717 }
718 }
719}
720
721#[derive(Clone, Copy)]
722pub(super) enum FunctionOrigin {
723 Handle(Handle<crate::Function>),
724 EntryPoint(proc::EntryPointIndex),
725}
726
727pub(super) trait NameKeyExt {
728 fn local(origin: FunctionOrigin, local_handle: Handle<crate::LocalVariable>) -> NameKey {
729 match origin {
730 FunctionOrigin::Handle(handle) => NameKey::FunctionLocal(handle, local_handle),
731 FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointLocal(idx, local_handle),
732 }
733 }
734
735 fn oob_local_for_type(origin: FunctionOrigin, ty: Handle<crate::Type>) -> NameKey {
740 match origin {
741 FunctionOrigin::Handle(handle) => NameKey::FunctionOobLocal(handle, ty),
742 FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointOobLocal(idx, ty),
743 }
744 }
745}
746
747impl NameKeyExt for NameKey {}
748
749#[derive(Clone, Copy)]
759enum LevelOfDetail {
760 Direct(Handle<crate::Expression>),
761 Restricted(Handle<crate::Expression>),
762}
763
764struct TexelAddress {
774 coordinate: Handle<crate::Expression>,
775 array_index: Option<Handle<crate::Expression>>,
776 sample: Option<Handle<crate::Expression>>,
777 level: Option<LevelOfDetail>,
778}
779
780pub(super) struct ExpressionContext<'a> {
781 pub(super) function: &'a crate::Function,
782 pub(super) origin: FunctionOrigin,
783 pub(super) info: &'a valid::FunctionInfo,
784 pub(super) module: &'a crate::Module,
785 pub(super) mod_info: &'a valid::ModuleInfo,
786 pub(super) pipeline_options: &'a PipelineOptions,
787 pub(super) lang_version: (u8, u8),
788 pub(super) policies: index::BoundsCheckPolicies,
789
790 pub(super) guarded_indices: HandleSet<crate::Expression>,
794 pub(super) force_loop_bounding: bool,
796 emit_int_div_checks: bool,
798 pub(super) ray_query_initialization_tracking: bool,
799}
800
801impl<'a> ExpressionContext<'a> {
802 fn resolve_type(&self, handle: Handle<crate::Expression>) -> &'a crate::TypeInner {
803 self.info[handle].ty.inner_with(&self.module.types)
804 }
805
806 fn binding_array_index_from_chain(
809 &self,
810 mut expr: Handle<crate::Expression>,
811 global: Handle<crate::GlobalVariable>,
812 ) -> Option<index::GuardedIndex> {
813 let expressions = &self.function.expressions;
814 loop {
815 match expressions[expr] {
816 crate::Expression::Load { pointer } => expr = pointer,
817 crate::Expression::Access { base, index } => {
818 if matches!(
819 expressions[base],
820 crate::Expression::GlobalVariable(g) if g == global
821 ) {
822 return Some(index::GuardedIndex::Expression(index));
823 }
824 expr = base;
825 }
826 crate::Expression::AccessIndex { base, index } => {
827 if matches!(
828 expressions[base],
829 crate::Expression::GlobalVariable(g) if g == global
830 ) {
831 return Some(index::GuardedIndex::Known(index));
832 }
833 expr = base;
834 }
835 crate::Expression::GlobalVariable(_) => return None,
836 _ => return None,
837 }
838 }
839 }
840
841 fn is_global_access_chain(&self, expr: Handle<crate::Expression>) -> bool {
843 let expressions = &self.function.expressions;
844 match expressions[expr] {
845 crate::Expression::Access { base, .. } => match expressions[base] {
846 crate::Expression::GlobalVariable(_) => true,
847 crate::Expression::Access { .. } => self.is_global_access_chain(base),
848 _ => false,
849 },
850 crate::Expression::AccessIndex { base, .. } => {
851 matches!(expressions[base], crate::Expression::GlobalVariable(_))
852 }
853 _ => false,
854 }
855 }
856
857 fn struct_member_needs_arrow(
858 &self,
859 base: Handle<crate::Expression>,
860 originating_global_ty: impl FnOnce(&crate::TypeInner) -> bool,
861 ) -> bool {
862 let originating_matches = match self.function.originating_global(base) {
863 Some(gv) => {
864 originating_global_ty(&self.module.types[self.module.global_variables[gv].ty].inner)
865 }
866 None => false,
867 };
868 originating_matches && self.is_global_access_chain(base)
869 }
870
871 fn image_needs_lod(&self, image: Handle<crate::Expression>) -> bool {
878 let image_ty = self.resolve_type(image);
879 if let crate::TypeInner::Image { dim, class, .. } = *image_ty {
880 class.is_mipmapped() && dim != crate::ImageDimension::D1
881 } else {
882 false
883 }
884 }
885
886 fn choose_bounds_check_policy(
887 &self,
888 pointer: Handle<crate::Expression>,
889 ) -> index::BoundsCheckPolicy {
890 self.policies
891 .choose_policy(pointer, &self.module.types, self.info)
892 }
893
894 fn access_needs_check(
896 &self,
897 base: Handle<crate::Expression>,
898 index: index::GuardedIndex,
899 ) -> Option<index::IndexableLength> {
900 index::access_needs_check(
901 base,
902 index,
903 self.module,
904 &self.function.expressions,
905 self.info,
906 )
907 }
908
909 fn bounds_check_iter(
911 &self,
912 chain: Handle<crate::Expression>,
913 ) -> impl Iterator<Item = BoundsCheck> + '_ {
914 index::bounds_check_iter(chain, self.module, self.function, self.info)
915 }
916
917 fn oob_local_types(&self) -> FastHashSet<Handle<crate::Type>> {
919 index::oob_local_types(self.module, self.function, self.info, self.policies)
920 }
921
922 fn get_packed_vec_kind(&self, expr_handle: Handle<crate::Expression>) -> Option<crate::Scalar> {
923 match self.function.expressions[expr_handle] {
924 crate::Expression::AccessIndex { base, index } => {
925 let ty = match *self.resolve_type(base) {
926 crate::TypeInner::Pointer { base, .. } => &self.module.types[base].inner,
927 ref ty => ty,
928 };
929 match *ty {
930 crate::TypeInner::Struct {
931 ref members, span, ..
932 } => should_pack_struct_member(members, span, index as usize, self.module),
933 _ => None,
934 }
935 }
936 _ => None,
937 }
938 }
939}
940
941pub(super) struct StatementContext<'a> {
942 pub(super) expression: ExpressionContext<'a>,
943 pub(super) result_struct: Option<&'a str>,
944}
945
946impl<W: Write> Writer<W> {
947 pub fn new(out: W) -> Self {
949 Writer {
950 out,
951 names: FastHashMap::default(),
952 named_expressions: Default::default(),
953 need_bake_expressions: Default::default(),
954 namer: proc::Namer::default(),
955 wrapped_functions: FastHashSet::default(),
956 emit_int_div_checks: true,
957 struct_member_pads: FastHashSet::default(),
958 needs_object_memory_barriers: false,
959 }
960 }
961
962 pub fn finish(self) -> W {
965 self.out
966 }
967
968 fn gen_force_bounded_loop_statements(
1075 &mut self,
1076 level: back::Level,
1077 context: &StatementContext,
1078 ) -> Option<(String, String)> {
1079 if !context.expression.force_loop_bounding {
1080 return None;
1081 }
1082
1083 let loop_bound_name = self.namer.call("loop_bound");
1084 let decl = format!("{level}uint2 {loop_bound_name} = uint2({}u);", u32::MAX);
1087 let level = level.next();
1088 let break_and_inc = format!(
1089 "{level}if ({NAMESPACE}::all({loop_bound_name} == uint2(0u))) {{ break; }}
1090{level}{loop_bound_name} -= uint2({loop_bound_name}.y == 0u, 1u);"
1091 );
1092
1093 Some((decl, break_and_inc))
1094 }
1095
1096 fn put_call_parameters(
1097 &mut self,
1098 parameters: impl Iterator<Item = Handle<crate::Expression>>,
1099 context: &ExpressionContext,
1100 ) -> BackendResult {
1101 self.put_call_parameters_impl(parameters, context, |writer, context, expr| {
1102 writer.put_expression(expr, context, true)
1103 })
1104 }
1105
1106 fn put_call_parameters_impl<C, E>(
1107 &mut self,
1108 parameters: impl Iterator<Item = Handle<crate::Expression>>,
1109 ctx: &C,
1110 put_expression: E,
1111 ) -> BackendResult
1112 where
1113 E: Fn(&mut Self, &C, Handle<crate::Expression>) -> BackendResult,
1114 {
1115 write!(self.out, "(")?;
1116 for (i, handle) in parameters.enumerate() {
1117 if i != 0 {
1118 write!(self.out, ", ")?;
1119 }
1120 put_expression(self, ctx, handle)?;
1121 }
1122 write!(self.out, ")")?;
1123 Ok(())
1124 }
1125
1126 fn put_locals(&mut self, context: &ExpressionContext) -> BackendResult {
1132 let oob_local_types = context.oob_local_types();
1133 for &ty in oob_local_types.iter() {
1134 let name_key = NameKey::oob_local_for_type(context.origin, ty);
1135 self.names.insert(name_key, self.namer.call("oob"));
1136 }
1137
1138 for (name_key, ty, init) in context
1139 .function
1140 .local_variables
1141 .iter()
1142 .map(|(local_handle, local)| {
1143 let name_key = NameKey::local(context.origin, local_handle);
1144 (name_key, local.ty, local.init)
1145 })
1146 .chain(oob_local_types.iter().map(|&ty| {
1147 let name_key = NameKey::oob_local_for_type(context.origin, ty);
1148 (name_key, ty, None)
1149 }))
1150 {
1151 let ty_name = TypeContext {
1152 handle: ty,
1153 gctx: context.module.to_ctx(),
1154 names: &self.names,
1155 access: crate::StorageAccess::empty(),
1156 first_time: false,
1157 };
1158 write!(
1159 self.out,
1160 "{}{} {}",
1161 back::INDENT,
1162 ty_name,
1163 self.names[&name_key]
1164 )?;
1165 match init {
1166 Some(value) => {
1167 write!(self.out, " = ")?;
1168 self.put_expression(value, context, true)?;
1169 }
1170 None => {
1171 write!(self.out, " = {{}}")?;
1172 }
1173 };
1174 writeln!(self.out, ";")?;
1175
1176 if context.ray_query_initialization_tracking {
1178 if let crate::TypeInner::RayQuery { .. } = context.module.types[ty].inner {
1179 writeln!(
1180 self.out,
1181 "{}uint {}{} = 0u;",
1182 back::INDENT,
1183 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
1184 self.names[&name_key]
1185 )?;
1186
1187 writeln!(
1188 self.out,
1189 "{}float {}{} = 0.0;",
1190 back::INDENT,
1191 super::ray::RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX,
1192 self.names[&name_key]
1193 )?;
1194 }
1195 }
1196 }
1197 Ok(())
1198 }
1199
1200 fn put_level_of_detail(
1201 &mut self,
1202 level: LevelOfDetail,
1203 context: &ExpressionContext,
1204 ) -> BackendResult {
1205 match level {
1206 LevelOfDetail::Direct(expr) => self.put_expression(expr, context, true)?,
1207 LevelOfDetail::Restricted(load) => write!(self.out, "{}", ClampedLod(load))?,
1208 }
1209 Ok(())
1210 }
1211
1212 fn put_image_query(
1213 &mut self,
1214 image: Handle<crate::Expression>,
1215 query: &str,
1216 level: Option<LevelOfDetail>,
1217 context: &ExpressionContext,
1218 ) -> BackendResult {
1219 self.put_expression(image, context, false)?;
1220 write!(self.out, ".get_{query}(")?;
1221 if let Some(level) = level {
1222 self.put_level_of_detail(level, context)?;
1223 }
1224 write!(self.out, ")")?;
1225 Ok(())
1226 }
1227
1228 fn put_image_size_query(
1229 &mut self,
1230 image: Handle<crate::Expression>,
1231 level: Option<LevelOfDetail>,
1232 kind: crate::ScalarKind,
1233 context: &ExpressionContext,
1234 ) -> BackendResult {
1235 if let crate::TypeInner::Image {
1236 class: crate::ImageClass::External,
1237 ..
1238 } = *context.resolve_type(image)
1239 {
1240 write!(self.out, "{IMAGE_SIZE_EXTERNAL_FUNCTION}(")?;
1241 self.put_expression(image, context, true)?;
1242 write!(self.out, ")")?;
1243 return Ok(());
1244 }
1245
1246 let dim = match *context.resolve_type(image) {
1249 crate::TypeInner::Image { dim, .. } => dim,
1250 ref other => unreachable!("Unexpected type {:?}", other),
1251 };
1252 let scalar = crate::Scalar { kind, width: 4 };
1253 let coordinate_type = scalar.to_msl_name();
1254 match dim {
1255 crate::ImageDimension::D1 => {
1256 if kind == crate::ScalarKind::Uint {
1260 self.put_image_query(image, "width", None, context)?;
1262 } else {
1263 write!(self.out, "int(")?;
1265 self.put_image_query(image, "width", None, context)?;
1266 write!(self.out, ")")?;
1267 }
1268 }
1269 crate::ImageDimension::D2 => {
1270 write!(self.out, "{NAMESPACE}::{coordinate_type}2(")?;
1271 self.put_image_query(image, "width", level, context)?;
1272 write!(self.out, ", ")?;
1273 self.put_image_query(image, "height", level, context)?;
1274 write!(self.out, ")")?;
1275 }
1276 crate::ImageDimension::D3 => {
1277 write!(self.out, "{NAMESPACE}::{coordinate_type}3(")?;
1278 self.put_image_query(image, "width", level, context)?;
1279 write!(self.out, ", ")?;
1280 self.put_image_query(image, "height", level, context)?;
1281 write!(self.out, ", ")?;
1282 self.put_image_query(image, "depth", level, context)?;
1283 write!(self.out, ")")?;
1284 }
1285 crate::ImageDimension::Cube => {
1286 write!(self.out, "{NAMESPACE}::{coordinate_type}2(")?;
1287 self.put_image_query(image, "width", level, context)?;
1288 write!(self.out, ")")?;
1289 }
1290 }
1291 Ok(())
1292 }
1293
1294 fn put_cast_to_uint_scalar_or_vector(
1295 &mut self,
1296 expr: Handle<crate::Expression>,
1297 context: &ExpressionContext,
1298 ) -> BackendResult {
1299 match *context.resolve_type(expr) {
1301 crate::TypeInner::Scalar(_) => {
1302 put_numeric_type(&mut self.out, crate::Scalar::U32, &[])?
1303 }
1304 crate::TypeInner::Vector { size, .. } => {
1305 put_numeric_type(&mut self.out, crate::Scalar::U32, &[size])?
1306 }
1307 _ => {
1308 return Err(Error::GenericValidation(
1309 "Invalid type for image coordinate".into(),
1310 ))
1311 }
1312 };
1313
1314 write!(self.out, "(")?;
1315 self.put_expression(expr, context, true)?;
1316 write!(self.out, ")")?;
1317 Ok(())
1318 }
1319
1320 fn put_image_sample_level(
1321 &mut self,
1322 image: Handle<crate::Expression>,
1323 level: crate::SampleLevel,
1324 context: &ExpressionContext,
1325 ) -> BackendResult {
1326 let has_levels = context.image_needs_lod(image);
1327 match level {
1328 crate::SampleLevel::Auto => {}
1329 crate::SampleLevel::Zero => {
1330 }
1332 _ if !has_levels => {
1333 log::warn!("1D image can't be sampled with level {level:?}");
1334 }
1335 crate::SampleLevel::Exact(h) => {
1336 write!(self.out, ", {NAMESPACE}::level(")?;
1337 self.put_expression(h, context, true)?;
1338 write!(self.out, ")")?;
1339 }
1340 crate::SampleLevel::Bias(h) => {
1341 write!(self.out, ", {NAMESPACE}::bias(")?;
1342 self.put_expression(h, context, true)?;
1343 write!(self.out, ")")?;
1344 }
1345 crate::SampleLevel::Gradient { x, y } => {
1346 write!(self.out, ", {NAMESPACE}::gradient2d(")?;
1347 self.put_expression(x, context, true)?;
1348 write!(self.out, ", ")?;
1349 self.put_expression(y, context, true)?;
1350 write!(self.out, ")")?;
1351 }
1352 }
1353 Ok(())
1354 }
1355
1356 fn put_image_coordinate_limits(
1357 &mut self,
1358 image: Handle<crate::Expression>,
1359 level: Option<LevelOfDetail>,
1360 context: &ExpressionContext,
1361 ) -> BackendResult {
1362 self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
1363 write!(self.out, " - 1")?;
1364 Ok(())
1365 }
1366
1367 fn put_restricted_scalar_image_index(
1385 &mut self,
1386 image: Handle<crate::Expression>,
1387 index: Handle<crate::Expression>,
1388 limit_method: &str,
1389 context: &ExpressionContext,
1390 ) -> BackendResult {
1391 write!(self.out, "{NAMESPACE}::min(uint(")?;
1392 self.put_expression(index, context, true)?;
1393 write!(self.out, "), ")?;
1394 self.put_expression(image, context, false)?;
1395 write!(self.out, ".{limit_method}() - 1)")?;
1396 Ok(())
1397 }
1398
1399 fn put_restricted_texel_address(
1400 &mut self,
1401 image: Handle<crate::Expression>,
1402 address: &TexelAddress,
1403 context: &ExpressionContext,
1404 ) -> BackendResult {
1405 write!(self.out, "{NAMESPACE}::min(")?;
1407 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1408 write!(self.out, ", ")?;
1409 self.put_image_coordinate_limits(image, address.level, context)?;
1410 write!(self.out, ")")?;
1411
1412 if let Some(array_index) = address.array_index {
1414 write!(self.out, ", ")?;
1415 self.put_restricted_scalar_image_index(image, array_index, "get_array_size", context)?;
1416 }
1417
1418 if let Some(sample) = address.sample {
1420 write!(self.out, ", ")?;
1421 self.put_restricted_scalar_image_index(image, sample, "get_num_samples", context)?;
1422 }
1423
1424 if let Some(level) = address.level {
1427 write!(self.out, ", ")?;
1428 self.put_level_of_detail(level, context)?;
1429 }
1430
1431 Ok(())
1432 }
1433
1434 fn put_image_access_bounds_check(
1436 &mut self,
1437 image: Handle<crate::Expression>,
1438 address: &TexelAddress,
1439 context: &ExpressionContext,
1440 ) -> BackendResult {
1441 let mut conjunction = "";
1442
1443 let level = if let Some(level) = address.level {
1446 write!(self.out, "uint(")?;
1447 self.put_level_of_detail(level, context)?;
1448 write!(self.out, ") < ")?;
1449 self.put_expression(image, context, true)?;
1450 write!(self.out, ".get_num_mip_levels()")?;
1451 conjunction = " && ";
1452 Some(level)
1453 } else {
1454 None
1455 };
1456
1457 if let Some(sample) = address.sample {
1459 write!(self.out, "uint(")?;
1460 self.put_expression(sample, context, true)?;
1461 write!(self.out, ") < ")?;
1462 self.put_expression(image, context, true)?;
1463 write!(self.out, ".get_num_samples()")?;
1464 conjunction = " && ";
1465 }
1466
1467 if let Some(array_index) = address.array_index {
1469 write!(self.out, "{conjunction}uint(")?;
1470 self.put_expression(array_index, context, true)?;
1471 write!(self.out, ") < ")?;
1472 self.put_expression(image, context, true)?;
1473 write!(self.out, ".get_array_size()")?;
1474 conjunction = " && ";
1475 }
1476
1477 let coord_is_vector = match *context.resolve_type(address.coordinate) {
1479 crate::TypeInner::Vector { .. } => true,
1480 _ => false,
1481 };
1482 write!(self.out, "{conjunction}")?;
1483 if coord_is_vector {
1484 write!(self.out, "{NAMESPACE}::all(")?;
1485 }
1486 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1487 write!(self.out, " < ")?;
1488 self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
1489 if coord_is_vector {
1490 write!(self.out, ")")?;
1491 }
1492
1493 Ok(())
1494 }
1495
1496 fn put_image_load(
1497 &mut self,
1498 load: Handle<crate::Expression>,
1499 image: Handle<crate::Expression>,
1500 mut address: TexelAddress,
1501 context: &ExpressionContext,
1502 ) -> BackendResult {
1503 if let crate::TypeInner::Image {
1504 class: crate::ImageClass::External,
1505 ..
1506 } = *context.resolve_type(image)
1507 {
1508 write!(self.out, "{IMAGE_LOAD_EXTERNAL_FUNCTION}(")?;
1509 self.put_expression(image, context, true)?;
1510 write!(self.out, ", ")?;
1511 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1512 write!(self.out, ")")?;
1513 return Ok(());
1514 }
1515
1516 match context.policies.image_load {
1517 proc::BoundsCheckPolicy::Restrict => {
1518 if address.level.is_some() {
1521 address.level = if context.image_needs_lod(image) {
1522 Some(LevelOfDetail::Restricted(load))
1523 } else {
1524 None
1525 }
1526 }
1527
1528 self.put_expression(image, context, false)?;
1529 write!(self.out, ".read(")?;
1530 self.put_restricted_texel_address(image, &address, context)?;
1531 write!(self.out, ")")?;
1532 }
1533 proc::BoundsCheckPolicy::ReadZeroSkipWrite => {
1534 write!(self.out, "(")?;
1535 self.put_image_access_bounds_check(image, &address, context)?;
1536 write!(self.out, " ? ")?;
1537 self.put_unchecked_image_load(image, &address, context)?;
1538 write!(self.out, ": DefaultConstructible())")?;
1539 }
1540 proc::BoundsCheckPolicy::Unchecked => {
1541 self.put_unchecked_image_load(image, &address, context)?;
1542 }
1543 }
1544
1545 Ok(())
1546 }
1547
1548 fn put_unchecked_image_load(
1549 &mut self,
1550 image: Handle<crate::Expression>,
1551 address: &TexelAddress,
1552 context: &ExpressionContext,
1553 ) -> BackendResult {
1554 self.put_expression(image, context, false)?;
1555 write!(self.out, ".read(")?;
1556 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1558 if let Some(expr) = address.array_index {
1559 write!(self.out, ", ")?;
1560 self.put_expression(expr, context, true)?;
1561 }
1562 if let Some(sample) = address.sample {
1563 write!(self.out, ", ")?;
1564 self.put_expression(sample, context, true)?;
1565 }
1566 if let Some(level) = address.level {
1567 if context.image_needs_lod(image) {
1568 write!(self.out, ", ")?;
1569 self.put_level_of_detail(level, context)?;
1570 }
1571 }
1572 write!(self.out, ")")?;
1573
1574 Ok(())
1575 }
1576
1577 fn put_image_atomic(
1578 &mut self,
1579 level: back::Level,
1580 image: Handle<crate::Expression>,
1581 address: &TexelAddress,
1582 fun: crate::AtomicFunction,
1583 value: Handle<crate::Expression>,
1584 context: &StatementContext,
1585 ) -> BackendResult {
1586 write!(self.out, "{level}")?;
1587 self.put_expression(image, &context.expression, false)?;
1588 let op = if context.expression.resolve_type(value).scalar_width() == Some(8) {
1589 fun.to_msl_64_bit()?
1590 } else {
1591 fun.to_msl()
1592 };
1593 write!(self.out, ".atomic_{op}(")?;
1594 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1596 write!(self.out, ", ")?;
1597 self.put_expression(value, &context.expression, true)?;
1598 writeln!(self.out, ");")?;
1599
1600 let value_ty = context.expression.resolve_type(value);
1606 let zero_value = match (value_ty.scalar_kind(), value_ty.scalar_width()) {
1607 (Some(crate::ScalarKind::Sint), _) => "int4(0)",
1608 (_, Some(8)) => "ulong4(0uL)",
1609 _ => "uint4(0u)",
1610 };
1611 let coord_ty = context.expression.resolve_type(address.coordinate);
1612 let x = if matches!(coord_ty, crate::TypeInner::Scalar(_)) {
1613 ""
1614 } else {
1615 ".x"
1616 };
1617 write!(self.out, "{level}if (")?;
1618 self.put_expression(address.coordinate, &context.expression, true)?;
1619 write!(self.out, "{x} == -99999) {{ ")?;
1620 self.put_expression(image, &context.expression, false)?;
1621 write!(self.out, ".write({zero_value}, ")?;
1622 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1623 if let Some(array_index) = address.array_index {
1624 write!(self.out, ", ")?;
1625 self.put_expression(array_index, &context.expression, true)?;
1626 }
1627 writeln!(self.out, "); }}")?;
1628
1629 Ok(())
1630 }
1631
1632 fn put_image_store(
1633 &mut self,
1634 level: back::Level,
1635 image: Handle<crate::Expression>,
1636 address: &TexelAddress,
1637 value: Handle<crate::Expression>,
1638 context: &StatementContext,
1639 ) -> BackendResult {
1640 write!(self.out, "{level}")?;
1641 self.put_expression(image, &context.expression, false)?;
1642 write!(self.out, ".write(")?;
1643 self.put_expression(value, &context.expression, true)?;
1644 write!(self.out, ", ")?;
1645 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1647 if let Some(expr) = address.array_index {
1648 write!(self.out, ", ")?;
1649 self.put_expression(expr, &context.expression, true)?;
1650 }
1651 writeln!(self.out, ");")?;
1652
1653 Ok(())
1654 }
1655
1656 fn binding_array_layout_count(
1671 module: &crate::Module,
1672 pipeline_options: &PipelineOptions,
1673 global: Handle<crate::GlobalVariable>,
1674 ) -> u32 {
1675 let var = &module.global_variables[global];
1676 let crate::TypeInner::BindingArray { size, .. } = module.types[var.ty].inner else {
1677 unreachable!("binding_array_layout_count called on non-binding-array global");
1678 };
1679 let from_shader = match size {
1680 crate::ArraySize::Constant(n) => n.get(),
1681 crate::ArraySize::Pending(_) | crate::ArraySize::Dynamic => 0,
1682 };
1683 let from_layout = var
1684 .binding
1685 .and_then(|br| pipeline_options.binding_array_length_map.get(&br))
1686 .copied()
1687 .unwrap_or(0);
1688 from_shader.max(from_layout).max(1)
1689 }
1690
1691 fn put_binding_array_size_member_index(
1692 &mut self,
1693 index: index::GuardedIndex,
1694 context: &ExpressionContext,
1695 ) -> BackendResult {
1696 match index {
1697 index::GuardedIndex::Expression(expr) => {
1698 write!(self.out, "unsigned(")?;
1699 self.put_expression(expr, context, true)?;
1700 write!(self.out, ")")?;
1701 }
1702 index::GuardedIndex::Known(value) => write!(self.out, "{value}u")?,
1703 }
1704 Ok(())
1705 }
1706
1707 fn put_dynamic_array_max_index(
1708 &mut self,
1709 handle: Handle<crate::GlobalVariable>,
1710 chain_expr: Handle<crate::Expression>,
1711 context: &ExpressionContext,
1712 ) -> BackendResult {
1713 let global = &context.module.global_variables[handle];
1714 let (offset, array_ty) = match context.module.types[global.ty].inner {
1715 crate::TypeInner::Struct { ref members, .. } => match members.last() {
1716 Some(&crate::StructMember { offset, ty, .. }) => (offset, ty),
1717 None => return Err(Error::GenericValidation("Struct has no members".into())),
1718 },
1719 crate::TypeInner::BindingArray { base, .. } => match context.module.types[base].inner {
1720 crate::TypeInner::Struct { ref members, .. } => match members.last() {
1721 Some(&crate::StructMember { offset, ty, .. }) => (offset, ty),
1722 None => return Err(Error::GenericValidation("Struct has no members".into())),
1723 },
1724 _ => {
1725 return Err(Error::GenericValidation(
1726 "binding_array element must be a struct with a runtime-sized array field"
1727 .into(),
1728 ))
1729 }
1730 },
1731 crate::TypeInner::Array {
1732 size: crate::ArraySize::Dynamic,
1733 ..
1734 } => (0, global.ty),
1735 ref ty => {
1736 return Err(Error::GenericValidation(format!(
1737 "Expected type with dynamic array, got {ty:?}"
1738 )))
1739 }
1740 };
1741
1742 let (size, stride) = match context.module.types[array_ty].inner {
1743 crate::TypeInner::Array { base, stride, .. } => (
1744 context.module.types[base]
1745 .inner
1746 .size(context.module.to_ctx()),
1747 stride,
1748 ),
1749 ref ty => {
1750 return Err(Error::GenericValidation(format!(
1751 "Expected array type, got {ty:?}"
1752 )))
1753 }
1754 };
1755
1756 write!(
1769 self.out,
1770 "(_buffer_sizes.{member}",
1771 member = ArraySizeMember(handle),
1772 )?;
1773 if let crate::TypeInner::BindingArray { .. } = context.module.types[global.ty].inner {
1774 let Some(array_index) = context.binding_array_index_from_chain(chain_expr, handle)
1775 else {
1776 return Err(Error::GenericValidation(
1777 "Could not find binding_array index for buffer size".into(),
1778 ));
1779 };
1780 write!(self.out, "[")?;
1781 match array_index {
1782 index::GuardedIndex::Expression(expr) => {
1783 write!(self.out, "unsigned(")?;
1784 self.put_expression(expr, context, true)?;
1785 write!(self.out, ")")?;
1786 }
1787 index::GuardedIndex::Known(i) => {
1788 write!(self.out, "{i}u")?;
1789 }
1790 }
1791 write!(self.out, "]")?;
1792 }
1793 write!(
1794 self.out,
1795 " - {offset} - {size}) / {stride}",
1796 offset = offset,
1797 size = size,
1798 stride = stride,
1799 )?;
1800 Ok(())
1801 }
1802
1803 fn put_dot_product<T: Copy>(
1808 &mut self,
1809 arg: T,
1810 arg1: T,
1811 size: usize,
1812 extractor: impl Fn(&mut Self, T, usize) -> BackendResult,
1813 ) -> BackendResult {
1814 write!(self.out, "(")?;
1817
1818 for index in 0..size {
1820 write!(self.out, " + ")?;
1823 extractor(self, arg, index)?;
1824 write!(self.out, " * ")?;
1825 extractor(self, arg1, index)?;
1826 }
1827
1828 write!(self.out, ")")?;
1829 Ok(())
1830 }
1831
1832 fn put_pack4x8(
1834 &mut self,
1835 arg: Handle<crate::Expression>,
1836 context: &ExpressionContext<'_>,
1837 was_signed: bool,
1838 clamp_bounds: Option<(&str, &str)>,
1839 ) -> Result<(), Error> {
1840 let write_arg = |this: &mut Self| -> BackendResult {
1841 if let Some((min, max)) = clamp_bounds {
1842 write!(this.out, "{NAMESPACE}::clamp(")?;
1844 this.put_expression(arg, context, true)?;
1845 write!(this.out, ", {min}, {max})")?;
1846 } else {
1847 this.put_expression(arg, context, true)?;
1848 }
1849 Ok(())
1850 };
1851
1852 if context.lang_version >= (2, 1) {
1853 let packed_type = if was_signed {
1854 "packed_char4"
1855 } else {
1856 "packed_uchar4"
1857 };
1858 write!(self.out, "as_type<uint>({packed_type}(")?;
1860 write_arg(self)?;
1861 write!(self.out, "))")?;
1862 } else {
1863 if was_signed {
1865 write!(self.out, "uint(")?;
1866 }
1867 write!(self.out, "(")?;
1868 write_arg(self)?;
1869 write!(self.out, "[0] & 0xFF) | ((")?;
1870 write_arg(self)?;
1871 write!(self.out, "[1] & 0xFF) << 8) | ((")?;
1872 write_arg(self)?;
1873 write!(self.out, "[2] & 0xFF) << 16) | ((")?;
1874 write_arg(self)?;
1875 write!(self.out, "[3] & 0xFF) << 24)")?;
1876 if was_signed {
1877 write!(self.out, ")")?;
1878 }
1879 }
1880
1881 Ok(())
1882 }
1883
1884 fn put_isign(
1887 &mut self,
1888 arg: Handle<crate::Expression>,
1889 context: &ExpressionContext,
1890 ) -> BackendResult {
1891 write!(self.out, "{NAMESPACE}::select({NAMESPACE}::select(")?;
1892 let scalar = context
1893 .resolve_type(arg)
1894 .scalar()
1895 .expect("put_isign should only be called for args which have an integer scalar type")
1896 .to_msl_name();
1897 match context.resolve_type(arg) {
1898 &crate::TypeInner::Vector { size, .. } => {
1899 let size = common::vector_size_str(size);
1900 write!(self.out, "{scalar}{size}(-1), {scalar}{size}(1)")?;
1901 }
1902 _ => {
1903 write!(self.out, "{scalar}(-1), {scalar}(1)")?;
1904 }
1905 }
1906 write!(self.out, ", (")?;
1907 self.put_expression(arg, context, true)?;
1908 write!(self.out, " > 0)), {scalar}(0), (")?;
1909 self.put_expression(arg, context, true)?;
1910 write!(self.out, " == 0))")?;
1911 Ok(())
1912 }
1913
1914 pub(super) fn put_const_expression(
1915 &mut self,
1916 expr_handle: Handle<crate::Expression>,
1917 module: &crate::Module,
1918 mod_info: &valid::ModuleInfo,
1919 arena: &crate::Arena<crate::Expression>,
1920 ) -> BackendResult {
1921 self.put_possibly_const_expression(
1922 expr_handle,
1923 arena,
1924 module,
1925 mod_info,
1926 &(module, mod_info),
1927 |&(_, mod_info), expr| &mod_info[expr],
1928 |writer, &(module, _), expr| writer.put_const_expression(expr, module, mod_info, arena),
1929 )
1930 }
1931
1932 fn put_literal(&mut self, literal: crate::Literal) -> BackendResult {
1933 match literal {
1934 crate::Literal::F64(_) => {
1935 return Err(Error::CapabilityNotSupported(valid::Capabilities::FLOAT64))
1936 }
1937 crate::Literal::F16(value) => {
1938 if value.is_infinite() {
1939 let sign = if value.is_sign_negative() { "-" } else { "" };
1940 write!(self.out, "{sign}INFINITY")?;
1941 } else if value.is_nan() {
1942 write!(self.out, "NAN")?;
1943 } else {
1944 let suffix = if value.fract() == f16::from_f32(0.0) {
1945 ".0h"
1946 } else {
1947 "h"
1948 };
1949 write!(self.out, "{value}{suffix}")?;
1950 }
1951 }
1952 crate::Literal::F32(value) => {
1953 if value.is_infinite() {
1954 let sign = if value.is_sign_negative() { "-" } else { "" };
1955 write!(self.out, "{sign}INFINITY")?;
1956 } else if value.is_nan() {
1957 write!(self.out, "NAN")?;
1958 } else {
1959 let suffix = if value.fract() == 0.0 { ".0" } else { "" };
1960 write!(self.out, "{value}{suffix}")?;
1961 }
1962 }
1963 crate::Literal::U16(value) => {
1964 write!(self.out, "static_cast<ushort>({value})")?;
1965 }
1966 crate::Literal::I16(value) => {
1967 write!(self.out, "static_cast<short>({value})")?;
1968 }
1969 crate::Literal::U32(value) => {
1970 write!(self.out, "{value}u")?;
1971 }
1972 crate::Literal::I32(value) => {
1973 if value == i32::MIN {
1978 write!(self.out, "({} - 1)", value + 1)?;
1979 } else {
1980 write!(self.out, "{value}")?;
1981 }
1982 }
1983 crate::Literal::U64(value) => {
1984 write!(self.out, "{value}uL")?;
1985 }
1986 crate::Literal::I64(value) => {
1987 if value == i64::MIN {
1994 write!(self.out, "({}L - 1L)", value + 1)?;
1995 } else {
1996 write!(self.out, "{value}L")?;
1997 }
1998 }
1999 crate::Literal::Bool(value) => {
2000 write!(self.out, "{value}")?;
2001 }
2002 crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {
2003 return Err(Error::GenericValidation(
2004 "Unsupported abstract literal".into(),
2005 ));
2006 }
2007 }
2008 Ok(())
2009 }
2010
2011 #[allow(clippy::too_many_arguments)]
2012 fn put_possibly_const_expression<C, I, E>(
2013 &mut self,
2014 expr_handle: Handle<crate::Expression>,
2015 expressions: &crate::Arena<crate::Expression>,
2016 module: &crate::Module,
2017 mod_info: &valid::ModuleInfo,
2018 ctx: &C,
2019 get_expr_ty: I,
2020 put_expression: E,
2021 ) -> BackendResult
2022 where
2023 I: Fn(&C, Handle<crate::Expression>) -> &TypeResolution,
2024 E: Fn(&mut Self, &C, Handle<crate::Expression>) -> BackendResult,
2025 {
2026 match expressions[expr_handle] {
2027 crate::Expression::Literal(literal) => {
2028 self.put_literal(literal)?;
2029 }
2030 crate::Expression::Constant(handle) => {
2031 let constant = &module.constants[handle];
2032 if constant.name.is_some() {
2033 write!(self.out, "{}", self.names[&NameKey::Constant(handle)])?;
2034 } else {
2035 self.put_const_expression(
2036 constant.init,
2037 module,
2038 mod_info,
2039 &module.global_expressions,
2040 )?;
2041 }
2042 }
2043 crate::Expression::ZeroValue(ty) => {
2044 let ty_name = TypeContext {
2045 handle: ty,
2046 gctx: module.to_ctx(),
2047 names: &self.names,
2048 access: crate::StorageAccess::empty(),
2049 first_time: false,
2050 };
2051 write!(self.out, "{ty_name} {{}}")?;
2052 }
2053 crate::Expression::Compose { ty, ref components } => {
2054 let ty_name = TypeContext {
2055 handle: ty,
2056 gctx: module.to_ctx(),
2057 names: &self.names,
2058 access: crate::StorageAccess::empty(),
2059 first_time: false,
2060 };
2061 write!(self.out, "{ty_name}")?;
2062 match module.types[ty].inner {
2063 crate::TypeInner::Scalar(_)
2064 | crate::TypeInner::Vector { .. }
2065 | crate::TypeInner::Matrix { .. } => {
2066 self.put_call_parameters_impl(
2067 components.iter().copied(),
2068 ctx,
2069 put_expression,
2070 )?;
2071 }
2072 crate::TypeInner::Array { .. } => {
2073 write!(self.out, " {{{{")?;
2076 for (index, &component) in components.iter().enumerate() {
2077 if index != 0 {
2078 write!(self.out, ", ")?;
2079 }
2080 put_expression(self, ctx, component)?;
2081 }
2082 write!(self.out, "}}}}")?;
2083 }
2084 crate::TypeInner::Struct { .. } => {
2085 write!(self.out, " {{")?;
2086 for (index, &component) in components.iter().enumerate() {
2087 if index != 0 {
2088 write!(self.out, ", ")?;
2089 }
2090 if self.struct_member_pads.contains(&(ty, index as u32)) {
2092 write!(self.out, "{{}}, ")?;
2093 }
2094 put_expression(self, ctx, component)?;
2095 }
2096 write!(self.out, "}}")?;
2097 }
2098 _ => return Err(Error::UnsupportedCompose(ty)),
2099 }
2100 }
2101 crate::Expression::Splat { size, value } => {
2102 let scalar = match *get_expr_ty(ctx, value).inner_with(&module.types) {
2103 crate::TypeInner::Scalar(scalar) => scalar,
2104 ref ty => {
2105 return Err(Error::GenericValidation(format!(
2106 "Expected splat value type must be a scalar, got {ty:?}",
2107 )))
2108 }
2109 };
2110 put_numeric_type(&mut self.out, scalar, &[size])?;
2111 write!(self.out, "(")?;
2112 put_expression(self, ctx, value)?;
2113 write!(self.out, ")")?;
2114 }
2115 _ => {
2116 return Err(Error::Override);
2117 }
2118 }
2119
2120 Ok(())
2121 }
2122
2123 pub(super) fn put_expression(
2135 &mut self,
2136 expr_handle: Handle<crate::Expression>,
2137 context: &ExpressionContext,
2138 is_scoped: bool,
2139 ) -> BackendResult {
2140 if let Some(name) = self.named_expressions.get(&expr_handle) {
2141 write!(self.out, "{name}")?;
2142 return Ok(());
2143 }
2144
2145 let expression = &context.function.expressions[expr_handle];
2146 match *expression {
2147 crate::Expression::Literal(_)
2148 | crate::Expression::Constant(_)
2149 | crate::Expression::ZeroValue(_)
2150 | crate::Expression::Compose { .. }
2151 | crate::Expression::Splat { .. } => {
2152 self.put_possibly_const_expression(
2153 expr_handle,
2154 &context.function.expressions,
2155 context.module,
2156 context.mod_info,
2157 context,
2158 |context, expr: Handle<crate::Expression>| &context.info[expr].ty,
2159 |writer, context, expr| writer.put_expression(expr, context, true),
2160 )?;
2161 }
2162 crate::Expression::Override(_) => return Err(Error::Override),
2163 crate::Expression::Access { base, .. }
2164 | crate::Expression::AccessIndex { base, .. } => {
2165 let policy = context.choose_bounds_check_policy(base);
2171 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
2172 && self.put_bounds_checks(
2173 expr_handle,
2174 context,
2175 back::Level(0),
2176 if is_scoped { "" } else { "(" },
2177 )?
2178 {
2179 write!(self.out, " ? ")?;
2180 self.put_access_chain(expr_handle, policy, context)?;
2181 write!(self.out, " : ")?;
2182
2183 if context.resolve_type(base).pointer_space().is_some() {
2184 let result_ty = context.info[expr_handle]
2188 .ty
2189 .inner_with(&context.module.types)
2190 .pointer_base_type();
2191 let result_ty_handle = match result_ty {
2192 Some(TypeResolution::Handle(handle)) => handle,
2193 Some(TypeResolution::Value(_)) => {
2194 unreachable!(
2202 "Expected type {result_ty:?} of access through pointer type {base:?} to be in the arena",
2203 );
2204 }
2205 None => {
2206 unreachable!(
2207 "Expected access through pointer type {base:?} to return a pointer, but got {result_ty:?}",
2208 )
2209 }
2210 };
2211 let name_key =
2212 NameKey::oob_local_for_type(context.origin, result_ty_handle);
2213 self.out.write_str(&self.names[&name_key])?;
2214 } else {
2215 write!(self.out, "DefaultConstructible()")?;
2216 }
2217
2218 if !is_scoped {
2219 write!(self.out, ")")?;
2220 }
2221 } else {
2222 self.put_access_chain(expr_handle, policy, context)?;
2223 }
2224 }
2225 crate::Expression::Swizzle {
2226 size,
2227 vector,
2228 pattern,
2229 } => {
2230 self.put_wrapped_expression_for_packed_vec3_access(
2231 vector,
2232 context,
2233 false,
2234 &Self::put_expression,
2235 )?;
2236 write!(self.out, ".")?;
2237 for &sc in pattern[..size as usize].iter() {
2238 write!(self.out, "{}", back::COMPONENTS[sc as usize])?;
2239 }
2240 }
2241 crate::Expression::FunctionArgument(index) => {
2242 let name_key = match context.origin {
2243 FunctionOrigin::Handle(handle) => NameKey::FunctionArgument(handle, index),
2244 FunctionOrigin::EntryPoint(ep_index) => {
2245 NameKey::EntryPointArgument(ep_index, index)
2246 }
2247 };
2248 let name = &self.names[&name_key];
2249 write!(self.out, "{name}")?;
2250 }
2251 crate::Expression::GlobalVariable(handle) => {
2252 let name = &self.names[&NameKey::GlobalVariable(handle)];
2253 write!(self.out, "{name}")?;
2254 }
2255 crate::Expression::LocalVariable(handle) => {
2256 let name_key = NameKey::local(context.origin, handle);
2257 let name = &self.names[&name_key];
2258 write!(self.out, "{name}")?;
2259 }
2260 crate::Expression::Load { pointer } => self.put_load(pointer, context, is_scoped)?,
2261 crate::Expression::ImageSample {
2262 coordinate,
2263 image,
2264 sampler,
2265 clamp_to_edge: true,
2266 gather: None,
2267 array_index: None,
2268 offset: None,
2269 level: crate::SampleLevel::Zero,
2270 depth_ref: None,
2271 } => {
2272 write!(self.out, "{IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION}(")?;
2273 self.put_expression(image, context, true)?;
2274 write!(self.out, ", ")?;
2275 self.put_expression(sampler, context, true)?;
2276 write!(self.out, ", ")?;
2277 self.put_expression(coordinate, context, true)?;
2278 write!(self.out, ")")?;
2279 }
2280 crate::Expression::ImageSample {
2281 image,
2282 sampler,
2283 gather,
2284 coordinate,
2285 array_index,
2286 offset,
2287 level,
2288 depth_ref,
2289 clamp_to_edge,
2290 } => {
2291 if clamp_to_edge {
2292 return Err(Error::GenericValidation(
2293 "ImageSample::clamp_to_edge should have been validated out".to_string(),
2294 ));
2295 }
2296
2297 let main_op = match gather {
2298 Some(_) => "gather",
2299 None => "sample",
2300 };
2301 let comparison_op = match depth_ref {
2302 Some(_) => "_compare",
2303 None => "",
2304 };
2305 self.put_expression(image, context, false)?;
2306 write!(self.out, ".{main_op}{comparison_op}(")?;
2307 self.put_expression(sampler, context, true)?;
2308 write!(self.out, ", ")?;
2309 self.put_expression(coordinate, context, true)?;
2310 if let Some(expr) = array_index {
2311 write!(self.out, ", ")?;
2312 self.put_expression(expr, context, true)?;
2313 }
2314 if let Some(dref) = depth_ref {
2315 write!(self.out, ", ")?;
2316 self.put_expression(dref, context, true)?;
2317 }
2318
2319 self.put_image_sample_level(image, level, context)?;
2320
2321 if let Some(offset) = offset {
2322 write!(self.out, ", ")?;
2323 self.put_expression(offset, context, true)?;
2324 }
2325
2326 match gather {
2327 None | Some(crate::SwizzleComponent::X) => {}
2328 Some(component) => {
2329 let is_cube_map = match *context.resolve_type(image) {
2330 crate::TypeInner::Image {
2331 dim: crate::ImageDimension::Cube,
2332 ..
2333 } => true,
2334 _ => false,
2335 };
2336 if offset.is_none() && !is_cube_map {
2339 write!(self.out, ", {NAMESPACE}::int2(0)")?;
2340 }
2341 let letter = back::COMPONENTS[component as usize];
2342 write!(self.out, ", {NAMESPACE}::component::{letter}")?;
2343 }
2344 }
2345 write!(self.out, ")")?;
2346 }
2347 crate::Expression::ImageLoad {
2348 image,
2349 coordinate,
2350 array_index,
2351 sample,
2352 level,
2353 } => {
2354 let address = TexelAddress {
2355 coordinate,
2356 array_index,
2357 sample,
2358 level: level.map(LevelOfDetail::Direct),
2359 };
2360 self.put_image_load(expr_handle, image, address, context)?;
2361 }
2362 crate::Expression::ImageQuery { image, query } => match query {
2365 crate::ImageQuery::Size { level } => {
2366 self.put_image_size_query(
2367 image,
2368 level.map(LevelOfDetail::Direct),
2369 crate::ScalarKind::Uint,
2370 context,
2371 )?;
2372 }
2373 crate::ImageQuery::NumLevels => {
2374 self.put_expression(image, context, false)?;
2375 write!(self.out, ".get_num_mip_levels()")?;
2376 }
2377 crate::ImageQuery::NumLayers => {
2378 self.put_expression(image, context, false)?;
2379 write!(self.out, ".get_array_size()")?;
2380 }
2381 crate::ImageQuery::NumSamples => {
2382 self.put_expression(image, context, false)?;
2383 write!(self.out, ".get_num_samples()")?;
2384 }
2385 },
2386 crate::Expression::Unary { op, expr } => {
2387 let op_str = match op {
2388 crate::UnaryOperator::Negate => {
2389 match context.resolve_type(expr).scalar_kind() {
2390 Some(crate::ScalarKind::Sint) => NEG_FUNCTION,
2391 _ => "-",
2392 }
2393 }
2394 crate::UnaryOperator::LogicalNot => "!",
2395 crate::UnaryOperator::BitwiseNot => "~",
2396 };
2397 write!(self.out, "{op_str}(")?;
2398 self.put_expression(expr, context, false)?;
2399 write!(self.out, ")")?;
2400 }
2401 crate::Expression::Binary { op, left, right } => {
2402 let kind = context
2403 .resolve_type(left)
2404 .scalar_kind()
2405 .ok_or(Error::UnsupportedBinaryOp(op))?;
2406
2407 if op == crate::BinaryOperator::Divide
2408 && (kind == crate::ScalarKind::Sint || kind == crate::ScalarKind::Uint)
2409 && context.emit_int_div_checks
2410 {
2411 write!(self.out, "{DIV_FUNCTION}(")?;
2412 self.put_expression(left, context, true)?;
2413 write!(self.out, ", ")?;
2414 self.put_expression(right, context, true)?;
2415 write!(self.out, ")")?;
2416 } else if op == crate::BinaryOperator::Modulo
2417 && (kind == crate::ScalarKind::Sint || kind == crate::ScalarKind::Uint)
2418 && context.emit_int_div_checks
2419 {
2420 write!(self.out, "{MOD_FUNCTION}(")?;
2421 self.put_expression(left, context, true)?;
2422 write!(self.out, ", ")?;
2423 self.put_expression(right, context, true)?;
2424 write!(self.out, ")")?;
2425 } else if op == crate::BinaryOperator::Modulo && kind == crate::ScalarKind::Float {
2426 write!(self.out, "{NAMESPACE}::fmod(")?;
2431 self.put_expression(left, context, true)?;
2432 write!(self.out, ", ")?;
2433 self.put_expression(right, context, true)?;
2434 write!(self.out, ")")?;
2435 } else if (op == crate::BinaryOperator::Add
2436 || op == crate::BinaryOperator::Subtract
2437 || op == crate::BinaryOperator::Multiply)
2438 && kind == crate::ScalarKind::Sint
2439 {
2440 let to_unsigned = |ty: &crate::TypeInner| match *ty {
2441 crate::TypeInner::Scalar(scalar) => {
2442 Ok(crate::TypeInner::Scalar(crate::Scalar {
2443 kind: crate::ScalarKind::Uint,
2444 ..scalar
2445 }))
2446 }
2447 crate::TypeInner::Vector { size, scalar } => Ok(crate::TypeInner::Vector {
2448 size,
2449 scalar: crate::Scalar {
2450 kind: crate::ScalarKind::Uint,
2451 ..scalar
2452 },
2453 }),
2454 _ => Err(Error::UnsupportedBitCast(ty.clone())),
2455 };
2456
2457 self.put_bitcasted_expression(
2462 context.resolve_type(expr_handle),
2463 expr_handle,
2464 context,
2465 &|writer, context, is_scoped| {
2466 writer.put_binop(
2467 op,
2468 left,
2469 right,
2470 context,
2471 is_scoped,
2472 &|writer, expr, context, _is_scoped| {
2473 writer.put_bitcasted_expression(
2474 &to_unsigned(context.resolve_type(expr))?,
2475 expr,
2476 context,
2477 &|writer, context, is_scoped| {
2478 writer.put_expression(expr, context, is_scoped)
2479 },
2480 )
2481 },
2482 )
2483 },
2484 )?;
2485 } else {
2486 self.put_binop(op, left, right, context, is_scoped, &Self::put_expression)?;
2487 }
2488 }
2489 crate::Expression::Select {
2490 condition,
2491 accept,
2492 reject,
2493 } => match *context.resolve_type(condition) {
2494 crate::TypeInner::Scalar(crate::Scalar {
2495 kind: crate::ScalarKind::Bool,
2496 ..
2497 }) => {
2498 if !is_scoped {
2499 write!(self.out, "(")?;
2500 }
2501 self.put_expression(condition, context, false)?;
2502 write!(self.out, " ? ")?;
2503 self.put_expression(accept, context, false)?;
2504 write!(self.out, " : ")?;
2505 self.put_expression(reject, context, false)?;
2506 if !is_scoped {
2507 write!(self.out, ")")?;
2508 }
2509 }
2510 crate::TypeInner::Vector {
2511 scalar:
2512 crate::Scalar {
2513 kind: crate::ScalarKind::Bool,
2514 ..
2515 },
2516 ..
2517 } => {
2518 write!(self.out, "{NAMESPACE}::select(")?;
2519 self.put_expression(reject, context, true)?;
2520 write!(self.out, ", ")?;
2521 self.put_expression(accept, context, true)?;
2522 write!(self.out, ", ")?;
2523 self.put_expression(condition, context, true)?;
2524 write!(self.out, ")")?;
2525 }
2526 ref ty => {
2527 return Err(Error::GenericValidation(format!(
2528 "Expected select condition to be a non-bool type, got {ty:?}",
2529 )))
2530 }
2531 },
2532 crate::Expression::Derivative { axis, expr, .. } => {
2533 use crate::DerivativeAxis as Axis;
2534 let op = match axis {
2535 Axis::X => "dfdx",
2536 Axis::Y => "dfdy",
2537 Axis::Width => "fwidth",
2538 };
2539 write!(self.out, "{NAMESPACE}::{op}")?;
2540 self.put_call_parameters(iter::once(expr), context)?;
2541 }
2542 crate::Expression::Relational { fun, argument } => {
2543 let op = match fun {
2544 crate::RelationalFunction::Any => "any",
2545 crate::RelationalFunction::All => "all",
2546 crate::RelationalFunction::IsNan => "isnan",
2547 crate::RelationalFunction::IsInf => "isinf",
2548 };
2549 write!(self.out, "{NAMESPACE}::{op}")?;
2550 self.put_call_parameters(iter::once(argument), context)?;
2551 }
2552 crate::Expression::Math {
2553 fun,
2554 arg,
2555 arg1,
2556 arg2,
2557 arg3,
2558 } => {
2559 use crate::MathFunction as Mf;
2560
2561 let arg_type = context.resolve_type(arg);
2562 let scalar_argument = match arg_type {
2563 &crate::TypeInner::Scalar(_) => true,
2564 _ => false,
2565 };
2566
2567 let fun_name = match fun {
2568 Mf::Abs => "abs",
2570 Mf::Min => "min",
2571 Mf::Max => "max",
2572 Mf::Clamp => "clamp",
2573 Mf::Saturate => "saturate",
2574 Mf::Cos => "cos",
2576 Mf::Cosh => "cosh",
2577 Mf::Sin => "sin",
2578 Mf::Sinh => "sinh",
2579 Mf::Tan => "tan",
2580 Mf::Tanh => "tanh",
2581 Mf::Acos => "acos",
2582 Mf::Asin => "asin",
2583 Mf::Atan => "atan",
2584 Mf::Atan2 => "atan2",
2585 Mf::Asinh => "asinh",
2586 Mf::Acosh => "acosh",
2587 Mf::Atanh => "atanh",
2588 Mf::Radians => "",
2589 Mf::Degrees => "",
2590 Mf::Ceil => "ceil",
2592 Mf::Floor => "floor",
2593 Mf::Round => "rint",
2594 Mf::Fract => "fract",
2595 Mf::Trunc => "trunc",
2596 Mf::Modf => MODF_FUNCTION,
2597 Mf::Frexp => FREXP_FUNCTION,
2598 Mf::Ldexp => "ldexp",
2599 Mf::Exp => "exp",
2601 Mf::Exp2 => "exp2",
2602 Mf::Log => "log",
2603 Mf::Log2 => "log2",
2604 Mf::Pow => "pow",
2605 Mf::Dot => match *context.resolve_type(arg) {
2607 crate::TypeInner::Vector {
2608 scalar:
2609 crate::Scalar {
2610 kind: crate::ScalarKind::Float,
2612 ..
2613 },
2614 ..
2615 } => "dot",
2616 crate::TypeInner::Vector {
2617 size,
2618 scalar:
2619 scalar @ crate::Scalar {
2620 kind: crate::ScalarKind::Sint | crate::ScalarKind::Uint,
2621 ..
2622 },
2623 } => {
2624 let fun_name = self.get_dot_wrapper_function_helper_name(scalar, size);
2626 write!(self.out, "{fun_name}(")?;
2627 self.put_expression(arg, context, true)?;
2628 write!(self.out, ", ")?;
2629 self.put_expression(arg1.unwrap(), context, true)?;
2630 write!(self.out, ")")?;
2631 return Ok(());
2632 }
2633 _ => unreachable!(
2634 "Correct TypeInner for dot product should be already validated"
2635 ),
2636 },
2637 fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed) => {
2638 if context.lang_version >= (2, 1) {
2639 let packed_type = match fun {
2643 Mf::Dot4I8Packed => "packed_char4",
2644 Mf::Dot4U8Packed => "packed_uchar4",
2645 _ => unreachable!(),
2646 };
2647
2648 return self.put_dot_product(
2649 Reinterpreted::new(packed_type, arg),
2650 Reinterpreted::new(packed_type, arg1.unwrap()),
2651 4,
2652 |writer, arg, index| {
2653 write!(writer.out, "{arg}[{index}]")?;
2656 Ok(())
2657 },
2658 );
2659 } else {
2660 let conversion = match fun {
2664 Mf::Dot4I8Packed => "int",
2665 Mf::Dot4U8Packed => "",
2666 _ => unreachable!(),
2667 };
2668
2669 return self.put_dot_product(
2670 arg,
2671 arg1.unwrap(),
2672 4,
2673 |writer, arg, index| {
2674 write!(writer.out, "({conversion}(")?;
2675 writer.put_expression(arg, context, true)?;
2676 if index == 3 {
2677 write!(writer.out, ") >> 24)")?;
2678 } else {
2679 write!(writer.out, ") << {} >> 24)", (3 - index) * 8)?;
2680 }
2681 Ok(())
2682 },
2683 );
2684 }
2685 }
2686 Mf::Outer => return Err(Error::UnsupportedCall(format!("{fun:?}"))),
2687 Mf::Cross => "cross",
2688 Mf::Distance => "distance",
2689 Mf::Length if scalar_argument => "abs",
2690 Mf::Length => "length",
2691 Mf::Normalize => "normalize",
2692 Mf::FaceForward => "faceforward",
2693 Mf::Reflect => "reflect",
2694 Mf::Refract => "refract",
2695 Mf::Sign => match arg_type.scalar_kind() {
2697 Some(crate::ScalarKind::Sint) => {
2698 return self.put_isign(arg, context);
2699 }
2700 _ => "sign",
2701 },
2702 Mf::Fma => "fma",
2703 Mf::Mix => "mix",
2704 Mf::Step => "step",
2705 Mf::SmoothStep => "smoothstep",
2706 Mf::Sqrt => "sqrt",
2707 Mf::InverseSqrt => "rsqrt",
2708 Mf::Inverse => return Err(Error::UnsupportedCall(format!("{fun:?}"))),
2709 Mf::Transpose => "transpose",
2710 Mf::Determinant => "determinant",
2711 Mf::QuantizeToF16 => "",
2712 Mf::CountTrailingZeros => "ctz",
2714 Mf::CountLeadingZeros => "clz",
2715 Mf::CountOneBits => "popcount",
2716 Mf::ReverseBits => "reverse_bits",
2717 Mf::ExtractBits => "",
2718 Mf::InsertBits => "",
2719 Mf::FirstTrailingBit => "",
2720 Mf::FirstLeadingBit => "",
2721 Mf::Pack4x8snorm => "pack_float_to_snorm4x8",
2723 Mf::Pack4x8unorm => "pack_float_to_unorm4x8",
2724 Mf::Pack2x16snorm => "pack_float_to_snorm2x16",
2725 Mf::Pack2x16unorm => "pack_float_to_unorm2x16",
2726 Mf::Pack2x16float => "",
2727 Mf::Pack4xI8 => "",
2728 Mf::Pack4xU8 => "",
2729 Mf::Pack4xI8Clamp => "",
2730 Mf::Pack4xU8Clamp => "",
2731 Mf::Unpack4x8snorm => "unpack_snorm4x8_to_float",
2733 Mf::Unpack4x8unorm => "unpack_unorm4x8_to_float",
2734 Mf::Unpack2x16snorm => "unpack_snorm2x16_to_float",
2735 Mf::Unpack2x16unorm => "unpack_unorm2x16_to_float",
2736 Mf::Unpack2x16float => "",
2737 Mf::Unpack4xI8 => "",
2738 Mf::Unpack4xU8 => "",
2739 };
2740
2741 match fun {
2742 Mf::ReverseBits | Mf::ExtractBits | Mf::InsertBits => {
2743 if context.lang_version < (1, 2) {
2752 return Err(Error::UnsupportedFunction(fun_name.to_string()));
2753 }
2754 }
2755 _ => {}
2756 }
2757
2758 match fun {
2759 Mf::Abs if arg_type.scalar_kind() == Some(crate::ScalarKind::Sint) => {
2760 write!(self.out, "{ABS_FUNCTION}(")?;
2761 self.put_expression(arg, context, true)?;
2762 write!(self.out, ")")?;
2763 }
2764 Mf::Distance if scalar_argument => {
2765 write!(self.out, "{NAMESPACE}::abs(")?;
2766 self.put_expression(arg, context, false)?;
2767 write!(self.out, " - ")?;
2768 self.put_expression(arg1.unwrap(), context, false)?;
2769 write!(self.out, ")")?;
2770 }
2771 Mf::FirstTrailingBit => {
2772 let scalar = context.resolve_type(arg).scalar().unwrap();
2773 let constant = scalar.width * 8 + 1;
2774
2775 write!(self.out, "((({NAMESPACE}::ctz(")?;
2776 self.put_expression(arg, context, true)?;
2777 write!(self.out, ") + 1) % {constant}) - 1)")?;
2778 }
2779 Mf::FirstLeadingBit => {
2780 let inner = context.resolve_type(arg);
2781 let scalar = inner.scalar().unwrap();
2782 let constant = scalar.width * 8 - 1;
2783
2784 write!(
2785 self.out,
2786 "{NAMESPACE}::select({constant} - {NAMESPACE}::clz("
2787 )?;
2788
2789 if scalar.kind == crate::ScalarKind::Sint {
2790 write!(self.out, "{NAMESPACE}::select(")?;
2791 self.put_expression(arg, context, true)?;
2792 write!(self.out, ", ~")?;
2793 self.put_expression(arg, context, true)?;
2794 write!(self.out, ", ")?;
2795 self.put_expression(arg, context, true)?;
2796 write!(self.out, " < 0)")?;
2797 } else {
2798 self.put_expression(arg, context, true)?;
2799 }
2800
2801 write!(self.out, "), ")?;
2802
2803 match *inner {
2805 crate::TypeInner::Vector { size, scalar } => {
2806 let size = common::vector_size_str(size);
2807 let name = scalar.to_msl_name();
2808 write!(self.out, "{name}{size}")?;
2809 }
2810 crate::TypeInner::Scalar(scalar) => {
2811 let name = scalar.to_msl_name();
2812 write!(self.out, "{name}")?;
2813 }
2814 _ => (),
2815 }
2816
2817 write!(self.out, "(-1), ")?;
2818 self.put_expression(arg, context, true)?;
2819 write!(self.out, " == 0")?;
2820 if scalar.kind == crate::ScalarKind::Sint {
2821 write!(self.out, " || ")?;
2822 self.put_expression(arg, context, true)?;
2823 write!(self.out, " == -1")?;
2824 }
2825 write!(self.out, ")")?;
2826 }
2827 Mf::Unpack2x16float => {
2828 write!(self.out, "float2(as_type<half2>(")?;
2829 self.put_expression(arg, context, false)?;
2830 write!(self.out, "))")?;
2831 }
2832 Mf::Pack2x16float => {
2833 write!(self.out, "as_type<uint>(half2(")?;
2834 self.put_expression(arg, context, false)?;
2835 write!(self.out, "))")?;
2836 }
2837 Mf::ExtractBits => {
2838 let scalar_bits = context.resolve_type(arg).scalar_width().unwrap() * 8;
2855
2856 write!(self.out, "{NAMESPACE}::extract_bits(")?;
2857 self.put_expression(arg, context, true)?;
2858 write!(self.out, ", {NAMESPACE}::min(")?;
2859 self.put_expression(arg1.unwrap(), context, true)?;
2860 write!(self.out, ", {scalar_bits}u), {NAMESPACE}::min(")?;
2861 self.put_expression(arg2.unwrap(), context, true)?;
2862 write!(self.out, ", {scalar_bits}u - {NAMESPACE}::min(")?;
2863 self.put_expression(arg1.unwrap(), context, true)?;
2864 write!(self.out, ", {scalar_bits}u)))")?;
2865 }
2866 Mf::InsertBits => {
2867 let scalar_bits = context.resolve_type(arg).scalar_width().unwrap() * 8;
2872
2873 write!(self.out, "{NAMESPACE}::insert_bits(")?;
2874 self.put_expression(arg, context, true)?;
2875 write!(self.out, ", ")?;
2876 self.put_expression(arg1.unwrap(), context, true)?;
2877 write!(self.out, ", {NAMESPACE}::min(")?;
2878 self.put_expression(arg2.unwrap(), context, true)?;
2879 write!(self.out, ", {scalar_bits}u), {NAMESPACE}::min(")?;
2880 self.put_expression(arg3.unwrap(), context, true)?;
2881 write!(self.out, ", {scalar_bits}u - {NAMESPACE}::min(")?;
2882 self.put_expression(arg2.unwrap(), context, true)?;
2883 write!(self.out, ", {scalar_bits}u)))")?;
2884 }
2885 Mf::Radians => {
2886 write!(self.out, "((")?;
2887 self.put_expression(arg, context, false)?;
2888 write!(self.out, ") * 0.017453292519943295474)")?;
2889 }
2890 Mf::Degrees => {
2891 write!(self.out, "((")?;
2892 self.put_expression(arg, context, false)?;
2893 write!(self.out, ") * 57.295779513082322865)")?;
2894 }
2895 Mf::Modf | Mf::Frexp => {
2896 write!(self.out, "{fun_name}")?;
2897 self.put_call_parameters(iter::once(arg), context)?;
2898 }
2899 Mf::Pack4xI8 => self.put_pack4x8(arg, context, true, None)?,
2900 Mf::Pack4xU8 => self.put_pack4x8(arg, context, false, None)?,
2901 Mf::Pack4xI8Clamp => {
2902 self.put_pack4x8(arg, context, true, Some(("-128", "127")))?
2903 }
2904 Mf::Pack4xU8Clamp => {
2905 self.put_pack4x8(arg, context, false, Some(("0", "255")))?
2906 }
2907 fun @ (Mf::Unpack4xI8 | Mf::Unpack4xU8) => {
2908 let sign_prefix = if matches!(fun, Mf::Unpack4xU8) {
2909 "u"
2910 } else {
2911 ""
2912 };
2913
2914 if context.lang_version >= (2, 1) {
2915 write!(
2917 self.out,
2918 "{sign_prefix}int4(as_type<packed_{sign_prefix}char4>("
2919 )?;
2920 self.put_expression(arg, context, true)?;
2921 write!(self.out, "))")?;
2922 } else {
2923 write!(self.out, "({sign_prefix}int4(")?;
2925 self.put_expression(arg, context, true)?;
2926 write!(self.out, ", ")?;
2927 self.put_expression(arg, context, true)?;
2928 write!(self.out, " >> 8, ")?;
2929 self.put_expression(arg, context, true)?;
2930 write!(self.out, " >> 16, ")?;
2931 self.put_expression(arg, context, true)?;
2932 write!(self.out, " >> 24) << 24 >> 24)")?;
2933 }
2934 }
2935 Mf::QuantizeToF16 => {
2936 match *context.resolve_type(arg) {
2937 crate::TypeInner::Scalar { .. } => write!(self.out, "float(half(")?,
2938 crate::TypeInner::Vector { size, .. } => write!(
2939 self.out,
2940 "{NAMESPACE}::float{size}({NAMESPACE}::half{size}(",
2941 size = common::vector_size_str(size),
2942 )?,
2943 _ => unreachable!(
2944 "Correct TypeInner for QuantizeToF16 should be already validated"
2945 ),
2946 };
2947
2948 self.put_expression(arg, context, true)?;
2949 write!(self.out, "))")?;
2950 }
2951 _ => {
2952 write!(self.out, "{NAMESPACE}::{fun_name}")?;
2953 self.put_call_parameters(
2954 iter::once(arg).chain(arg1).chain(arg2).chain(arg3),
2955 context,
2956 )?;
2957 }
2958 }
2959 }
2960 crate::Expression::As {
2961 expr,
2962 kind,
2963 convert,
2964 } => match *context.resolve_type(expr) {
2965 crate::TypeInner::Scalar(src) | crate::TypeInner::Vector { scalar: src, .. } => {
2966 if src.kind == crate::ScalarKind::Float
2967 && (kind == crate::ScalarKind::Sint || kind == crate::ScalarKind::Uint)
2968 && convert.is_some()
2969 {
2970 let fun_name = match (kind, convert) {
2974 (crate::ScalarKind::Sint, Some(4)) => F2I32_FUNCTION,
2975 (crate::ScalarKind::Uint, Some(4)) => F2U32_FUNCTION,
2976 (crate::ScalarKind::Sint, Some(8)) => F2I64_FUNCTION,
2977 (crate::ScalarKind::Uint, Some(8)) => F2U64_FUNCTION,
2978 _ => unreachable!(),
2979 };
2980 write!(self.out, "{fun_name}(")?;
2981 self.put_expression(expr, context, true)?;
2982 write!(self.out, ")")?;
2983 } else {
2984 let target_scalar = crate::Scalar {
2985 kind,
2986 width: convert.unwrap_or(src.width),
2987 };
2988 let op = match convert {
2989 Some(_) => "static_cast",
2990 None => "as_type",
2991 };
2992 write!(self.out, "{op}<")?;
2993 match *context.resolve_type(expr) {
2994 crate::TypeInner::Vector { size, .. } => {
2995 put_numeric_type(&mut self.out, target_scalar, &[size])?
2996 }
2997 _ => put_numeric_type(&mut self.out, target_scalar, &[])?,
2998 };
2999 write!(self.out, ">(")?;
3000 self.put_expression(expr, context, true)?;
3001 write!(self.out, ")")?;
3002 }
3003 }
3004 crate::TypeInner::Matrix {
3005 columns,
3006 rows,
3007 scalar,
3008 } => {
3009 let target_scalar = crate::Scalar {
3010 kind,
3011 width: convert.unwrap_or(scalar.width),
3012 };
3013 put_numeric_type(&mut self.out, target_scalar, &[rows, columns])?;
3014 write!(self.out, "(")?;
3015 self.put_expression(expr, context, true)?;
3016 write!(self.out, ")")?;
3017 }
3018 ref ty => {
3019 return Err(Error::GenericValidation(format!(
3020 "Unsupported type for As: {ty:?}"
3021 )))
3022 }
3023 },
3024 crate::Expression::CallResult(_)
3026 | crate::Expression::AtomicResult { .. }
3027 | crate::Expression::WorkGroupUniformLoadResult { .. }
3028 | crate::Expression::SubgroupBallotResult
3029 | crate::Expression::SubgroupOperationResult { .. }
3030 | crate::Expression::RayQueryProceedResult => {
3031 unreachable!()
3032 }
3033 crate::Expression::ArrayLength(expr) => {
3034 let global = context.function.originating_global(expr).ok_or_else(|| {
3035 Error::GenericValidation(format!(
3036 "Could not find global variable for ArrayLength operand {:?}",
3037 context.function.expressions[expr]
3038 ))
3039 })?;
3040
3041 if !is_scoped {
3042 write!(self.out, "(")?;
3043 }
3044 write!(self.out, "1 + ")?;
3045 self.put_dynamic_array_max_index(global, expr, context)?;
3046 if !is_scoped {
3047 write!(self.out, ")")?;
3048 }
3049 }
3050 crate::Expression::RayQueryVertexPositions { .. } => {
3051 unimplemented!()
3052 }
3053 crate::Expression::RayQueryGetIntersection { query, committed } => {
3054 if context.lang_version < (2, 4) {
3055 return Err(Error::UnsupportedRayTracing);
3056 }
3057
3058 let crate::Expression::LocalVariable(query_var) =
3060 context.function.expressions[query]
3061 else {
3062 unreachable!()
3063 };
3064
3065 let tracker_expr_name = format!(
3066 "{}{}",
3067 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
3068 self.names[&NameKey::local(context.origin, query_var)]
3069 );
3070
3071 write!(
3072 self.out,
3073 "{}_{committed}(",
3074 super::ray::INTERSECTION_FUNCTION_NAME
3075 )?;
3076 self.put_expression(query, context, true)?;
3077 if context.ray_query_initialization_tracking {
3078 write!(self.out, ", {tracker_expr_name}")?;
3079 }
3080 write!(self.out, ")")?;
3081 }
3082 crate::Expression::CooperativeLoad { ref data, .. } => {
3083 if context.lang_version < (2, 3) {
3084 return Err(Error::UnsupportedCooperativeMatrix);
3085 }
3086 write!(self.out, "{COOPERATIVE_LOAD_FUNCTION}(")?;
3087 write!(self.out, "&")?;
3088 self.put_access_chain(data.pointer, context.policies.index, context)?;
3089 write!(self.out, ", ")?;
3090 self.put_expression(data.stride, context, true)?;
3091 write!(self.out, ", {})", !data.row_major)?;
3098 }
3099 crate::Expression::CooperativeMultiplyAdd { a, b, c } => {
3100 if context.lang_version < (2, 3) {
3101 return Err(Error::UnsupportedCooperativeMatrix);
3102 }
3103 write!(self.out, "{COOPERATIVE_MULTIPLY_ADD_FUNCTION}(")?;
3104 self.put_expression(a, context, true)?;
3105 write!(self.out, ", ")?;
3106 self.put_expression(b, context, true)?;
3107 write!(self.out, ", ")?;
3108 self.put_expression(c, context, true)?;
3109 write!(self.out, ")")?;
3110 }
3111 }
3112 Ok(())
3113 }
3114
3115 fn put_binop<F>(
3118 &mut self,
3119 op: crate::BinaryOperator,
3120 left: Handle<crate::Expression>,
3121 right: Handle<crate::Expression>,
3122 context: &ExpressionContext,
3123 is_scoped: bool,
3124 put_expression: &F,
3125 ) -> BackendResult
3126 where
3127 F: Fn(&mut Self, Handle<crate::Expression>, &ExpressionContext, bool) -> BackendResult,
3128 {
3129 let op_str = back::binary_operation_str(op);
3130
3131 if !is_scoped {
3132 write!(self.out, "(")?;
3133 }
3134
3135 if op == crate::BinaryOperator::Multiply
3138 && matches!(
3139 context.resolve_type(right),
3140 &crate::TypeInner::Matrix { .. }
3141 )
3142 {
3143 self.put_wrapped_expression_for_packed_vec3_access(
3144 left,
3145 context,
3146 false,
3147 put_expression,
3148 )?;
3149 } else {
3150 put_expression(self, left, context, false)?;
3151 }
3152
3153 write!(self.out, " {op_str} ")?;
3154
3155 if op == crate::BinaryOperator::Multiply
3157 && matches!(context.resolve_type(left), &crate::TypeInner::Matrix { .. })
3158 {
3159 self.put_wrapped_expression_for_packed_vec3_access(
3160 right,
3161 context,
3162 false,
3163 put_expression,
3164 )?;
3165 } else {
3166 put_expression(self, right, context, false)?;
3167 }
3168
3169 if !is_scoped {
3170 write!(self.out, ")")?;
3171 }
3172
3173 Ok(())
3174 }
3175
3176 fn put_wrapped_expression_for_packed_vec3_access<F>(
3178 &mut self,
3179 expr_handle: Handle<crate::Expression>,
3180 context: &ExpressionContext,
3181 is_scoped: bool,
3182 put_expression: &F,
3183 ) -> BackendResult
3184 where
3185 F: Fn(&mut Self, Handle<crate::Expression>, &ExpressionContext, bool) -> BackendResult,
3186 {
3187 if let Some(scalar) = context.get_packed_vec_kind(expr_handle) {
3188 write!(self.out, "{}::{}3(", NAMESPACE, scalar.to_msl_name())?;
3189 put_expression(self, expr_handle, context, is_scoped)?;
3190 write!(self.out, ")")?;
3191 } else {
3192 put_expression(self, expr_handle, context, is_scoped)?;
3193 }
3194 Ok(())
3195 }
3196
3197 fn put_bitcasted_expression<F>(
3200 &mut self,
3201 cast_to: &crate::TypeInner,
3202 inner_expr: Handle<crate::Expression>,
3203 context: &ExpressionContext,
3204 put_expression: &F,
3205 ) -> BackendResult
3206 where
3207 F: Fn(&mut Self, &ExpressionContext, bool) -> BackendResult,
3208 {
3209 let needs_truncation = match *cast_to {
3214 crate::TypeInner::Scalar(scalar) => scalar.width < 4,
3215 crate::TypeInner::Vector { scalar, .. } => scalar.width < 4,
3216 _ => false,
3217 };
3218
3219 write!(self.out, "as_type<")?;
3220 match *cast_to {
3221 crate::TypeInner::Scalar(scalar) => put_numeric_type(&mut self.out, scalar, &[])?,
3222 crate::TypeInner::Vector { size, scalar } => {
3223 put_numeric_type(&mut self.out, scalar, &[size])?
3224 }
3225 _ => return Err(Error::UnsupportedBitCast(cast_to.clone())),
3226 };
3227 write!(self.out, ">(")?;
3228
3229 if needs_truncation {
3230 write!(self.out, "static_cast<")?;
3231 let unsigned_scalar = match *cast_to {
3233 crate::TypeInner::Scalar(scalar) => crate::Scalar {
3234 kind: crate::ScalarKind::Uint,
3235 ..scalar
3236 },
3237 crate::TypeInner::Vector { scalar, .. } => crate::Scalar {
3238 kind: crate::ScalarKind::Uint,
3239 ..scalar
3240 },
3241 _ => unreachable!(),
3242 };
3243 match *cast_to {
3244 crate::TypeInner::Scalar(_) => {
3245 put_numeric_type(&mut self.out, unsigned_scalar, &[])?
3246 }
3247 crate::TypeInner::Vector { size, .. } => {
3248 put_numeric_type(&mut self.out, unsigned_scalar, &[size])?
3249 }
3250 _ => unreachable!(),
3251 };
3252 write!(self.out, ">(")?;
3253 }
3254
3255 if let Some(scalar) = context.get_packed_vec_kind(inner_expr) {
3257 put_numeric_type(&mut self.out, scalar, &[crate::VectorSize::Tri])?;
3258 write!(self.out, "(")?;
3259 put_expression(self, context, true)?;
3260 write!(self.out, ")")?;
3261 } else {
3262 put_expression(self, context, true)?;
3263 }
3264
3265 if needs_truncation {
3266 write!(self.out, ")")?;
3267 }
3268
3269 write!(self.out, ")")?;
3270 Ok(())
3271 }
3272
3273 fn put_index(
3275 &mut self,
3276 index: index::GuardedIndex,
3277 context: &ExpressionContext,
3278 is_scoped: bool,
3279 ) -> BackendResult {
3280 match index {
3281 index::GuardedIndex::Expression(expr) => {
3282 self.put_expression(expr, context, is_scoped)?
3283 }
3284 index::GuardedIndex::Known(value) => write!(self.out, "{value}")?,
3285 }
3286 Ok(())
3287 }
3288
3289 fn put_bounds_checks(
3319 &mut self,
3320 chain: Handle<crate::Expression>,
3321 context: &ExpressionContext,
3322 level: back::Level,
3323 prefix: &'static str,
3324 ) -> Result<bool, Error> {
3325 let mut check_written = false;
3326
3327 for item in context.bounds_check_iter(chain) {
3329 let BoundsCheck {
3330 base,
3331 index,
3332 length,
3333 } = item;
3334
3335 if check_written {
3336 write!(self.out, " && ")?;
3337 } else {
3338 write!(self.out, "{level}{prefix}")?;
3339 check_written = true;
3340 }
3341
3342 write!(self.out, "uint(")?;
3346 self.put_index(index, context, true)?;
3347 self.out.write_str(") < ")?;
3348 match length {
3349 index::IndexableLength::Known(value) => write!(self.out, "{value}")?,
3350 index::IndexableLength::Dynamic => {
3351 let global = context.function.originating_global(base).ok_or_else(|| {
3352 Error::GenericValidation("Could not find originating global".into())
3353 })?;
3354 if matches!(
3355 context.module.types[context.module.global_variables[global].ty].inner,
3356 crate::TypeInner::BindingArray { .. }
3357 ) {
3358 write!(
3359 self.out,
3360 "{} && _buffer_sizes.{}[",
3361 Self::binding_array_layout_count(
3362 context.module,
3363 context.pipeline_options,
3364 global,
3365 ),
3366 ArraySizeMember(global),
3367 )?;
3368 self.put_binding_array_size_member_index(index, context)?;
3369 write!(self.out, "] != 0u")?;
3370 } else {
3371 write!(self.out, "1 + ")?;
3372 self.put_dynamic_array_max_index(global, base, context)?
3373 }
3374 }
3375 }
3376 }
3377
3378 Ok(check_written)
3379 }
3380
3381 fn put_access_chain(
3401 &mut self,
3402 chain: Handle<crate::Expression>,
3403 policy: index::BoundsCheckPolicy,
3404 context: &ExpressionContext,
3405 ) -> BackendResult {
3406 match context.function.expressions[chain] {
3407 crate::Expression::Access { base, index } => {
3408 let mut base_ty = context.resolve_type(base);
3409
3410 if let crate::TypeInner::Pointer { base, space: _ } = *base_ty {
3412 base_ty = &context.module.types[base].inner;
3413 }
3414
3415 self.put_subscripted_access_chain(
3416 base,
3417 base_ty,
3418 index::GuardedIndex::Expression(index),
3419 policy,
3420 context,
3421 )?;
3422 }
3423 crate::Expression::AccessIndex { base, index } => {
3424 let base_resolution = &context.info[base].ty;
3425 let mut base_ty = base_resolution.inner_with(&context.module.types);
3426 let mut base_ty_handle = base_resolution.handle();
3427
3428 if let crate::TypeInner::Pointer { base, space: _ } = *base_ty {
3430 base_ty = &context.module.types[base].inner;
3431 base_ty_handle = Some(base);
3432 }
3433
3434 match *base_ty {
3438 crate::TypeInner::Struct { .. } => {
3439 let base_ty = base_ty_handle.unwrap();
3440 self.put_access_chain(base, policy, context)?;
3441 let name = &self.names[&NameKey::StructMember(base_ty, index)];
3442 write!(
3443 self.out,
3444 "{}{name}",
3445 if context.struct_member_needs_arrow(base, |ty| {
3446 matches!(ty, crate::TypeInner::BindingArray { .. })
3447 }) {
3448 "->"
3449 } else {
3450 "."
3451 },
3452 )?;
3453 }
3454 crate::TypeInner::ValuePointer { .. } | crate::TypeInner::Vector { .. } => {
3455 self.put_access_chain(base, policy, context)?;
3456 if context.get_packed_vec_kind(base).is_some() {
3459 write!(self.out, "[{index}]")?;
3460 } else {
3461 write!(self.out, ".{}", back::COMPONENTS[index as usize])?;
3462 }
3463 }
3464 _ => {
3465 self.put_subscripted_access_chain(
3466 base,
3467 base_ty,
3468 index::GuardedIndex::Known(index),
3469 policy,
3470 context,
3471 )?;
3472 }
3473 }
3474 }
3475 _ => self.put_expression(chain, context, false)?,
3476 }
3477
3478 Ok(())
3479 }
3480
3481 fn put_subscripted_access_chain(
3498 &mut self,
3499 base: Handle<crate::Expression>,
3500 base_ty: &crate::TypeInner,
3501 index: index::GuardedIndex,
3502 policy: index::BoundsCheckPolicy,
3503 context: &ExpressionContext,
3504 ) -> BackendResult {
3505 let accessing_wrapped_array = match *base_ty {
3506 crate::TypeInner::Array {
3507 size: crate::ArraySize::Constant(_) | crate::ArraySize::Pending(_),
3508 ..
3509 } => true,
3510 _ => false,
3511 };
3512 let accessing_wrapped_binding_array =
3513 matches!(*base_ty, crate::TypeInner::BindingArray { .. });
3514
3515 self.put_access_chain(base, policy, context)?;
3516 if accessing_wrapped_array {
3517 write!(self.out, ".{WRAPPED_ARRAY_FIELD}")?;
3518 }
3519 write!(self.out, "[")?;
3520
3521 let restriction_needed = if policy == index::BoundsCheckPolicy::Restrict {
3523 context.access_needs_check(base, index)
3524 } else {
3525 None
3526 };
3527 if let Some(limit) = restriction_needed {
3528 write!(self.out, "{NAMESPACE}::min(unsigned(")?;
3529 self.put_index(index, context, true)?;
3530 write!(self.out, "), ")?;
3531 match limit {
3532 index::IndexableLength::Known(limit) => {
3533 write!(self.out, "{}u", limit - 1)?;
3534 }
3535 index::IndexableLength::Dynamic => {
3536 let global = context.function.originating_global(base).ok_or_else(|| {
3537 Error::GenericValidation("Could not find originating global".into())
3538 })?;
3539 self.put_dynamic_array_max_index(global, base, context)?;
3540 }
3541 }
3542 write!(self.out, ")")?;
3543 } else {
3544 self.put_index(index, context, true)?;
3545 }
3546
3547 write!(self.out, "]")?;
3548
3549 if accessing_wrapped_binding_array {
3550 write!(self.out, ".{WRAPPED_ARRAY_FIELD}")?;
3551 }
3552
3553 Ok(())
3554 }
3555
3556 fn put_load(
3557 &mut self,
3558 pointer: Handle<crate::Expression>,
3559 context: &ExpressionContext,
3560 is_scoped: bool,
3561 ) -> BackendResult {
3562 let policy = context.choose_bounds_check_policy(pointer);
3565 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
3566 && self.put_bounds_checks(
3567 pointer,
3568 context,
3569 back::Level(0),
3570 if is_scoped { "" } else { "(" },
3571 )?
3572 {
3573 write!(self.out, " ? ")?;
3574 self.put_unchecked_load(pointer, policy, context)?;
3575 write!(self.out, " : DefaultConstructible()")?;
3576
3577 if !is_scoped {
3578 write!(self.out, ")")?;
3579 }
3580 } else {
3581 self.put_unchecked_load(pointer, policy, context)?;
3582 }
3583
3584 Ok(())
3585 }
3586
3587 fn put_unchecked_load(
3588 &mut self,
3589 pointer: Handle<crate::Expression>,
3590 policy: index::BoundsCheckPolicy,
3591 context: &ExpressionContext,
3592 ) -> BackendResult {
3593 let is_atomic_pointer = context
3594 .resolve_type(pointer)
3595 .is_atomic_pointer(&context.module.types);
3596
3597 if is_atomic_pointer {
3598 write!(
3599 self.out,
3600 "{NAMESPACE}::atomic_load_explicit({ATOMIC_REFERENCE}"
3601 )?;
3602 self.put_access_chain(pointer, policy, context)?;
3603 write!(self.out, ", {NAMESPACE}::memory_order_relaxed)")?;
3604 } else {
3605 self.put_access_chain(pointer, policy, context)?;
3609 }
3610
3611 Ok(())
3612 }
3613
3614 fn put_return_value(
3615 &mut self,
3616 level: back::Level,
3617 expr_handle: Handle<crate::Expression>,
3618 result_struct: Option<&str>,
3619 context: &ExpressionContext,
3620 ) -> BackendResult {
3621 match result_struct {
3622 Some(struct_name) => {
3623 let mut has_point_size = false;
3624 let result_ty = context.function.result.as_ref().unwrap().ty;
3625 match context.module.types[result_ty].inner {
3626 crate::TypeInner::Struct { ref members, .. } => {
3627 let tmp = self.namer.call("_tmp");
3628 write!(self.out, "{level}const auto {tmp} = ")?;
3629 self.put_expression(expr_handle, context, true)?;
3630 writeln!(self.out, ";")?;
3631 write!(self.out, "{level}return {struct_name} {{")?;
3632
3633 let mut is_first = true;
3634
3635 for (index, member) in members.iter().enumerate() {
3636 if let Some(crate::Binding::BuiltIn(crate::BuiltIn::PointSize)) =
3637 member.binding
3638 {
3639 has_point_size = true;
3640 if !context.pipeline_options.allow_and_force_point_size {
3641 continue;
3642 }
3643 }
3644
3645 let comma = if is_first { "" } else { "," };
3646 is_first = false;
3647 let name = &self.names[&NameKey::StructMember(result_ty, index as u32)];
3648 if let crate::TypeInner::Array {
3652 size: crate::ArraySize::Constant(size),
3653 ..
3654 } = context.module.types[member.ty].inner
3655 {
3656 write!(self.out, "{comma} {{")?;
3657 for j in 0..size.get() {
3658 if j != 0 {
3659 write!(self.out, ",")?;
3660 }
3661 write!(self.out, "{tmp}.{name}.{WRAPPED_ARRAY_FIELD}[{j}]")?;
3662 }
3663 write!(self.out, "}}")?;
3664 } else {
3665 write!(self.out, "{comma} {tmp}.{name}")?;
3666 }
3667 }
3668 }
3669 _ => {
3670 write!(self.out, "{level}return {struct_name} {{ ")?;
3671 self.put_expression(expr_handle, context, true)?;
3672 }
3673 }
3674
3675 if let FunctionOrigin::EntryPoint(ep_index) = context.origin {
3676 let stage = context.module.entry_points[ep_index as usize].stage;
3677 if context.pipeline_options.allow_and_force_point_size
3678 && stage == crate::ShaderStage::Vertex
3679 && !has_point_size
3680 {
3681 write!(self.out, ", 1.0")?;
3683 }
3684 }
3685 write!(self.out, " }}")?;
3686 }
3687 None => {
3688 write!(self.out, "{level}return ")?;
3689 self.put_expression(expr_handle, context, true)?;
3690 }
3691 }
3692 writeln!(self.out, ";")?;
3693 Ok(())
3694 }
3695
3696 fn update_expressions_to_bake(
3701 &mut self,
3702 func: &crate::Function,
3703 info: &valid::FunctionInfo,
3704 context: &ExpressionContext,
3705 ) {
3706 use crate::Expression;
3707 self.need_bake_expressions.clear();
3708
3709 for (expr_handle, expr) in func.expressions.iter() {
3710 let expr_info = &info[expr_handle];
3713 let min_ref_count = func.expressions[expr_handle].bake_ref_count();
3714 if min_ref_count <= expr_info.ref_count {
3715 self.need_bake_expressions.insert(expr_handle);
3716 } else {
3717 match expr_info.ty {
3718 TypeResolution::Handle(h)
3720 if Some(h) == context.module.special_types.ray_desc =>
3721 {
3722 self.need_bake_expressions.insert(expr_handle);
3723 }
3724 _ => {}
3725 }
3726 }
3727
3728 if let Expression::Math {
3729 fun,
3730 arg,
3731 arg1,
3732 arg2,
3733 ..
3734 } = *expr
3735 {
3736 match fun {
3737 crate::MathFunction::Dot4U8Packed | crate::MathFunction::Dot4I8Packed => {
3747 self.need_bake_expressions.insert(arg);
3748 self.need_bake_expressions.insert(arg1.unwrap());
3749 }
3750 crate::MathFunction::FirstLeadingBit => {
3751 self.need_bake_expressions.insert(arg);
3752 }
3753 crate::MathFunction::Pack4xI8
3754 | crate::MathFunction::Pack4xU8
3755 | crate::MathFunction::Pack4xI8Clamp
3756 | crate::MathFunction::Pack4xU8Clamp
3757 | crate::MathFunction::Unpack4xI8
3758 | crate::MathFunction::Unpack4xU8 => {
3759 if context.lang_version < (2, 1) {
3762 self.need_bake_expressions.insert(arg);
3763 }
3764 }
3765 crate::MathFunction::ExtractBits => {
3766 self.need_bake_expressions.insert(arg1.unwrap());
3768 }
3769 crate::MathFunction::InsertBits => {
3770 self.need_bake_expressions.insert(arg2.unwrap());
3772 }
3773 crate::MathFunction::Sign => {
3774 let inner = context.resolve_type(expr_handle);
3779 if inner.scalar_kind() == Some(crate::ScalarKind::Sint) {
3780 self.need_bake_expressions.insert(arg);
3781 }
3782 }
3783 _ => {}
3784 }
3785 }
3786 }
3787 }
3788
3789 pub(super) fn start_baking_expression(
3790 &mut self,
3791 handle: Handle<crate::Expression>,
3792 context: &ExpressionContext,
3793 name: &str,
3794 ) -> BackendResult {
3795 match context.info[handle].ty {
3796 TypeResolution::Handle(ty_handle) => {
3797 let ty_name = TypeContext {
3798 handle: ty_handle,
3799 gctx: context.module.to_ctx(),
3800 names: &self.names,
3801 access: crate::StorageAccess::empty(),
3802 first_time: false,
3803 };
3804 write!(self.out, "{ty_name}")?;
3805 }
3806 TypeResolution::Value(crate::TypeInner::Scalar(scalar)) => {
3807 put_numeric_type(&mut self.out, scalar, &[])?;
3808 }
3809 TypeResolution::Value(crate::TypeInner::Vector { size, scalar }) => {
3810 put_numeric_type(&mut self.out, scalar, &[size])?;
3811 }
3812 TypeResolution::Value(crate::TypeInner::Matrix {
3813 columns,
3814 rows,
3815 scalar,
3816 }) => {
3817 put_numeric_type(&mut self.out, scalar, &[rows, columns])?;
3818 }
3819 TypeResolution::Value(crate::TypeInner::CooperativeMatrix {
3820 columns,
3821 rows,
3822 scalar,
3823 role: _,
3824 }) => {
3825 write!(
3826 self.out,
3827 "{}::simdgroup_{}{}x{}",
3828 NAMESPACE,
3829 scalar.to_msl_name(),
3830 columns as u32,
3831 rows as u32,
3832 )?;
3833 }
3834 TypeResolution::Value(ref other) => {
3835 log::warn!("Type {other:?} isn't a known local");
3836 return Err(Error::FeatureNotImplemented("weird local type".to_string()));
3837 }
3838 }
3839
3840 write!(self.out, " {name} = ")?;
3842
3843 Ok(())
3844 }
3845
3846 fn put_cache_restricted_level(
3859 &mut self,
3860 load: Handle<crate::Expression>,
3861 image: Handle<crate::Expression>,
3862 mip_level: Option<Handle<crate::Expression>>,
3863 indent: back::Level,
3864 context: &StatementContext,
3865 ) -> BackendResult {
3866 let level_of_detail = match mip_level {
3869 Some(level) => level,
3870 None => return Ok(()),
3871 };
3872
3873 if context.expression.policies.image_load != index::BoundsCheckPolicy::Restrict
3874 || !context.expression.image_needs_lod(image)
3875 {
3876 return Ok(());
3877 }
3878
3879 write!(self.out, "{}uint {} = ", indent, ClampedLod(load),)?;
3880 self.put_restricted_scalar_image_index(
3881 image,
3882 level_of_detail,
3883 "get_num_mip_levels",
3884 &context.expression,
3885 )?;
3886 writeln!(self.out, ";")?;
3887
3888 Ok(())
3889 }
3890
3891 fn put_casting_to_packed_chars(
3897 &mut self,
3898 fun: crate::MathFunction,
3899 arg0: Handle<crate::Expression>,
3900 arg1: Handle<crate::Expression>,
3901 indent: back::Level,
3902 context: &StatementContext<'_>,
3903 ) -> Result<(), Error> {
3904 let packed_type = match fun {
3905 crate::MathFunction::Dot4I8Packed => "packed_char4",
3906 crate::MathFunction::Dot4U8Packed => "packed_uchar4",
3907 _ => unreachable!(),
3908 };
3909
3910 for arg in [arg0, arg1] {
3911 write!(
3912 self.out,
3913 "{indent}{packed_type} {0} = as_type<{packed_type}>(",
3914 Reinterpreted::new(packed_type, arg)
3915 )?;
3916 self.put_expression(arg, &context.expression, true)?;
3917 writeln!(self.out, ");")?;
3918 }
3919
3920 Ok(())
3921 }
3922
3923 fn put_block(
3924 &mut self,
3925 level: back::Level,
3926 statements: &[crate::Statement],
3927 context: &StatementContext,
3928 ) -> BackendResult {
3929 for statement in statements {
3930 log::trace!("statement[{}] {:?}", level.0, statement);
3931 match *statement {
3932 crate::Statement::Emit(ref range) => {
3933 for handle in range.clone() {
3934 use crate::MathFunction as Mf;
3935
3936 match context.expression.function.expressions[handle] {
3937 crate::Expression::ImageLoad {
3940 image,
3941 level: mip_level,
3942 ..
3943 } => {
3944 self.put_cache_restricted_level(
3945 handle, image, mip_level, level, context,
3946 )?;
3947 }
3948
3949 crate::Expression::Math {
3958 fun: fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed),
3959 arg,
3960 arg1,
3961 ..
3962 } if context.expression.lang_version >= (2, 1) => {
3963 self.put_casting_to_packed_chars(
3964 fun,
3965 arg,
3966 arg1.unwrap(),
3967 level,
3968 context,
3969 )?;
3970 }
3971
3972 _ => (),
3973 }
3974
3975 let ptr_class = context.expression.resolve_type(handle).pointer_space();
3976 let expr_name = if ptr_class.is_some() {
3977 None } else if let Some(name) =
3979 context.expression.function.named_expressions.get(&handle)
3980 {
3981 Some(self.namer.call(name))
3991 } else {
3992 let bake = if context.expression.guarded_indices.contains(handle) {
3996 true
3997 } else {
3998 self.need_bake_expressions.contains(&handle)
3999 };
4000
4001 if bake {
4002 Some(Baked(handle).to_string())
4003 } else {
4004 None
4005 }
4006 };
4007
4008 if let Some(name) = expr_name {
4009 write!(self.out, "{level}")?;
4010 self.start_baking_expression(handle, &context.expression, &name)?;
4011 self.put_expression(handle, &context.expression, true)?;
4012 self.named_expressions.insert(handle, name);
4013 writeln!(self.out, ";")?;
4014 }
4015 }
4016 }
4017 crate::Statement::Block(ref block) => {
4018 if !block.is_empty() {
4019 writeln!(self.out, "{level}{{")?;
4020 self.put_block(level.next(), block, context)?;
4021 writeln!(self.out, "{level}}}")?;
4022 }
4023 }
4024 crate::Statement::If {
4025 condition,
4026 ref accept,
4027 ref reject,
4028 } => {
4029 write!(self.out, "{level}if (")?;
4030 self.put_expression(condition, &context.expression, true)?;
4031 writeln!(self.out, ") {{")?;
4032 self.put_block(level.next(), accept, context)?;
4033 if !reject.is_empty() {
4034 writeln!(self.out, "{level}}} else {{")?;
4035 self.put_block(level.next(), reject, context)?;
4036 }
4037 writeln!(self.out, "{level}}}")?;
4038 }
4039 crate::Statement::Switch {
4040 selector,
4041 ref cases,
4042 } => {
4043 write!(self.out, "{level}switch(")?;
4044 self.put_expression(selector, &context.expression, true)?;
4045 writeln!(self.out, ") {{")?;
4046 let lcase = level.next();
4047 for case in cases.iter() {
4048 match case.value {
4049 crate::SwitchValue::I32(value) => {
4050 write!(self.out, "{lcase}case {value}:")?;
4051 }
4052 crate::SwitchValue::U32(value) => {
4053 write!(self.out, "{lcase}case {value}u:")?;
4054 }
4055 crate::SwitchValue::Default => {
4056 write!(self.out, "{lcase}default:")?;
4057 }
4058 }
4059
4060 let write_block_braces = !(case.fall_through && case.body.is_empty());
4061 if write_block_braces {
4062 writeln!(self.out, " {{")?;
4063 } else {
4064 writeln!(self.out)?;
4065 }
4066
4067 self.put_block(lcase.next(), &case.body, context)?;
4068 if !case.fall_through && case.body.last().is_none_or(|s| !s.is_terminator())
4069 {
4070 writeln!(self.out, "{}break;", lcase.next())?;
4071 }
4072
4073 if write_block_braces {
4074 writeln!(self.out, "{lcase}}}")?;
4075 }
4076 }
4077 writeln!(self.out, "{level}}}")?;
4078 }
4079 crate::Statement::Loop {
4080 ref body,
4081 ref continuing,
4082 break_if,
4083 } => {
4084 let force_loop_bound_statements =
4085 self.gen_force_bounded_loop_statements(level, context);
4086 let gate_name = (!continuing.is_empty() || break_if.is_some())
4087 .then(|| self.namer.call("loop_init"));
4088
4089 if let Some((ref decl, _)) = force_loop_bound_statements {
4090 writeln!(self.out, "{decl}")?;
4091 }
4092 if let Some(ref gate_name) = gate_name {
4093 writeln!(self.out, "{level}bool {gate_name} = true;")?;
4094 }
4095
4096 writeln!(self.out, "{level}while(true) {{",)?;
4097 if let Some((_, ref break_and_inc)) = force_loop_bound_statements {
4098 writeln!(self.out, "{break_and_inc}")?;
4099 }
4100 if let Some(ref gate_name) = gate_name {
4101 let lif = level.next();
4102 let lcontinuing = lif.next();
4103 writeln!(self.out, "{lif}if (!{gate_name}) {{")?;
4104 self.put_block(lcontinuing, continuing, context)?;
4105 if let Some(condition) = break_if {
4106 write!(self.out, "{lcontinuing}if (")?;
4107 self.put_expression(condition, &context.expression, true)?;
4108 writeln!(self.out, ") {{")?;
4109 writeln!(self.out, "{}break;", lcontinuing.next())?;
4110 writeln!(self.out, "{lcontinuing}}}")?;
4111 }
4112 writeln!(self.out, "{lif}}}")?;
4113 writeln!(self.out, "{lif}{gate_name} = false;")?;
4114 }
4115 self.put_block(level.next(), body, context)?;
4116
4117 writeln!(self.out, "{level}}}")?;
4118 }
4119 crate::Statement::Break => {
4120 writeln!(self.out, "{level}break;")?;
4121 }
4122 crate::Statement::Continue => {
4123 writeln!(self.out, "{level}continue;")?;
4124 }
4125 crate::Statement::Return {
4126 value: Some(expr_handle),
4127 } => {
4128 self.put_return_value(
4129 level,
4130 expr_handle,
4131 context.result_struct,
4132 &context.expression,
4133 )?;
4134 }
4135 crate::Statement::Return { value: None } => {
4136 writeln!(self.out, "{level}return;")?;
4137 }
4138 crate::Statement::Kill => {
4139 writeln!(self.out, "{level}{NAMESPACE}::discard_fragment();")?;
4140 }
4141 crate::Statement::ControlBarrier(flags)
4142 | crate::Statement::MemoryBarrier(flags) => {
4143 self.write_barrier(flags, level)?;
4144 }
4145 crate::Statement::Store { pointer, value } => {
4146 self.put_store(pointer, value, level, context)?
4147 }
4148 crate::Statement::ImageStore {
4149 image,
4150 coordinate,
4151 array_index,
4152 value,
4153 } => {
4154 let address = TexelAddress {
4155 coordinate,
4156 array_index,
4157 sample: None,
4158 level: None,
4159 };
4160 self.put_image_store(level, image, &address, value, context)?
4161 }
4162 crate::Statement::Call {
4163 function,
4164 ref arguments,
4165 result,
4166 } => {
4167 write!(self.out, "{level}")?;
4168 if let Some(expr) = result {
4169 let name = Baked(expr).to_string();
4170 self.start_baking_expression(expr, &context.expression, &name)?;
4171 self.named_expressions.insert(expr, name);
4172 }
4173 let fun_name = &self.names[&NameKey::Function(function)];
4174 write!(self.out, "{fun_name}(")?;
4175 for (i, &handle) in arguments.iter().enumerate() {
4177 if i != 0 {
4178 write!(self.out, ", ")?;
4179 }
4180 self.put_expression(handle, &context.expression, true)?;
4181 }
4182 let mut separate = !arguments.is_empty();
4184 let fun_info = &context.expression.mod_info[function];
4185 let mut needs_buffer_sizes = false;
4186 for (handle, var) in context.expression.module.global_variables.iter() {
4187 if fun_info[handle].is_empty() {
4188 continue;
4189 }
4190 if var.space.needs_pass_through() {
4191 let name = &self.names[&NameKey::GlobalVariable(handle)];
4192 if separate {
4193 write!(self.out, ", ")?;
4194 } else {
4195 separate = true;
4196 }
4197 write!(self.out, "{name}")?;
4198 }
4199 needs_buffer_sizes |= context.expression.module.types[var.ty]
4200 .inner
4201 .needs_host_buffer_byte_size(&context.expression.module.types);
4202 }
4203 if needs_buffer_sizes {
4204 if separate {
4205 write!(self.out, ", ")?;
4206 }
4207 write!(self.out, "_buffer_sizes")?;
4208 }
4209
4210 writeln!(self.out, ");")?;
4212 }
4213 crate::Statement::Atomic {
4214 pointer,
4215 ref fun,
4216 value,
4217 result,
4218 } => {
4219 let context = &context.expression;
4220
4221 write!(self.out, "{level}")?;
4226 let fun_key = if let Some(result) = result {
4227 let res_name = Baked(result).to_string();
4228 self.start_baking_expression(result, context, &res_name)?;
4229 self.named_expressions.insert(result, res_name);
4230 fun.to_msl()
4231 } else if context.resolve_type(value).scalar_width() == Some(8) {
4232 fun.to_msl_64_bit()?
4233 } else {
4234 fun.to_msl()
4235 };
4236
4237 let policy = context.choose_bounds_check_policy(pointer);
4241 let checked = policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
4242 && self.put_bounds_checks(pointer, context, back::Level(0), "")?;
4243
4244 if checked {
4246 write!(self.out, " ? ")?;
4247 }
4248
4249 match *fun {
4251 crate::AtomicFunction::Exchange { compare: Some(cmp) } => {
4252 write!(self.out, "{ATOMIC_COMP_EXCH_FUNCTION}({ATOMIC_REFERENCE}")?;
4253 self.put_access_chain(pointer, policy, context)?;
4254 write!(self.out, ", ")?;
4255 self.put_expression(cmp, context, true)?;
4256 write!(self.out, ", ")?;
4257 self.put_expression(value, context, true)?;
4258 write!(self.out, ")")?;
4259 }
4260 _ => {
4261 write!(
4262 self.out,
4263 "{NAMESPACE}::atomic_{fun_key}_explicit({ATOMIC_REFERENCE}"
4264 )?;
4265 self.put_access_chain(pointer, policy, context)?;
4266 write!(self.out, ", ")?;
4267 self.put_expression(value, context, true)?;
4268 write!(self.out, ", {NAMESPACE}::memory_order_relaxed)")?;
4269 }
4270 }
4271
4272 if checked {
4274 write!(self.out, " : DefaultConstructible()")?;
4275 }
4276
4277 writeln!(self.out, ";")?;
4279 }
4280 crate::Statement::ImageAtomic {
4281 image,
4282 coordinate,
4283 array_index,
4284 fun,
4285 value,
4286 } => {
4287 let address = TexelAddress {
4288 coordinate,
4289 array_index,
4290 sample: None,
4291 level: None,
4292 };
4293 self.put_image_atomic(level, image, &address, fun, value, context)?
4294 }
4295 crate::Statement::WorkGroupUniformLoad { pointer, result } => {
4296 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
4297
4298 write!(self.out, "{level}")?;
4299 let name = self.namer.call("");
4300 self.start_baking_expression(result, &context.expression, &name)?;
4301 self.put_load(pointer, &context.expression, true)?;
4302 self.named_expressions.insert(result, name);
4303
4304 writeln!(self.out, ";")?;
4305 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
4306 }
4307 crate::Statement::RayQuery { query, ref fun } => {
4308 self.write_ray_query_stmt(level, context, query, fun)?;
4309 }
4310 crate::Statement::SubgroupBallot { result, predicate } => {
4311 write!(self.out, "{level}")?;
4312 let name = self.namer.call("");
4313 self.start_baking_expression(result, &context.expression, &name)?;
4314 self.named_expressions.insert(result, name);
4315 write!(
4316 self.out,
4317 "{NAMESPACE}::uint4((uint64_t){NAMESPACE}::simd_ballot("
4318 )?;
4319 if let Some(predicate) = predicate {
4320 self.put_expression(predicate, &context.expression, true)?;
4321 } else {
4322 write!(self.out, "true")?;
4323 }
4324 writeln!(self.out, "), 0, 0, 0);")?;
4325 }
4326 crate::Statement::SubgroupCollectiveOperation {
4327 op,
4328 collective_op,
4329 argument,
4330 result,
4331 } => {
4332 write!(self.out, "{level}")?;
4333 let name = self.namer.call("");
4334 self.start_baking_expression(result, &context.expression, &name)?;
4335 self.named_expressions.insert(result, name);
4336 match (collective_op, op) {
4337 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::All) => {
4338 write!(self.out, "{NAMESPACE}::simd_all(")?
4339 }
4340 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Any) => {
4341 write!(self.out, "{NAMESPACE}::simd_any(")?
4342 }
4343 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Add) => {
4344 write!(self.out, "{NAMESPACE}::simd_sum(")?
4345 }
4346 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Mul) => {
4347 write!(self.out, "{NAMESPACE}::simd_product(")?
4348 }
4349 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Max) => {
4350 write!(self.out, "{NAMESPACE}::simd_max(")?
4351 }
4352 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Min) => {
4353 write!(self.out, "{NAMESPACE}::simd_min(")?
4354 }
4355 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::And) => {
4356 write!(self.out, "{NAMESPACE}::simd_and(")?
4357 }
4358 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Or) => {
4359 write!(self.out, "{NAMESPACE}::simd_or(")?
4360 }
4361 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Xor) => {
4362 write!(self.out, "{NAMESPACE}::simd_xor(")?
4363 }
4364 (
4365 crate::CollectiveOperation::ExclusiveScan,
4366 crate::SubgroupOperation::Add,
4367 ) => write!(self.out, "{NAMESPACE}::simd_prefix_exclusive_sum(")?,
4368 (
4369 crate::CollectiveOperation::ExclusiveScan,
4370 crate::SubgroupOperation::Mul,
4371 ) => write!(self.out, "{NAMESPACE}::simd_prefix_exclusive_product(")?,
4372 (
4373 crate::CollectiveOperation::InclusiveScan,
4374 crate::SubgroupOperation::Add,
4375 ) => write!(self.out, "{NAMESPACE}::simd_prefix_inclusive_sum(")?,
4376 (
4377 crate::CollectiveOperation::InclusiveScan,
4378 crate::SubgroupOperation::Mul,
4379 ) => write!(self.out, "{NAMESPACE}::simd_prefix_inclusive_product(")?,
4380 _ => unimplemented!(),
4381 }
4382 self.put_expression(argument, &context.expression, true)?;
4383 writeln!(self.out, ");")?;
4384 }
4385 crate::Statement::SubgroupGather {
4386 mode,
4387 argument,
4388 result,
4389 } => {
4390 write!(self.out, "{level}")?;
4391 let name = self.namer.call("");
4392 self.start_baking_expression(result, &context.expression, &name)?;
4393 self.named_expressions.insert(result, name);
4394 match mode {
4395 crate::GatherMode::BroadcastFirst => {
4396 write!(self.out, "{NAMESPACE}::simd_broadcast_first(")?;
4397 }
4398 crate::GatherMode::Broadcast(_) => {
4399 write!(self.out, "{NAMESPACE}::simd_broadcast(")?;
4400 }
4401 crate::GatherMode::Shuffle(_) => {
4402 write!(self.out, "{NAMESPACE}::simd_shuffle(")?;
4403 }
4404 crate::GatherMode::ShuffleDown(_) => {
4405 write!(self.out, "{NAMESPACE}::simd_shuffle_down(")?;
4406 }
4407 crate::GatherMode::ShuffleUp(_) => {
4408 write!(self.out, "{NAMESPACE}::simd_shuffle_up(")?;
4409 }
4410 crate::GatherMode::ShuffleXor(_) => {
4411 write!(self.out, "{NAMESPACE}::simd_shuffle_xor(")?;
4412 }
4413 crate::GatherMode::QuadBroadcast(_) => {
4414 write!(self.out, "{NAMESPACE}::quad_broadcast(")?;
4415 }
4416 crate::GatherMode::QuadSwap(_) => {
4417 write!(self.out, "{NAMESPACE}::quad_shuffle_xor(")?;
4418 }
4419 }
4420 self.put_expression(argument, &context.expression, true)?;
4421 match mode {
4422 crate::GatherMode::BroadcastFirst => {}
4423 crate::GatherMode::Broadcast(index)
4424 | crate::GatherMode::Shuffle(index)
4425 | crate::GatherMode::ShuffleDown(index)
4426 | crate::GatherMode::ShuffleUp(index)
4427 | crate::GatherMode::ShuffleXor(index)
4428 | crate::GatherMode::QuadBroadcast(index) => {
4429 write!(self.out, ", ")?;
4430 self.put_expression(index, &context.expression, true)?;
4431 }
4432 crate::GatherMode::QuadSwap(direction) => {
4433 write!(self.out, ", ")?;
4434 match direction {
4435 crate::Direction::X => {
4436 write!(self.out, "1u")?;
4437 }
4438 crate::Direction::Y => {
4439 write!(self.out, "2u")?;
4440 }
4441 crate::Direction::Diagonal => {
4442 write!(self.out, "3u")?;
4443 }
4444 }
4445 }
4446 }
4447 writeln!(self.out, ");")?;
4448 }
4449 crate::Statement::CooperativeStore { target, ref data } => {
4450 write!(self.out, "{level}simdgroup_store(")?;
4451 self.put_expression(target, &context.expression, true)?;
4452 write!(self.out, ", &")?;
4453 self.put_access_chain(
4454 data.pointer,
4455 context.expression.policies.index,
4456 &context.expression,
4457 )?;
4458 write!(self.out, ", ")?;
4459 self.put_expression(data.stride, &context.expression, true)?;
4460 if !data.row_major {
4465 let matrix_origin = "0";
4466 let transpose = true;
4467 write!(self.out, ", {matrix_origin}, {transpose}")?;
4468 }
4469 writeln!(self.out, ");")?;
4470 }
4471 crate::Statement::RayPipelineFunction(_) => unreachable!(),
4472 }
4473 }
4474
4475 for statement in statements {
4478 if let crate::Statement::Emit(ref range) = *statement {
4479 for handle in range.clone() {
4480 self.named_expressions.shift_remove(&handle);
4481 }
4482 }
4483 }
4484 Ok(())
4485 }
4486
4487 fn put_store(
4488 &mut self,
4489 pointer: Handle<crate::Expression>,
4490 value: Handle<crate::Expression>,
4491 level: back::Level,
4492 context: &StatementContext,
4493 ) -> BackendResult {
4494 let policy = context.expression.choose_bounds_check_policy(pointer);
4495 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
4496 && self.put_bounds_checks(pointer, &context.expression, level, "if (")?
4497 {
4498 writeln!(self.out, ") {{")?;
4499 self.put_unchecked_store(pointer, value, policy, level.next(), context)?;
4500 writeln!(self.out, "{level}}}")?;
4501 } else {
4502 self.put_unchecked_store(pointer, value, policy, level, context)?;
4503 }
4504
4505 Ok(())
4506 }
4507
4508 fn put_unchecked_store(
4509 &mut self,
4510 pointer: Handle<crate::Expression>,
4511 value: Handle<crate::Expression>,
4512 policy: index::BoundsCheckPolicy,
4513 level: back::Level,
4514 context: &StatementContext,
4515 ) -> BackendResult {
4516 let is_atomic_pointer = context
4517 .expression
4518 .resolve_type(pointer)
4519 .is_atomic_pointer(&context.expression.module.types);
4520
4521 if is_atomic_pointer {
4522 write!(
4523 self.out,
4524 "{level}{NAMESPACE}::atomic_store_explicit({ATOMIC_REFERENCE}"
4525 )?;
4526 self.put_access_chain(pointer, policy, &context.expression)?;
4527 write!(self.out, ", ")?;
4528 self.put_expression(value, &context.expression, true)?;
4529 writeln!(self.out, ", {NAMESPACE}::memory_order_relaxed);")?;
4530 } else {
4531 write!(self.out, "{level}")?;
4532 self.put_access_chain(pointer, policy, &context.expression)?;
4533 write!(self.out, " = ")?;
4534 self.put_expression(value, &context.expression, true)?;
4535 writeln!(self.out, ";")?;
4536 }
4537
4538 Ok(())
4539 }
4540
4541 pub fn write(
4542 &mut self,
4543 module: &crate::Module,
4544 info: &valid::ModuleInfo,
4545 options: &Options,
4546 pipeline_options: &PipelineOptions,
4547 ) -> Result<TranslationInfo, Error> {
4548 self.emit_int_div_checks = options.emit_int_div_checks;
4549 self.names.clear();
4550 self.namer.reset(
4551 module,
4552 &super::keywords::RESERVED_SET,
4553 proc::KeywordSet::empty(),
4554 proc::CaseInsensitiveKeywordSet::empty(),
4555 &[
4556 CLAMPED_LOD_LOAD_PREFIX,
4557 super::ray::INTERSECTION_FUNCTION_NAME,
4558 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
4559 super::ray::RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX,
4560 ],
4561 &mut self.names,
4562 );
4563 self.wrapped_functions.clear();
4564 self.struct_member_pads.clear();
4565
4566 writeln!(
4567 self.out,
4568 "// language: metal{}.{}",
4569 options.lang_version.0, options.lang_version.1
4570 )?;
4571 writeln!(self.out, "#include <metal_stdlib>")?;
4572 writeln!(self.out, "#include <simd/simd.h>")?;
4573 writeln!(self.out)?;
4574 writeln!(self.out, "using {NAMESPACE}::uint;")?;
4576
4577 if module.uses_mesh_shaders() && options.lang_version < (3, 0) {
4578 return Err(Error::UnsupportedMeshShader);
4579 }
4580 self.needs_object_memory_barriers = module
4581 .entry_points
4582 .iter()
4583 .any(|e| e.stage == crate::ShaderStage::Task && e.task_payload.is_some());
4584
4585 if module.special_types.ray_desc.is_some()
4586 || module.special_types.ray_intersection.is_some()
4587 {
4588 if options.lang_version < (2, 4) {
4589 return Err(Error::UnsupportedRayTracing);
4590 }
4591 }
4592
4593 if options
4594 .bounds_check_policies
4595 .contains(index::BoundsCheckPolicy::ReadZeroSkipWrite)
4596 {
4597 self.put_default_constructible()?;
4598 }
4599 writeln!(self.out)?;
4600
4601 {
4602 let globals: Vec<Handle<crate::GlobalVariable>> = module
4605 .global_variables
4606 .iter()
4607 .filter(|&(_, var)| {
4608 module.types[var.ty]
4609 .inner
4610 .needs_host_buffer_byte_size(&module.types)
4611 })
4612 .map(|(handle, _)| handle)
4613 .collect();
4614
4615 let mut buffer_indices = vec![];
4616 for vbm in &pipeline_options.vertex_buffer_mappings {
4617 buffer_indices.push(vbm.id);
4618 }
4619
4620 if !globals.is_empty() || !buffer_indices.is_empty() {
4621 writeln!(self.out, "struct _mslBufferSizes {{")?;
4622
4623 for global in globals {
4624 let var = &module.global_variables[global];
4625 let var_ty = var.ty;
4626 match module.types[var_ty].inner {
4627 crate::TypeInner::BindingArray { .. } => {
4628 let n =
4629 Self::binding_array_layout_count(module, pipeline_options, global);
4630 writeln!(
4631 self.out,
4632 "{}uint {}[{n}];",
4633 back::INDENT,
4634 ArraySizeMember(global),
4635 )?;
4636 }
4637 _ => writeln!(
4638 self.out,
4639 "{}uint {};",
4640 back::INDENT,
4641 ArraySizeMember(global)
4642 )?,
4643 }
4644 }
4645
4646 for idx in buffer_indices {
4647 writeln!(self.out, "{}uint buffer_size{};", back::INDENT, idx)?;
4648 }
4649
4650 writeln!(self.out, "}};")?;
4651 writeln!(self.out)?;
4652 }
4653 };
4654
4655 self.write_type_defs(module)?;
4656 self.write_global_constants(module, info)?;
4657 self.write_functions(module, info, options, pipeline_options)
4658 }
4659
4660 fn put_default_constructible(&mut self) -> BackendResult {
4673 let tab = back::INDENT;
4674 writeln!(self.out, "struct DefaultConstructible {{")?;
4675 writeln!(self.out, "{tab}template<typename T>")?;
4676 writeln!(self.out, "{tab}operator T() && {{")?;
4677 writeln!(self.out, "{tab}{tab}return T {{}};")?;
4678 writeln!(self.out, "{tab}}}")?;
4679 writeln!(self.out, "}};")?;
4680 Ok(())
4681 }
4682
4683 fn write_type_defs(&mut self, module: &crate::Module) -> BackendResult {
4684 let mut generated_argument_buffer_wrapper = false;
4685 let mut generated_external_texture_wrapper = false;
4686 for (handle, ty) in module.types.iter() {
4687 match ty.inner {
4688 crate::TypeInner::BindingArray { .. } if !generated_argument_buffer_wrapper => {
4689 writeln!(self.out, "template <typename T>")?;
4690 writeln!(self.out, "struct {ARGUMENT_BUFFER_WRAPPER_STRUCT} {{")?;
4691 writeln!(self.out, "{}T {WRAPPED_ARRAY_FIELD};", back::INDENT)?;
4692 writeln!(self.out, "}};")?;
4693 generated_argument_buffer_wrapper = true;
4694 }
4695 crate::TypeInner::Image {
4696 class: crate::ImageClass::External,
4697 ..
4698 } if !generated_external_texture_wrapper => {
4699 let params_ty_name = &self.names
4700 [&NameKey::Type(module.special_types.external_texture_params.unwrap())];
4701 writeln!(self.out, "struct {EXTERNAL_TEXTURE_WRAPPER_STRUCT} {{")?;
4702 writeln!(
4703 self.out,
4704 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane0;",
4705 back::INDENT
4706 )?;
4707 writeln!(
4708 self.out,
4709 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane1;",
4710 back::INDENT
4711 )?;
4712 writeln!(
4713 self.out,
4714 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane2;",
4715 back::INDENT
4716 )?;
4717 writeln!(self.out, "{}{params_ty_name} params;", back::INDENT)?;
4718 writeln!(self.out, "}};")?;
4719 generated_external_texture_wrapper = true;
4720 }
4721 _ => {}
4722 }
4723
4724 if !ty.needs_alias() {
4725 continue;
4726 }
4727 let name = &self.names[&NameKey::Type(handle)];
4728 match ty.inner {
4729 crate::TypeInner::Array {
4743 base,
4744 size,
4745 stride: _,
4746 } => {
4747 let base_name = TypeContext {
4748 handle: base,
4749 gctx: module.to_ctx(),
4750 names: &self.names,
4751 access: crate::StorageAccess::empty(),
4752 first_time: false,
4753 };
4754
4755 match size.resolve(module.to_ctx())? {
4756 proc::IndexableLength::Known(size) => {
4757 writeln!(self.out, "struct {name} {{")?;
4758 writeln!(
4759 self.out,
4760 "{}{} {}[{}];",
4761 back::INDENT,
4762 base_name,
4763 WRAPPED_ARRAY_FIELD,
4764 size
4765 )?;
4766 writeln!(self.out, "}};")?;
4767 }
4768 proc::IndexableLength::Dynamic => {
4769 writeln!(self.out, "typedef {base_name} {name}[1];")?;
4770 }
4771 }
4772 }
4773 crate::TypeInner::Struct {
4774 ref members, span, ..
4775 } => {
4776 writeln!(self.out, "struct {name} {{")?;
4777 let mut last_offset = 0;
4778 for (index, member) in members.iter().enumerate() {
4779 if member.offset > last_offset {
4780 self.struct_member_pads.insert((handle, index as u32));
4781 let pad = member.offset - last_offset;
4782 writeln!(self.out, "{}char _pad{}[{}];", back::INDENT, index, pad)?;
4783 }
4784 let ty_inner = &module.types[member.ty].inner;
4785 last_offset = member.offset + ty_inner.size(module.to_ctx());
4786
4787 let member_name = &self.names[&NameKey::StructMember(handle, index as u32)];
4788
4789 match should_pack_struct_member(members, span, index, module) {
4791 Some(scalar) => {
4792 writeln!(
4793 self.out,
4794 "{}{}::packed_{}3 {};",
4795 back::INDENT,
4796 NAMESPACE,
4797 scalar.to_msl_name(),
4798 member_name
4799 )?;
4800 }
4801 None => {
4802 let base_name = TypeContext {
4803 handle: member.ty,
4804 gctx: module.to_ctx(),
4805 names: &self.names,
4806 access: crate::StorageAccess::empty(),
4807 first_time: false,
4808 };
4809 writeln!(
4810 self.out,
4811 "{}{} {};",
4812 back::INDENT,
4813 base_name,
4814 member_name
4815 )?;
4816
4817 if let crate::TypeInner::Vector {
4819 size: crate::VectorSize::Tri,
4820 scalar,
4821 } = *ty_inner
4822 {
4823 last_offset += scalar.width as u32;
4824 }
4825 }
4826 }
4827 }
4828 if last_offset < span {
4829 let pad = span - last_offset;
4830 writeln!(
4831 self.out,
4832 "{}char _pad{}[{}];",
4833 back::INDENT,
4834 members.len(),
4835 pad
4836 )?;
4837 }
4838 writeln!(self.out, "}};")?;
4839 }
4840 _ => {
4841 let ty_name = TypeContext {
4842 handle,
4843 gctx: module.to_ctx(),
4844 names: &self.names,
4845 access: crate::StorageAccess::empty(),
4846 first_time: true,
4847 };
4848 writeln!(self.out, "typedef {ty_name} {name};")?;
4849 }
4850 }
4851 }
4852
4853 for (type_key, struct_ty) in module.special_types.predeclared_types.iter() {
4855 match type_key {
4856 &crate::PredeclaredType::ModfResult { size, scalar }
4857 | &crate::PredeclaredType::FrexpResult { size, scalar } => {
4858 let arg_type_name_owner;
4859 let arg_type_name = if let Some(size) = size {
4860 arg_type_name_owner = format!(
4861 "{NAMESPACE}::{}{}",
4862 if scalar.width == 8 { "double" } else { "float" },
4863 size as u8
4864 );
4865 &arg_type_name_owner
4866 } else if scalar.width == 8 {
4867 "double"
4868 } else {
4869 "float"
4870 };
4871
4872 let other_type_name_owner;
4873 let (defined_func_name, called_func_name, other_type_name) =
4874 if matches!(type_key, &crate::PredeclaredType::ModfResult { .. }) {
4875 (MODF_FUNCTION, "modf", arg_type_name)
4876 } else {
4877 let other_type_name = if let Some(size) = size {
4878 other_type_name_owner = format!("int{}", size as u8);
4879 &other_type_name_owner
4880 } else {
4881 "int"
4882 };
4883 (FREXP_FUNCTION, "frexp", other_type_name)
4884 };
4885
4886 let struct_name = &self.names[&NameKey::Type(*struct_ty)];
4887
4888 writeln!(self.out)?;
4889 writeln!(
4890 self.out,
4891 "{struct_name} {defined_func_name}({arg_type_name} arg) {{
4892 {other_type_name} other;
4893 {arg_type_name} fract = {NAMESPACE}::{called_func_name}(arg, other);
4894 return {struct_name}{{ fract, other }};
4895}}"
4896 )?;
4897 }
4898 &crate::PredeclaredType::AtomicCompareExchangeWeakResult(scalar) => {
4899 let arg_type_name = scalar.to_msl_name();
4900 let called_func_name = "atomic_compare_exchange_weak_explicit";
4901 let defined_func_name = ATOMIC_COMP_EXCH_FUNCTION;
4902 let struct_name = &self.names[&NameKey::Type(*struct_ty)];
4903
4904 writeln!(self.out)?;
4905
4906 for address_space_name in ["device", "threadgroup"] {
4907 writeln!(
4908 self.out,
4909 "\
4910template <typename A>
4911{struct_name} {defined_func_name}(
4912 {address_space_name} A *atomic_ptr,
4913 {arg_type_name} cmp,
4914 {arg_type_name} v
4915) {{
4916 bool swapped = {NAMESPACE}::{called_func_name}(
4917 atomic_ptr, &cmp, v,
4918 metal::memory_order_relaxed, metal::memory_order_relaxed
4919 );
4920 return {struct_name}{{cmp, swapped}};
4921}}"
4922 )?;
4923 }
4924 }
4925 }
4926 }
4927
4928 Ok(())
4929 }
4930
4931 fn write_global_constants(
4933 &mut self,
4934 module: &crate::Module,
4935 mod_info: &valid::ModuleInfo,
4936 ) -> BackendResult {
4937 let constants = module.constants.iter().filter(|&(_, c)| c.name.is_some());
4938
4939 for (handle, constant) in constants {
4940 let ty_name = TypeContext {
4941 handle: constant.ty,
4942 gctx: module.to_ctx(),
4943 names: &self.names,
4944 access: crate::StorageAccess::empty(),
4945 first_time: false,
4946 };
4947 let name = &self.names[&NameKey::Constant(handle)];
4948 write!(self.out, "constant {ty_name} {name} = ")?;
4949 self.put_const_expression(constant.init, module, mod_info, &module.global_expressions)?;
4950 writeln!(self.out, ";")?;
4951 }
4952
4953 Ok(())
4954 }
4955
4956 fn put_inline_sampler_properties(
4957 &mut self,
4958 level: back::Level,
4959 sampler: &sm::InlineSampler,
4960 options: &Options,
4961 ) -> BackendResult {
4962 for (&letter, address) in ['s', 't', 'r'].iter().zip(sampler.address.iter()) {
4963 writeln!(
4964 self.out,
4965 "{}{}::{}_address::{},",
4966 level,
4967 NAMESPACE,
4968 letter,
4969 address.as_str(),
4970 )?;
4971 }
4972 writeln!(
4973 self.out,
4974 "{}{}::mag_filter::{},",
4975 level,
4976 NAMESPACE,
4977 sampler.mag_filter.as_str(),
4978 )?;
4979 writeln!(
4980 self.out,
4981 "{}{}::min_filter::{},",
4982 level,
4983 NAMESPACE,
4984 sampler.min_filter.as_str(),
4985 )?;
4986 if let Some(filter) = sampler.mip_filter {
4987 writeln!(
4988 self.out,
4989 "{}{}::mip_filter::{},",
4990 level,
4991 NAMESPACE,
4992 filter.as_str(),
4993 )?;
4994 }
4995 if sampler.border_color != sm::BorderColor::TransparentBlack {
4997 writeln!(
4998 self.out,
4999 "{}{}::border_color::{},",
5000 level,
5001 NAMESPACE,
5002 sampler.border_color.as_str(),
5003 )?;
5004 }
5005
5006 if options.lang_version >= (1, 2) {
5007 if let Some(ref lod) = sampler.lod_clamp {
5008 writeln!(
5009 self.out,
5010 "{}{}::lod_clamp({},{}),",
5011 level, NAMESPACE, lod.start, lod.end,
5012 )?;
5013 }
5014 if let Some(aniso) = sampler.max_anisotropy {
5015 writeln!(
5016 self.out,
5017 "{}{}::max_anisotropy({}),",
5018 level,
5019 NAMESPACE,
5020 aniso.get(),
5021 )?;
5022 }
5023 }
5024
5025 if sampler.compare_func != sm::CompareFunc::Never {
5026 writeln!(
5027 self.out,
5028 "{}{}::compare_func::{},",
5029 level,
5030 NAMESPACE,
5031 sampler.compare_func.as_str(),
5032 )?;
5033 }
5034 writeln!(
5035 self.out,
5036 "{}{}::coord::{}",
5037 level,
5038 NAMESPACE,
5039 sampler.coord.as_str()
5040 )?;
5041 Ok(())
5042 }
5043
5044 fn write_unpacking_function(
5045 &mut self,
5046 format: nt::VertexFormat,
5047 ) -> Result<(String, u32, Option<crate::VectorSize>, crate::Scalar), Error> {
5048 use crate::{Scalar, VectorSize};
5049 use nt::VertexFormat::*;
5050 match format {
5051 Uint8 => {
5052 let name = self.namer.call("unpackUint8");
5053 writeln!(self.out, "uint {name}(metal::uchar b0) {{")?;
5054 writeln!(self.out, "{}return uint(b0);", back::INDENT)?;
5055 writeln!(self.out, "}}")?;
5056 Ok((name, 1, None, Scalar::U32))
5057 }
5058 Uint8x2 => {
5059 let name = self.namer.call("unpackUint8x2");
5060 writeln!(
5061 self.out,
5062 "metal::uint2 {name}(metal::uchar b0, \
5063 metal::uchar b1) {{"
5064 )?;
5065 writeln!(self.out, "{}return metal::uint2(b0, b1);", back::INDENT)?;
5066 writeln!(self.out, "}}")?;
5067 Ok((name, 2, Some(VectorSize::Bi), Scalar::U32))
5068 }
5069 Uint8x4 => {
5070 let name = self.namer.call("unpackUint8x4");
5071 writeln!(
5072 self.out,
5073 "metal::uint4 {name}(metal::uchar b0, \
5074 metal::uchar b1, \
5075 metal::uchar b2, \
5076 metal::uchar b3) {{"
5077 )?;
5078 writeln!(
5079 self.out,
5080 "{}return metal::uint4(b0, b1, b2, b3);",
5081 back::INDENT
5082 )?;
5083 writeln!(self.out, "}}")?;
5084 Ok((name, 4, Some(VectorSize::Quad), Scalar::U32))
5085 }
5086 Sint8 => {
5087 let name = self.namer.call("unpackSint8");
5088 writeln!(self.out, "int {name}(metal::uchar b0) {{")?;
5089 writeln!(self.out, "{}return int(as_type<char>(b0));", back::INDENT)?;
5090 writeln!(self.out, "}}")?;
5091 Ok((name, 1, None, Scalar::I32))
5092 }
5093 Sint8x2 => {
5094 let name = self.namer.call("unpackSint8x2");
5095 writeln!(
5096 self.out,
5097 "metal::int2 {name}(metal::uchar b0, \
5098 metal::uchar b1) {{"
5099 )?;
5100 writeln!(
5101 self.out,
5102 "{}return metal::int2(as_type<char>(b0), \
5103 as_type<char>(b1));",
5104 back::INDENT
5105 )?;
5106 writeln!(self.out, "}}")?;
5107 Ok((name, 2, Some(VectorSize::Bi), Scalar::I32))
5108 }
5109 Sint8x4 => {
5110 let name = self.namer.call("unpackSint8x4");
5111 writeln!(
5112 self.out,
5113 "metal::int4 {name}(metal::uchar b0, \
5114 metal::uchar b1, \
5115 metal::uchar b2, \
5116 metal::uchar b3) {{"
5117 )?;
5118 writeln!(
5119 self.out,
5120 "{}return metal::int4(as_type<char>(b0), \
5121 as_type<char>(b1), \
5122 as_type<char>(b2), \
5123 as_type<char>(b3));",
5124 back::INDENT
5125 )?;
5126 writeln!(self.out, "}}")?;
5127 Ok((name, 4, Some(VectorSize::Quad), Scalar::I32))
5128 }
5129 Unorm8 => {
5130 let name = self.namer.call("unpackUnorm8");
5131 writeln!(self.out, "float {name}(metal::uchar b0) {{")?;
5132 writeln!(
5133 self.out,
5134 "{}return float(float(b0) / 255.0f);",
5135 back::INDENT
5136 )?;
5137 writeln!(self.out, "}}")?;
5138 Ok((name, 1, None, Scalar::F32))
5139 }
5140 Unorm8x2 => {
5141 let name = self.namer.call("unpackUnorm8x2");
5142 writeln!(
5143 self.out,
5144 "metal::float2 {name}(metal::uchar b0, \
5145 metal::uchar b1) {{"
5146 )?;
5147 writeln!(
5148 self.out,
5149 "{}return metal::float2(float(b0) / 255.0f, \
5150 float(b1) / 255.0f);",
5151 back::INDENT
5152 )?;
5153 writeln!(self.out, "}}")?;
5154 Ok((name, 2, Some(VectorSize::Bi), Scalar::F32))
5155 }
5156 Unorm8x4 => {
5157 let name = self.namer.call("unpackUnorm8x4");
5158 writeln!(
5159 self.out,
5160 "metal::float4 {name}(metal::uchar b0, \
5161 metal::uchar b1, \
5162 metal::uchar b2, \
5163 metal::uchar b3) {{"
5164 )?;
5165 writeln!(
5166 self.out,
5167 "{}return metal::float4(float(b0) / 255.0f, \
5168 float(b1) / 255.0f, \
5169 float(b2) / 255.0f, \
5170 float(b3) / 255.0f);",
5171 back::INDENT
5172 )?;
5173 writeln!(self.out, "}}")?;
5174 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5175 }
5176 Snorm8 => {
5177 let name = self.namer.call("unpackSnorm8");
5178 writeln!(self.out, "float {name}(metal::uchar b0) {{")?;
5179 writeln!(
5180 self.out,
5181 "{}return float(metal::max(-1.0f, as_type<char>(b0) / 127.0f));",
5182 back::INDENT
5183 )?;
5184 writeln!(self.out, "}}")?;
5185 Ok((name, 1, None, Scalar::F32))
5186 }
5187 Snorm8x2 => {
5188 let name = self.namer.call("unpackSnorm8x2");
5189 writeln!(
5190 self.out,
5191 "metal::float2 {name}(metal::uchar b0, \
5192 metal::uchar b1) {{"
5193 )?;
5194 writeln!(
5195 self.out,
5196 "{}return metal::float2(metal::max(-1.0f, as_type<char>(b0) / 127.0f), \
5197 metal::max(-1.0f, as_type<char>(b1) / 127.0f));",
5198 back::INDENT
5199 )?;
5200 writeln!(self.out, "}}")?;
5201 Ok((name, 2, Some(VectorSize::Bi), Scalar::F32))
5202 }
5203 Snorm8x4 => {
5204 let name = self.namer.call("unpackSnorm8x4");
5205 writeln!(
5206 self.out,
5207 "metal::float4 {name}(metal::uchar b0, \
5208 metal::uchar b1, \
5209 metal::uchar b2, \
5210 metal::uchar b3) {{"
5211 )?;
5212 writeln!(
5213 self.out,
5214 "{}return metal::float4(metal::max(-1.0f, as_type<char>(b0) / 127.0f), \
5215 metal::max(-1.0f, as_type<char>(b1) / 127.0f), \
5216 metal::max(-1.0f, as_type<char>(b2) / 127.0f), \
5217 metal::max(-1.0f, as_type<char>(b3) / 127.0f));",
5218 back::INDENT
5219 )?;
5220 writeln!(self.out, "}}")?;
5221 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5222 }
5223 Uint16 => {
5224 let name = self.namer.call("unpackUint16");
5225 writeln!(
5226 self.out,
5227 "metal::uint {name}(metal::uint b0, \
5228 metal::uint b1) {{"
5229 )?;
5230 writeln!(
5231 self.out,
5232 "{}return metal::uint(b1 << 8 | b0);",
5233 back::INDENT
5234 )?;
5235 writeln!(self.out, "}}")?;
5236 Ok((name, 2, None, Scalar::U32))
5237 }
5238 Uint16x2 => {
5239 let name = self.namer.call("unpackUint16x2");
5240 writeln!(
5241 self.out,
5242 "metal::uint2 {name}(metal::uint b0, \
5243 metal::uint b1, \
5244 metal::uint b2, \
5245 metal::uint b3) {{"
5246 )?;
5247 writeln!(
5248 self.out,
5249 "{}return metal::uint2(b1 << 8 | b0, \
5250 b3 << 8 | b2);",
5251 back::INDENT
5252 )?;
5253 writeln!(self.out, "}}")?;
5254 Ok((name, 4, Some(VectorSize::Bi), Scalar::U32))
5255 }
5256 Uint16x4 => {
5257 let name = self.namer.call("unpackUint16x4");
5258 writeln!(
5259 self.out,
5260 "metal::uint4 {name}(metal::uint b0, \
5261 metal::uint b1, \
5262 metal::uint b2, \
5263 metal::uint b3, \
5264 metal::uint b4, \
5265 metal::uint b5, \
5266 metal::uint b6, \
5267 metal::uint b7) {{"
5268 )?;
5269 writeln!(
5270 self.out,
5271 "{}return metal::uint4(b1 << 8 | b0, \
5272 b3 << 8 | b2, \
5273 b5 << 8 | b4, \
5274 b7 << 8 | b6);",
5275 back::INDENT
5276 )?;
5277 writeln!(self.out, "}}")?;
5278 Ok((name, 8, Some(VectorSize::Quad), Scalar::U32))
5279 }
5280 Sint16 => {
5281 let name = self.namer.call("unpackSint16");
5282 writeln!(
5283 self.out,
5284 "int {name}(metal::ushort b0, \
5285 metal::ushort b1) {{"
5286 )?;
5287 writeln!(
5288 self.out,
5289 "{}return int(as_type<short>(metal::ushort(b1 << 8 | b0)));",
5290 back::INDENT
5291 )?;
5292 writeln!(self.out, "}}")?;
5293 Ok((name, 2, None, Scalar::I32))
5294 }
5295 Sint16x2 => {
5296 let name = self.namer.call("unpackSint16x2");
5297 writeln!(
5298 self.out,
5299 "metal::int2 {name}(metal::ushort b0, \
5300 metal::ushort b1, \
5301 metal::ushort b2, \
5302 metal::ushort b3) {{"
5303 )?;
5304 writeln!(
5305 self.out,
5306 "{}return metal::int2(as_type<short>(metal::ushort(b1 << 8 | b0)), \
5307 as_type<short>(metal::ushort(b3 << 8 | b2)));",
5308 back::INDENT
5309 )?;
5310 writeln!(self.out, "}}")?;
5311 Ok((name, 4, Some(VectorSize::Bi), Scalar::I32))
5312 }
5313 Sint16x4 => {
5314 let name = self.namer.call("unpackSint16x4");
5315 writeln!(
5316 self.out,
5317 "metal::int4 {name}(metal::ushort b0, \
5318 metal::ushort b1, \
5319 metal::ushort b2, \
5320 metal::ushort b3, \
5321 metal::ushort b4, \
5322 metal::ushort b5, \
5323 metal::ushort b6, \
5324 metal::ushort b7) {{"
5325 )?;
5326 writeln!(
5327 self.out,
5328 "{}return metal::int4(as_type<short>(metal::ushort(b1 << 8 | b0)), \
5329 as_type<short>(metal::ushort(b3 << 8 | b2)), \
5330 as_type<short>(metal::ushort(b5 << 8 | b4)), \
5331 as_type<short>(metal::ushort(b7 << 8 | b6)));",
5332 back::INDENT
5333 )?;
5334 writeln!(self.out, "}}")?;
5335 Ok((name, 8, Some(VectorSize::Quad), Scalar::I32))
5336 }
5337 Unorm16 => {
5338 let name = self.namer.call("unpackUnorm16");
5339 writeln!(
5340 self.out,
5341 "float {name}(metal::ushort b0, \
5342 metal::ushort b1) {{"
5343 )?;
5344 writeln!(
5345 self.out,
5346 "{}return float(float(b1 << 8 | b0) / 65535.0f);",
5347 back::INDENT
5348 )?;
5349 writeln!(self.out, "}}")?;
5350 Ok((name, 2, None, Scalar::F32))
5351 }
5352 Unorm16x2 => {
5353 let name = self.namer.call("unpackUnorm16x2");
5354 writeln!(
5355 self.out,
5356 "metal::float2 {name}(metal::ushort b0, \
5357 metal::ushort b1, \
5358 metal::ushort b2, \
5359 metal::ushort b3) {{"
5360 )?;
5361 writeln!(
5362 self.out,
5363 "{}return metal::float2(float(b1 << 8 | b0) / 65535.0f, \
5364 float(b3 << 8 | b2) / 65535.0f);",
5365 back::INDENT
5366 )?;
5367 writeln!(self.out, "}}")?;
5368 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5369 }
5370 Unorm16x4 => {
5371 let name = self.namer.call("unpackUnorm16x4");
5372 writeln!(
5373 self.out,
5374 "metal::float4 {name}(metal::ushort b0, \
5375 metal::ushort b1, \
5376 metal::ushort b2, \
5377 metal::ushort b3, \
5378 metal::ushort b4, \
5379 metal::ushort b5, \
5380 metal::ushort b6, \
5381 metal::ushort b7) {{"
5382 )?;
5383 writeln!(
5384 self.out,
5385 "{}return metal::float4(float(b1 << 8 | b0) / 65535.0f, \
5386 float(b3 << 8 | b2) / 65535.0f, \
5387 float(b5 << 8 | b4) / 65535.0f, \
5388 float(b7 << 8 | b6) / 65535.0f);",
5389 back::INDENT
5390 )?;
5391 writeln!(self.out, "}}")?;
5392 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5393 }
5394 Snorm16 => {
5395 let name = self.namer.call("unpackSnorm16");
5396 writeln!(
5397 self.out,
5398 "float {name}(metal::ushort b0, \
5399 metal::ushort b1) {{"
5400 )?;
5401 writeln!(
5402 self.out,
5403 "{}return metal::unpack_snorm2x16_to_float(b1 << 8 | b0).x;",
5404 back::INDENT
5405 )?;
5406 writeln!(self.out, "}}")?;
5407 Ok((name, 2, None, Scalar::F32))
5408 }
5409 Snorm16x2 => {
5410 let name = self.namer.call("unpackSnorm16x2");
5411 writeln!(
5412 self.out,
5413 "metal::float2 {name}(uint b0, \
5414 uint b1, \
5415 uint b2, \
5416 uint b3) {{"
5417 )?;
5418 writeln!(
5419 self.out,
5420 "{}return metal::unpack_snorm2x16_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5421 back::INDENT
5422 )?;
5423 writeln!(self.out, "}}")?;
5424 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5425 }
5426 Snorm16x4 => {
5427 let name = self.namer.call("unpackSnorm16x4");
5428 writeln!(
5429 self.out,
5430 "metal::float4 {name}(uint b0, \
5431 uint b1, \
5432 uint b2, \
5433 uint b3, \
5434 uint b4, \
5435 uint b5, \
5436 uint b6, \
5437 uint b7) {{"
5438 )?;
5439 writeln!(
5440 self.out,
5441 "{}return metal::float4(metal::unpack_snorm2x16_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5442 metal::unpack_snorm2x16_to_float(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5443 back::INDENT
5444 )?;
5445 writeln!(self.out, "}}")?;
5446 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5447 }
5448 Float16 => {
5449 let name = self.namer.call("unpackFloat16");
5450 writeln!(
5451 self.out,
5452 "float {name}(metal::ushort b0, \
5453 metal::ushort b1) {{"
5454 )?;
5455 writeln!(
5456 self.out,
5457 "{}return float(as_type<half>(metal::ushort(b1 << 8 | b0)));",
5458 back::INDENT
5459 )?;
5460 writeln!(self.out, "}}")?;
5461 Ok((name, 2, None, Scalar::F32))
5462 }
5463 Float16x2 => {
5464 let name = self.namer.call("unpackFloat16x2");
5465 writeln!(
5466 self.out,
5467 "metal::float2 {name}(metal::ushort b0, \
5468 metal::ushort b1, \
5469 metal::ushort b2, \
5470 metal::ushort b3) {{"
5471 )?;
5472 writeln!(
5473 self.out,
5474 "{}return metal::float2(as_type<half>(metal::ushort(b1 << 8 | b0)), \
5475 as_type<half>(metal::ushort(b3 << 8 | b2)));",
5476 back::INDENT
5477 )?;
5478 writeln!(self.out, "}}")?;
5479 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5480 }
5481 Float16x4 => {
5482 let name = self.namer.call("unpackFloat16x4");
5483 writeln!(
5484 self.out,
5485 "metal::float4 {name}(metal::ushort b0, \
5486 metal::ushort b1, \
5487 metal::ushort b2, \
5488 metal::ushort b3, \
5489 metal::ushort b4, \
5490 metal::ushort b5, \
5491 metal::ushort b6, \
5492 metal::ushort b7) {{"
5493 )?;
5494 writeln!(
5495 self.out,
5496 "{}return metal::float4(as_type<half>(metal::ushort(b1 << 8 | b0)), \
5497 as_type<half>(metal::ushort(b3 << 8 | b2)), \
5498 as_type<half>(metal::ushort(b5 << 8 | b4)), \
5499 as_type<half>(metal::ushort(b7 << 8 | b6)));",
5500 back::INDENT
5501 )?;
5502 writeln!(self.out, "}}")?;
5503 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5504 }
5505 Float32 => {
5506 let name = self.namer.call("unpackFloat32");
5507 writeln!(
5508 self.out,
5509 "float {name}(uint b0, \
5510 uint b1, \
5511 uint b2, \
5512 uint b3) {{"
5513 )?;
5514 writeln!(
5515 self.out,
5516 "{}return as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5517 back::INDENT
5518 )?;
5519 writeln!(self.out, "}}")?;
5520 Ok((name, 4, None, Scalar::F32))
5521 }
5522 Float32x2 => {
5523 let name = self.namer.call("unpackFloat32x2");
5524 writeln!(
5525 self.out,
5526 "metal::float2 {name}(uint b0, \
5527 uint b1, \
5528 uint b2, \
5529 uint b3, \
5530 uint b4, \
5531 uint b5, \
5532 uint b6, \
5533 uint b7) {{"
5534 )?;
5535 writeln!(
5536 self.out,
5537 "{}return metal::float2(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5538 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5539 back::INDENT
5540 )?;
5541 writeln!(self.out, "}}")?;
5542 Ok((name, 8, Some(VectorSize::Bi), Scalar::F32))
5543 }
5544 Float32x3 => {
5545 let name = self.namer.call("unpackFloat32x3");
5546 writeln!(
5547 self.out,
5548 "metal::float3 {name}(uint b0, \
5549 uint b1, \
5550 uint b2, \
5551 uint b3, \
5552 uint b4, \
5553 uint b5, \
5554 uint b6, \
5555 uint b7, \
5556 uint b8, \
5557 uint b9, \
5558 uint b10, \
5559 uint b11) {{"
5560 )?;
5561 writeln!(
5562 self.out,
5563 "{}return metal::float3(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5564 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5565 as_type<float>(b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5566 back::INDENT
5567 )?;
5568 writeln!(self.out, "}}")?;
5569 Ok((name, 12, Some(VectorSize::Tri), Scalar::F32))
5570 }
5571 Float32x4 => {
5572 let name = self.namer.call("unpackFloat32x4");
5573 writeln!(
5574 self.out,
5575 "metal::float4 {name}(uint b0, \
5576 uint b1, \
5577 uint b2, \
5578 uint b3, \
5579 uint b4, \
5580 uint b5, \
5581 uint b6, \
5582 uint b7, \
5583 uint b8, \
5584 uint b9, \
5585 uint b10, \
5586 uint b11, \
5587 uint b12, \
5588 uint b13, \
5589 uint b14, \
5590 uint b15) {{"
5591 )?;
5592 writeln!(
5593 self.out,
5594 "{}return metal::float4(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5595 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5596 as_type<float>(b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5597 as_type<float>(b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5598 back::INDENT
5599 )?;
5600 writeln!(self.out, "}}")?;
5601 Ok((name, 16, Some(VectorSize::Quad), Scalar::F32))
5602 }
5603 Uint32 => {
5604 let name = self.namer.call("unpackUint32");
5605 writeln!(
5606 self.out,
5607 "uint {name}(uint b0, \
5608 uint b1, \
5609 uint b2, \
5610 uint b3) {{"
5611 )?;
5612 writeln!(
5613 self.out,
5614 "{}return (b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5615 back::INDENT
5616 )?;
5617 writeln!(self.out, "}}")?;
5618 Ok((name, 4, None, Scalar::U32))
5619 }
5620 Uint32x2 => {
5621 let name = self.namer.call("unpackUint32x2");
5622 writeln!(
5623 self.out,
5624 "uint2 {name}(uint b0, \
5625 uint b1, \
5626 uint b2, \
5627 uint b3, \
5628 uint b4, \
5629 uint b5, \
5630 uint b6, \
5631 uint b7) {{"
5632 )?;
5633 writeln!(
5634 self.out,
5635 "{}return uint2((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5636 (b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5637 back::INDENT
5638 )?;
5639 writeln!(self.out, "}}")?;
5640 Ok((name, 8, Some(VectorSize::Bi), Scalar::U32))
5641 }
5642 Uint32x3 => {
5643 let name = self.namer.call("unpackUint32x3");
5644 writeln!(
5645 self.out,
5646 "uint3 {name}(uint b0, \
5647 uint b1, \
5648 uint b2, \
5649 uint b3, \
5650 uint b4, \
5651 uint b5, \
5652 uint b6, \
5653 uint b7, \
5654 uint b8, \
5655 uint b9, \
5656 uint b10, \
5657 uint b11) {{"
5658 )?;
5659 writeln!(
5660 self.out,
5661 "{}return uint3((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5662 (b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5663 (b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5664 back::INDENT
5665 )?;
5666 writeln!(self.out, "}}")?;
5667 Ok((name, 12, Some(VectorSize::Tri), Scalar::U32))
5668 }
5669 Uint32x4 => {
5670 let name = self.namer.call("unpackUint32x4");
5671 writeln!(
5672 self.out,
5673 "{NAMESPACE}::uint4 {name}(uint b0, \
5674 uint b1, \
5675 uint b2, \
5676 uint b3, \
5677 uint b4, \
5678 uint b5, \
5679 uint b6, \
5680 uint b7, \
5681 uint b8, \
5682 uint b9, \
5683 uint b10, \
5684 uint b11, \
5685 uint b12, \
5686 uint b13, \
5687 uint b14, \
5688 uint b15) {{"
5689 )?;
5690 writeln!(
5691 self.out,
5692 "{}return {NAMESPACE}::uint4((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5693 (b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5694 (b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5695 (b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5696 back::INDENT
5697 )?;
5698 writeln!(self.out, "}}")?;
5699 Ok((name, 16, Some(VectorSize::Quad), Scalar::U32))
5700 }
5701 Sint32 => {
5702 let name = self.namer.call("unpackSint32");
5703 writeln!(
5704 self.out,
5705 "int {name}(uint b0, \
5706 uint b1, \
5707 uint b2, \
5708 uint b3) {{"
5709 )?;
5710 writeln!(
5711 self.out,
5712 "{}return as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5713 back::INDENT
5714 )?;
5715 writeln!(self.out, "}}")?;
5716 Ok((name, 4, None, Scalar::I32))
5717 }
5718 Sint32x2 => {
5719 let name = self.namer.call("unpackSint32x2");
5720 writeln!(
5721 self.out,
5722 "metal::int2 {name}(uint b0, \
5723 uint b1, \
5724 uint b2, \
5725 uint b3, \
5726 uint b4, \
5727 uint b5, \
5728 uint b6, \
5729 uint b7) {{"
5730 )?;
5731 writeln!(
5732 self.out,
5733 "{}return metal::int2(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5734 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5735 back::INDENT
5736 )?;
5737 writeln!(self.out, "}}")?;
5738 Ok((name, 8, Some(VectorSize::Bi), Scalar::I32))
5739 }
5740 Sint32x3 => {
5741 let name = self.namer.call("unpackSint32x3");
5742 writeln!(
5743 self.out,
5744 "metal::int3 {name}(uint b0, \
5745 uint b1, \
5746 uint b2, \
5747 uint b3, \
5748 uint b4, \
5749 uint b5, \
5750 uint b6, \
5751 uint b7, \
5752 uint b8, \
5753 uint b9, \
5754 uint b10, \
5755 uint b11) {{"
5756 )?;
5757 writeln!(
5758 self.out,
5759 "{}return metal::int3(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5760 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5761 as_type<int>(b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5762 back::INDENT
5763 )?;
5764 writeln!(self.out, "}}")?;
5765 Ok((name, 12, Some(VectorSize::Tri), Scalar::I32))
5766 }
5767 Sint32x4 => {
5768 let name = self.namer.call("unpackSint32x4");
5769 writeln!(
5770 self.out,
5771 "metal::int4 {name}(uint b0, \
5772 uint b1, \
5773 uint b2, \
5774 uint b3, \
5775 uint b4, \
5776 uint b5, \
5777 uint b6, \
5778 uint b7, \
5779 uint b8, \
5780 uint b9, \
5781 uint b10, \
5782 uint b11, \
5783 uint b12, \
5784 uint b13, \
5785 uint b14, \
5786 uint b15) {{"
5787 )?;
5788 writeln!(
5789 self.out,
5790 "{}return metal::int4(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5791 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5792 as_type<int>(b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5793 as_type<int>(b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5794 back::INDENT
5795 )?;
5796 writeln!(self.out, "}}")?;
5797 Ok((name, 16, Some(VectorSize::Quad), Scalar::I32))
5798 }
5799 Unorm10_10_10_2 => {
5800 let name = self.namer.call("unpackUnorm10_10_10_2");
5801 writeln!(
5802 self.out,
5803 "metal::float4 {name}(uint b0, \
5804 uint b1, \
5805 uint b2, \
5806 uint b3) {{"
5807 )?;
5808 writeln!(
5809 self.out,
5810 "{}return metal::unpack_unorm10a2_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5822 back::INDENT
5823 )?;
5824 writeln!(self.out, "}}")?;
5825 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5826 }
5827 Unorm8x4Bgra => {
5828 let name = self.namer.call("unpackUnorm8x4Bgra");
5829 writeln!(
5830 self.out,
5831 "metal::float4 {name}(metal::uchar b0, \
5832 metal::uchar b1, \
5833 metal::uchar b2, \
5834 metal::uchar b3) {{"
5835 )?;
5836 writeln!(
5837 self.out,
5838 "{}return metal::float4(float(b2) / 255.0f, \
5839 float(b1) / 255.0f, \
5840 float(b0) / 255.0f, \
5841 float(b3) / 255.0f);",
5842 back::INDENT
5843 )?;
5844 writeln!(self.out, "}}")?;
5845 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5846 }
5847 Float64 | Float64x2 | Float64x3 | Float64x4 => unreachable!(),
5848 }
5849 }
5850
5851 fn write_wrapped_unary_op(
5852 &mut self,
5853 module: &crate::Module,
5854 func_ctx: &back::FunctionCtx,
5855 op: crate::UnaryOperator,
5856 operand: Handle<crate::Expression>,
5857 ) -> BackendResult {
5858 let operand_ty = func_ctx.resolve_type(operand, &module.types);
5859 match op {
5860 crate::UnaryOperator::Negate
5867 if operand_ty.scalar_kind() == Some(crate::ScalarKind::Sint) =>
5868 {
5869 let Some((vector_size, scalar)) = operand_ty.vector_size_and_scalar() else {
5870 return Ok(());
5871 };
5872 let wrapped = WrappedFunction::UnaryOp {
5873 op,
5874 ty: (vector_size, scalar),
5875 };
5876 if !self.wrapped_functions.insert(wrapped) {
5877 return Ok(());
5878 }
5879
5880 let unsigned_scalar = crate::Scalar {
5881 kind: crate::ScalarKind::Uint,
5882 ..scalar
5883 };
5884 let mut type_name = String::new();
5885 let mut unsigned_type_name = String::new();
5886 match vector_size {
5887 None => {
5888 put_numeric_type(&mut type_name, scalar, &[])?;
5889 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[])?
5890 }
5891 Some(size) => {
5892 put_numeric_type(&mut type_name, scalar, &[size])?;
5893 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[size])?;
5894 }
5895 };
5896
5897 writeln!(self.out, "{type_name} {NEG_FUNCTION}({type_name} val) {{")?;
5898 let level = back::Level(1);
5899 if scalar.width < 4 {
5903 writeln!(
5904 self.out,
5905 "{level}return as_type<{type_name}>(static_cast<{unsigned_type_name}>(-as_type<{unsigned_type_name}>(val)));"
5906 )?;
5907 } else {
5908 writeln!(
5909 self.out,
5910 "{level}return as_type<{type_name}>(-as_type<{unsigned_type_name}>(val));"
5911 )?;
5912 }
5913 writeln!(self.out, "}}")?;
5914 writeln!(self.out)?;
5915 }
5916 _ => {}
5917 }
5918 Ok(())
5919 }
5920
5921 fn write_wrapped_binary_op(
5922 &mut self,
5923 module: &crate::Module,
5924 func_ctx: &back::FunctionCtx,
5925 expr: Handle<crate::Expression>,
5926 op: crate::BinaryOperator,
5927 left: Handle<crate::Expression>,
5928 right: Handle<crate::Expression>,
5929 ) -> BackendResult {
5930 let expr_ty = func_ctx.resolve_type(expr, &module.types);
5931 let left_ty = func_ctx.resolve_type(left, &module.types);
5932 let right_ty = func_ctx.resolve_type(right, &module.types);
5933 match (op, expr_ty.scalar_kind()) {
5934 (
5941 crate::BinaryOperator::Divide,
5942 Some(crate::ScalarKind::Sint | crate::ScalarKind::Uint),
5943 ) if self.emit_int_div_checks => {
5944 let Some(left_wrapped_ty) = left_ty.vector_size_and_scalar() else {
5945 return Ok(());
5946 };
5947 let Some(right_wrapped_ty) = right_ty.vector_size_and_scalar() else {
5948 return Ok(());
5949 };
5950 let wrapped = WrappedFunction::BinaryOp {
5951 op,
5952 left_ty: left_wrapped_ty,
5953 right_ty: right_wrapped_ty,
5954 };
5955 if !self.wrapped_functions.insert(wrapped) {
5956 return Ok(());
5957 }
5958
5959 let Some((vector_size, scalar)) = expr_ty.vector_size_and_scalar() else {
5960 return Ok(());
5961 };
5962 let mut type_name = String::new();
5963 match vector_size {
5964 None => put_numeric_type(&mut type_name, scalar, &[])?,
5965 Some(size) => put_numeric_type(&mut type_name, scalar, &[size])?,
5966 };
5967 writeln!(
5968 self.out,
5969 "{type_name} {DIV_FUNCTION}({type_name} lhs, {type_name} rhs) {{"
5970 )?;
5971 let level = back::Level(1);
5972 let (lp, rp) = if scalar.width < 4 {
5976 (format!("{type_name}("), ")".to_string())
5977 } else {
5978 (String::new(), String::new())
5979 };
5980 match scalar.kind {
5981 crate::ScalarKind::Sint => {
5982 let min_val = match scalar.width {
5983 2 => crate::Literal::I16(i16::MIN),
5984 4 => crate::Literal::I32(i32::MIN),
5985 8 => crate::Literal::I64(i64::MIN),
5986 _ => {
5987 return Err(Error::GenericValidation(format!(
5988 "Unexpected width for scalar {scalar:?}"
5989 )));
5990 }
5991 };
5992 write!(
5993 self.out,
5994 "{level}return lhs / metal::select(rhs, {lp}1{rp}, (lhs == "
5995 )?;
5996 self.put_literal(min_val)?;
5997 writeln!(self.out, " & rhs == {lp}-1{rp}) | (rhs == {lp}0{rp}));")?
5998 }
5999 crate::ScalarKind::Uint => {
6000 let suffix = if scalar.width < 4 { "" } else { "u" };
6001 writeln!(
6002 self.out,
6003 "{level}return lhs / metal::select(rhs, {lp}1{suffix}{rp}, rhs == {lp}0{suffix}{rp});"
6004 )?
6005 }
6006 _ => unreachable!(),
6007 }
6008 writeln!(self.out, "}}")?;
6009 writeln!(self.out)?;
6010 }
6011 (
6024 crate::BinaryOperator::Modulo,
6025 Some(crate::ScalarKind::Sint | crate::ScalarKind::Uint),
6026 ) if self.emit_int_div_checks => {
6027 let Some(left_wrapped_ty) = left_ty.vector_size_and_scalar() else {
6028 return Ok(());
6029 };
6030 let Some((right_vector_size, right_scalar)) = right_ty.vector_size_and_scalar()
6031 else {
6032 return Ok(());
6033 };
6034 let wrapped = WrappedFunction::BinaryOp {
6035 op,
6036 left_ty: left_wrapped_ty,
6037 right_ty: (right_vector_size, right_scalar),
6038 };
6039 if !self.wrapped_functions.insert(wrapped) {
6040 return Ok(());
6041 }
6042
6043 let Some((vector_size, scalar)) = expr_ty.vector_size_and_scalar() else {
6044 return Ok(());
6045 };
6046 let mut type_name = String::new();
6047 match vector_size {
6048 None => put_numeric_type(&mut type_name, scalar, &[])?,
6049 Some(size) => put_numeric_type(&mut type_name, scalar, &[size])?,
6050 };
6051 let mut rhs_type_name = String::new();
6052 match right_vector_size {
6053 None => put_numeric_type(&mut rhs_type_name, right_scalar, &[])?,
6054 Some(size) => put_numeric_type(&mut rhs_type_name, right_scalar, &[size])?,
6055 };
6056
6057 writeln!(
6058 self.out,
6059 "{type_name} {MOD_FUNCTION}({type_name} lhs, {type_name} rhs) {{"
6060 )?;
6061 let level = back::Level(1);
6062 let (lp, rp) = if scalar.width < 4 {
6063 (format!("{type_name}("), ")".to_string())
6064 } else {
6065 (String::new(), String::new())
6066 };
6067 match scalar.kind {
6068 crate::ScalarKind::Sint => {
6069 let min_val = match scalar.width {
6070 2 => crate::Literal::I16(i16::MIN),
6071 4 => crate::Literal::I32(i32::MIN),
6072 8 => crate::Literal::I64(i64::MIN),
6073 _ => {
6074 return Err(Error::GenericValidation(format!(
6075 "Unexpected width for scalar {scalar:?}"
6076 )));
6077 }
6078 };
6079 write!(
6080 self.out,
6081 "{level}{rhs_type_name} divisor = metal::select(rhs, {lp}1{rp}, (lhs == "
6082 )?;
6083 self.put_literal(min_val)?;
6084 writeln!(self.out, " & rhs == {lp}-1{rp}) | (rhs == {lp}0{rp}));")?;
6085 writeln!(self.out, "{level}return lhs - (lhs / divisor) * divisor;")?
6086 }
6087 crate::ScalarKind::Uint => {
6088 let suffix = if scalar.width < 4 { "" } else { "u" };
6089 writeln!(
6090 self.out,
6091 "{level}return lhs % metal::select(rhs, {lp}1{suffix}{rp}, rhs == {lp}0{suffix}{rp});"
6092 )?
6093 }
6094 _ => unreachable!(),
6095 }
6096 writeln!(self.out, "}}")?;
6097 writeln!(self.out)?;
6098 }
6099 _ => {}
6100 }
6101 Ok(())
6102 }
6103
6104 fn get_dot_wrapper_function_helper_name(
6110 &self,
6111 scalar: crate::Scalar,
6112 size: crate::VectorSize,
6113 ) -> String {
6114 debug_assert!(concrete_int_scalars().any(|s| s == scalar));
6116
6117 let type_name = scalar.to_msl_name();
6118 let size_suffix = common::vector_size_str(size);
6119 format!("{DOT_FUNCTION_PREFIX}_{type_name}{size_suffix}")
6120 }
6121
6122 #[allow(clippy::too_many_arguments)]
6123 fn write_wrapped_math_function(
6124 &mut self,
6125 module: &crate::Module,
6126 func_ctx: &back::FunctionCtx,
6127 fun: crate::MathFunction,
6128 arg: Handle<crate::Expression>,
6129 _arg1: Option<Handle<crate::Expression>>,
6130 _arg2: Option<Handle<crate::Expression>>,
6131 _arg3: Option<Handle<crate::Expression>>,
6132 ) -> BackendResult {
6133 let arg_ty = func_ctx.resolve_type(arg, &module.types);
6134 match fun {
6135 crate::MathFunction::Abs if arg_ty.scalar_kind() == Some(crate::ScalarKind::Sint) => {
6143 let Some((vector_size, scalar)) = arg_ty.vector_size_and_scalar() else {
6144 return Ok(());
6145 };
6146 let wrapped = WrappedFunction::Math {
6147 fun,
6148 arg_ty: (vector_size, scalar),
6149 };
6150 if !self.wrapped_functions.insert(wrapped) {
6151 return Ok(());
6152 }
6153
6154 let unsigned_scalar = crate::Scalar {
6155 kind: crate::ScalarKind::Uint,
6156 ..scalar
6157 };
6158 let mut type_name = String::new();
6159 let mut unsigned_type_name = String::new();
6160 match vector_size {
6161 None => {
6162 put_numeric_type(&mut type_name, scalar, &[])?;
6163 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[])?
6164 }
6165 Some(size) => {
6166 put_numeric_type(&mut type_name, scalar, &[size])?;
6167 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[size])?;
6168 }
6169 };
6170
6171 writeln!(self.out, "{type_name} {ABS_FUNCTION}({type_name} val) {{")?;
6172 let level = back::Level(1);
6173 let zero = if scalar.width < 4 {
6174 format!("{type_name}(0)")
6175 } else {
6176 "0".to_string()
6177 };
6178 let neg_expr = if scalar.width < 4 {
6179 format!(
6180 "static_cast<{unsigned_type_name}>(-as_type<{unsigned_type_name}>(val))"
6181 )
6182 } else {
6183 format!("-as_type<{unsigned_type_name}>(val)")
6184 };
6185 writeln!(self.out, "{level}return metal::select(as_type<{type_name}>({neg_expr}), val, val >= {zero});")?;
6186 writeln!(self.out, "}}")?;
6187 writeln!(self.out)?;
6188 }
6189
6190 crate::MathFunction::Dot => match *arg_ty {
6191 crate::TypeInner::Vector { size, scalar }
6192 if matches!(
6193 scalar.kind,
6194 crate::ScalarKind::Sint | crate::ScalarKind::Uint
6195 ) =>
6196 {
6197 let wrapped = WrappedFunction::Math {
6199 fun,
6200 arg_ty: (Some(size), scalar),
6201 };
6202 if !self.wrapped_functions.insert(wrapped) {
6203 return Ok(());
6204 }
6205
6206 let mut vec_ty = String::new();
6207 put_numeric_type(&mut vec_ty, scalar, &[size])?;
6208 let mut ret_ty = String::new();
6209 put_numeric_type(&mut ret_ty, scalar, &[])?;
6210
6211 let fun_name = self.get_dot_wrapper_function_helper_name(scalar, size);
6212
6213 writeln!(self.out, "{ret_ty} {fun_name}({vec_ty} a, {vec_ty} b) {{")?;
6215 let level = back::Level(1);
6216 write!(self.out, "{level}return ")?;
6217 self.put_dot_product("a", "b", size as usize, |writer, name, index| {
6218 write!(writer.out, "{name}.{}", back::COMPONENTS[index])?;
6219 Ok(())
6220 })?;
6221 writeln!(self.out, ";")?;
6222 writeln!(self.out, "}}")?;
6223 writeln!(self.out)?;
6224 }
6225 _ => {}
6226 },
6227
6228 _ => {}
6229 }
6230 Ok(())
6231 }
6232
6233 fn write_wrapped_cast(
6234 &mut self,
6235 module: &crate::Module,
6236 func_ctx: &back::FunctionCtx,
6237 expr: Handle<crate::Expression>,
6238 kind: crate::ScalarKind,
6239 convert: Option<crate::Bytes>,
6240 ) -> BackendResult {
6241 let src_ty = func_ctx.resolve_type(expr, &module.types);
6252 let Some(width) = convert else {
6253 return Ok(());
6254 };
6255 let Some((vector_size, src_scalar)) = src_ty.vector_size_and_scalar() else {
6256 return Ok(());
6257 };
6258 let dst_scalar = crate::Scalar { kind, width };
6259 if src_scalar.kind != crate::ScalarKind::Float
6260 || (dst_scalar.kind != crate::ScalarKind::Sint
6261 && dst_scalar.kind != crate::ScalarKind::Uint)
6262 {
6263 return Ok(());
6264 }
6265 let wrapped = WrappedFunction::Cast {
6266 src_scalar,
6267 vector_size,
6268 dst_scalar,
6269 };
6270 if !self.wrapped_functions.insert(wrapped) {
6271 return Ok(());
6272 }
6273 let (min, max) = proc::min_max_float_representable_by(src_scalar, dst_scalar);
6274
6275 let mut src_type_name = String::new();
6276 match vector_size {
6277 None => put_numeric_type(&mut src_type_name, src_scalar, &[])?,
6278 Some(size) => put_numeric_type(&mut src_type_name, src_scalar, &[size])?,
6279 };
6280 let mut dst_type_name = String::new();
6281 match vector_size {
6282 None => put_numeric_type(&mut dst_type_name, dst_scalar, &[])?,
6283 Some(size) => put_numeric_type(&mut dst_type_name, dst_scalar, &[size])?,
6284 };
6285 let fun_name = match dst_scalar {
6286 crate::Scalar::I32 => F2I32_FUNCTION,
6287 crate::Scalar::U32 => F2U32_FUNCTION,
6288 crate::Scalar::I64 => F2I64_FUNCTION,
6289 crate::Scalar::U64 => F2U64_FUNCTION,
6290 _ => unreachable!(),
6291 };
6292
6293 writeln!(
6294 self.out,
6295 "{dst_type_name} {fun_name}({src_type_name} value) {{"
6296 )?;
6297 let level = back::Level(1);
6298 write!(
6299 self.out,
6300 "{level}return static_cast<{dst_type_name}>({NAMESPACE}::clamp(value, "
6301 )?;
6302 self.put_literal(min)?;
6303 write!(self.out, ", ")?;
6304 self.put_literal(max)?;
6305 writeln!(self.out, "));")?;
6306 writeln!(self.out, "}}")?;
6307 writeln!(self.out)?;
6308 Ok(())
6309 }
6310
6311 fn write_convert_yuv_to_rgb_and_return(
6319 &mut self,
6320 level: back::Level,
6321 y: &str,
6322 uv: &str,
6323 params: &str,
6324 ) -> BackendResult {
6325 let l1 = level;
6326 let l2 = l1.next();
6327
6328 writeln!(
6330 self.out,
6331 "{l1}float3 srcGammaRgb = ({params}.yuv_conversion_matrix * float4({y}, {uv}, 1.0)).rgb;"
6332 )?;
6333
6334 writeln!(self.out, "{l1}float3 srcLinearRgb = {NAMESPACE}::select(")?;
6337 writeln!(self.out, "{l2}{NAMESPACE}::pow((srcGammaRgb + {params}.src_tf.a - 1.0) / {params}.src_tf.a, {params}.src_tf.g),")?;
6338 writeln!(self.out, "{l2}srcGammaRgb / {params}.src_tf.k,")?;
6339 writeln!(
6340 self.out,
6341 "{l2}srcGammaRgb < {params}.src_tf.k * {params}.src_tf.b);"
6342 )?;
6343
6344 writeln!(
6347 self.out,
6348 "{l1}float3 dstLinearRgb = {params}.gamut_conversion_matrix * srcLinearRgb;"
6349 )?;
6350
6351 writeln!(self.out, "{l1}float3 dstGammaRgb = {NAMESPACE}::select(")?;
6354 writeln!(self.out, "{l2}{params}.dst_tf.a * {NAMESPACE}::pow(dstLinearRgb, 1.0 / {params}.dst_tf.g) - ({params}.dst_tf.a - 1),")?;
6355 writeln!(self.out, "{l2}{params}.dst_tf.k * dstLinearRgb,")?;
6356 writeln!(self.out, "{l2}dstLinearRgb < {params}.dst_tf.b);")?;
6357
6358 writeln!(self.out, "{l1}return float4(dstGammaRgb, 1.0);")?;
6359 Ok(())
6360 }
6361
6362 #[allow(clippy::too_many_arguments)]
6363 fn write_wrapped_image_load(
6364 &mut self,
6365 module: &crate::Module,
6366 func_ctx: &back::FunctionCtx,
6367 image: Handle<crate::Expression>,
6368 _coordinate: Handle<crate::Expression>,
6369 _array_index: Option<Handle<crate::Expression>>,
6370 _sample: Option<Handle<crate::Expression>>,
6371 _level: Option<Handle<crate::Expression>>,
6372 ) -> BackendResult {
6373 let class = match *func_ctx.resolve_type(image, &module.types) {
6375 crate::TypeInner::Image { class, .. } => class,
6376 _ => unreachable!(),
6377 };
6378 if class != crate::ImageClass::External {
6379 return Ok(());
6380 }
6381 let wrapped = WrappedFunction::ImageLoad { class };
6382 if !self.wrapped_functions.insert(wrapped) {
6383 return Ok(());
6384 }
6385
6386 writeln!(self.out, "float4 {IMAGE_LOAD_EXTERNAL_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex, uint2 coords) {{")?;
6387 let l1 = back::Level(1);
6388 let l2 = l1.next();
6389 let l3 = l2.next();
6390 writeln!(
6391 self.out,
6392 "{l1}uint2 plane0_size = uint2(tex.plane0.get_width(), tex.plane0.get_height());"
6393 )?;
6394 writeln!(
6398 self.out,
6399 "{l1}uint2 cropped_size = {NAMESPACE}::any(tex.params.size != 0) ? tex.params.size : plane0_size;"
6400 )?;
6401 writeln!(
6402 self.out,
6403 "{l1}coords = {NAMESPACE}::min(coords, cropped_size - 1);"
6404 )?;
6405
6406 writeln!(self.out, "{l1}uint2 plane0_coords = uint2({NAMESPACE}::round(tex.params.load_transform * float3(float2(coords), 1.0)));")?;
6408 writeln!(self.out, "{l1}if (tex.params.num_planes == 1u) {{")?;
6409 writeln!(self.out, "{l2}return tex.plane0.read(plane0_coords);")?;
6411 writeln!(self.out, "{l1}}} else {{")?;
6412
6413 writeln!(
6415 self.out,
6416 "{l2}uint2 plane1_size = uint2(tex.plane1.get_width(), tex.plane1.get_height());"
6417 )?;
6418 writeln!(self.out, "{l2}uint2 plane1_coords = uint2({NAMESPACE}::floor(float2(plane0_coords) * float2(plane1_size) / float2(plane0_size)));")?;
6419
6420 writeln!(self.out, "{l2}float y = tex.plane0.read(plane0_coords).x;")?;
6422
6423 writeln!(self.out, "{l2}float2 uv;")?;
6424 writeln!(self.out, "{l2}if (tex.params.num_planes == 2u) {{")?;
6425 writeln!(self.out, "{l3}uv = tex.plane1.read(plane1_coords).xy;")?;
6427 writeln!(self.out, "{l2}}} else {{")?;
6428 writeln!(
6430 self.out,
6431 "{l2}uint2 plane2_size = uint2(tex.plane2.get_width(), tex.plane2.get_height());"
6432 )?;
6433 writeln!(self.out, "{l2}uint2 plane2_coords = uint2({NAMESPACE}::floor(float2(plane0_coords) * float2(plane2_size) / float2(plane0_size)));")?;
6434 writeln!(
6435 self.out,
6436 "{l3}uv = float2(tex.plane1.read(plane1_coords).x, tex.plane2.read(plane2_coords).x);"
6437 )?;
6438 writeln!(self.out, "{l2}}}")?;
6439
6440 self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "tex.params")?;
6441
6442 writeln!(self.out, "{l1}}}")?;
6443 writeln!(self.out, "}}")?;
6444 writeln!(self.out)?;
6445 Ok(())
6446 }
6447
6448 #[allow(clippy::too_many_arguments)]
6449 fn write_wrapped_image_sample(
6450 &mut self,
6451 module: &crate::Module,
6452 func_ctx: &back::FunctionCtx,
6453 image: Handle<crate::Expression>,
6454 _sampler: Handle<crate::Expression>,
6455 _gather: Option<crate::SwizzleComponent>,
6456 _coordinate: Handle<crate::Expression>,
6457 _array_index: Option<Handle<crate::Expression>>,
6458 _offset: Option<Handle<crate::Expression>>,
6459 _level: crate::SampleLevel,
6460 _depth_ref: Option<Handle<crate::Expression>>,
6461 clamp_to_edge: bool,
6462 ) -> BackendResult {
6463 if !clamp_to_edge {
6466 return Ok(());
6467 }
6468 let class = match *func_ctx.resolve_type(image, &module.types) {
6469 crate::TypeInner::Image { class, .. } => class,
6470 _ => unreachable!(),
6471 };
6472 let wrapped = WrappedFunction::ImageSample {
6473 class,
6474 clamp_to_edge: true,
6475 };
6476 if !self.wrapped_functions.insert(wrapped) {
6477 return Ok(());
6478 }
6479 match class {
6480 crate::ImageClass::External => {
6481 writeln!(self.out, "float4 {IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex, {NAMESPACE}::sampler samp, float2 coords) {{")?;
6482 let l1 = back::Level(1);
6483 let l2 = l1.next();
6484 let l3 = l2.next();
6485 writeln!(self.out, "{l1}uint2 plane0_size = uint2(tex.plane0.get_width(), tex.plane0.get_height());")?;
6486 writeln!(
6487 self.out,
6488 "{l1}coords = tex.params.sample_transform * float3(coords, 1.0);"
6489 )?;
6490
6491 writeln!(
6499 self.out,
6500 "{l1}float2 bounds_min = tex.params.sample_transform * float3(0.0, 0.0, 1.0);"
6501 )?;
6502 writeln!(
6503 self.out,
6504 "{l1}float2 bounds_max = tex.params.sample_transform * float3(1.0, 1.0, 1.0);"
6505 )?;
6506 writeln!(self.out, "{l1}float4 bounds = float4({NAMESPACE}::min(bounds_min, bounds_max), {NAMESPACE}::max(bounds_min, bounds_max));")?;
6507 writeln!(
6508 self.out,
6509 "{l1}float2 plane0_half_texel = float2(0.5, 0.5) / float2(plane0_size);"
6510 )?;
6511 writeln!(
6512 self.out,
6513 "{l1}float2 plane0_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane0_half_texel, bounds.zw - plane0_half_texel);"
6514 )?;
6515 writeln!(self.out, "{l1}if (tex.params.num_planes == 1u) {{")?;
6516 writeln!(
6518 self.out,
6519 "{l2}return tex.plane0.sample(samp, plane0_coords, {NAMESPACE}::level(0.0f));"
6520 )?;
6521 writeln!(self.out, "{l1}}} else {{")?;
6522 writeln!(self.out, "{l2}uint2 plane1_size = uint2(tex.plane1.get_width(), tex.plane1.get_height());")?;
6523 writeln!(
6524 self.out,
6525 "{l2}float2 plane1_half_texel = float2(0.5, 0.5) / float2(plane1_size);"
6526 )?;
6527 writeln!(
6528 self.out,
6529 "{l2}float2 plane1_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane1_half_texel, bounds.zw - plane1_half_texel);"
6530 )?;
6531
6532 writeln!(
6534 self.out,
6535 "{l2}float y = tex.plane0.sample(samp, plane0_coords, {NAMESPACE}::level(0.0f)).r;"
6536 )?;
6537 writeln!(self.out, "{l2}float2 uv = float2(0.0, 0.0);")?;
6538 writeln!(self.out, "{l2}if (tex.params.num_planes == 2u) {{")?;
6539 writeln!(
6541 self.out,
6542 "{l3}uv = tex.plane1.sample(samp, plane1_coords, {NAMESPACE}::level(0.0f)).xy;"
6543 )?;
6544 writeln!(self.out, "{l2}}} else {{")?;
6545 writeln!(self.out, "{l3}uint2 plane2_size = uint2(tex.plane2.get_width(), tex.plane2.get_height());")?;
6547 writeln!(
6548 self.out,
6549 "{l3}float2 plane2_half_texel = float2(0.5, 0.5) / float2(plane2_size);"
6550 )?;
6551 writeln!(
6552 self.out,
6553 "{l3}float2 plane2_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane2_half_texel, bounds.zw - plane1_half_texel);"
6554 )?;
6555 writeln!(self.out, "{l3}uv.x = tex.plane1.sample(samp, plane1_coords, {NAMESPACE}::level(0.0f)).x;")?;
6556 writeln!(self.out, "{l3}uv.y = tex.plane2.sample(samp, plane2_coords, {NAMESPACE}::level(0.0f)).x;")?;
6557 writeln!(self.out, "{l2}}}")?;
6558
6559 self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "tex.params")?;
6560
6561 writeln!(self.out, "{l1}}}")?;
6562 writeln!(self.out, "}}")?;
6563 writeln!(self.out)?;
6564 }
6565 _ => {
6566 writeln!(self.out, "{NAMESPACE}::float4 {IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION}({NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> tex, {NAMESPACE}::sampler samp, {NAMESPACE}::float2 coords) {{")?;
6567 let l1 = back::Level(1);
6568 writeln!(self.out, "{l1}{NAMESPACE}::float2 half_texel = 0.5 / {NAMESPACE}::float2(tex.get_width(0u), tex.get_height(0u));")?;
6569 writeln!(
6570 self.out,
6571 "{l1}return tex.sample(samp, {NAMESPACE}::clamp(coords, half_texel, 1.0 - half_texel), {NAMESPACE}::level(0.0));"
6572 )?;
6573 writeln!(self.out, "}}")?;
6574 writeln!(self.out)?;
6575 }
6576 }
6577 Ok(())
6578 }
6579
6580 fn write_wrapped_image_query(
6581 &mut self,
6582 module: &crate::Module,
6583 func_ctx: &back::FunctionCtx,
6584 image: Handle<crate::Expression>,
6585 query: crate::ImageQuery,
6586 ) -> BackendResult {
6587 if !matches!(query, crate::ImageQuery::Size { .. }) {
6589 return Ok(());
6590 }
6591 let class = match *func_ctx.resolve_type(image, &module.types) {
6592 crate::TypeInner::Image { class, .. } => class,
6593 _ => unreachable!(),
6594 };
6595 if class != crate::ImageClass::External {
6596 return Ok(());
6597 }
6598 let wrapped = WrappedFunction::ImageQuerySize { class };
6599 if !self.wrapped_functions.insert(wrapped) {
6600 return Ok(());
6601 }
6602 writeln!(
6603 self.out,
6604 "uint2 {IMAGE_SIZE_EXTERNAL_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex) {{"
6605 )?;
6606 let l1 = back::Level(1);
6607 let l2 = l1.next();
6608 writeln!(
6609 self.out,
6610 "{l1}if ({NAMESPACE}::any(tex.params.size != uint2(0u))) {{"
6611 )?;
6612 writeln!(self.out, "{l2}return tex.params.size;")?;
6613 writeln!(self.out, "{l1}}} else {{")?;
6614 writeln!(
6616 self.out,
6617 "{l2}return uint2(tex.plane0.get_width(), tex.plane0.get_height());"
6618 )?;
6619 writeln!(self.out, "{l1}}}")?;
6620 writeln!(self.out, "}}")?;
6621 writeln!(self.out)?;
6622 Ok(())
6623 }
6624
6625 fn write_wrapped_cooperative_load(
6626 &mut self,
6627 module: &crate::Module,
6628 func_ctx: &back::FunctionCtx,
6629 columns: crate::CooperativeSize,
6630 rows: crate::CooperativeSize,
6631 pointer: Handle<crate::Expression>,
6632 ) -> BackendResult {
6633 let ptr_ty = func_ctx.resolve_type(pointer, &module.types);
6634 let space = ptr_ty.pointer_space().unwrap();
6635 let space_name = space.to_msl_name().unwrap_or_default();
6636 let scalar = ptr_ty
6637 .pointer_base_type()
6638 .unwrap()
6639 .inner_with(&module.types)
6640 .scalar()
6641 .unwrap();
6642 let wrapped = WrappedFunction::CooperativeLoad {
6643 space_name,
6644 columns,
6645 rows,
6646 scalar,
6647 };
6648 if !self.wrapped_functions.insert(wrapped) {
6649 return Ok(());
6650 }
6651 let scalar_name = scalar.to_msl_name();
6652 writeln!(
6653 self.out,
6654 "{NAMESPACE}::simdgroup_{scalar_name}{}x{} {COOPERATIVE_LOAD_FUNCTION}(const {space_name} {scalar_name}* ptr, int stride, bool is_row_major) {{",
6655 columns as u32, rows as u32,
6656 )?;
6657 let l1 = back::Level(1);
6658 writeln!(
6659 self.out,
6660 "{l1}{NAMESPACE}::simdgroup_{scalar_name}{}x{} m;",
6661 columns as u32, rows as u32
6662 )?;
6663 let matrix_origin = "0";
6664 writeln!(
6665 self.out,
6666 "{l1}simdgroup_load(m, ptr, stride, {matrix_origin}, is_row_major);"
6667 )?;
6668 writeln!(self.out, "{l1}return m;")?;
6669 writeln!(self.out, "}}")?;
6670 writeln!(self.out)?;
6671 Ok(())
6672 }
6673
6674 fn write_wrapped_cooperative_multiply_add(
6675 &mut self,
6676 module: &crate::Module,
6677 func_ctx: &back::FunctionCtx,
6678 space: crate::AddressSpace,
6679 a: Handle<crate::Expression>,
6680 b: Handle<crate::Expression>,
6681 c: Handle<crate::Expression>,
6682 ) -> BackendResult {
6683 let space_name = space.to_msl_name().unwrap_or_default();
6684 let (a_c, a_r, ab_scalar) = match *func_ctx.resolve_type(a, &module.types) {
6685 crate::TypeInner::CooperativeMatrix {
6686 columns,
6687 rows,
6688 scalar,
6689 ..
6690 } => (columns, rows, scalar),
6691 _ => unreachable!(),
6692 };
6693 let (b_c, b_r) = match *func_ctx.resolve_type(b, &module.types) {
6694 crate::TypeInner::CooperativeMatrix { columns, rows, .. } => (columns, rows),
6695 _ => unreachable!(),
6696 };
6697 let c_scalar = match *func_ctx.resolve_type(c, &module.types) {
6698 crate::TypeInner::CooperativeMatrix { scalar, .. } => scalar,
6699 _ => unreachable!(),
6700 };
6701 let wrapped = WrappedFunction::CooperativeMultiplyAdd {
6702 space_name,
6703 columns: b_c,
6704 rows: a_r,
6705 intermediate: a_c,
6706 ab_scalar,
6707 c_scalar,
6708 };
6709 if !self.wrapped_functions.insert(wrapped) {
6710 return Ok(());
6711 }
6712 let ab_scalar_name = ab_scalar.to_msl_name();
6713 let c_scalar_name = c_scalar.to_msl_name();
6714 writeln!(
6715 self.out,
6716 "{NAMESPACE}::simdgroup_{c_scalar_name}{}x{} {COOPERATIVE_MULTIPLY_ADD_FUNCTION}(const {space_name} {NAMESPACE}::simdgroup_{ab_scalar_name}{}x{}& a, const {space_name} {NAMESPACE}::simdgroup_{ab_scalar_name}{}x{}& b, const {space_name} {NAMESPACE}::simdgroup_{c_scalar_name}{}x{}& c) {{",
6717 b_c as u32, a_r as u32, a_c as u32, a_r as u32, b_c as u32, b_r as u32, b_c as u32, a_r as u32,
6718 )?;
6719 let l1 = back::Level(1);
6720 writeln!(
6721 self.out,
6722 "{l1}{NAMESPACE}::simdgroup_{c_scalar_name}{}x{} d;",
6723 b_c as u32, a_r as u32
6724 )?;
6725 writeln!(self.out, "{l1}simdgroup_multiply_accumulate(d,a,b,c);")?;
6726 writeln!(self.out, "{l1}return d;")?;
6727 writeln!(self.out, "}}")?;
6728 writeln!(self.out)?;
6729 Ok(())
6730 }
6731
6732 pub(super) fn write_wrapped_functions(
6733 &mut self,
6734 module: &crate::Module,
6735 func_ctx: &back::FunctionCtx,
6736 options: &Options,
6737 ) -> BackendResult {
6738 for (expr_handle, expr) in func_ctx.expressions.iter() {
6739 match *expr {
6740 crate::Expression::Unary { op, expr: operand } => {
6741 self.write_wrapped_unary_op(module, func_ctx, op, operand)?;
6742 }
6743 crate::Expression::Binary { op, left, right } => {
6744 self.write_wrapped_binary_op(module, func_ctx, expr_handle, op, left, right)?;
6745 }
6746 crate::Expression::Math {
6747 fun,
6748 arg,
6749 arg1,
6750 arg2,
6751 arg3,
6752 } => {
6753 self.write_wrapped_math_function(module, func_ctx, fun, arg, arg1, arg2, arg3)?;
6754 }
6755 crate::Expression::As {
6756 expr,
6757 kind,
6758 convert,
6759 } => {
6760 self.write_wrapped_cast(module, func_ctx, expr, kind, convert)?;
6761 }
6762 crate::Expression::ImageLoad {
6763 image,
6764 coordinate,
6765 array_index,
6766 sample,
6767 level,
6768 } => {
6769 self.write_wrapped_image_load(
6770 module,
6771 func_ctx,
6772 image,
6773 coordinate,
6774 array_index,
6775 sample,
6776 level,
6777 )?;
6778 }
6779 crate::Expression::ImageSample {
6780 image,
6781 sampler,
6782 gather,
6783 coordinate,
6784 array_index,
6785 offset,
6786 level,
6787 depth_ref,
6788 clamp_to_edge,
6789 } => {
6790 self.write_wrapped_image_sample(
6791 module,
6792 func_ctx,
6793 image,
6794 sampler,
6795 gather,
6796 coordinate,
6797 array_index,
6798 offset,
6799 level,
6800 depth_ref,
6801 clamp_to_edge,
6802 )?;
6803 }
6804 crate::Expression::ImageQuery { image, query } => {
6805 self.write_wrapped_image_query(module, func_ctx, image, query)?;
6806 }
6807 crate::Expression::CooperativeLoad {
6808 columns,
6809 rows,
6810 role: _,
6811 ref data,
6812 } => {
6813 self.write_wrapped_cooperative_load(
6814 module,
6815 func_ctx,
6816 columns,
6817 rows,
6818 data.pointer,
6819 )?;
6820 }
6821 crate::Expression::CooperativeMultiplyAdd { a, b, c } => {
6822 let space = crate::AddressSpace::Private;
6823 self.write_wrapped_cooperative_multiply_add(module, func_ctx, space, a, b, c)?;
6824 }
6825 crate::Expression::RayQueryGetIntersection { committed, .. } => {
6826 self.write_rq_get_intersection_function(module, committed, options)?;
6827 }
6828 _ => {}
6829 }
6830 }
6831
6832 Ok(())
6833 }
6834
6835 fn write_functions(
6837 &mut self,
6838 module: &crate::Module,
6839 mod_info: &valid::ModuleInfo,
6840 options: &Options,
6841 pipeline_options: &PipelineOptions,
6842 ) -> Result<TranslationInfo, Error> {
6843 use nt::VertexFormat;
6844
6845 struct AttributeMappingResolved {
6848 ty_name: String,
6849 dimension: Option<crate::VectorSize>,
6850 scalar: crate::Scalar,
6851 name: String,
6852 }
6853 let mut am_resolved = FastHashMap::<u32, AttributeMappingResolved>::default();
6854
6855 struct VertexBufferMappingResolved<'a> {
6856 id: u32,
6857 stride: u32,
6858 step_mode: back::msl::VertexBufferStepMode,
6859 ty_name: String,
6860 param_name: String,
6861 elem_name: String,
6862 attributes: &'a Vec<back::msl::AttributeMapping>,
6863 }
6864 let mut vbm_resolved = Vec::<VertexBufferMappingResolved>::new();
6865
6866 struct UnpackingFunction {
6868 name: String,
6869 byte_count: u32,
6870 dimension: Option<crate::VectorSize>,
6871 scalar: crate::Scalar,
6872 }
6873 let mut unpacking_functions = FastHashMap::<VertexFormat, UnpackingFunction>::default();
6874
6875 let mut needs_vertex_id = false;
6881 let v_id = self.namer.call("v_id");
6882
6883 let mut needs_instance_id = false;
6884 let i_id = self.namer.call("i_id");
6885 if pipeline_options.vertex_pulling_transform {
6886 for vbm in &pipeline_options.vertex_buffer_mappings {
6887 let buffer_id = vbm.id;
6888 let buffer_stride = vbm.stride;
6889
6890 assert!(
6891 buffer_stride > 0,
6892 "Vertex pulling requires a non-zero buffer stride."
6893 );
6894
6895 match vbm.step_mode {
6896 back::msl::VertexBufferStepMode::Constant => {}
6897 back::msl::VertexBufferStepMode::ByVertex => {
6898 needs_vertex_id = true;
6899 }
6900 back::msl::VertexBufferStepMode::ByInstance => {
6901 needs_instance_id = true;
6902 }
6903 }
6904
6905 let buffer_ty = self.namer.call(format!("vb_{buffer_id}_type").as_str());
6906 let buffer_param = self.namer.call(format!("vb_{buffer_id}_in").as_str());
6907 let buffer_elem = self.namer.call(format!("vb_{buffer_id}_elem").as_str());
6908
6909 vbm_resolved.push(VertexBufferMappingResolved {
6910 id: buffer_id,
6911 stride: buffer_stride,
6912 step_mode: vbm.step_mode,
6913 ty_name: buffer_ty,
6914 param_name: buffer_param,
6915 elem_name: buffer_elem,
6916 attributes: &vbm.attributes,
6917 });
6918
6919 for attribute in &vbm.attributes {
6921 if unpacking_functions.contains_key(&attribute.format) {
6922 continue;
6923 }
6924 let (name, byte_count, dimension, scalar) =
6925 match self.write_unpacking_function(attribute.format) {
6926 Ok((name, byte_count, dimension, scalar)) => {
6927 (name, byte_count, dimension, scalar)
6928 }
6929 _ => {
6930 continue;
6931 }
6932 };
6933 unpacking_functions.insert(
6934 attribute.format,
6935 UnpackingFunction {
6936 name,
6937 byte_count,
6938 dimension,
6939 scalar,
6940 },
6941 );
6942 }
6943 }
6944 }
6945
6946 let mut pass_through_globals = Vec::new();
6947 for (fun_handle, fun) in module.functions.iter() {
6948 log::trace!(
6949 "function {:?}, handle {:?}",
6950 fun.name.as_deref().unwrap_or("(anonymous)"),
6951 fun_handle
6952 );
6953
6954 let ctx = back::FunctionCtx {
6955 ty: back::FunctionType::Function(fun_handle),
6956 info: &mod_info[fun_handle],
6957 expressions: &fun.expressions,
6958 named_expressions: &fun.named_expressions,
6959 };
6960
6961 writeln!(self.out)?;
6962 self.write_wrapped_functions(module, &ctx, options)?;
6963
6964 let fun_info = &mod_info[fun_handle];
6965 pass_through_globals.clear();
6966 let mut needs_buffer_sizes = false;
6967 for (handle, var) in module.global_variables.iter() {
6968 if !fun_info[handle].is_empty() {
6969 if var.space.needs_pass_through() {
6970 pass_through_globals.push(handle);
6971 }
6972 needs_buffer_sizes |= module.types[var.ty]
6973 .inner
6974 .needs_host_buffer_byte_size(&module.types);
6975 }
6976 }
6977
6978 let fun_name = &self.names[&NameKey::Function(fun_handle)];
6979 match fun.result {
6980 Some(ref result) => {
6981 let ty_name = TypeContext {
6982 handle: result.ty,
6983 gctx: module.to_ctx(),
6984 names: &self.names,
6985 access: crate::StorageAccess::empty(),
6986 first_time: false,
6987 };
6988 write!(self.out, "{ty_name}")?;
6989 }
6990 None => {
6991 write!(self.out, "void")?;
6992 }
6993 }
6994 writeln!(self.out, " {fun_name}(")?;
6995
6996 for (index, arg) in fun.arguments.iter().enumerate() {
6997 let name = &self.names[&NameKey::FunctionArgument(fun_handle, index as u32)];
6998 let param_type_name = TypeContext {
6999 handle: arg.ty,
7000 gctx: module.to_ctx(),
7001 names: &self.names,
7002 access: match module.types[arg.ty].inner {
7003 crate::TypeInner::Image {
7004 class: crate::ImageClass::Storage { access, .. },
7005 ..
7006 } => access,
7007 _ => crate::StorageAccess::empty(),
7008 },
7009 first_time: false,
7010 };
7011 let separator = separate(
7012 !pass_through_globals.is_empty()
7013 || index + 1 != fun.arguments.len()
7014 || needs_buffer_sizes,
7015 );
7016 writeln!(
7017 self.out,
7018 "{}{} {}{}",
7019 back::INDENT,
7020 param_type_name,
7021 name,
7022 separator
7023 )?;
7024 }
7025 for (index, &handle) in pass_through_globals.iter().enumerate() {
7026 let tyvar = TypedGlobalVariable {
7027 module,
7028 names: &self.names,
7029 handle,
7030 usage: fun_info[handle],
7031 reference: true,
7032 };
7033 let separator =
7034 separate(index + 1 != pass_through_globals.len() || needs_buffer_sizes);
7035 write!(self.out, "{}", back::INDENT)?;
7036 tyvar.try_fmt(&mut self.out)?;
7037 writeln!(self.out, "{separator}")?;
7038 }
7039
7040 if needs_buffer_sizes {
7041 writeln!(
7042 self.out,
7043 "{}constant _mslBufferSizes& _buffer_sizes",
7044 back::INDENT
7045 )?;
7046 }
7047
7048 writeln!(self.out, ") {{")?;
7049
7050 let guarded_indices =
7051 index::find_checked_indexes(module, fun, fun_info, options.bounds_check_policies);
7052
7053 let context = StatementContext {
7054 expression: ExpressionContext {
7055 function: fun,
7056 origin: FunctionOrigin::Handle(fun_handle),
7057 info: fun_info,
7058 lang_version: options.lang_version,
7059 policies: options.bounds_check_policies,
7060 guarded_indices,
7061 module,
7062 mod_info,
7063 pipeline_options,
7064 force_loop_bounding: options.force_loop_bounding,
7065 emit_int_div_checks: options.emit_int_div_checks,
7066 ray_query_initialization_tracking: options.ray_query_initialization_tracking,
7067 },
7068 result_struct: None,
7069 };
7070
7071 self.put_locals(&context.expression)?;
7072 self.update_expressions_to_bake(fun, fun_info, &context.expression);
7073 self.put_block(back::Level(1), &fun.body, &context)?;
7074 writeln!(self.out, "}}")?;
7075 self.named_expressions.clear();
7076 }
7077
7078 let ep_range = get_entry_points(module, pipeline_options.entry_point.as_ref())
7079 .map_err(|(stage, name)| Error::EntryPointNotFound(stage, name))?;
7080
7081 let mut info = TranslationInfo {
7082 entry_point_names: Vec::with_capacity(ep_range.len()),
7083 };
7084
7085 for ep_index in ep_range {
7086 let ep = &module.entry_points[ep_index];
7087 let fun = &ep.function;
7088 let fun_info = mod_info.get_entry_point(ep_index);
7089 let mut ep_error = None;
7090
7091 let mut v_existing_id = None;
7095 let mut i_existing_id = None;
7096
7097 log::trace!(
7098 "entry point {:?}, index {:?}",
7099 fun.name.as_deref().unwrap_or("(anonymous)"),
7100 ep_index
7101 );
7102
7103 let ctx = back::FunctionCtx {
7104 ty: back::FunctionType::EntryPoint(ep_index as u16),
7105 info: fun_info,
7106 expressions: &fun.expressions,
7107 named_expressions: &fun.named_expressions,
7108 };
7109
7110 self.write_wrapped_functions(module, &ctx, options)?;
7111
7112 let (em_str, in_mode, out_mode, can_vertex_pull) = match ep.stage {
7113 crate::ShaderStage::Vertex => (
7114 Some("vertex"),
7115 LocationMode::VertexInput,
7116 LocationMode::VertexOutput,
7117 true,
7118 ),
7119 crate::ShaderStage::Fragment => (
7120 Some("fragment"),
7121 LocationMode::FragmentInput,
7122 LocationMode::FragmentOutput,
7123 false,
7124 ),
7125 crate::ShaderStage::Compute => (
7126 Some("kernel"),
7127 LocationMode::Uniform,
7128 LocationMode::Uniform,
7129 false,
7130 ),
7131 crate::ShaderStage::Task => {
7132 (None, LocationMode::Uniform, LocationMode::Uniform, false)
7133 }
7134 crate::ShaderStage::Mesh => {
7135 (None, LocationMode::Uniform, LocationMode::MeshOutput, false)
7136 }
7137 crate::ShaderStage::RayGeneration
7138 | crate::ShaderStage::AnyHit
7139 | crate::ShaderStage::ClosestHit
7140 | crate::ShaderStage::Miss => unimplemented!(),
7141 };
7142
7143 let do_vertex_pulling = can_vertex_pull
7145 && pipeline_options.vertex_pulling_transform
7146 && !pipeline_options.vertex_buffer_mappings.is_empty();
7147
7148 let needs_buffer_sizes = do_vertex_pulling
7150 || module
7151 .global_variables
7152 .iter()
7153 .filter(|&(handle, _)| !fun_info[handle].is_empty())
7154 .any(|(_, var)| {
7155 module.types[var.ty]
7156 .inner
7157 .needs_host_buffer_byte_size(&module.types)
7158 });
7159
7160 if !options.fake_missing_bindings {
7163 for (var_handle, var) in module.global_variables.iter() {
7164 if fun_info[var_handle].is_empty() {
7165 continue;
7166 }
7167 match var.space {
7168 crate::AddressSpace::Uniform
7169 | crate::AddressSpace::Storage { .. }
7170 | crate::AddressSpace::Handle => {
7171 let br = match var.binding {
7172 Some(ref br) => br,
7173 None => {
7174 let var_name = var.name.clone().unwrap_or_default();
7175 ep_error =
7176 Some(super::EntryPointError::MissingBinding(var_name));
7177 break;
7178 }
7179 };
7180 let target = options.get_resource_binding_target(ep, br);
7181 let good = match target {
7182 Some(target) => {
7183 match module.types[var.ty].inner {
7187 crate::TypeInner::Image {
7188 class: crate::ImageClass::External,
7189 ..
7190 } => target.external_texture.is_some(),
7191 crate::TypeInner::Image { .. } => target.texture.is_some(),
7192 crate::TypeInner::Sampler { .. } => {
7193 target.sampler.is_some()
7194 }
7195 _ => target.buffer.is_some(),
7196 }
7197 }
7198 None => false,
7199 };
7200 if !good {
7201 ep_error = Some(super::EntryPointError::MissingBindTarget(*br));
7202 break;
7203 }
7204 }
7205 crate::AddressSpace::Immediate => {
7206 if let Err(e) = options.resolve_immediates(ep) {
7207 ep_error = Some(e);
7208 break;
7209 }
7210 }
7211 crate::AddressSpace::Function
7212 | crate::AddressSpace::Private
7213 | crate::AddressSpace::WorkGroup
7214 | crate::AddressSpace::TaskPayload => {}
7215 crate::AddressSpace::RayPayload
7216 | crate::AddressSpace::IncomingRayPayload => unimplemented!(),
7217 }
7218 }
7219 if needs_buffer_sizes {
7220 if let Err(err) = options.resolve_sizes_buffer(ep) {
7221 ep_error = Some(err);
7222 }
7223 }
7224 }
7225
7226 if let Some(err) = ep_error {
7227 info.entry_point_names.push(Err(err));
7228 continue;
7229 }
7230 let fun_name = self.names[&NameKey::EntryPoint(ep_index as _)].clone();
7231 info.entry_point_names.push(Ok(fun_name.clone()));
7232
7233 writeln!(self.out)?;
7234
7235 let mut flattened_member_names = FastHashMap::default();
7241 let mut varyings_namer = proc::Namer::default();
7243
7244 let mut empty_names = FastHashMap::default(); varyings_namer.reset(
7246 module,
7247 &super::keywords::RESERVED_SET,
7248 proc::KeywordSet::empty(),
7249 proc::CaseInsensitiveKeywordSet::empty(),
7250 &[CLAMPED_LOD_LOAD_PREFIX],
7251 &mut empty_names,
7252 );
7253
7254 let mut flattened_arguments = Vec::new();
7259 for (arg_index, arg) in fun.arguments.iter().enumerate() {
7260 match module.types[arg.ty].inner {
7261 crate::TypeInner::Struct { ref members, .. } => {
7262 for (member_index, member) in members.iter().enumerate() {
7263 let member_index = member_index as u32;
7264 flattened_arguments.push((
7265 NameKey::StructMember(arg.ty, member_index),
7266 member.ty,
7267 member.binding.as_ref(),
7268 ));
7269 let name_key = NameKey::StructMember(arg.ty, member_index);
7270 let name = match member.binding {
7271 Some(crate::Binding::Location { .. }) => {
7272 if do_vertex_pulling {
7273 self.namer.call(&self.names[&name_key])
7274 } else {
7275 varyings_namer.call(&self.names[&name_key])
7276 }
7277 }
7278 _ => self.namer.call(&self.names[&name_key]),
7279 };
7280 flattened_member_names.insert(name_key, name);
7281 }
7282 }
7283 _ => flattened_arguments.push((
7284 NameKey::EntryPointArgument(ep_index as _, arg_index as u32),
7285 arg.ty,
7286 arg.binding.as_ref(),
7287 )),
7288 }
7289 }
7290
7291 let stage_in_name = self.namer.call(&format!("{fun_name}Input"));
7296 let varyings_member_name = self.namer.call("varyings");
7297 let mut has_varyings = false;
7298
7299 if !flattened_arguments.is_empty() {
7300 if !do_vertex_pulling {
7301 writeln!(self.out, "struct {stage_in_name} {{")?;
7302 }
7303 for &(ref name_key, ty, binding) in flattened_arguments.iter() {
7304 let Some(binding) = binding else {
7305 continue;
7306 };
7307 let name = match *name_key {
7308 NameKey::StructMember(..) => &flattened_member_names[name_key],
7309 _ => &self.names[name_key],
7310 };
7311 let ty_name = TypeContext {
7312 handle: ty,
7313 gctx: module.to_ctx(),
7314 names: &self.names,
7315 access: crate::StorageAccess::empty(),
7316 first_time: false,
7317 };
7318 let resolved = options.resolve_local_binding(binding, in_mode)?;
7319 let location = match *binding {
7320 crate::Binding::Location { location, .. } => Some(location),
7321 crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. }) => None,
7322 crate::Binding::BuiltIn(_) => continue,
7323 };
7324 if do_vertex_pulling {
7325 let Some(location) = location else {
7326 continue;
7327 };
7328 am_resolved.insert(
7330 location,
7331 AttributeMappingResolved {
7332 ty_name: ty_name.to_string(),
7333 dimension: ty_name.vector_size(),
7334 scalar: ty_name.scalar().unwrap(),
7335 name: name.to_string(),
7336 },
7337 );
7338 } else {
7339 has_varyings = true;
7340 if let super::ResolvedBinding::User {
7341 prefix,
7342 index,
7343 interpolation: Some(super::ResolvedInterpolation::PerVertex),
7344 } = resolved
7345 {
7346 if options.lang_version < (4, 0) {
7347 return Err(Error::PerVertexNotSupported);
7348 }
7349 write!(
7350 self.out,
7351 "{}{NAMESPACE}::vertex_value<{}> {name} [[user({prefix}{index})]]",
7352 back::INDENT,
7353 ty_name.unwrap_array()
7354 )?;
7355 } else {
7356 write!(self.out, "{}{} {}", back::INDENT, ty_name, name)?;
7357 resolved.try_fmt(&mut self.out)?;
7358 }
7359 writeln!(self.out, ";")?;
7360 }
7361 }
7362 if !do_vertex_pulling {
7363 writeln!(self.out, "}};")?;
7364 }
7365 }
7366
7367 let stage_out_name = self.namer.call(&format!("{fun_name}Output"));
7370 let result_member_name = self.namer.call("member");
7371 let result_type_name = match fun.result {
7372 Some(ref result) if ep.stage != crate::ShaderStage::Task => {
7373 let mut result_members = Vec::new();
7374 if let crate::TypeInner::Struct { ref members, .. } =
7375 module.types[result.ty].inner
7376 {
7377 for (member_index, member) in members.iter().enumerate() {
7378 result_members.push((
7379 &self.names[&NameKey::StructMember(result.ty, member_index as u32)],
7380 member.ty,
7381 member.binding.as_ref(),
7382 ));
7383 }
7384 } else {
7385 result_members.push((
7386 &result_member_name,
7387 result.ty,
7388 result.binding.as_ref(),
7389 ));
7390 }
7391
7392 writeln!(self.out, "struct {stage_out_name} {{")?;
7393 let mut has_point_size = false;
7394 for (name, ty, binding) in result_members {
7395 let ty_name = TypeContext {
7396 handle: ty,
7397 gctx: module.to_ctx(),
7398 names: &self.names,
7399 access: crate::StorageAccess::empty(),
7400 first_time: true,
7401 };
7402 let binding = binding.ok_or_else(|| {
7403 Error::GenericValidation("Expected binding, got None".into())
7404 })?;
7405
7406 if let crate::Binding::BuiltIn(crate::BuiltIn::PointSize) = *binding {
7407 has_point_size = true;
7408 if !pipeline_options.allow_and_force_point_size {
7409 continue;
7410 }
7411 }
7412
7413 let array_len = match module.types[ty].inner {
7414 crate::TypeInner::Array {
7415 size: crate::ArraySize::Constant(size),
7416 ..
7417 } => Some(size),
7418 _ => None,
7419 };
7420 let resolved = options.resolve_local_binding(binding, out_mode)?;
7421 write!(self.out, "{}{} {}", back::INDENT, ty_name, name)?;
7422 resolved.try_fmt(&mut self.out)?;
7423 if let Some(array_len) = array_len {
7424 write!(self.out, " [{array_len}]")?;
7425 }
7426 writeln!(self.out, ";")?;
7427 }
7428
7429 if pipeline_options.allow_and_force_point_size
7430 && ep.stage == crate::ShaderStage::Vertex
7431 && !has_point_size
7432 {
7433 writeln!(
7435 self.out,
7436 "{}float _point_size [[point_size]];",
7437 back::INDENT
7438 )?;
7439 }
7440 writeln!(self.out, "}};")?;
7441 &stage_out_name
7442 }
7443 Some(ref result) if ep.stage == crate::ShaderStage::Task => {
7444 assert_eq!(
7445 module.types[result.ty].inner,
7446 crate::TypeInner::Vector {
7447 size: crate::VectorSize::Tri,
7448 scalar: crate::Scalar::U32
7449 }
7450 );
7451
7452 "metal::uint3"
7453 }
7454 _ => "void",
7455 };
7456
7457 let out_mesh_info = if let Some(ref mesh_info) = ep.mesh_info {
7458 Some(self.write_mesh_output_types(
7459 mesh_info,
7460 &fun_name,
7461 module,
7462 pipeline_options.allow_and_force_point_size,
7463 options,
7464 )?)
7465 } else {
7466 None
7467 };
7468
7469 if do_vertex_pulling {
7472 for vbm in &vbm_resolved {
7473 let buffer_stride = vbm.stride;
7474 let buffer_ty = &vbm.ty_name;
7475
7476 writeln!(
7480 self.out,
7481 "struct {buffer_ty} {{ metal::uchar data[{buffer_stride}]; }};"
7482 )?;
7483 }
7484 }
7485
7486 let is_wrapped = matches!(
7487 ep.stage,
7488 crate::ShaderStage::Task | crate::ShaderStage::Mesh
7489 );
7490 let fun_name = fun_name.clone();
7491 let nested_fun_name = if is_wrapped {
7492 self.namer.call(&format!("_{fun_name}"))
7493 } else {
7494 fun_name.clone()
7495 };
7496
7497 if ep.stage == crate::ShaderStage::Compute && options.lang_version >= (2, 1) {
7499 let total_threads =
7500 ep.workgroup_size[0] * ep.workgroup_size[1] * ep.workgroup_size[2];
7501 write!(
7502 self.out,
7503 "[[max_total_threads_per_threadgroup({total_threads})]] "
7504 )?;
7505 }
7506
7507 if let Some(em_str) = em_str {
7509 write!(self.out, "{em_str} ")?;
7510 }
7511 writeln!(self.out, "{result_type_name} {nested_fun_name}(")?;
7512
7513 let mut args = Vec::new();
7514
7515 if has_varyings {
7518 args.push(EntryPointArgument {
7519 ty_name: stage_in_name,
7520 name: varyings_member_name.clone(),
7521 binding: " [[stage_in]]".to_string(),
7522 init: None,
7523 });
7524 }
7525
7526 let mut local_invocation_index = None;
7527
7528 for &(ref name_key, ty, binding) in flattened_arguments.iter() {
7531 let binding = match binding {
7532 Some(&crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. })) => continue,
7533 Some(binding @ &crate::Binding::BuiltIn { .. }) => binding,
7534 _ => continue,
7535 };
7536 let name = match *name_key {
7537 NameKey::StructMember(..) => &flattened_member_names[name_key],
7538 _ => &self.names[name_key],
7539 };
7540
7541 if binding == &crate::Binding::BuiltIn(crate::BuiltIn::LocalInvocationIndex) {
7542 local_invocation_index = Some(name_key);
7543 }
7544
7545 let ty_name = TypeContext {
7546 handle: ty,
7547 gctx: module.to_ctx(),
7548 names: &self.names,
7549 access: crate::StorageAccess::empty(),
7550 first_time: false,
7551 };
7552
7553 match *binding {
7554 crate::Binding::BuiltIn(crate::BuiltIn::VertexIndex) => {
7555 v_existing_id = Some(name.clone());
7556 }
7557 crate::Binding::BuiltIn(crate::BuiltIn::InstanceIndex) => {
7558 i_existing_id = Some(name.clone());
7559 }
7560 _ => {}
7561 };
7562
7563 let resolved = options.resolve_local_binding(binding, in_mode)?;
7564 let mut binding = String::new();
7565 resolved.try_fmt(&mut binding)?;
7566
7567 args.push(EntryPointArgument {
7568 ty_name: format!("{ty_name}"),
7569 name: name.clone(),
7570 binding,
7571 init: None,
7572 });
7573 }
7574
7575 let need_workgroup_variables_initialization =
7576 self.need_workgroup_variables_initialization(options, ep, module, fun_info);
7577
7578 if local_invocation_index.is_none()
7579 && (need_workgroup_variables_initialization
7580 || ep.stage == crate::ShaderStage::Task
7581 || ep.stage == crate::ShaderStage::Mesh)
7582 {
7583 args.push(EntryPointArgument {
7584 ty_name: "uint".to_string(),
7585 name: "__local_invocation_index".to_string(),
7586 binding: " [[thread_index_in_threadgroup]]".to_string(),
7587 init: None,
7588 });
7589 }
7590
7591 for (handle, var) in module.global_variables.iter() {
7596 let usage = fun_info[handle];
7597 if usage.is_empty() || var.space == crate::AddressSpace::Private {
7598 continue;
7599 }
7600
7601 if options.lang_version < (1, 2) {
7602 match var.space {
7603 crate::AddressSpace::Storage { access }
7613 if access.contains(crate::StorageAccess::STORE)
7614 && ep.stage == crate::ShaderStage::Fragment =>
7615 {
7616 return Err(Error::UnsupportedWritableStorageBuffer)
7617 }
7618 crate::AddressSpace::Handle => {
7619 match module.types[var.ty].inner {
7620 crate::TypeInner::Image {
7621 class: crate::ImageClass::Storage { access, .. },
7622 ..
7623 } => {
7624 if access.contains(crate::StorageAccess::STORE)
7634 && (ep.stage == crate::ShaderStage::Vertex
7635 || ep.stage == crate::ShaderStage::Fragment)
7636 {
7637 return Err(Error::UnsupportedWritableStorageTexture(
7638 ep.stage,
7639 ));
7640 }
7641
7642 if access.contains(
7643 crate::StorageAccess::LOAD | crate::StorageAccess::STORE,
7644 ) {
7645 return Err(Error::UnsupportedRWStorageTexture);
7646 }
7647 }
7648 _ => {}
7649 }
7650 }
7651 _ => {}
7652 }
7653 }
7654
7655 match var.space {
7657 crate::AddressSpace::Handle => match module.types[var.ty].inner {
7658 crate::TypeInner::BindingArray { base, .. } => {
7659 match module.types[base].inner {
7660 crate::TypeInner::Sampler { .. } => {
7661 if options.lang_version < (2, 0) {
7662 return Err(Error::UnsupportedArrayOf(
7663 "samplers".to_string(),
7664 ));
7665 }
7666 }
7667 crate::TypeInner::Image { class, .. } => match class {
7668 crate::ImageClass::Sampled { .. }
7669 | crate::ImageClass::Depth { .. }
7670 | crate::ImageClass::Storage {
7671 access: crate::StorageAccess::LOAD,
7672 ..
7673 } => {
7674 if options.lang_version < (2, 0) {
7679 return Err(Error::UnsupportedArrayOf(
7680 "textures".to_string(),
7681 ));
7682 }
7683 }
7684 crate::ImageClass::Storage {
7685 access: crate::StorageAccess::STORE,
7686 ..
7687 } => {
7688 if options.lang_version < (2, 0) {
7693 return Err(Error::UnsupportedArrayOf(
7694 "write-only textures".to_string(),
7695 ));
7696 }
7697 }
7698 crate::ImageClass::Storage { .. } => {
7699 if options.lang_version < (3, 0) {
7700 return Err(Error::UnsupportedArrayOf(
7701 "read-write textures".to_string(),
7702 ));
7703 }
7704 }
7705 crate::ImageClass::External => {
7706 return Err(Error::UnsupportedArrayOf(
7707 "external textures".to_string(),
7708 ));
7709 }
7710 },
7711 _ => {
7712 return Err(Error::UnsupportedArrayOfType(base));
7713 }
7714 }
7715 }
7716 _ => {}
7717 },
7718 _ => {}
7719 }
7720
7721 let resolved = match var.space {
7723 crate::AddressSpace::Immediate => options.resolve_immediates(ep).ok(),
7724 crate::AddressSpace::WorkGroup => None,
7725 crate::AddressSpace::TaskPayload => Some(back::msl::ResolvedBinding::Payload),
7726 _ => options
7727 .resolve_resource_binding(ep, var.binding.as_ref().unwrap())
7728 .ok(),
7729 };
7730 if let Some(ref resolved) = resolved {
7731 if resolved.as_inline_sampler(options).is_some() {
7733 continue;
7734 }
7735 }
7736
7737 match module.types[var.ty].inner {
7738 crate::TypeInner::Image {
7739 class: crate::ImageClass::External,
7740 ..
7741 } => {
7742 let target = match resolved {
7746 Some(back::msl::ResolvedBinding::Resource(target)) => {
7747 target.external_texture
7748 }
7749 _ => None,
7750 };
7751
7752 for i in 0..3 {
7753 let plane_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
7754 handle,
7755 ExternalTextureNameKey::Plane(i),
7756 )];
7757 let ty_name = format!(
7758 "{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample>"
7759 );
7760 let name = plane_name.clone();
7761 let binding = if let Some(ref target) = target {
7762 format!(" [[texture({})]]", target.planes[i])
7763 } else {
7764 String::new()
7765 };
7766 args.push(EntryPointArgument {
7767 ty_name,
7768 name,
7769 binding,
7770 init: None,
7771 });
7772 }
7773 let params_ty_name = &self.names
7774 [&NameKey::Type(module.special_types.external_texture_params.unwrap())];
7775 let params_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
7776 handle,
7777 ExternalTextureNameKey::Params,
7778 )];
7779 let binding = if let Some(ref target) = target {
7780 format!(" [[buffer({})]]", target.params)
7781 } else {
7782 String::new()
7783 };
7784
7785 args.push(EntryPointArgument {
7786 ty_name: format!("constant {params_ty_name}&"),
7787 name: params_name.clone(),
7788 binding,
7789 init: None,
7790 });
7791 }
7792 _ => {
7793 if var.space == crate::AddressSpace::WorkGroup
7794 && ep.stage == crate::ShaderStage::Mesh
7795 {
7796 continue;
7797 }
7798 let tyvar = TypedGlobalVariable {
7799 module,
7800 names: &self.names,
7801 handle,
7802 usage,
7803 reference: true,
7804 };
7805 let parts = tyvar.to_parts()?;
7806 let mut binding = String::new();
7807 if let Some(resolved) = resolved {
7808 resolved.try_fmt(&mut binding)?;
7809 }
7810 args.push(EntryPointArgument {
7811 ty_name: parts.ty_name,
7812 name: parts.var_name,
7813 binding,
7814 init: var.init,
7815 });
7816 }
7817 }
7818 }
7819
7820 if do_vertex_pulling {
7821 if needs_vertex_id && v_existing_id.is_none() {
7822 args.push(EntryPointArgument {
7824 ty_name: "uint".to_string(),
7825 name: v_id.clone(),
7826 binding: " [[vertex_id]]".to_string(),
7827 init: None,
7828 });
7829 }
7830
7831 if needs_instance_id && i_existing_id.is_none() {
7832 args.push(EntryPointArgument {
7833 ty_name: "uint".to_string(),
7834 name: i_id.clone(),
7835 binding: " [[instance_id]]".to_string(),
7836 init: None,
7837 });
7838 }
7839
7840 for vbm in &vbm_resolved {
7843 let id = &vbm.id;
7844 let ty_name = &vbm.ty_name;
7845 let param_name = &vbm.param_name;
7846 args.push(EntryPointArgument {
7847 ty_name: format!("const device {ty_name}*"),
7848 name: param_name.clone(),
7849 binding: format!(" [[buffer({id})]]"),
7850 init: None,
7851 });
7852 }
7853 }
7854
7855 if needs_buffer_sizes {
7858 let resolved = options.resolve_sizes_buffer(ep).unwrap();
7860 let mut binding = String::new();
7861 resolved.try_fmt(&mut binding)?;
7862 args.push(EntryPointArgument {
7863 ty_name: "constant _mslBufferSizes&".to_string(),
7864 name: "_buffer_sizes".to_string(),
7865 binding,
7866 init: None,
7867 });
7868 }
7869
7870 let mut is_first_arg = true;
7871 for arg in &args {
7872 if is_first_arg {
7873 write!(self.out, " ")?;
7874 } else {
7875 write!(self.out, ", ")?;
7876 }
7877 is_first_arg = false;
7878 write!(self.out, "{} {}", arg.ty_name, arg.name)?;
7879 if !is_wrapped {
7880 write!(self.out, "{}", arg.binding)?;
7881 if let Some(init) = arg.init {
7882 write!(self.out, " = ")?;
7883 self.put_const_expression(
7884 init,
7885 module,
7886 mod_info,
7887 &module.global_expressions,
7888 )?;
7889 }
7890 }
7891 writeln!(self.out)?;
7892 }
7893 if ep.stage == crate::ShaderStage::Mesh {
7894 for (handle, var) in module.global_variables.iter() {
7895 if var.space != crate::AddressSpace::WorkGroup || fun_info[handle].is_empty() {
7896 continue;
7897 }
7898 if is_first_arg {
7899 write!(self.out, " ")?;
7900 } else {
7901 write!(self.out, ", ")?;
7902 }
7903 let ty_context = TypeContext {
7904 handle: module.global_variables[handle].ty,
7905 gctx: module.to_ctx(),
7906 names: &self.names,
7907 access: crate::StorageAccess::empty(),
7908 first_time: false,
7909 };
7910 writeln!(
7911 self.out,
7912 "threadgroup {ty_context}& {}",
7913 self.names[&NameKey::GlobalVariable(handle)]
7914 )?;
7915 }
7916 }
7917
7918 writeln!(self.out, ") {{")?;
7920
7921 if do_vertex_pulling {
7923 for vbm in &vbm_resolved {
7926 for attribute in vbm.attributes {
7927 let location = attribute.shader_location;
7928 let am_option = am_resolved.get(&location);
7929 if am_option.is_none() {
7930 continue;
7933 }
7934 let am = am_option.unwrap();
7935 let attribute_ty_name = &am.ty_name;
7936 let attribute_name = &am.name;
7937
7938 writeln!(
7939 self.out,
7940 "{}{attribute_ty_name} {attribute_name} = {{}};",
7941 back::Level(1)
7942 )?;
7943 }
7944
7945 write!(self.out, "{}if (", back::Level(1))?;
7948
7949 let idx = &vbm.id;
7950 let stride = &vbm.stride;
7951 let index_name = match vbm.step_mode {
7952 back::msl::VertexBufferStepMode::Constant => "0",
7953 back::msl::VertexBufferStepMode::ByVertex => {
7954 if let Some(ref name) = v_existing_id {
7955 name
7956 } else {
7957 &v_id
7958 }
7959 }
7960 back::msl::VertexBufferStepMode::ByInstance => {
7961 if let Some(ref name) = i_existing_id {
7962 name
7963 } else {
7964 &i_id
7965 }
7966 }
7967 };
7968 write!(
7969 self.out,
7970 "{index_name} < (_buffer_sizes.buffer_size{idx} / {stride})"
7971 )?;
7972
7973 writeln!(self.out, ") {{")?;
7974
7975 let ty_name = &vbm.ty_name;
7977 let elem_name = &vbm.elem_name;
7978 let param_name = &vbm.param_name;
7979
7980 writeln!(
7981 self.out,
7982 "{}const {ty_name} {elem_name} = {param_name}[{index_name}];",
7983 back::Level(2),
7984 )?;
7985
7986 for attribute in vbm.attributes {
7989 let location = attribute.shader_location;
7990 let Some(am) = am_resolved.get(&location) else {
7991 continue;
7995 };
7996 let attribute_name = &am.name;
7997 let attribute_ty_name = &am.ty_name;
7998
7999 let offset = attribute.offset;
8000 let func = unpacking_functions
8001 .get(&attribute.format)
8002 .expect("Should have generated this unpacking function earlier.");
8003 let func_name = &func.name;
8004
8005 let needs_padding_or_truncation = am.dimension.cmp(&func.dimension);
8011
8012 let needs_conversion = am.scalar != func.scalar;
8015
8016 if needs_padding_or_truncation != Ordering::Equal {
8017 writeln!(
8020 self.out,
8021 "{}// {attribute_ty_name} <- {:?}",
8022 back::Level(2),
8023 attribute.format
8024 )?;
8025 }
8026
8027 write!(self.out, "{}{attribute_name} = ", back::Level(2),)?;
8028
8029 if needs_padding_or_truncation == Ordering::Greater {
8030 write!(self.out, "{attribute_ty_name}(")?;
8032 }
8033
8034 if needs_conversion {
8036 put_numeric_type(&mut self.out, am.scalar, func.dimension.as_slice())?;
8037 write!(self.out, "(")?;
8038 }
8039 write!(self.out, "{func_name}({elem_name}.data[{offset}]")?;
8040 for i in (offset + 1)..(offset + func.byte_count) {
8041 write!(self.out, ", {elem_name}.data[{i}]")?;
8042 }
8043 write!(self.out, ")")?;
8044 if needs_conversion {
8045 write!(self.out, ")")?;
8046 }
8047
8048 match needs_padding_or_truncation {
8049 Ordering::Greater => {
8050 let ty_is_int = scalar_is_int(am.scalar);
8052 let zero_value = if ty_is_int { "0" } else { "0.0" };
8053 let one_value = if ty_is_int { "1" } else { "1.0" };
8054 for i in func.dimension.map_or(1, u8::from)
8055 ..am.dimension.map_or(1, u8::from)
8056 {
8057 write!(
8058 self.out,
8059 ", {}",
8060 if i == 3 { one_value } else { zero_value }
8061 )?;
8062 }
8063 }
8064 Ordering::Less => {
8065 write!(
8067 self.out,
8068 ".{}",
8069 &"xyzw"[0..usize::from(am.dimension.map_or(1, u8::from))]
8070 )?;
8071 }
8072 Ordering::Equal => {}
8073 }
8074
8075 if needs_padding_or_truncation == Ordering::Greater {
8076 write!(self.out, ")")?;
8077 }
8078
8079 writeln!(self.out, ";")?;
8080 }
8081
8082 writeln!(self.out, "{}}}", back::Level(1))?;
8084 }
8085 }
8086
8087 for (handle, var) in module.global_variables.iter() {
8090 let usage = fun_info[handle];
8091 if usage.is_empty() {
8092 continue;
8093 }
8094 if var.space == crate::AddressSpace::Private {
8095 let tyvar = TypedGlobalVariable {
8096 module,
8097 names: &self.names,
8098 handle,
8099 usage,
8100
8101 reference: false,
8102 };
8103 write!(self.out, "{}", back::INDENT)?;
8104 tyvar.try_fmt(&mut self.out)?;
8105 match var.init {
8106 Some(value) => {
8107 write!(self.out, " = ")?;
8108 self.put_const_expression(
8109 value,
8110 module,
8111 mod_info,
8112 &module.global_expressions,
8113 )?;
8114 writeln!(self.out, ";")?;
8115 }
8116 None => {
8117 writeln!(self.out, " = {{}};")?;
8118 }
8119 };
8120 } else if let Some(ref binding) = var.binding {
8121 let resolved = options.resolve_resource_binding(ep, binding).unwrap();
8122 if let Some(sampler) = resolved.as_inline_sampler(options) {
8123 let name = &self.names[&NameKey::GlobalVariable(handle)];
8125 writeln!(
8126 self.out,
8127 "{}constexpr {}::sampler {}(",
8128 back::INDENT,
8129 NAMESPACE,
8130 name
8131 )?;
8132 self.put_inline_sampler_properties(back::Level(2), sampler, options)?;
8133 writeln!(self.out, "{});", back::INDENT)?;
8134 } else if let crate::TypeInner::Image {
8135 class: crate::ImageClass::External,
8136 ..
8137 } = module.types[var.ty].inner
8138 {
8139 let wrapper_name = &self.names[&NameKey::GlobalVariable(handle)];
8142 let l1 = back::Level(1);
8143 let l2 = l1.next();
8144 writeln!(
8145 self.out,
8146 "{l1}const {EXTERNAL_TEXTURE_WRAPPER_STRUCT} {wrapper_name} {{"
8147 )?;
8148 for i in 0..3 {
8149 let plane_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
8150 handle,
8151 ExternalTextureNameKey::Plane(i),
8152 )];
8153 writeln!(self.out, "{l2}.plane{i} = {plane_name},")?;
8154 }
8155 let params_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
8156 handle,
8157 ExternalTextureNameKey::Params,
8158 )];
8159 writeln!(self.out, "{l2}.params = {params_name},")?;
8160 writeln!(self.out, "{l1}}};")?;
8161 }
8162 }
8163 }
8164
8165 if need_workgroup_variables_initialization {
8166 self.write_workgroup_variables_initialization(
8167 module,
8168 mod_info,
8169 fun_info,
8170 local_invocation_index,
8171 ep.stage,
8172 )?;
8173 }
8174
8175 for (arg_index, arg) in fun.arguments.iter().enumerate() {
8186 let arg_name =
8187 &self.names[&NameKey::EntryPointArgument(ep_index as _, arg_index as u32)];
8188 match module.types[arg.ty].inner {
8189 crate::TypeInner::Struct { ref members, .. } => {
8190 let struct_name = &self.names[&NameKey::Type(arg.ty)];
8191 write!(
8192 self.out,
8193 "{}const {} {} = {{ ",
8194 back::INDENT,
8195 struct_name,
8196 arg_name
8197 )?;
8198 for (member_index, member) in members.iter().enumerate() {
8199 let key = NameKey::StructMember(arg.ty, member_index as u32);
8200 let name = &flattened_member_names[&key];
8201 if member_index != 0 {
8202 write!(self.out, ", ")?;
8203 }
8204 if self
8206 .struct_member_pads
8207 .contains(&(arg.ty, member_index as u32))
8208 {
8209 write!(self.out, "{{}}, ")?;
8210 }
8211 match member.binding {
8212 Some(crate::Binding::Location {
8213 interpolation: Some(crate::Interpolation::PerVertex),
8214 ..
8215 }) => {
8216 writeln!(
8217 self.out,
8218 "{0}{{ {1}.{2}.get({NAMESPACE}::vertex_index::first), {1}.{2}.get({NAMESPACE}::vertex_index::second), {1}.{2}.get({NAMESPACE}::vertex_index::third) }}",
8219 back::INDENT,
8220 varyings_member_name,
8221 arg_name,
8222 )?;
8223 continue;
8224 }
8225 Some(crate::Binding::Location { .. }) => {
8226 if has_varyings {
8227 write!(self.out, "{varyings_member_name}.")?;
8228 }
8229 }
8230 _ => (),
8231 }
8232 write!(self.out, "{name}")?;
8233 }
8234 writeln!(self.out, " }};")?;
8235 }
8236 _ => match arg.binding {
8237 Some(crate::Binding::Location {
8238 interpolation: Some(crate::Interpolation::PerVertex),
8239 ..
8240 }) => {
8241 let ty_name = TypeContext {
8242 handle: arg.ty,
8243 gctx: module.to_ctx(),
8244 names: &self.names,
8245 access: crate::StorageAccess::empty(),
8246 first_time: false,
8247 };
8248 writeln!(
8249 self.out,
8250 "{0}const {ty_name} {arg_name} = {{ {1}.{2}.get({NAMESPACE}::vertex_index::first), {1}.{2}.get({NAMESPACE}::vertex_index::second), {1}.{2}.get({NAMESPACE}::vertex_index::third) }};",
8251 back::INDENT,
8252 varyings_member_name,
8253 arg_name,
8254 )?;
8255 }
8256 Some(crate::Binding::Location { .. })
8257 | Some(crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. })) => {
8258 if has_varyings {
8259 writeln!(
8260 self.out,
8261 "{}const auto {} = {}.{};",
8262 back::INDENT,
8263 arg_name,
8264 varyings_member_name,
8265 arg_name
8266 )?;
8267 }
8268 }
8269 _ => {}
8270 },
8271 }
8272 }
8273
8274 let guarded_indices =
8275 index::find_checked_indexes(module, fun, fun_info, options.bounds_check_policies);
8276
8277 let context = StatementContext {
8278 expression: ExpressionContext {
8279 function: fun,
8280 origin: FunctionOrigin::EntryPoint(ep_index as _),
8281 info: fun_info,
8282 lang_version: options.lang_version,
8283 policies: options.bounds_check_policies,
8284 guarded_indices,
8285 module,
8286 mod_info,
8287 pipeline_options,
8288 force_loop_bounding: options.force_loop_bounding,
8289 emit_int_div_checks: options.emit_int_div_checks,
8290 ray_query_initialization_tracking: options.ray_query_initialization_tracking,
8291 },
8292 result_struct: if ep.stage == crate::ShaderStage::Task {
8293 None
8294 } else {
8295 Some(&stage_out_name)
8296 },
8297 };
8298
8299 self.put_locals(&context.expression)?;
8302 self.update_expressions_to_bake(fun, fun_info, &context.expression);
8303 self.put_block(back::Level(1), &fun.body, &context)?;
8304 writeln!(self.out, "}}")?;
8305 if ep_index + 1 != module.entry_points.len() {
8306 writeln!(self.out)?;
8307 }
8308 self.named_expressions.clear();
8309
8310 if is_wrapped {
8311 self.write_wrapper_function(NestedFunctionInfo {
8312 options,
8313 ep,
8314 module,
8315 mod_info,
8316 fun_info,
8317 args,
8318 local_invocation_index,
8319 nested_name: &nested_fun_name,
8320 outer_name: &fun_name,
8321 out_mesh_info,
8322 })?;
8323 }
8324 }
8325
8326 Ok(info)
8327 }
8328
8329 pub(super) fn write_barrier(
8330 &mut self,
8331 flags: crate::Barrier,
8332 level: back::Level,
8333 ) -> BackendResult {
8334 if flags.is_empty() {
8337 writeln!(
8338 self.out,
8339 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_none);",
8340 )?;
8341 }
8342 if flags.contains(crate::Barrier::STORAGE) {
8343 writeln!(
8344 self.out,
8345 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_device);",
8346 )?;
8347 }
8348 if flags.contains(crate::Barrier::WORK_GROUP) {
8349 writeln!(
8350 self.out,
8351 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_threadgroup);",
8352 )?;
8353 if self.needs_object_memory_barriers {
8354 writeln!(
8355 self.out,
8356 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_object_data);",
8357 )?;
8358 }
8359 }
8360 if flags.contains(crate::Barrier::SUB_GROUP) {
8361 writeln!(
8362 self.out,
8363 "{level}{NAMESPACE}::simdgroup_barrier({NAMESPACE}::mem_flags::mem_threadgroup);",
8364 )?;
8365 }
8366 if flags.contains(crate::Barrier::TEXTURE) {
8367 writeln!(
8368 self.out,
8369 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_texture);",
8370 )?;
8371 }
8372 Ok(())
8373 }
8374}
8375
8376mod workgroup_mem_init {
8379 use crate::EntryPoint;
8380
8381 use super::*;
8382
8383 enum Access {
8384 GlobalVariable(Handle<crate::GlobalVariable>),
8385 StructMember(Handle<crate::Type>, u32),
8386 Array(usize),
8387 }
8388
8389 impl Access {
8390 fn write<W: Write>(
8391 &self,
8392 writer: &mut W,
8393 names: &FastHashMap<NameKey, String>,
8394 ) -> Result<(), core::fmt::Error> {
8395 match *self {
8396 Access::GlobalVariable(handle) => {
8397 write!(writer, "{}", &names[&NameKey::GlobalVariable(handle)])
8398 }
8399 Access::StructMember(handle, index) => {
8400 write!(writer, ".{}", &names[&NameKey::StructMember(handle, index)])
8401 }
8402 Access::Array(depth) => write!(writer, ".{WRAPPED_ARRAY_FIELD}[__i{depth}]"),
8403 }
8404 }
8405 }
8406
8407 struct AccessStack {
8408 stack: Vec<Access>,
8409 array_depth: usize,
8410 }
8411
8412 impl AccessStack {
8413 const fn new() -> Self {
8414 Self {
8415 stack: Vec::new(),
8416 array_depth: 0,
8417 }
8418 }
8419
8420 fn enter_array<R>(&mut self, cb: impl FnOnce(&mut Self, usize) -> R) -> R {
8421 let array_depth = self.array_depth;
8422 self.stack.push(Access::Array(array_depth));
8423 self.array_depth += 1;
8424 let res = cb(self, array_depth);
8425 self.stack.pop();
8426 self.array_depth -= 1;
8427 res
8428 }
8429
8430 fn enter<R>(&mut self, new: Access, cb: impl FnOnce(&mut Self) -> R) -> R {
8431 self.stack.push(new);
8432 let res = cb(self);
8433 self.stack.pop();
8434 res
8435 }
8436
8437 fn write<W: Write>(
8438 &self,
8439 writer: &mut W,
8440 names: &FastHashMap<NameKey, String>,
8441 ) -> Result<(), core::fmt::Error> {
8442 for next in self.stack.iter() {
8443 next.write(writer, names)?;
8444 }
8445 Ok(())
8446 }
8447 }
8448
8449 impl<W: Write> Writer<W> {
8450 pub(super) fn need_workgroup_variables_initialization(
8451 &mut self,
8452 options: &Options,
8453 ep: &EntryPoint,
8454 module: &crate::Module,
8455 fun_info: &valid::FunctionInfo,
8456 ) -> bool {
8457 let is_task = ep.stage == crate::ShaderStage::Task;
8458 options.zero_initialize_workgroup_memory
8459 && ep.stage.compute_like()
8460 && module.global_variables.iter().any(|(handle, var)| {
8461 let is_right_address_space = var.space == crate::AddressSpace::WorkGroup
8462 || (var.space == crate::AddressSpace::TaskPayload && is_task);
8463 !fun_info[handle].is_empty() && is_right_address_space
8464 })
8465 }
8466
8467 pub fn write_workgroup_variables_initialization(
8468 &mut self,
8469 module: &crate::Module,
8470 module_info: &valid::ModuleInfo,
8471 fun_info: &valid::FunctionInfo,
8472 local_invocation_index: Option<&NameKey>,
8473 stage: crate::ShaderStage,
8474 ) -> BackendResult {
8475 let level = back::Level(1);
8476
8477 writeln!(
8478 self.out,
8479 "{}if ({} == 0u) {{",
8480 level,
8481 local_invocation_index
8482 .map(|name_key| self.names[name_key].as_str())
8483 .unwrap_or("__local_invocation_index"),
8484 )?;
8485
8486 let mut access_stack = AccessStack::new();
8487
8488 let is_task = stage == crate::ShaderStage::Task;
8489 let vars = module.global_variables.iter().filter(|&(handle, var)| {
8490 let is_right_address_space = var.space == crate::AddressSpace::WorkGroup
8491 || (var.space == crate::AddressSpace::TaskPayload && is_task);
8492 !fun_info[handle].is_empty() && is_right_address_space
8493 });
8494
8495 for (handle, var) in vars {
8496 access_stack.enter(Access::GlobalVariable(handle), |access_stack| {
8497 self.write_workgroup_variable_initialization(
8498 module,
8499 module_info,
8500 var.ty,
8501 access_stack,
8502 level.next(),
8503 )
8504 })?;
8505 }
8506
8507 writeln!(self.out, "{level}}}")?;
8508 self.write_barrier(crate::Barrier::WORK_GROUP, level)
8509 }
8510
8511 fn write_workgroup_variable_initialization(
8512 &mut self,
8513 module: &crate::Module,
8514 module_info: &valid::ModuleInfo,
8515 ty: Handle<crate::Type>,
8516 access_stack: &mut AccessStack,
8517 level: back::Level,
8518 ) -> BackendResult {
8519 if module_info[ty].contains(valid::TypeFlags::CONSTRUCTIBLE) {
8520 write!(self.out, "{level}")?;
8521 access_stack.write(&mut self.out, &self.names)?;
8522 writeln!(self.out, " = {{}};")?;
8523 } else {
8524 match module.types[ty].inner {
8525 crate::TypeInner::Atomic { .. } => {
8526 write!(
8527 self.out,
8528 "{level}{NAMESPACE}::atomic_store_explicit({ATOMIC_REFERENCE}"
8529 )?;
8530 access_stack.write(&mut self.out, &self.names)?;
8531 writeln!(self.out, ", 0, {NAMESPACE}::memory_order_relaxed);")?;
8532 }
8533 crate::TypeInner::Array { base, size, .. } => {
8534 let count = match size.resolve(module.to_ctx())? {
8535 proc::IndexableLength::Known(count) => count,
8536 proc::IndexableLength::Dynamic => unreachable!(),
8537 };
8538
8539 access_stack.enter_array(|access_stack, array_depth| {
8540 writeln!(
8541 self.out,
8542 "{level}for (int __i{array_depth} = 0; __i{array_depth} < {count}; __i{array_depth}++) {{"
8543 )?;
8544 self.write_workgroup_variable_initialization(
8545 module,
8546 module_info,
8547 base,
8548 access_stack,
8549 level.next(),
8550 )?;
8551 writeln!(self.out, "{level}}}")?;
8552 BackendResult::Ok(())
8553 })?;
8554 }
8555 crate::TypeInner::Struct { ref members, .. } => {
8556 for (index, member) in members.iter().enumerate() {
8557 access_stack.enter(
8558 Access::StructMember(ty, index as u32),
8559 |access_stack| {
8560 self.write_workgroup_variable_initialization(
8561 module,
8562 module_info,
8563 member.ty,
8564 access_stack,
8565 level,
8566 )
8567 },
8568 )?;
8569 }
8570 }
8571 _ => unreachable!(),
8572 }
8573 }
8574
8575 Ok(())
8576 }
8577 }
8578}
8579
8580impl crate::AtomicFunction {
8581 const fn to_msl(self) -> &'static str {
8582 match self {
8583 Self::Add => "fetch_add",
8584 Self::Subtract => "fetch_sub",
8585 Self::And => "fetch_and",
8586 Self::InclusiveOr => "fetch_or",
8587 Self::ExclusiveOr => "fetch_xor",
8588 Self::Min => "fetch_min",
8589 Self::Max => "fetch_max",
8590 Self::Exchange { compare: None } => "exchange",
8591 Self::Exchange { compare: Some(_) } => ATOMIC_COMP_EXCH_FUNCTION,
8592 }
8593 }
8594
8595 fn to_msl_64_bit(self) -> Result<&'static str, Error> {
8596 Ok(match self {
8597 Self::Min => "min",
8598 Self::Max => "max",
8599 _ => Err(Error::FeatureNotImplemented(
8600 "64-bit atomic operation other than min/max".to_string(),
8601 ))?,
8602 })
8603 }
8604}