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,
557 pub(super) hoisted_loop_expressions: FastHashMap<Handle<crate::Expression>, String>,
573 need_bake_expressions: back::NeedBakeExpressions,
575 pub(super) namer: proc::Namer,
576 pub(super) wrapped_functions: FastHashSet<WrappedFunction>,
577 emit_int_div_checks: bool,
578 struct_member_pads: FastHashSet<(Handle<crate::Type>, u32)>,
581 needs_object_memory_barriers: bool,
582}
583
584impl crate::Scalar {
585 pub(super) fn to_msl_name(self) -> &'static str {
586 use crate::ScalarKind as Sk;
587 match self {
588 Self {
589 kind: Sk::Float,
590 width: 4,
591 } => "float",
592 Self {
593 kind: Sk::Float,
594 width: 2,
595 } => "half",
596 Self {
597 kind: Sk::Sint,
598 width: 2,
599 } => "short",
600 Self {
601 kind: Sk::Uint,
602 width: 2,
603 } => "ushort",
604 Self {
605 kind: Sk::Sint,
606 width: 4,
607 } => "int",
608 Self {
609 kind: Sk::Uint,
610 width: 4,
611 } => "uint",
612 Self {
613 kind: Sk::Sint,
614 width: 8,
615 } => "long",
616 Self {
617 kind: Sk::Uint,
618 width: 8,
619 } => "ulong",
620 Self {
621 kind: Sk::Bool,
622 width: _,
623 } => "bool",
624 Self {
625 kind: Sk::AbstractInt | Sk::AbstractFloat,
626 width: _,
627 } => unreachable!("Found Abstract scalar kind"),
628 _ => unreachable!("Unsupported scalar kind: {:?}", self),
629 }
630 }
631}
632
633const fn is_frag_depth(binding: Option<&crate::Binding>) -> bool {
635 matches!(
636 binding,
637 Some(crate::Binding::BuiltIn(crate::BuiltIn::FragDepth))
638 )
639}
640
641fn only_writes_frag_depth(module: &crate::ir::Module, entry_point: &crate::ir::EntryPoint) -> bool {
644 entry_point
645 .function
646 .result
647 .as_ref()
648 .is_some_and(|result| match module.types[result.ty].inner {
649 crate::TypeInner::Struct { ref members, .. } if result.binding.is_none() => members
650 .iter()
651 .all(|member| is_frag_depth(member.binding.as_ref())),
652 _ => is_frag_depth(result.binding.as_ref()),
653 })
654}
655
656fn discards_frag_depth(entry_point: &crate::ir::EntryPoint) -> bool {
662 entry_point.stage == crate::ShaderStage::Fragment
663 && matches!(
664 entry_point.early_depth_test,
665 Some(
666 crate::EarlyDepthTest::Force
667 | crate::EarlyDepthTest::Allow {
668 conservative: crate::ConservativeDepth::Unchanged
669 }
670 )
671 )
672}
673
674fn only_writes_discarded_depth(
676 module: &crate::ir::Module,
677 entry_point: &crate::ir::EntryPoint,
678) -> bool {
679 discards_frag_depth(entry_point) && only_writes_frag_depth(module, entry_point)
680}
681
682const fn separate(need_separator: bool) -> &'static str {
683 if need_separator {
684 ","
685 } else {
686 ""
687 }
688}
689
690fn should_pack_struct_member(
691 members: &[crate::StructMember],
692 span: u32,
693 index: usize,
694 module: &crate::Module,
695) -> Option<crate::Scalar> {
696 let member = &members[index];
697
698 let ty_inner = &module.types[member.ty].inner;
699 let last_offset = member.offset + ty_inner.size(module.to_ctx());
700 let next_offset = match members.get(index + 1) {
701 Some(next) => next.offset,
702 None => span,
703 };
704 let is_tight = next_offset == last_offset;
705
706 match *ty_inner {
707 crate::TypeInner::Vector {
708 size: crate::VectorSize::Tri,
709 scalar: scalar @ crate::Scalar { width: 4 | 2, .. },
710 } if is_tight => Some(scalar),
711 _ => None,
712 }
713}
714
715impl crate::AddressSpace {
716 const fn needs_pass_through(&self) -> bool {
720 match *self {
721 Self::Uniform
722 | Self::Storage { .. }
723 | Self::Private
724 | Self::WorkGroup
725 | Self::Immediate
726 | Self::Handle
727 | Self::TaskPayload => true,
728 Self::Function => false,
729 Self::RayPayload | Self::IncomingRayPayload => unreachable!(),
730 }
731 }
732
733 const fn needs_access_qualifier(&self) -> bool {
735 match *self {
736 Self::Storage { .. } => true,
741 Self::TaskPayload => true,
742 Self::RayPayload | Self::IncomingRayPayload => unimplemented!(),
743 Self::Private | Self::WorkGroup => false,
745 Self::Uniform | Self::Immediate => false,
747 Self::Handle | Self::Function => false,
749 }
750 }
751
752 const fn to_msl_name(self) -> Option<&'static str> {
753 match self {
754 Self::Handle => None,
755 Self::Uniform | Self::Immediate => Some("constant"),
756 Self::Storage { .. } => Some("device"),
757 Self::Private | Self::Function | Self::RayPayload => Some("thread"),
761 Self::WorkGroup => Some("threadgroup"),
762 Self::TaskPayload => Some("object_data"),
763 Self::IncomingRayPayload => Some("ray_data"),
764 }
765 }
766}
767
768impl crate::Type {
769 const fn needs_alias(&self) -> bool {
771 use crate::TypeInner as Ti;
772
773 match self.inner {
774 Ti::Scalar(_)
776 | Ti::Vector { .. }
777 | Ti::Matrix { .. }
778 | Ti::CooperativeMatrix { .. }
779 | Ti::Atomic(_)
780 | Ti::Pointer { .. }
781 | Ti::ValuePointer { .. } => self.name.is_some(),
782 Ti::Struct { .. } | Ti::Array { .. } => true,
784 Ti::Image { .. }
786 | Ti::Sampler { .. }
787 | Ti::AccelerationStructure { .. }
788 | Ti::RayQuery { .. }
789 | Ti::BindingArray { .. } => false,
790 }
791 }
792}
793
794#[derive(Clone, Copy)]
795pub(super) enum FunctionOrigin {
796 Handle(Handle<crate::Function>),
797 EntryPoint(proc::EntryPointIndex),
798}
799
800pub(super) trait NameKeyExt {
801 fn local(origin: FunctionOrigin, local_handle: Handle<crate::LocalVariable>) -> NameKey {
802 match origin {
803 FunctionOrigin::Handle(handle) => NameKey::FunctionLocal(handle, local_handle),
804 FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointLocal(idx, local_handle),
805 }
806 }
807
808 fn oob_local_for_type(origin: FunctionOrigin, ty: Handle<crate::Type>) -> NameKey {
813 match origin {
814 FunctionOrigin::Handle(handle) => NameKey::FunctionOobLocal(handle, ty),
815 FunctionOrigin::EntryPoint(idx) => NameKey::EntryPointOobLocal(idx, ty),
816 }
817 }
818}
819
820impl NameKeyExt for NameKey {}
821
822#[derive(Clone, Copy)]
832enum LevelOfDetail {
833 Direct(Handle<crate::Expression>),
834 Restricted(Handle<crate::Expression>),
835}
836
837struct TexelAddress {
847 coordinate: Handle<crate::Expression>,
848 array_index: Option<Handle<crate::Expression>>,
849 sample: Option<Handle<crate::Expression>>,
850 level: Option<LevelOfDetail>,
851}
852
853pub(super) struct ExpressionContext<'a> {
854 pub(super) function: &'a crate::Function,
855 pub(super) origin: FunctionOrigin,
856 pub(super) info: &'a valid::FunctionInfo,
857 pub(super) module: &'a crate::Module,
858 pub(super) mod_info: &'a valid::ModuleInfo,
859 pub(super) pipeline_options: &'a PipelineOptions,
860 pub(super) lang_version: (u8, u8),
861 pub(super) policies: index::BoundsCheckPolicies,
862
863 pub(super) guarded_indices: HandleSet<crate::Expression>,
867 pub(super) force_loop_bounding: bool,
869 emit_int_div_checks: bool,
871 pub(super) ray_query_initialization_tracking: bool,
872}
873
874impl<'a> ExpressionContext<'a> {
875 fn resolve_type(&self, handle: Handle<crate::Expression>) -> &'a crate::TypeInner {
876 self.info[handle].ty.inner_with(&self.module.types)
877 }
878
879 fn binding_array_index_from_chain(
882 &self,
883 mut expr: Handle<crate::Expression>,
884 global: Handle<crate::GlobalVariable>,
885 ) -> Option<index::GuardedIndex> {
886 let expressions = &self.function.expressions;
887 loop {
888 match expressions[expr] {
889 crate::Expression::Load { pointer } => expr = pointer,
890 crate::Expression::Access { base, index } => {
891 if matches!(
892 expressions[base],
893 crate::Expression::GlobalVariable(g) if g == global
894 ) {
895 return Some(index::GuardedIndex::Expression(index));
896 }
897 expr = base;
898 }
899 crate::Expression::AccessIndex { base, index } => {
900 if matches!(
901 expressions[base],
902 crate::Expression::GlobalVariable(g) if g == global
903 ) {
904 return Some(index::GuardedIndex::Known(index));
905 }
906 expr = base;
907 }
908 crate::Expression::GlobalVariable(_) => return None,
909 _ => return None,
910 }
911 }
912 }
913
914 fn is_global_access_chain(&self, expr: Handle<crate::Expression>) -> bool {
916 let expressions = &self.function.expressions;
917 match expressions[expr] {
918 crate::Expression::Access { base, .. } => match expressions[base] {
919 crate::Expression::GlobalVariable(_) => true,
920 crate::Expression::Access { .. } => self.is_global_access_chain(base),
921 _ => false,
922 },
923 crate::Expression::AccessIndex { base, .. } => {
924 matches!(expressions[base], crate::Expression::GlobalVariable(_))
925 }
926 _ => false,
927 }
928 }
929
930 fn struct_member_needs_arrow(
931 &self,
932 base: Handle<crate::Expression>,
933 originating_global_ty: impl FnOnce(&crate::TypeInner) -> bool,
934 ) -> bool {
935 let originating_matches = match self.function.originating_global(base) {
936 Some(gv) => {
937 originating_global_ty(&self.module.types[self.module.global_variables[gv].ty].inner)
938 }
939 None => false,
940 };
941 originating_matches && self.is_global_access_chain(base)
942 }
943
944 fn image_needs_lod(&self, image: Handle<crate::Expression>) -> bool {
951 let image_ty = self.resolve_type(image);
952 if let crate::TypeInner::Image { dim, class, .. } = *image_ty {
953 class.is_mipmapped() && dim != crate::ImageDimension::D1
954 } else {
955 false
956 }
957 }
958
959 fn choose_bounds_check_policy(
960 &self,
961 pointer: Handle<crate::Expression>,
962 ) -> index::BoundsCheckPolicy {
963 self.policies
964 .choose_policy(pointer, &self.module.types, self.info)
965 }
966
967 fn access_needs_check(
969 &self,
970 base: Handle<crate::Expression>,
971 index: index::GuardedIndex,
972 ) -> Option<index::IndexableLength> {
973 index::access_needs_check(
974 base,
975 index,
976 self.module,
977 &self.function.expressions,
978 self.info,
979 )
980 }
981
982 fn bounds_check_iter(
984 &self,
985 chain: Handle<crate::Expression>,
986 ) -> impl Iterator<Item = BoundsCheck> + '_ {
987 index::bounds_check_iter(chain, self.module, self.function, self.info)
988 }
989
990 fn oob_local_types(&self) -> FastHashSet<Handle<crate::Type>> {
992 index::oob_local_types(self.module, self.function, self.info, self.policies)
993 }
994
995 fn get_packed_vec_kind(&self, expr_handle: Handle<crate::Expression>) -> Option<crate::Scalar> {
996 match self.function.expressions[expr_handle] {
997 crate::Expression::AccessIndex { base, index } => {
998 let ty = match *self.resolve_type(base) {
999 crate::TypeInner::Pointer { base, .. } => &self.module.types[base].inner,
1000 ref ty => ty,
1001 };
1002 match *ty {
1003 crate::TypeInner::Struct {
1004 ref members, span, ..
1005 } => should_pack_struct_member(members, span, index as usize, self.module),
1006 _ => None,
1007 }
1008 }
1009 _ => None,
1010 }
1011 }
1012}
1013
1014pub(super) struct StatementContext<'a> {
1015 pub(super) expression: ExpressionContext<'a>,
1016 pub(super) result_struct: Option<&'a str>,
1017}
1018
1019impl<W: Write> Writer<W> {
1020 pub fn new(out: W) -> Self {
1022 Writer {
1023 out,
1024 names: FastHashMap::default(),
1025 named_expressions: Default::default(),
1026 hoisted_loop_expressions: FastHashMap::default(),
1027 need_bake_expressions: Default::default(),
1028 namer: proc::Namer::default(),
1029 wrapped_functions: FastHashSet::default(),
1030 emit_int_div_checks: true,
1031 struct_member_pads: FastHashSet::default(),
1032 needs_object_memory_barriers: false,
1033 }
1034 }
1035
1036 pub fn finish(self) -> W {
1039 self.out
1040 }
1041
1042 fn gen_force_bounded_loop_statements(
1149 &mut self,
1150 level: back::Level,
1151 context: &StatementContext,
1152 ) -> Option<(String, String)> {
1153 if !context.expression.force_loop_bounding {
1154 return None;
1155 }
1156
1157 let loop_bound_name = self.namer.call("loop_bound");
1158 let decl = format!("{level}uint2 {loop_bound_name} = uint2({}u);", u32::MAX);
1161 let level = level.next();
1162 let break_and_inc = format!(
1163 "{level}if ({NAMESPACE}::all({loop_bound_name} == uint2(0u))) {{ break; }}
1164{level}{loop_bound_name} -= uint2({loop_bound_name}.y == 0u, 1u);"
1165 );
1166
1167 Some((decl, break_and_inc))
1168 }
1169
1170 fn put_call_parameters(
1171 &mut self,
1172 parameters: impl Iterator<Item = Handle<crate::Expression>>,
1173 context: &ExpressionContext,
1174 ) -> BackendResult {
1175 self.put_call_parameters_impl(parameters, context, |writer, context, expr| {
1176 writer.put_expression(expr, context, true)
1177 })
1178 }
1179
1180 fn put_call_parameters_impl<C, E>(
1181 &mut self,
1182 parameters: impl Iterator<Item = Handle<crate::Expression>>,
1183 ctx: &C,
1184 put_expression: E,
1185 ) -> BackendResult
1186 where
1187 E: Fn(&mut Self, &C, Handle<crate::Expression>) -> BackendResult,
1188 {
1189 write!(self.out, "(")?;
1190 for (i, handle) in parameters.enumerate() {
1191 if i != 0 {
1192 write!(self.out, ", ")?;
1193 }
1194 put_expression(self, ctx, handle)?;
1195 }
1196 write!(self.out, ")")?;
1197 Ok(())
1198 }
1199
1200 fn put_locals(&mut self, context: &ExpressionContext) -> BackendResult {
1206 let oob_local_types = context.oob_local_types();
1207 for &ty in oob_local_types.iter() {
1208 let name_key = NameKey::oob_local_for_type(context.origin, ty);
1209 self.names.insert(name_key, self.namer.call("oob"));
1210 }
1211
1212 for (name_key, ty, init) in context
1213 .function
1214 .local_variables
1215 .iter()
1216 .map(|(local_handle, local)| {
1217 let name_key = NameKey::local(context.origin, local_handle);
1218 (name_key, local.ty, local.init)
1219 })
1220 .chain(oob_local_types.iter().map(|&ty| {
1221 let name_key = NameKey::oob_local_for_type(context.origin, ty);
1222 (name_key, ty, None)
1223 }))
1224 {
1225 let ty_name = TypeContext {
1226 handle: ty,
1227 gctx: context.module.to_ctx(),
1228 names: &self.names,
1229 access: crate::StorageAccess::empty(),
1230 first_time: false,
1231 };
1232 write!(
1233 self.out,
1234 "{}{} {}",
1235 back::INDENT,
1236 ty_name,
1237 self.names[&name_key]
1238 )?;
1239 match init {
1240 Some(value) => {
1241 write!(self.out, " = ")?;
1242 self.put_expression(value, context, true)?;
1243 }
1244 None => {
1245 write!(self.out, " = {{}}")?;
1246 }
1247 };
1248 writeln!(self.out, ";")?;
1249
1250 if context.ray_query_initialization_tracking {
1252 if let crate::TypeInner::RayQuery { .. } = context.module.types[ty].inner {
1253 writeln!(
1254 self.out,
1255 "{}uint {}{} = 0u;",
1256 back::INDENT,
1257 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
1258 self.names[&name_key]
1259 )?;
1260
1261 writeln!(
1262 self.out,
1263 "{}float {}{} = 0.0;",
1264 back::INDENT,
1265 super::ray::RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX,
1266 self.names[&name_key]
1267 )?;
1268 }
1269 }
1270 }
1271 Ok(())
1272 }
1273
1274 fn put_level_of_detail(
1275 &mut self,
1276 level: LevelOfDetail,
1277 context: &ExpressionContext,
1278 ) -> BackendResult {
1279 match level {
1280 LevelOfDetail::Direct(expr) => self.put_expression(expr, context, true)?,
1281 LevelOfDetail::Restricted(load) => write!(self.out, "{}", ClampedLod(load))?,
1282 }
1283 Ok(())
1284 }
1285
1286 fn put_image_query(
1287 &mut self,
1288 image: Handle<crate::Expression>,
1289 query: &str,
1290 level: Option<LevelOfDetail>,
1291 context: &ExpressionContext,
1292 ) -> BackendResult {
1293 self.put_expression(image, context, false)?;
1294 write!(self.out, ".get_{query}(")?;
1295 if let Some(level) = level {
1296 self.put_level_of_detail(level, context)?;
1297 }
1298 write!(self.out, ")")?;
1299 Ok(())
1300 }
1301
1302 fn put_image_size_query(
1303 &mut self,
1304 image: Handle<crate::Expression>,
1305 level: Option<LevelOfDetail>,
1306 kind: crate::ScalarKind,
1307 context: &ExpressionContext,
1308 ) -> BackendResult {
1309 if context.resolve_type(image).is_external_image() {
1310 write!(self.out, "{IMAGE_SIZE_EXTERNAL_FUNCTION}(")?;
1311 self.put_expression(image, context, true)?;
1312 write!(self.out, ")")?;
1313 return Ok(());
1314 }
1315
1316 let dim = match *context.resolve_type(image) {
1319 crate::TypeInner::Image { dim, .. } => dim,
1320 ref other => unreachable!("Unexpected type {:?}", other),
1321 };
1322 let scalar = crate::Scalar { kind, width: 4 };
1323 let coordinate_type = scalar.to_msl_name();
1324 match dim {
1325 crate::ImageDimension::D1 => {
1326 if kind == crate::ScalarKind::Uint {
1330 self.put_image_query(image, "width", None, context)?;
1332 } else {
1333 write!(self.out, "int(")?;
1335 self.put_image_query(image, "width", None, context)?;
1336 write!(self.out, ")")?;
1337 }
1338 }
1339 crate::ImageDimension::D2 => {
1340 write!(self.out, "{NAMESPACE}::{coordinate_type}2(")?;
1341 self.put_image_query(image, "width", level, context)?;
1342 write!(self.out, ", ")?;
1343 self.put_image_query(image, "height", level, context)?;
1344 write!(self.out, ")")?;
1345 }
1346 crate::ImageDimension::D3 => {
1347 write!(self.out, "{NAMESPACE}::{coordinate_type}3(")?;
1348 self.put_image_query(image, "width", level, context)?;
1349 write!(self.out, ", ")?;
1350 self.put_image_query(image, "height", level, context)?;
1351 write!(self.out, ", ")?;
1352 self.put_image_query(image, "depth", level, context)?;
1353 write!(self.out, ")")?;
1354 }
1355 crate::ImageDimension::Cube => {
1356 write!(self.out, "{NAMESPACE}::{coordinate_type}2(")?;
1357 self.put_image_query(image, "width", level, context)?;
1358 write!(self.out, ")")?;
1359 }
1360 }
1361 Ok(())
1362 }
1363
1364 fn put_cast_to_uint_scalar_or_vector(
1365 &mut self,
1366 expr: Handle<crate::Expression>,
1367 context: &ExpressionContext,
1368 ) -> BackendResult {
1369 match *context.resolve_type(expr) {
1371 crate::TypeInner::Scalar(_) => {
1372 put_numeric_type(&mut self.out, crate::Scalar::U32, &[])?
1373 }
1374 crate::TypeInner::Vector { size, .. } => {
1375 put_numeric_type(&mut self.out, crate::Scalar::U32, &[size])?
1376 }
1377 _ => {
1378 return Err(Error::GenericValidation(
1379 "Invalid type for image coordinate".into(),
1380 ))
1381 }
1382 };
1383
1384 write!(self.out, "(")?;
1385 self.put_expression(expr, context, true)?;
1386 write!(self.out, ")")?;
1387 Ok(())
1388 }
1389
1390 fn put_image_sample_level(
1391 &mut self,
1392 image: Handle<crate::Expression>,
1393 level: crate::SampleLevel,
1394 context: &ExpressionContext,
1395 ) -> BackendResult {
1396 let has_levels = context.image_needs_lod(image);
1397 match level {
1398 crate::SampleLevel::Auto => {}
1399 crate::SampleLevel::Zero => {
1400 }
1402 _ if !has_levels => {
1403 log::warn!("1D image can't be sampled with level {level:?}");
1404 }
1405 crate::SampleLevel::Exact(h) => {
1406 write!(self.out, ", {NAMESPACE}::level(")?;
1407 self.put_expression(h, context, true)?;
1408 write!(self.out, ")")?;
1409 }
1410 crate::SampleLevel::Bias(h) => {
1411 write!(self.out, ", {NAMESPACE}::bias(")?;
1412 self.put_expression(h, context, true)?;
1413 write!(self.out, ")")?;
1414 }
1415 crate::SampleLevel::Gradient { x, y } => {
1416 let dim = match *context.resolve_type(image) {
1417 crate::TypeInner::Image { dim, .. } => dim,
1418 ref other => unreachable!("Unexpected type {other:?}"),
1419 };
1420 let gradient = match dim {
1421 crate::ImageDimension::D2 => "gradient2d",
1422 crate::ImageDimension::D3 => "gradient3d",
1423 crate::ImageDimension::Cube => "gradientcube",
1424 crate::ImageDimension::D1 => {
1427 unreachable!("1D image can't be sampled with a gradient")
1428 }
1429 };
1430 write!(self.out, ", {NAMESPACE}::{gradient}(")?;
1431 self.put_expression(x, context, true)?;
1432 write!(self.out, ", ")?;
1433 self.put_expression(y, context, true)?;
1434 write!(self.out, ")")?;
1435 }
1436 }
1437 Ok(())
1438 }
1439
1440 fn put_image_coordinate_limits(
1441 &mut self,
1442 image: Handle<crate::Expression>,
1443 level: Option<LevelOfDetail>,
1444 context: &ExpressionContext,
1445 ) -> BackendResult {
1446 self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
1447 write!(self.out, " - 1")?;
1448 Ok(())
1449 }
1450
1451 fn put_restricted_scalar_image_index(
1469 &mut self,
1470 image: Handle<crate::Expression>,
1471 index: Handle<crate::Expression>,
1472 limit_method: &str,
1473 context: &ExpressionContext,
1474 ) -> BackendResult {
1475 write!(self.out, "{NAMESPACE}::min(uint(")?;
1476 self.put_expression(index, context, true)?;
1477 write!(self.out, "), ")?;
1478 self.put_expression(image, context, false)?;
1479 write!(self.out, ".{limit_method}() - 1)")?;
1480 Ok(())
1481 }
1482
1483 fn put_restricted_texel_address(
1484 &mut self,
1485 image: Handle<crate::Expression>,
1486 address: &TexelAddress,
1487 context: &ExpressionContext,
1488 ) -> BackendResult {
1489 write!(self.out, "{NAMESPACE}::min(")?;
1491 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1492 write!(self.out, ", ")?;
1493 self.put_image_coordinate_limits(image, address.level, context)?;
1494 write!(self.out, ")")?;
1495
1496 if let Some(array_index) = address.array_index {
1498 write!(self.out, ", ")?;
1499 self.put_restricted_scalar_image_index(image, array_index, "get_array_size", context)?;
1500 }
1501
1502 if let Some(sample) = address.sample {
1504 write!(self.out, ", ")?;
1505 self.put_restricted_scalar_image_index(image, sample, "get_num_samples", context)?;
1506 }
1507
1508 if let Some(level) = address.level {
1511 write!(self.out, ", ")?;
1512 self.put_level_of_detail(level, context)?;
1513 }
1514
1515 Ok(())
1516 }
1517
1518 fn put_image_access_bounds_check(
1520 &mut self,
1521 image: Handle<crate::Expression>,
1522 address: &TexelAddress,
1523 context: &ExpressionContext,
1524 ) -> BackendResult {
1525 let mut conjunction = "";
1526
1527 let level = if let Some(level) = address.level {
1530 write!(self.out, "uint(")?;
1531 self.put_level_of_detail(level, context)?;
1532 write!(self.out, ") < ")?;
1533 self.put_expression(image, context, true)?;
1534 write!(self.out, ".get_num_mip_levels()")?;
1535 conjunction = " && ";
1536 Some(level)
1537 } else {
1538 None
1539 };
1540
1541 if let Some(sample) = address.sample {
1543 write!(self.out, "uint(")?;
1544 self.put_expression(sample, context, true)?;
1545 write!(self.out, ") < ")?;
1546 self.put_expression(image, context, true)?;
1547 write!(self.out, ".get_num_samples()")?;
1548 conjunction = " && ";
1549 }
1550
1551 if let Some(array_index) = address.array_index {
1553 write!(self.out, "{conjunction}uint(")?;
1554 self.put_expression(array_index, context, true)?;
1555 write!(self.out, ") < ")?;
1556 self.put_expression(image, context, true)?;
1557 write!(self.out, ".get_array_size()")?;
1558 conjunction = " && ";
1559 }
1560
1561 let coord_is_vector = match *context.resolve_type(address.coordinate) {
1563 crate::TypeInner::Vector { .. } => true,
1564 _ => false,
1565 };
1566 write!(self.out, "{conjunction}")?;
1567 if coord_is_vector {
1568 write!(self.out, "{NAMESPACE}::all(")?;
1569 }
1570 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1571 write!(self.out, " < ")?;
1572 self.put_image_size_query(image, level, crate::ScalarKind::Uint, context)?;
1573 if coord_is_vector {
1574 write!(self.out, ")")?;
1575 }
1576
1577 Ok(())
1578 }
1579
1580 fn put_image_load(
1581 &mut self,
1582 load: Handle<crate::Expression>,
1583 image: Handle<crate::Expression>,
1584 mut address: TexelAddress,
1585 context: &ExpressionContext,
1586 ) -> BackendResult {
1587 if context.resolve_type(image).is_external_image() {
1588 write!(self.out, "{IMAGE_LOAD_EXTERNAL_FUNCTION}(")?;
1589 self.put_expression(image, context, true)?;
1590 write!(self.out, ", ")?;
1591 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1592 write!(self.out, ")")?;
1593 return Ok(());
1594 }
1595
1596 match context.policies.image_load {
1597 proc::BoundsCheckPolicy::Restrict => {
1598 if address.level.is_some() {
1601 address.level = if context.image_needs_lod(image) {
1602 Some(LevelOfDetail::Restricted(load))
1603 } else {
1604 None
1605 }
1606 }
1607
1608 self.put_expression(image, context, false)?;
1609 write!(self.out, ".read(")?;
1610 self.put_restricted_texel_address(image, &address, context)?;
1611 write!(self.out, ")")?;
1612 }
1613 proc::BoundsCheckPolicy::ReadZeroSkipWrite => {
1614 write!(self.out, "(")?;
1615 self.put_image_access_bounds_check(image, &address, context)?;
1616 write!(self.out, " ? ")?;
1617 self.put_unchecked_image_load(image, &address, context)?;
1618 write!(self.out, ": DefaultConstructible())")?;
1619 }
1620 proc::BoundsCheckPolicy::Unchecked => {
1621 self.put_unchecked_image_load(image, &address, context)?;
1622 }
1623 }
1624
1625 Ok(())
1626 }
1627
1628 fn put_unchecked_image_load(
1629 &mut self,
1630 image: Handle<crate::Expression>,
1631 address: &TexelAddress,
1632 context: &ExpressionContext,
1633 ) -> BackendResult {
1634 self.put_expression(image, context, false)?;
1635 write!(self.out, ".read(")?;
1636 self.put_cast_to_uint_scalar_or_vector(address.coordinate, context)?;
1638 if let Some(expr) = address.array_index {
1639 write!(self.out, ", ")?;
1640 self.put_expression(expr, context, true)?;
1641 }
1642 if let Some(sample) = address.sample {
1643 write!(self.out, ", ")?;
1644 self.put_expression(sample, context, true)?;
1645 }
1646 if let Some(level) = address.level {
1647 if context.image_needs_lod(image) {
1648 write!(self.out, ", ")?;
1649 self.put_level_of_detail(level, context)?;
1650 }
1651 }
1652 write!(self.out, ")")?;
1653
1654 Ok(())
1655 }
1656
1657 fn put_image_atomic(
1658 &mut self,
1659 level: back::Level,
1660 image: Handle<crate::Expression>,
1661 address: &TexelAddress,
1662 fun: crate::AtomicFunction,
1663 value: Handle<crate::Expression>,
1664 context: &StatementContext,
1665 ) -> BackendResult {
1666 write!(self.out, "{level}")?;
1667 self.put_expression(image, &context.expression, false)?;
1668 let op = if context.expression.resolve_type(value).scalar_width() == Some(8) {
1669 fun.to_msl_64_bit()?
1670 } else {
1671 fun.to_msl()
1672 };
1673 write!(self.out, ".atomic_{op}(")?;
1674 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1676 write!(self.out, ", ")?;
1677 self.put_expression(value, &context.expression, true)?;
1678 writeln!(self.out, ");")?;
1679
1680 let value_ty = context.expression.resolve_type(value);
1686 let zero_value = match (value_ty.scalar_kind(), value_ty.scalar_width()) {
1687 (Some(crate::ScalarKind::Sint), _) => "int4(0)",
1688 (_, Some(8)) => "ulong4(0uL)",
1689 _ => "uint4(0u)",
1690 };
1691 let coord_ty = context.expression.resolve_type(address.coordinate);
1692 let x = if matches!(coord_ty, crate::TypeInner::Scalar(_)) {
1693 ""
1694 } else {
1695 ".x"
1696 };
1697 write!(self.out, "{level}if (")?;
1698 self.put_expression(address.coordinate, &context.expression, true)?;
1699 write!(self.out, "{x} == -99999) {{ ")?;
1700 self.put_expression(image, &context.expression, false)?;
1701 write!(self.out, ".write({zero_value}, ")?;
1702 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1703 if let Some(array_index) = address.array_index {
1704 write!(self.out, ", ")?;
1705 self.put_expression(array_index, &context.expression, true)?;
1706 }
1707 writeln!(self.out, "); }}")?;
1708
1709 Ok(())
1710 }
1711
1712 fn put_image_store(
1713 &mut self,
1714 level: back::Level,
1715 image: Handle<crate::Expression>,
1716 address: &TexelAddress,
1717 value: Handle<crate::Expression>,
1718 context: &StatementContext,
1719 ) -> BackendResult {
1720 write!(self.out, "{level}")?;
1721 self.put_expression(image, &context.expression, false)?;
1722 write!(self.out, ".write(")?;
1723 self.put_expression(value, &context.expression, true)?;
1724 write!(self.out, ", ")?;
1725 self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;
1727 if let Some(expr) = address.array_index {
1728 write!(self.out, ", ")?;
1729 self.put_expression(expr, &context.expression, true)?;
1730 }
1731 writeln!(self.out, ");")?;
1732
1733 Ok(())
1734 }
1735
1736 fn binding_array_layout_count(
1751 module: &crate::Module,
1752 pipeline_options: &PipelineOptions,
1753 global: Handle<crate::GlobalVariable>,
1754 ) -> u32 {
1755 let var = &module.global_variables[global];
1756 let crate::TypeInner::BindingArray { size, .. } = module.types[var.ty].inner else {
1757 unreachable!("binding_array_layout_count called on non-binding-array global");
1758 };
1759 let from_shader = match size {
1760 crate::ArraySize::Constant(n) => n.get(),
1761 crate::ArraySize::Pending(_) | crate::ArraySize::Dynamic => 0,
1762 };
1763 let from_layout = var
1764 .binding
1765 .and_then(|br| pipeline_options.binding_array_length_map.get(&br))
1766 .copied()
1767 .unwrap_or(0);
1768 from_shader.max(from_layout).max(1)
1769 }
1770
1771 fn put_binding_array_size_member_index(
1772 &mut self,
1773 index: index::GuardedIndex,
1774 context: &ExpressionContext,
1775 ) -> BackendResult {
1776 match index {
1777 index::GuardedIndex::Expression(expr) => {
1778 write!(self.out, "unsigned(")?;
1779 self.put_expression(expr, context, true)?;
1780 write!(self.out, ")")?;
1781 }
1782 index::GuardedIndex::Known(value) => write!(self.out, "{value}u")?,
1783 }
1784 Ok(())
1785 }
1786
1787 fn put_dynamic_array_max_index(
1788 &mut self,
1789 handle: Handle<crate::GlobalVariable>,
1790 chain_expr: Handle<crate::Expression>,
1791 context: &ExpressionContext,
1792 ) -> BackendResult {
1793 let global = &context.module.global_variables[handle];
1794 let (offset, array_ty) = match context.module.types[global.ty].inner {
1795 crate::TypeInner::Struct { ref members, .. } => match members.last() {
1796 Some(&crate::StructMember { offset, ty, .. }) => (offset, ty),
1797 None => return Err(Error::GenericValidation("Struct has no members".into())),
1798 },
1799 crate::TypeInner::BindingArray { base, .. } => match context.module.types[base].inner {
1800 crate::TypeInner::Struct { ref members, .. } => match members.last() {
1801 Some(&crate::StructMember { offset, ty, .. }) => (offset, ty),
1802 None => return Err(Error::GenericValidation("Struct has no members".into())),
1803 },
1804 _ => {
1805 return Err(Error::GenericValidation(
1806 "binding_array element must be a struct with a runtime-sized array field"
1807 .into(),
1808 ))
1809 }
1810 },
1811 crate::TypeInner::Array {
1812 size: crate::ArraySize::Dynamic,
1813 ..
1814 } => (0, global.ty),
1815 ref ty => {
1816 return Err(Error::GenericValidation(format!(
1817 "Expected type with dynamic array, got {ty:?}"
1818 )))
1819 }
1820 };
1821
1822 let (size, stride) = match context.module.types[array_ty].inner {
1823 crate::TypeInner::Array { base, stride, .. } => (
1824 context.module.types[base]
1825 .inner
1826 .size(context.module.to_ctx()),
1827 stride,
1828 ),
1829 ref ty => {
1830 return Err(Error::GenericValidation(format!(
1831 "Expected array type, got {ty:?}"
1832 )))
1833 }
1834 };
1835
1836 write!(
1849 self.out,
1850 "(_buffer_sizes.{member}",
1851 member = ArraySizeMember(handle),
1852 )?;
1853 if let crate::TypeInner::BindingArray { .. } = context.module.types[global.ty].inner {
1854 let Some(array_index) = context.binding_array_index_from_chain(chain_expr, handle)
1855 else {
1856 return Err(Error::GenericValidation(
1857 "Could not find binding_array index for buffer size".into(),
1858 ));
1859 };
1860 write!(self.out, "[")?;
1861 match array_index {
1862 index::GuardedIndex::Expression(expr) => {
1863 write!(self.out, "unsigned(")?;
1864 self.put_expression(expr, context, true)?;
1865 write!(self.out, ")")?;
1866 }
1867 index::GuardedIndex::Known(i) => {
1868 write!(self.out, "{i}u")?;
1869 }
1870 }
1871 write!(self.out, "]")?;
1872 }
1873 write!(
1874 self.out,
1875 " - {offset} - {size}) / {stride}",
1876 offset = offset,
1877 size = size,
1878 stride = stride,
1879 )?;
1880 Ok(())
1881 }
1882
1883 fn put_dot_product<T: Copy>(
1888 &mut self,
1889 arg: T,
1890 arg1: T,
1891 size: usize,
1892 extractor: impl Fn(&mut Self, T, usize) -> BackendResult,
1893 ) -> BackendResult {
1894 write!(self.out, "(")?;
1897
1898 for index in 0..size {
1900 write!(self.out, " + ")?;
1903 extractor(self, arg, index)?;
1904 write!(self.out, " * ")?;
1905 extractor(self, arg1, index)?;
1906 }
1907
1908 write!(self.out, ")")?;
1909 Ok(())
1910 }
1911
1912 fn put_pack4x8(
1914 &mut self,
1915 arg: Handle<crate::Expression>,
1916 context: &ExpressionContext<'_>,
1917 was_signed: bool,
1918 clamp_bounds: Option<(&str, &str)>,
1919 ) -> Result<(), Error> {
1920 let write_arg = |this: &mut Self| -> BackendResult {
1921 if let Some((min, max)) = clamp_bounds {
1922 write!(this.out, "{NAMESPACE}::clamp(")?;
1924 this.put_expression(arg, context, true)?;
1925 write!(this.out, ", {min}, {max})")?;
1926 } else {
1927 this.put_expression(arg, context, true)?;
1928 }
1929 Ok(())
1930 };
1931
1932 if context.lang_version >= (2, 1) {
1933 let packed_type = if was_signed {
1934 "packed_char4"
1935 } else {
1936 "packed_uchar4"
1937 };
1938 write!(self.out, "as_type<uint>({packed_type}(")?;
1940 write_arg(self)?;
1941 write!(self.out, "))")?;
1942 } else {
1943 if was_signed {
1945 write!(self.out, "uint(")?;
1946 }
1947 write!(self.out, "(")?;
1948 write_arg(self)?;
1949 write!(self.out, "[0] & 0xFF) | ((")?;
1950 write_arg(self)?;
1951 write!(self.out, "[1] & 0xFF) << 8) | ((")?;
1952 write_arg(self)?;
1953 write!(self.out, "[2] & 0xFF) << 16) | ((")?;
1954 write_arg(self)?;
1955 write!(self.out, "[3] & 0xFF) << 24)")?;
1956 if was_signed {
1957 write!(self.out, ")")?;
1958 }
1959 }
1960
1961 Ok(())
1962 }
1963
1964 fn put_isign(
1967 &mut self,
1968 arg: Handle<crate::Expression>,
1969 context: &ExpressionContext,
1970 ) -> BackendResult {
1971 write!(self.out, "{NAMESPACE}::select({NAMESPACE}::select(")?;
1972 let scalar = context
1973 .resolve_type(arg)
1974 .scalar()
1975 .expect("put_isign should only be called for args which have an integer scalar type")
1976 .to_msl_name();
1977 match context.resolve_type(arg) {
1978 &crate::TypeInner::Vector { size, .. } => {
1979 let size = common::vector_size_str(size);
1980 write!(self.out, "{scalar}{size}(-1), {scalar}{size}(1)")?;
1981 }
1982 _ => {
1983 write!(self.out, "{scalar}(-1), {scalar}(1)")?;
1984 }
1985 }
1986 write!(self.out, ", (")?;
1987 self.put_expression(arg, context, true)?;
1988 write!(self.out, " > 0)), {scalar}(0), (")?;
1989 self.put_expression(arg, context, true)?;
1990 write!(self.out, " == 0))")?;
1991 Ok(())
1992 }
1993
1994 pub(super) fn put_const_expression(
1995 &mut self,
1996 expr_handle: Handle<crate::Expression>,
1997 module: &crate::Module,
1998 mod_info: &valid::ModuleInfo,
1999 arena: &crate::Arena<crate::Expression>,
2000 ) -> BackendResult {
2001 self.put_possibly_const_expression(
2002 expr_handle,
2003 arena,
2004 module,
2005 mod_info,
2006 &(module, mod_info),
2007 |&(_, mod_info), expr| &mod_info[expr],
2008 |writer, &(module, _), expr| writer.put_const_expression(expr, module, mod_info, arena),
2009 )
2010 }
2011
2012 fn put_literal(&mut self, literal: crate::Literal) -> BackendResult {
2013 match literal {
2014 crate::Literal::F64(_) => {
2015 return Err(Error::CapabilityNotSupported(valid::Capabilities::FLOAT64))
2016 }
2017 crate::Literal::F16(value) => {
2018 if value.is_infinite() {
2019 let sign = if value.is_sign_negative() { "-" } else { "" };
2020 write!(self.out, "{sign}INFINITY")?;
2021 } else if value.is_nan() {
2022 write!(self.out, "NAN")?;
2023 } else {
2024 let suffix = if value.fract() == f16::from_f32(0.0) {
2025 ".0h"
2026 } else {
2027 "h"
2028 };
2029 write!(self.out, "{value}{suffix}")?;
2030 }
2031 }
2032 crate::Literal::F32(value) => {
2033 if value.is_infinite() {
2034 let sign = if value.is_sign_negative() { "-" } else { "" };
2035 write!(self.out, "{sign}INFINITY")?;
2036 } else if value.is_nan() {
2037 write!(self.out, "NAN")?;
2038 } else {
2039 let suffix = if value.fract() == 0.0 { ".0" } else { "" };
2040 write!(self.out, "{value}{suffix}")?;
2041 }
2042 }
2043 crate::Literal::U16(value) => {
2044 write!(self.out, "static_cast<ushort>({value})")?;
2045 }
2046 crate::Literal::I16(value) => {
2047 write!(self.out, "static_cast<short>({value})")?;
2048 }
2049 crate::Literal::U32(value) => {
2050 write!(self.out, "{value}u")?;
2051 }
2052 crate::Literal::I32(value) => {
2053 if value == i32::MIN {
2058 write!(self.out, "({} - 1)", value + 1)?;
2059 } else {
2060 write!(self.out, "{value}")?;
2061 }
2062 }
2063 crate::Literal::U64(value) => {
2064 write!(self.out, "{value}uL")?;
2065 }
2066 crate::Literal::I64(value) => {
2067 if value == i64::MIN {
2074 write!(self.out, "({}L - 1L)", value + 1)?;
2075 } else {
2076 write!(self.out, "{value}L")?;
2077 }
2078 }
2079 crate::Literal::Bool(value) => {
2080 write!(self.out, "{value}")?;
2081 }
2082 crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {
2083 return Err(Error::GenericValidation(
2084 "Unsupported abstract literal".into(),
2085 ));
2086 }
2087 }
2088 Ok(())
2089 }
2090
2091 #[allow(clippy::too_many_arguments)]
2092 fn put_possibly_const_expression<C, I, E>(
2093 &mut self,
2094 expr_handle: Handle<crate::Expression>,
2095 expressions: &crate::Arena<crate::Expression>,
2096 module: &crate::Module,
2097 mod_info: &valid::ModuleInfo,
2098 ctx: &C,
2099 get_expr_ty: I,
2100 put_expression: E,
2101 ) -> BackendResult
2102 where
2103 I: Fn(&C, Handle<crate::Expression>) -> &TypeResolution,
2104 E: Fn(&mut Self, &C, Handle<crate::Expression>) -> BackendResult,
2105 {
2106 match expressions[expr_handle] {
2107 crate::Expression::Literal(literal) => {
2108 self.put_literal(literal)?;
2109 }
2110 crate::Expression::Constant(handle) => {
2111 let constant = &module.constants[handle];
2112 if constant.name.is_some() {
2113 write!(self.out, "{}", self.names[&NameKey::Constant(handle)])?;
2114 } else {
2115 self.put_const_expression(
2116 constant.init,
2117 module,
2118 mod_info,
2119 &module.global_expressions,
2120 )?;
2121 }
2122 }
2123 crate::Expression::ZeroValue(ty) => {
2124 let ty_name = TypeContext {
2125 handle: ty,
2126 gctx: module.to_ctx(),
2127 names: &self.names,
2128 access: crate::StorageAccess::empty(),
2129 first_time: false,
2130 };
2131 write!(self.out, "{ty_name} {{}}")?;
2132 }
2133 crate::Expression::Compose { ty, ref components } => {
2134 let ty_name = TypeContext {
2135 handle: ty,
2136 gctx: module.to_ctx(),
2137 names: &self.names,
2138 access: crate::StorageAccess::empty(),
2139 first_time: false,
2140 };
2141 write!(self.out, "{ty_name}")?;
2142 match module.types[ty].inner {
2143 crate::TypeInner::Scalar(_)
2144 | crate::TypeInner::Vector { .. }
2145 | crate::TypeInner::Matrix { .. } => {
2146 self.put_call_parameters_impl(
2147 components.iter().copied(),
2148 ctx,
2149 put_expression,
2150 )?;
2151 }
2152 crate::TypeInner::Array { .. } => {
2153 write!(self.out, " {{{{")?;
2156 for (index, &component) in components.iter().enumerate() {
2157 if index != 0 {
2158 write!(self.out, ", ")?;
2159 }
2160 put_expression(self, ctx, component)?;
2161 }
2162 write!(self.out, "}}}}")?;
2163 }
2164 crate::TypeInner::Struct { .. } => {
2165 write!(self.out, " {{")?;
2166 for (index, &component) in components.iter().enumerate() {
2167 if index != 0 {
2168 write!(self.out, ", ")?;
2169 }
2170 if self.struct_member_pads.contains(&(ty, index as u32)) {
2172 write!(self.out, "{{}}, ")?;
2173 }
2174 put_expression(self, ctx, component)?;
2175 }
2176 write!(self.out, "}}")?;
2177 }
2178 _ => return Err(Error::UnsupportedCompose(ty)),
2179 }
2180 }
2181 crate::Expression::Splat { size, value } => {
2182 let scalar = match *get_expr_ty(ctx, value).inner_with(&module.types) {
2183 crate::TypeInner::Scalar(scalar) => scalar,
2184 ref ty => {
2185 return Err(Error::GenericValidation(format!(
2186 "Expected splat value type must be a scalar, got {ty:?}",
2187 )))
2188 }
2189 };
2190 put_numeric_type(&mut self.out, scalar, &[size])?;
2191 write!(self.out, "(")?;
2192 put_expression(self, ctx, value)?;
2193 write!(self.out, ")")?;
2194 }
2195 _ => {
2196 return Err(Error::Override);
2197 }
2198 }
2199
2200 Ok(())
2201 }
2202
2203 pub(super) fn put_expression(
2215 &mut self,
2216 expr_handle: Handle<crate::Expression>,
2217 context: &ExpressionContext,
2218 is_scoped: bool,
2219 ) -> BackendResult {
2220 let stored = self
2224 .named_expressions
2225 .get(&expr_handle)
2226 .or_else(|| self.hoisted_loop_expressions.get(&expr_handle));
2227 if let Some(name) = stored {
2228 write!(self.out, "{name}")?;
2229 return Ok(());
2230 }
2231
2232 let expression = &context.function.expressions[expr_handle];
2233 match *expression {
2234 crate::Expression::Literal(_)
2235 | crate::Expression::Constant(_)
2236 | crate::Expression::ZeroValue(_)
2237 | crate::Expression::Compose { .. }
2238 | crate::Expression::Splat { .. } => {
2239 self.put_possibly_const_expression(
2240 expr_handle,
2241 &context.function.expressions,
2242 context.module,
2243 context.mod_info,
2244 context,
2245 |context, expr: Handle<crate::Expression>| &context.info[expr].ty,
2246 |writer, context, expr| writer.put_expression(expr, context, true),
2247 )?;
2248 }
2249 crate::Expression::Override(_) => return Err(Error::Override),
2250 crate::Expression::Access { base, .. }
2251 | crate::Expression::AccessIndex { base, .. } => {
2252 let policy = context.choose_bounds_check_policy(base);
2258 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
2259 && self.put_bounds_checks(
2260 expr_handle,
2261 context,
2262 back::Level(0),
2263 if is_scoped { "" } else { "(" },
2264 )?
2265 {
2266 write!(self.out, " ? ")?;
2267 self.put_access_chain(expr_handle, policy, context)?;
2268 write!(self.out, " : ")?;
2269
2270 if context.resolve_type(base).pointer_space().is_some() {
2271 let result_ty = context.info[expr_handle]
2275 .ty
2276 .inner_with(&context.module.types)
2277 .pointer_base_type();
2278 let result_ty_handle = match result_ty {
2279 Some(TypeResolution::Handle(handle)) => handle,
2280 Some(TypeResolution::Value(_)) => {
2281 unreachable!(
2289 "Expected type {result_ty:?} of access through pointer type {base:?} to be in the arena",
2290 );
2291 }
2292 None => {
2293 unreachable!(
2294 "Expected access through pointer type {base:?} to return a pointer, but got {result_ty:?}",
2295 )
2296 }
2297 };
2298 let name_key =
2299 NameKey::oob_local_for_type(context.origin, result_ty_handle);
2300 self.out.write_str(&self.names[&name_key])?;
2301 } else {
2302 write!(self.out, "DefaultConstructible()")?;
2303 }
2304
2305 if !is_scoped {
2306 write!(self.out, ")")?;
2307 }
2308 } else {
2309 self.put_access_chain(expr_handle, policy, context)?;
2310 }
2311 }
2312 crate::Expression::Swizzle {
2313 size,
2314 vector,
2315 pattern,
2316 } => {
2317 self.put_wrapped_expression_for_packed_vec3_access(
2318 vector,
2319 context,
2320 false,
2321 &Self::put_expression,
2322 )?;
2323 write!(self.out, ".")?;
2324 for &sc in pattern[..size as usize].iter() {
2325 write!(self.out, "{}", back::COMPONENTS[sc as usize])?;
2326 }
2327 }
2328 crate::Expression::FunctionArgument(index) => {
2329 let name_key = match context.origin {
2330 FunctionOrigin::Handle(handle) => NameKey::FunctionArgument(handle, index),
2331 FunctionOrigin::EntryPoint(ep_index) => {
2332 NameKey::EntryPointArgument(ep_index, index)
2333 }
2334 };
2335 let name = &self.names[&name_key];
2336 write!(self.out, "{name}")?;
2337 }
2338 crate::Expression::GlobalVariable(handle) => {
2339 let name = &self.names[&NameKey::GlobalVariable(handle)];
2340 write!(self.out, "{name}")?;
2341 }
2342 crate::Expression::LocalVariable(handle) => {
2343 let name_key = NameKey::local(context.origin, handle);
2344 let name = &self.names[&name_key];
2345 write!(self.out, "{name}")?;
2346 }
2347 crate::Expression::Load { pointer } => self.put_load(pointer, context, is_scoped)?,
2348 crate::Expression::ImageSample {
2349 coordinate,
2350 image,
2351 sampler,
2352 clamp_to_edge: true,
2353 gather: None,
2354 array_index: None,
2355 offset: None,
2356 level: crate::SampleLevel::Zero,
2357 depth_ref: None,
2358 } => {
2359 write!(self.out, "{IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION}(")?;
2360 self.put_expression(image, context, true)?;
2361 write!(self.out, ", ")?;
2362 self.put_expression(sampler, context, true)?;
2363 write!(self.out, ", ")?;
2364 self.put_expression(coordinate, context, true)?;
2365 write!(self.out, ")")?;
2366 }
2367 crate::Expression::ImageSample {
2368 image,
2369 sampler,
2370 gather,
2371 coordinate,
2372 array_index,
2373 offset,
2374 level,
2375 depth_ref,
2376 clamp_to_edge,
2377 } => {
2378 if clamp_to_edge {
2379 return Err(Error::GenericValidation(
2380 "ImageSample::clamp_to_edge should have been validated out".to_string(),
2381 ));
2382 }
2383
2384 let main_op = match gather {
2385 Some(_) => "gather",
2386 None => "sample",
2387 };
2388 let comparison_op = match depth_ref {
2389 Some(_) => "_compare",
2390 None => "",
2391 };
2392 self.put_expression(image, context, false)?;
2393 write!(self.out, ".{main_op}{comparison_op}(")?;
2394 self.put_expression(sampler, context, true)?;
2395 write!(self.out, ", ")?;
2396 self.put_expression(coordinate, context, true)?;
2397 if let Some(expr) = array_index {
2398 write!(self.out, ", ")?;
2399 self.put_expression(expr, context, true)?;
2400 }
2401 if let Some(dref) = depth_ref {
2402 write!(self.out, ", ")?;
2403 self.put_expression(dref, context, true)?;
2404 }
2405
2406 self.put_image_sample_level(image, level, context)?;
2407
2408 if let Some(offset) = offset {
2409 write!(self.out, ", ")?;
2410 self.put_expression(offset, context, true)?;
2411 }
2412
2413 match gather {
2414 None | Some(crate::SwizzleComponent::X) => {}
2415 Some(component) => {
2416 let is_cube_map = match *context.resolve_type(image) {
2417 crate::TypeInner::Image {
2418 dim: crate::ImageDimension::Cube,
2419 ..
2420 } => true,
2421 _ => false,
2422 };
2423 if offset.is_none() && !is_cube_map {
2426 write!(self.out, ", {NAMESPACE}::int2(0)")?;
2427 }
2428 let letter = back::COMPONENTS[component as usize];
2429 write!(self.out, ", {NAMESPACE}::component::{letter}")?;
2430 }
2431 }
2432 write!(self.out, ")")?;
2433 }
2434 crate::Expression::ImageLoad {
2435 image,
2436 coordinate,
2437 array_index,
2438 sample,
2439 level,
2440 } => {
2441 let address = TexelAddress {
2442 coordinate,
2443 array_index,
2444 sample,
2445 level: level.map(LevelOfDetail::Direct),
2446 };
2447 self.put_image_load(expr_handle, image, address, context)?;
2448 }
2449 crate::Expression::ImageQuery { image, query } => match query {
2452 crate::ImageQuery::Size { level } => {
2453 self.put_image_size_query(
2454 image,
2455 level.map(LevelOfDetail::Direct),
2456 crate::ScalarKind::Uint,
2457 context,
2458 )?;
2459 }
2460 crate::ImageQuery::NumLevels => {
2461 self.put_expression(image, context, false)?;
2462 write!(self.out, ".get_num_mip_levels()")?;
2463 }
2464 crate::ImageQuery::NumLayers => {
2465 self.put_expression(image, context, false)?;
2466 write!(self.out, ".get_array_size()")?;
2467 }
2468 crate::ImageQuery::NumSamples => {
2469 self.put_expression(image, context, false)?;
2470 write!(self.out, ".get_num_samples()")?;
2471 }
2472 },
2473 crate::Expression::Unary { op, expr } => {
2474 let op_str = match op {
2475 crate::UnaryOperator::Negate => {
2476 match context.resolve_type(expr).scalar_kind() {
2477 Some(crate::ScalarKind::Sint) => NEG_FUNCTION,
2478 _ => "-",
2479 }
2480 }
2481 crate::UnaryOperator::LogicalNot => "!",
2482 crate::UnaryOperator::BitwiseNot => "~",
2483 };
2484 write!(self.out, "{op_str}(")?;
2485 self.put_expression(expr, context, false)?;
2486 write!(self.out, ")")?;
2487 }
2488 crate::Expression::Binary { op, left, right } => {
2489 let kind = context
2490 .resolve_type(left)
2491 .scalar_kind()
2492 .ok_or(Error::UnsupportedBinaryOp(op))?;
2493
2494 if op == crate::BinaryOperator::Divide
2495 && (kind == crate::ScalarKind::Sint || kind == crate::ScalarKind::Uint)
2496 && context.emit_int_div_checks
2497 {
2498 write!(self.out, "{DIV_FUNCTION}(")?;
2499 self.put_expression(left, context, true)?;
2500 write!(self.out, ", ")?;
2501 self.put_expression(right, context, true)?;
2502 write!(self.out, ")")?;
2503 } else if op == crate::BinaryOperator::Modulo
2504 && (kind == crate::ScalarKind::Sint || kind == crate::ScalarKind::Uint)
2505 && context.emit_int_div_checks
2506 {
2507 write!(self.out, "{MOD_FUNCTION}(")?;
2508 self.put_expression(left, context, true)?;
2509 write!(self.out, ", ")?;
2510 self.put_expression(right, context, true)?;
2511 write!(self.out, ")")?;
2512 } else if op == crate::BinaryOperator::Modulo && kind == crate::ScalarKind::Float {
2513 write!(self.out, "{NAMESPACE}::fmod(")?;
2518 self.put_expression(left, context, true)?;
2519 write!(self.out, ", ")?;
2520 self.put_expression(right, context, true)?;
2521 write!(self.out, ")")?;
2522 } else if (op == crate::BinaryOperator::Add
2523 || op == crate::BinaryOperator::Subtract
2524 || op == crate::BinaryOperator::Multiply)
2525 && kind == crate::ScalarKind::Sint
2526 {
2527 let to_unsigned = |ty: &crate::TypeInner| match *ty {
2528 crate::TypeInner::Scalar(scalar) => {
2529 Ok(crate::TypeInner::Scalar(crate::Scalar {
2530 kind: crate::ScalarKind::Uint,
2531 ..scalar
2532 }))
2533 }
2534 crate::TypeInner::Vector { size, scalar } => Ok(crate::TypeInner::Vector {
2535 size,
2536 scalar: crate::Scalar {
2537 kind: crate::ScalarKind::Uint,
2538 ..scalar
2539 },
2540 }),
2541 _ => Err(Error::UnsupportedBitCast(ty.clone())),
2542 };
2543
2544 self.put_bitcasted_expression(
2549 context.resolve_type(expr_handle),
2550 expr_handle,
2551 context,
2552 &|writer, context, is_scoped| {
2553 writer.put_binop(
2554 op,
2555 left,
2556 right,
2557 context,
2558 is_scoped,
2559 &|writer, expr, context, _is_scoped| {
2560 writer.put_bitcasted_expression(
2561 &to_unsigned(context.resolve_type(expr))?,
2562 expr,
2563 context,
2564 &|writer, context, is_scoped| {
2565 writer.put_expression(expr, context, is_scoped)
2566 },
2567 )
2568 },
2569 )
2570 },
2571 )?;
2572 } else {
2573 self.put_binop(op, left, right, context, is_scoped, &Self::put_expression)?;
2574 }
2575 }
2576 crate::Expression::Select {
2577 condition,
2578 accept,
2579 reject,
2580 } => match *context.resolve_type(condition) {
2581 crate::TypeInner::Scalar(crate::Scalar {
2582 kind: crate::ScalarKind::Bool,
2583 ..
2584 }) => {
2585 if !is_scoped {
2586 write!(self.out, "(")?;
2587 }
2588 self.put_expression(condition, context, false)?;
2589 write!(self.out, " ? ")?;
2590 self.put_expression(accept, context, false)?;
2591 write!(self.out, " : ")?;
2592 self.put_expression(reject, context, false)?;
2593 if !is_scoped {
2594 write!(self.out, ")")?;
2595 }
2596 }
2597 crate::TypeInner::Vector {
2598 scalar:
2599 crate::Scalar {
2600 kind: crate::ScalarKind::Bool,
2601 ..
2602 },
2603 ..
2604 } => {
2605 write!(self.out, "{NAMESPACE}::select(")?;
2606 self.put_expression(reject, context, true)?;
2607 write!(self.out, ", ")?;
2608 self.put_expression(accept, context, true)?;
2609 write!(self.out, ", ")?;
2610 self.put_expression(condition, context, true)?;
2611 write!(self.out, ")")?;
2612 }
2613 ref ty => {
2614 return Err(Error::GenericValidation(format!(
2615 "Expected select condition to be a non-bool type, got {ty:?}",
2616 )))
2617 }
2618 },
2619 crate::Expression::Derivative { axis, expr, .. } => {
2620 use crate::DerivativeAxis as Axis;
2621 let op = match axis {
2622 Axis::X => "dfdx",
2623 Axis::Y => "dfdy",
2624 Axis::Width => "fwidth",
2625 };
2626 write!(self.out, "{NAMESPACE}::{op}")?;
2627 self.put_call_parameters(iter::once(expr), context)?;
2628 }
2629 crate::Expression::Relational { fun, argument } => {
2630 let op = match fun {
2631 crate::RelationalFunction::Any => "any",
2632 crate::RelationalFunction::All => "all",
2633 crate::RelationalFunction::IsNan => "isnan",
2634 crate::RelationalFunction::IsInf => "isinf",
2635 };
2636 write!(self.out, "{NAMESPACE}::{op}")?;
2637 self.put_call_parameters(iter::once(argument), context)?;
2638 }
2639 crate::Expression::Math {
2640 fun,
2641 arg,
2642 arg1,
2643 arg2,
2644 arg3,
2645 } => {
2646 use crate::MathFunction as Mf;
2647
2648 let arg_type = context.resolve_type(arg);
2649 let scalar_argument = match arg_type {
2650 &crate::TypeInner::Scalar(_) => true,
2651 _ => false,
2652 };
2653
2654 let fun_name = match fun {
2655 Mf::Abs => "abs",
2657 Mf::Min => "min",
2658 Mf::Max => "max",
2659 Mf::Clamp => "clamp",
2660 Mf::Saturate => "saturate",
2661 Mf::Cos => "cos",
2663 Mf::Cosh => "cosh",
2664 Mf::Sin => "sin",
2665 Mf::Sinh => "sinh",
2666 Mf::Tan => "tan",
2667 Mf::Tanh => "tanh",
2668 Mf::Acos => "acos",
2669 Mf::Asin => "asin",
2670 Mf::Atan => "atan",
2671 Mf::Atan2 => "atan2",
2672 Mf::Asinh => "asinh",
2673 Mf::Acosh => "acosh",
2674 Mf::Atanh => "atanh",
2675 Mf::Radians => "",
2676 Mf::Degrees => "",
2677 Mf::Ceil => "ceil",
2679 Mf::Floor => "floor",
2680 Mf::Round => "rint",
2681 Mf::Fract => "fract",
2682 Mf::Trunc => "trunc",
2683 Mf::Modf => MODF_FUNCTION,
2684 Mf::Frexp => FREXP_FUNCTION,
2685 Mf::Ldexp => "ldexp",
2686 Mf::Exp => "exp",
2688 Mf::Exp2 => "exp2",
2689 Mf::Log => "log",
2690 Mf::Log2 => "log2",
2691 Mf::Pow => "pow",
2692 Mf::Dot => match *context.resolve_type(arg) {
2694 crate::TypeInner::Vector {
2695 scalar:
2696 crate::Scalar {
2697 kind: crate::ScalarKind::Float,
2699 ..
2700 },
2701 ..
2702 } => "dot",
2703 crate::TypeInner::Vector {
2704 size,
2705 scalar:
2706 scalar @ crate::Scalar {
2707 kind: crate::ScalarKind::Sint | crate::ScalarKind::Uint,
2708 ..
2709 },
2710 } => {
2711 let fun_name = self.get_dot_wrapper_function_helper_name(scalar, size);
2713 write!(self.out, "{fun_name}(")?;
2714 self.put_expression(arg, context, true)?;
2715 write!(self.out, ", ")?;
2716 self.put_expression(arg1.unwrap(), context, true)?;
2717 write!(self.out, ")")?;
2718 return Ok(());
2719 }
2720 _ => unreachable!(
2721 "Correct TypeInner for dot product should be already validated"
2722 ),
2723 },
2724 fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed) => {
2725 if context.lang_version >= (2, 1) {
2726 let packed_type = match fun {
2730 Mf::Dot4I8Packed => "packed_char4",
2731 Mf::Dot4U8Packed => "packed_uchar4",
2732 _ => unreachable!(),
2733 };
2734
2735 return self.put_dot_product(
2736 Reinterpreted::new(packed_type, arg),
2737 Reinterpreted::new(packed_type, arg1.unwrap()),
2738 4,
2739 |writer, arg, index| {
2740 write!(writer.out, "{arg}[{index}]")?;
2743 Ok(())
2744 },
2745 );
2746 } else {
2747 let conversion = match fun {
2751 Mf::Dot4I8Packed => "int",
2752 Mf::Dot4U8Packed => "",
2753 _ => unreachable!(),
2754 };
2755
2756 return self.put_dot_product(
2757 arg,
2758 arg1.unwrap(),
2759 4,
2760 |writer, arg, index| {
2761 write!(writer.out, "({conversion}(")?;
2762 writer.put_expression(arg, context, true)?;
2763 if index == 3 {
2764 write!(writer.out, ") >> 24)")?;
2765 } else {
2766 write!(writer.out, ") << {} >> 24)", (3 - index) * 8)?;
2767 }
2768 Ok(())
2769 },
2770 );
2771 }
2772 }
2773 Mf::Outer => return Err(Error::UnsupportedCall(format!("{fun:?}"))),
2774 Mf::Cross => "cross",
2775 Mf::Distance => "distance",
2776 Mf::Length if scalar_argument => "abs",
2777 Mf::Length => "length",
2778 Mf::Normalize => "normalize",
2779 Mf::FaceForward => "faceforward",
2780 Mf::Reflect => "reflect",
2781 Mf::Refract => "refract",
2782 Mf::Sign => match arg_type.scalar_kind() {
2784 Some(crate::ScalarKind::Sint) => {
2785 return self.put_isign(arg, context);
2786 }
2787 _ => "sign",
2788 },
2789 Mf::Fma => "fma",
2790 Mf::Mix => "mix",
2791 Mf::Step => "step",
2792 Mf::SmoothStep => "smoothstep",
2793 Mf::Sqrt => "sqrt",
2794 Mf::InverseSqrt => "rsqrt",
2795 Mf::Inverse => return Err(Error::UnsupportedCall(format!("{fun:?}"))),
2796 Mf::Transpose => "transpose",
2797 Mf::Determinant => "determinant",
2798 Mf::QuantizeToF16 => "",
2799 Mf::CountTrailingZeros => "ctz",
2801 Mf::CountLeadingZeros => "clz",
2802 Mf::CountOneBits => "popcount",
2803 Mf::ReverseBits => "reverse_bits",
2804 Mf::ExtractBits => "",
2805 Mf::InsertBits => "",
2806 Mf::FirstTrailingBit => "",
2807 Mf::FirstLeadingBit => "",
2808 Mf::Pack4x8snorm => "pack_float_to_snorm4x8",
2810 Mf::Pack4x8unorm => "pack_float_to_unorm4x8",
2811 Mf::Pack2x16snorm => "pack_float_to_snorm2x16",
2812 Mf::Pack2x16unorm => "pack_float_to_unorm2x16",
2813 Mf::Pack2x16float => "",
2814 Mf::Pack4xI8 => "",
2815 Mf::Pack4xU8 => "",
2816 Mf::Pack4xI8Clamp => "",
2817 Mf::Pack4xU8Clamp => "",
2818 Mf::Unpack4x8snorm => "unpack_snorm4x8_to_float",
2820 Mf::Unpack4x8unorm => "unpack_unorm4x8_to_float",
2821 Mf::Unpack2x16snorm => "unpack_snorm2x16_to_float",
2822 Mf::Unpack2x16unorm => "unpack_unorm2x16_to_float",
2823 Mf::Unpack2x16float => "",
2824 Mf::Unpack4xI8 => "",
2825 Mf::Unpack4xU8 => "",
2826 };
2827
2828 match fun {
2829 Mf::ReverseBits | Mf::ExtractBits | Mf::InsertBits
2838 if context.lang_version < (1, 2) =>
2839 {
2840 return Err(Error::UnsupportedFunction(fun_name.to_string()));
2841 }
2842 _ => {}
2843 }
2844
2845 match fun {
2846 Mf::Abs if arg_type.scalar_kind() == Some(crate::ScalarKind::Sint) => {
2847 write!(self.out, "{ABS_FUNCTION}(")?;
2848 self.put_expression(arg, context, true)?;
2849 write!(self.out, ")")?;
2850 }
2851 Mf::Distance if scalar_argument => {
2852 write!(self.out, "{NAMESPACE}::abs(")?;
2853 self.put_expression(arg, context, false)?;
2854 write!(self.out, " - ")?;
2855 self.put_expression(arg1.unwrap(), context, false)?;
2856 write!(self.out, ")")?;
2857 }
2858 Mf::FirstTrailingBit => {
2859 let scalar = context.resolve_type(arg).scalar().unwrap();
2860 let constant = scalar.width * 8 + 1;
2861
2862 write!(self.out, "((({NAMESPACE}::ctz(")?;
2863 self.put_expression(arg, context, true)?;
2864 write!(self.out, ") + 1) % {constant}) - 1)")?;
2865 }
2866 Mf::FirstLeadingBit => {
2867 let inner = context.resolve_type(arg);
2868 let scalar = inner.scalar().unwrap();
2869 let constant = scalar.width * 8 - 1;
2870
2871 write!(
2872 self.out,
2873 "{NAMESPACE}::select({constant} - {NAMESPACE}::clz("
2874 )?;
2875
2876 if scalar.kind == crate::ScalarKind::Sint {
2877 write!(self.out, "{NAMESPACE}::select(")?;
2878 self.put_expression(arg, context, true)?;
2879 write!(self.out, ", ~")?;
2880 self.put_expression(arg, context, true)?;
2881 write!(self.out, ", ")?;
2882 self.put_expression(arg, context, true)?;
2883 write!(self.out, " < 0)")?;
2884 } else {
2885 self.put_expression(arg, context, true)?;
2886 }
2887
2888 write!(self.out, "), ")?;
2889
2890 match *inner {
2892 crate::TypeInner::Vector { size, scalar } => {
2893 let size = common::vector_size_str(size);
2894 let name = scalar.to_msl_name();
2895 write!(self.out, "{name}{size}")?;
2896 }
2897 crate::TypeInner::Scalar(scalar) => {
2898 let name = scalar.to_msl_name();
2899 write!(self.out, "{name}")?;
2900 }
2901 _ => (),
2902 }
2903
2904 write!(self.out, "(-1), ")?;
2905 self.put_expression(arg, context, true)?;
2906 write!(self.out, " == 0")?;
2907 if scalar.kind == crate::ScalarKind::Sint {
2908 write!(self.out, " || ")?;
2909 self.put_expression(arg, context, true)?;
2910 write!(self.out, " == -1")?;
2911 }
2912 write!(self.out, ")")?;
2913 }
2914 Mf::Unpack2x16float => {
2915 write!(self.out, "float2(as_type<half2>(")?;
2916 self.put_expression(arg, context, false)?;
2917 write!(self.out, "))")?;
2918 }
2919 Mf::Pack2x16float => {
2920 write!(self.out, "as_type<uint>(half2(")?;
2921 self.put_expression(arg, context, false)?;
2922 write!(self.out, "))")?;
2923 }
2924 Mf::ExtractBits => {
2925 let scalar_bits = context.resolve_type(arg).scalar_width().unwrap() * 8;
2942
2943 write!(self.out, "{NAMESPACE}::extract_bits(")?;
2944 self.put_expression(arg, context, true)?;
2945 write!(self.out, ", {NAMESPACE}::min(")?;
2946 self.put_expression(arg1.unwrap(), context, true)?;
2947 write!(self.out, ", {scalar_bits}u), {NAMESPACE}::min(")?;
2948 self.put_expression(arg2.unwrap(), context, true)?;
2949 write!(self.out, ", {scalar_bits}u - {NAMESPACE}::min(")?;
2950 self.put_expression(arg1.unwrap(), context, true)?;
2951 write!(self.out, ", {scalar_bits}u)))")?;
2952 }
2953 Mf::InsertBits => {
2954 let scalar_bits = context.resolve_type(arg).scalar_width().unwrap() * 8;
2959
2960 write!(self.out, "{NAMESPACE}::insert_bits(")?;
2961 self.put_expression(arg, context, true)?;
2962 write!(self.out, ", ")?;
2963 self.put_expression(arg1.unwrap(), context, true)?;
2964 write!(self.out, ", {NAMESPACE}::min(")?;
2965 self.put_expression(arg2.unwrap(), context, true)?;
2966 write!(self.out, ", {scalar_bits}u), {NAMESPACE}::min(")?;
2967 self.put_expression(arg3.unwrap(), context, true)?;
2968 write!(self.out, ", {scalar_bits}u - {NAMESPACE}::min(")?;
2969 self.put_expression(arg2.unwrap(), context, true)?;
2970 write!(self.out, ", {scalar_bits}u)))")?;
2971 }
2972 Mf::Radians => {
2973 write!(self.out, "((")?;
2974 self.put_expression(arg, context, false)?;
2975 write!(self.out, ") * 0.017453292519943295474)")?;
2976 }
2977 Mf::Degrees => {
2978 write!(self.out, "((")?;
2979 self.put_expression(arg, context, false)?;
2980 write!(self.out, ") * 57.295779513082322865)")?;
2981 }
2982 Mf::Modf | Mf::Frexp => {
2983 write!(self.out, "{fun_name}")?;
2984 self.put_call_parameters(iter::once(arg), context)?;
2985 }
2986 Mf::Pack4xI8 => self.put_pack4x8(arg, context, true, None)?,
2987 Mf::Pack4xU8 => self.put_pack4x8(arg, context, false, None)?,
2988 Mf::Pack4xI8Clamp => {
2989 self.put_pack4x8(arg, context, true, Some(("-128", "127")))?
2990 }
2991 Mf::Pack4xU8Clamp => {
2992 self.put_pack4x8(arg, context, false, Some(("0", "255")))?
2993 }
2994 fun @ (Mf::Unpack4xI8 | Mf::Unpack4xU8) => {
2995 let sign_prefix = if matches!(fun, Mf::Unpack4xU8) {
2996 "u"
2997 } else {
2998 ""
2999 };
3000
3001 if context.lang_version >= (2, 1) {
3002 write!(
3004 self.out,
3005 "{sign_prefix}int4(as_type<packed_{sign_prefix}char4>("
3006 )?;
3007 self.put_expression(arg, context, true)?;
3008 write!(self.out, "))")?;
3009 } else {
3010 write!(self.out, "({sign_prefix}int4(")?;
3012 self.put_expression(arg, context, true)?;
3013 write!(self.out, ", ")?;
3014 self.put_expression(arg, context, true)?;
3015 write!(self.out, " >> 8, ")?;
3016 self.put_expression(arg, context, true)?;
3017 write!(self.out, " >> 16, ")?;
3018 self.put_expression(arg, context, true)?;
3019 write!(self.out, " >> 24) << 24 >> 24)")?;
3020 }
3021 }
3022 Mf::QuantizeToF16 => {
3023 match *context.resolve_type(arg) {
3024 crate::TypeInner::Scalar { .. } => write!(self.out, "float(half(")?,
3025 crate::TypeInner::Vector { size, .. } => write!(
3026 self.out,
3027 "{NAMESPACE}::float{size}({NAMESPACE}::half{size}(",
3028 size = common::vector_size_str(size),
3029 )?,
3030 _ => unreachable!(
3031 "Correct TypeInner for QuantizeToF16 should be already validated"
3032 ),
3033 };
3034
3035 self.put_expression(arg, context, true)?;
3036 write!(self.out, "))")?;
3037 }
3038 _ => {
3039 write!(self.out, "{NAMESPACE}::{fun_name}")?;
3040 self.put_call_parameters(
3041 iter::once(arg).chain(arg1).chain(arg2).chain(arg3),
3042 context,
3043 )?;
3044 }
3045 }
3046 }
3047 crate::Expression::As { expr, target } => match *context.resolve_type(expr) {
3048 crate::TypeInner::Scalar(src) | crate::TypeInner::Vector { scalar: src, .. } => {
3049 if src.kind == crate::ScalarKind::Float
3050 && (target.kind == crate::ScalarKind::Sint
3051 || target.kind == crate::ScalarKind::Uint)
3052 {
3053 let fun_name = match (target.kind, target.width) {
3057 (crate::ScalarKind::Sint, 4) => F2I32_FUNCTION,
3058 (crate::ScalarKind::Uint, 4) => F2U32_FUNCTION,
3059 (crate::ScalarKind::Sint, 8) => F2I64_FUNCTION,
3060 (crate::ScalarKind::Uint, 8) => F2U64_FUNCTION,
3061 _ => unreachable!(),
3062 };
3063 write!(self.out, "{fun_name}(")?;
3064 self.put_expression(expr, context, true)?;
3065 write!(self.out, ")")?;
3066 } else {
3067 write!(self.out, "static_cast<")?;
3068 match *context.resolve_type(expr) {
3069 crate::TypeInner::Vector { size, .. } => {
3070 put_numeric_type(&mut self.out, target, &[size])?
3071 }
3072 _ => put_numeric_type(&mut self.out, target, &[])?,
3073 };
3074 write!(self.out, ">(")?;
3075 self.put_expression(expr, context, true)?;
3076 write!(self.out, ")")?;
3077 }
3078 }
3079 crate::TypeInner::Matrix { columns, rows, .. } => {
3080 put_numeric_type(&mut self.out, target, &[rows, columns])?;
3081 write!(self.out, "(")?;
3082 self.put_expression(expr, context, true)?;
3083 write!(self.out, ")")?;
3084 }
3085 ref ty => {
3086 return Err(Error::GenericValidation(format!(
3087 "Unsupported type for As: {ty:?}"
3088 )))
3089 }
3090 },
3091 crate::Expression::Bitcast { expr, ty } => {
3092 write!(self.out, "as_type<")?;
3093 match context.module.types[ty].inner {
3094 crate::TypeInner::Scalar(scalar) => {
3095 put_numeric_type(&mut self.out, scalar, &[])?
3096 }
3097 crate::TypeInner::Vector { size, scalar } => {
3098 put_numeric_type(&mut self.out, scalar, &[size])?
3099 }
3100 ref other => {
3101 return Err(Error::GenericValidation(format!(
3102 "Unsupported type for Bitcast: {other:?}"
3103 )))
3104 }
3105 }
3106 write!(self.out, ">(")?;
3107 self.put_expression(expr, context, true)?;
3108 write!(self.out, ")")?;
3109 }
3110 crate::Expression::CallResult(_)
3117 | crate::Expression::AtomicResult { .. }
3118 | crate::Expression::WorkGroupUniformLoadResult { .. }
3119 | crate::Expression::SubgroupBallotResult
3120 | crate::Expression::SubgroupOperationResult { .. }
3121 | crate::Expression::RayQueryProceedResult => {
3122 unreachable!("statement result {expr_handle:?} has no temporary holding it")
3123 }
3124 crate::Expression::ArrayLength(expr) => {
3125 let global = context.function.originating_global(expr).ok_or_else(|| {
3126 Error::GenericValidation(format!(
3127 "Could not find global variable for ArrayLength operand {:?}",
3128 context.function.expressions[expr]
3129 ))
3130 })?;
3131
3132 if !is_scoped {
3133 write!(self.out, "(")?;
3134 }
3135 write!(self.out, "1 + ")?;
3136 self.put_dynamic_array_max_index(global, expr, context)?;
3137 if !is_scoped {
3138 write!(self.out, ")")?;
3139 }
3140 }
3141 crate::Expression::RayQueryVertexPositions { .. } => {
3142 unimplemented!()
3143 }
3144 crate::Expression::RayQueryGetIntersection { query, committed } => {
3145 if context.lang_version < (2, 4) {
3146 return Err(Error::UnsupportedRayTracing);
3147 }
3148
3149 let crate::Expression::LocalVariable(query_var) =
3151 context.function.expressions[query]
3152 else {
3153 unreachable!()
3154 };
3155
3156 let tracker_expr_name = format!(
3157 "{}{}",
3158 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
3159 self.names[&NameKey::local(context.origin, query_var)]
3160 );
3161
3162 write!(
3163 self.out,
3164 "{}_{committed}(",
3165 super::ray::INTERSECTION_FUNCTION_NAME
3166 )?;
3167 self.put_expression(query, context, true)?;
3168 if context.ray_query_initialization_tracking {
3169 write!(self.out, ", {tracker_expr_name}")?;
3170 }
3171 write!(self.out, ")")?;
3172 }
3173 crate::Expression::CooperativeLoad { ref data, .. } => {
3174 if context.lang_version < (2, 3) {
3175 return Err(Error::UnsupportedCooperativeMatrix);
3176 }
3177 write!(self.out, "{COOPERATIVE_LOAD_FUNCTION}(")?;
3178 write!(self.out, "&")?;
3179 self.put_access_chain(data.pointer, context.policies.index, context)?;
3180 write!(self.out, ", ")?;
3181 self.put_expression(data.stride, context, true)?;
3182 write!(self.out, ", {})", !data.row_major)?;
3189 }
3190 crate::Expression::CooperativeMultiplyAdd { a, b, c } => {
3191 if context.lang_version < (2, 3) {
3192 return Err(Error::UnsupportedCooperativeMatrix);
3193 }
3194 write!(self.out, "{COOPERATIVE_MULTIPLY_ADD_FUNCTION}(")?;
3195 self.put_expression(a, context, true)?;
3196 write!(self.out, ", ")?;
3197 self.put_expression(b, context, true)?;
3198 write!(self.out, ", ")?;
3199 self.put_expression(c, context, true)?;
3200 write!(self.out, ")")?;
3201 }
3202 }
3203 Ok(())
3204 }
3205
3206 fn put_binop<F>(
3209 &mut self,
3210 op: crate::BinaryOperator,
3211 left: Handle<crate::Expression>,
3212 right: Handle<crate::Expression>,
3213 context: &ExpressionContext,
3214 is_scoped: bool,
3215 put_expression: &F,
3216 ) -> BackendResult
3217 where
3218 F: Fn(&mut Self, Handle<crate::Expression>, &ExpressionContext, bool) -> BackendResult,
3219 {
3220 let op_str = back::binary_operation_str(op);
3221
3222 if !is_scoped {
3223 write!(self.out, "(")?;
3224 }
3225
3226 if op == crate::BinaryOperator::Multiply
3229 && matches!(
3230 context.resolve_type(right),
3231 &crate::TypeInner::Matrix { .. }
3232 )
3233 {
3234 self.put_wrapped_expression_for_packed_vec3_access(
3235 left,
3236 context,
3237 false,
3238 put_expression,
3239 )?;
3240 } else {
3241 put_expression(self, left, context, false)?;
3242 }
3243
3244 write!(self.out, " {op_str} ")?;
3245
3246 if op == crate::BinaryOperator::Multiply
3248 && matches!(context.resolve_type(left), &crate::TypeInner::Matrix { .. })
3249 {
3250 self.put_wrapped_expression_for_packed_vec3_access(
3251 right,
3252 context,
3253 false,
3254 put_expression,
3255 )?;
3256 } else {
3257 put_expression(self, right, context, false)?;
3258 }
3259
3260 if !is_scoped {
3261 write!(self.out, ")")?;
3262 }
3263
3264 Ok(())
3265 }
3266
3267 fn put_wrapped_expression_for_packed_vec3_access<F>(
3269 &mut self,
3270 expr_handle: Handle<crate::Expression>,
3271 context: &ExpressionContext,
3272 is_scoped: bool,
3273 put_expression: &F,
3274 ) -> BackendResult
3275 where
3276 F: Fn(&mut Self, Handle<crate::Expression>, &ExpressionContext, bool) -> BackendResult,
3277 {
3278 if let Some(scalar) = context.get_packed_vec_kind(expr_handle) {
3279 write!(self.out, "{}::{}3(", NAMESPACE, scalar.to_msl_name())?;
3280 put_expression(self, expr_handle, context, is_scoped)?;
3281 write!(self.out, ")")?;
3282 } else {
3283 put_expression(self, expr_handle, context, is_scoped)?;
3284 }
3285 Ok(())
3286 }
3287
3288 fn put_bitcasted_expression<F>(
3291 &mut self,
3292 cast_to: &crate::TypeInner,
3293 inner_expr: Handle<crate::Expression>,
3294 context: &ExpressionContext,
3295 put_expression: &F,
3296 ) -> BackendResult
3297 where
3298 F: Fn(&mut Self, &ExpressionContext, bool) -> BackendResult,
3299 {
3300 let needs_truncation = match *cast_to {
3305 crate::TypeInner::Scalar(scalar) => scalar.width < 4,
3306 crate::TypeInner::Vector { scalar, .. } => scalar.width < 4,
3307 _ => false,
3308 };
3309
3310 write!(self.out, "as_type<")?;
3311 match *cast_to {
3312 crate::TypeInner::Scalar(scalar) => put_numeric_type(&mut self.out, scalar, &[])?,
3313 crate::TypeInner::Vector { size, scalar } => {
3314 put_numeric_type(&mut self.out, scalar, &[size])?
3315 }
3316 _ => return Err(Error::UnsupportedBitCast(cast_to.clone())),
3317 };
3318 write!(self.out, ">(")?;
3319
3320 if needs_truncation {
3321 write!(self.out, "static_cast<")?;
3322 let unsigned_scalar = match *cast_to {
3324 crate::TypeInner::Scalar(scalar) => crate::Scalar {
3325 kind: crate::ScalarKind::Uint,
3326 ..scalar
3327 },
3328 crate::TypeInner::Vector { scalar, .. } => crate::Scalar {
3329 kind: crate::ScalarKind::Uint,
3330 ..scalar
3331 },
3332 _ => unreachable!(),
3333 };
3334 match *cast_to {
3335 crate::TypeInner::Scalar(_) => {
3336 put_numeric_type(&mut self.out, unsigned_scalar, &[])?
3337 }
3338 crate::TypeInner::Vector { size, .. } => {
3339 put_numeric_type(&mut self.out, unsigned_scalar, &[size])?
3340 }
3341 _ => unreachable!(),
3342 };
3343 write!(self.out, ">(")?;
3344 }
3345
3346 if let Some(scalar) = context.get_packed_vec_kind(inner_expr) {
3348 put_numeric_type(&mut self.out, scalar, &[crate::VectorSize::Tri])?;
3349 write!(self.out, "(")?;
3350 put_expression(self, context, true)?;
3351 write!(self.out, ")")?;
3352 } else {
3353 put_expression(self, context, true)?;
3354 }
3355
3356 if needs_truncation {
3357 write!(self.out, ")")?;
3358 }
3359
3360 write!(self.out, ")")?;
3361 Ok(())
3362 }
3363
3364 fn put_index(
3366 &mut self,
3367 index: index::GuardedIndex,
3368 context: &ExpressionContext,
3369 is_scoped: bool,
3370 ) -> BackendResult {
3371 match index {
3372 index::GuardedIndex::Expression(expr) => {
3373 self.put_expression(expr, context, is_scoped)?
3374 }
3375 index::GuardedIndex::Known(value) => write!(self.out, "{value}")?,
3376 }
3377 Ok(())
3378 }
3379
3380 fn put_bounds_checks(
3410 &mut self,
3411 chain: Handle<crate::Expression>,
3412 context: &ExpressionContext,
3413 level: back::Level,
3414 prefix: &'static str,
3415 ) -> Result<bool, Error> {
3416 let mut check_written = false;
3417
3418 for item in context.bounds_check_iter(chain) {
3420 let BoundsCheck {
3421 base,
3422 index,
3423 length,
3424 } = item;
3425
3426 if check_written {
3427 write!(self.out, " && ")?;
3428 } else {
3429 write!(self.out, "{level}{prefix}")?;
3430 check_written = true;
3431 }
3432
3433 write!(self.out, "uint(")?;
3437 self.put_index(index, context, true)?;
3438 self.out.write_str(") < ")?;
3439 match length {
3440 index::IndexableLength::Known(value) => write!(self.out, "{value}")?,
3441 index::IndexableLength::Dynamic => {
3442 let global = context.function.originating_global(base).ok_or_else(|| {
3443 Error::GenericValidation("Could not find originating global".into())
3444 })?;
3445 if matches!(
3446 context.module.types[context.module.global_variables[global].ty].inner,
3447 crate::TypeInner::BindingArray { .. }
3448 ) {
3449 write!(
3450 self.out,
3451 "{} && _buffer_sizes.{}[",
3452 Self::binding_array_layout_count(
3453 context.module,
3454 context.pipeline_options,
3455 global,
3456 ),
3457 ArraySizeMember(global),
3458 )?;
3459 self.put_binding_array_size_member_index(index, context)?;
3460 write!(self.out, "] != 0u")?;
3461 } else {
3462 write!(self.out, "1 + ")?;
3463 self.put_dynamic_array_max_index(global, base, context)?
3464 }
3465 }
3466 }
3467 }
3468
3469 Ok(check_written)
3470 }
3471
3472 fn put_access_chain(
3492 &mut self,
3493 chain: Handle<crate::Expression>,
3494 policy: index::BoundsCheckPolicy,
3495 context: &ExpressionContext,
3496 ) -> BackendResult {
3497 match context.function.expressions[chain] {
3498 crate::Expression::Access { base, index } => {
3499 let mut base_ty = context.resolve_type(base);
3500
3501 if let crate::TypeInner::Pointer { base, space: _ } = *base_ty {
3503 base_ty = &context.module.types[base].inner;
3504 }
3505
3506 self.put_subscripted_access_chain(
3507 base,
3508 base_ty,
3509 index::GuardedIndex::Expression(index),
3510 policy,
3511 context,
3512 )?;
3513 }
3514 crate::Expression::AccessIndex { base, index } => {
3515 let base_resolution = &context.info[base].ty;
3516 let mut base_ty = base_resolution.inner_with(&context.module.types);
3517 let mut base_ty_handle = base_resolution.handle();
3518
3519 if let crate::TypeInner::Pointer { base, space: _ } = *base_ty {
3521 base_ty = &context.module.types[base].inner;
3522 base_ty_handle = Some(base);
3523 }
3524
3525 match *base_ty {
3529 crate::TypeInner::Struct { .. } => {
3530 let base_ty = base_ty_handle.unwrap();
3531 self.put_access_chain(base, policy, context)?;
3532 let name = &self.names[&NameKey::StructMember(base_ty, index)];
3533 write!(
3534 self.out,
3535 "{}{name}",
3536 if context.struct_member_needs_arrow(base, |ty| {
3537 matches!(ty, crate::TypeInner::BindingArray { .. })
3538 }) {
3539 "->"
3540 } else {
3541 "."
3542 },
3543 )?;
3544 }
3545 crate::TypeInner::ValuePointer { .. } | crate::TypeInner::Vector { .. } => {
3546 self.put_access_chain(base, policy, context)?;
3547 if context.get_packed_vec_kind(base).is_some() {
3550 write!(self.out, "[{index}]")?;
3551 } else {
3552 write!(self.out, ".{}", back::COMPONENTS[index as usize])?;
3553 }
3554 }
3555 _ => {
3556 self.put_subscripted_access_chain(
3557 base,
3558 base_ty,
3559 index::GuardedIndex::Known(index),
3560 policy,
3561 context,
3562 )?;
3563 }
3564 }
3565 }
3566 _ => self.put_expression(chain, context, false)?,
3567 }
3568
3569 Ok(())
3570 }
3571
3572 fn put_subscripted_access_chain(
3589 &mut self,
3590 base: Handle<crate::Expression>,
3591 base_ty: &crate::TypeInner,
3592 index: index::GuardedIndex,
3593 policy: index::BoundsCheckPolicy,
3594 context: &ExpressionContext,
3595 ) -> BackendResult {
3596 let accessing_wrapped_array = match *base_ty {
3597 crate::TypeInner::Array {
3598 size: crate::ArraySize::Constant(_) | crate::ArraySize::Pending(_),
3599 ..
3600 } => true,
3601 _ => false,
3602 };
3603 let accessing_wrapped_binding_array =
3604 matches!(*base_ty, crate::TypeInner::BindingArray { .. });
3605
3606 self.put_access_chain(base, policy, context)?;
3607 if accessing_wrapped_array {
3608 write!(self.out, ".{WRAPPED_ARRAY_FIELD}")?;
3609 }
3610 write!(self.out, "[")?;
3611
3612 let restriction_needed = if policy == index::BoundsCheckPolicy::Restrict {
3614 context.access_needs_check(base, index)
3615 } else {
3616 None
3617 };
3618 if let Some(limit) = restriction_needed {
3619 write!(self.out, "{NAMESPACE}::min(unsigned(")?;
3620 self.put_index(index, context, true)?;
3621 write!(self.out, "), ")?;
3622 match limit {
3623 index::IndexableLength::Known(limit) => {
3624 write!(self.out, "{}u", limit - 1)?;
3625 }
3626 index::IndexableLength::Dynamic => {
3627 let global = context.function.originating_global(base).ok_or_else(|| {
3628 Error::GenericValidation("Could not find originating global".into())
3629 })?;
3630 self.put_dynamic_array_max_index(global, base, context)?;
3631 }
3632 }
3633 write!(self.out, ")")?;
3634 } else {
3635 self.put_index(index, context, true)?;
3636 }
3637
3638 write!(self.out, "]")?;
3639
3640 if accessing_wrapped_binding_array {
3641 write!(self.out, ".{WRAPPED_ARRAY_FIELD}")?;
3642 }
3643
3644 Ok(())
3645 }
3646
3647 fn put_load(
3648 &mut self,
3649 pointer: Handle<crate::Expression>,
3650 context: &ExpressionContext,
3651 is_scoped: bool,
3652 ) -> BackendResult {
3653 let policy = context.choose_bounds_check_policy(pointer);
3656 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
3657 && self.put_bounds_checks(
3658 pointer,
3659 context,
3660 back::Level(0),
3661 if is_scoped { "" } else { "(" },
3662 )?
3663 {
3664 write!(self.out, " ? ")?;
3665 self.put_unchecked_load(pointer, policy, context)?;
3666 write!(self.out, " : DefaultConstructible()")?;
3667
3668 if !is_scoped {
3669 write!(self.out, ")")?;
3670 }
3671 } else {
3672 self.put_unchecked_load(pointer, policy, context)?;
3673 }
3674
3675 Ok(())
3676 }
3677
3678 fn put_unchecked_load(
3679 &mut self,
3680 pointer: Handle<crate::Expression>,
3681 policy: index::BoundsCheckPolicy,
3682 context: &ExpressionContext,
3683 ) -> BackendResult {
3684 let is_atomic_pointer = context
3685 .resolve_type(pointer)
3686 .is_atomic_pointer(&context.module.types);
3687
3688 if is_atomic_pointer {
3689 write!(
3690 self.out,
3691 "{NAMESPACE}::atomic_load_explicit({ATOMIC_REFERENCE}"
3692 )?;
3693 self.put_access_chain(pointer, policy, context)?;
3694 write!(self.out, ", {NAMESPACE}::memory_order_relaxed)")?;
3695 } else {
3696 self.put_access_chain(pointer, policy, context)?;
3700 }
3701
3702 Ok(())
3703 }
3704
3705 fn put_return_value(
3706 &mut self,
3707 level: back::Level,
3708 expr_handle: Handle<crate::Expression>,
3709 result_struct: Option<&str>,
3710 context: &ExpressionContext,
3711 ) -> BackendResult {
3712 let entry_point = match context.origin {
3713 FunctionOrigin::EntryPoint(index) => Some(&context.module.entry_points[index as usize]),
3714 FunctionOrigin::Handle(_) => None,
3715 };
3716 if entry_point.is_some_and(|ep| only_writes_discarded_depth(context.module, ep)) {
3717 writeln!(self.out, "{level}return;")?;
3718 return Ok(());
3719 }
3720
3721 match result_struct {
3722 Some(struct_name) => {
3723 let mut has_point_size = false;
3724 let result_ty = context.function.result.as_ref().unwrap().ty;
3725 match context.module.types[result_ty].inner {
3726 crate::TypeInner::Struct { ref members, .. } => {
3727 let tmp = self.namer.call("_tmp");
3728 write!(self.out, "{level}const auto {tmp} = ")?;
3729 self.put_expression(expr_handle, context, true)?;
3730 writeln!(self.out, ";")?;
3731 write!(self.out, "{level}return {struct_name} {{")?;
3732
3733 let mut is_first = true;
3734
3735 for (index, member) in members.iter().enumerate() {
3736 if let Some(crate::Binding::BuiltIn(crate::BuiltIn::PointSize)) =
3737 member.binding
3738 {
3739 has_point_size = true;
3740 if !context.pipeline_options.allow_and_force_point_size {
3741 continue;
3742 }
3743 }
3744
3745 if entry_point.is_some_and(discards_frag_depth)
3746 && is_frag_depth(member.binding.as_ref())
3747 {
3748 continue;
3749 }
3750
3751 let comma = if is_first { "" } else { "," };
3752 is_first = false;
3753 let name = &self.names[&NameKey::StructMember(result_ty, index as u32)];
3754 if let crate::TypeInner::Array {
3758 size: crate::ArraySize::Constant(size),
3759 ..
3760 } = context.module.types[member.ty].inner
3761 {
3762 write!(self.out, "{comma} {{")?;
3763 for j in 0..size.get() {
3764 if j != 0 {
3765 write!(self.out, ",")?;
3766 }
3767 write!(self.out, "{tmp}.{name}.{WRAPPED_ARRAY_FIELD}[{j}]")?;
3768 }
3769 write!(self.out, "}}")?;
3770 } else {
3771 write!(self.out, "{comma} {tmp}.{name}")?;
3772 }
3773 }
3774 }
3775 _ => {
3776 write!(self.out, "{level}return {struct_name} {{ ")?;
3777 self.put_expression(expr_handle, context, true)?;
3778 }
3779 }
3780
3781 if let FunctionOrigin::EntryPoint(ep_index) = context.origin {
3782 let stage = context.module.entry_points[ep_index as usize].stage;
3783 if context.pipeline_options.allow_and_force_point_size
3784 && stage == crate::ShaderStage::Vertex
3785 && !has_point_size
3786 {
3787 write!(self.out, ", 1.0")?;
3789 }
3790 }
3791 write!(self.out, " }}")?;
3792 }
3793 None => {
3794 write!(self.out, "{level}return ")?;
3795 self.put_expression(expr_handle, context, true)?;
3796 }
3797 }
3798 writeln!(self.out, ";")?;
3799 Ok(())
3800 }
3801
3802 fn update_expressions_to_bake(
3807 &mut self,
3808 func: &crate::Function,
3809 info: &valid::FunctionInfo,
3810 context: &ExpressionContext,
3811 ) {
3812 use crate::Expression;
3813 self.need_bake_expressions.clear();
3814
3815 for (expr_handle, expr) in func.expressions.iter() {
3816 let expr_info = &info[expr_handle];
3819 let min_ref_count = func.expressions[expr_handle].bake_ref_count();
3820 if min_ref_count <= expr_info.ref_count {
3821 self.need_bake_expressions.insert(expr_handle);
3822 } else {
3823 match expr_info.ty {
3824 TypeResolution::Handle(h)
3826 if Some(h) == context.module.special_types.ray_desc =>
3827 {
3828 self.need_bake_expressions.insert(expr_handle);
3829 }
3830 _ => {}
3831 }
3832 }
3833
3834 if let Expression::Math {
3835 fun,
3836 arg,
3837 arg1,
3838 arg2,
3839 ..
3840 } = *expr
3841 {
3842 match fun {
3843 crate::MathFunction::Dot4U8Packed | crate::MathFunction::Dot4I8Packed => {
3853 self.need_bake_expressions.insert(arg);
3854 self.need_bake_expressions.insert(arg1.unwrap());
3855 }
3856 crate::MathFunction::FirstLeadingBit => {
3857 self.need_bake_expressions.insert(arg);
3858 }
3859 crate::MathFunction::Pack4xI8
3862 | crate::MathFunction::Pack4xU8
3863 | crate::MathFunction::Pack4xI8Clamp
3864 | crate::MathFunction::Pack4xU8Clamp
3865 | crate::MathFunction::Unpack4xI8
3866 | crate::MathFunction::Unpack4xU8
3867 if context.lang_version < (2, 1) =>
3868 {
3869 self.need_bake_expressions.insert(arg);
3870 }
3871 crate::MathFunction::ExtractBits => {
3872 self.need_bake_expressions.insert(arg1.unwrap());
3874 }
3875 crate::MathFunction::InsertBits => {
3876 self.need_bake_expressions.insert(arg2.unwrap());
3878 }
3879 crate::MathFunction::Sign => {
3880 let inner = context.resolve_type(expr_handle);
3885 if inner.scalar_kind() == Some(crate::ScalarKind::Sint) {
3886 self.need_bake_expressions.insert(arg);
3887 }
3888 }
3889 _ => {}
3890 }
3891 }
3892 }
3893 }
3894
3895 pub(super) fn put_baked_type(
3898 &mut self,
3899 handle: Handle<crate::Expression>,
3900 context: &ExpressionContext,
3901 ) -> BackendResult {
3902 match context.info[handle].ty {
3903 TypeResolution::Handle(ty_handle) => {
3904 let ty_name = TypeContext {
3905 handle: ty_handle,
3906 gctx: context.module.to_ctx(),
3907 names: &self.names,
3908 access: crate::StorageAccess::empty(),
3909 first_time: false,
3910 };
3911 write!(self.out, "{ty_name}")?;
3912 }
3913 TypeResolution::Value(crate::TypeInner::Scalar(scalar)) => {
3914 put_numeric_type(&mut self.out, scalar, &[])?;
3915 }
3916 TypeResolution::Value(crate::TypeInner::Vector { size, scalar }) => {
3917 put_numeric_type(&mut self.out, scalar, &[size])?;
3918 }
3919 TypeResolution::Value(crate::TypeInner::Matrix {
3920 columns,
3921 rows,
3922 scalar,
3923 }) => {
3924 put_numeric_type(&mut self.out, scalar, &[rows, columns])?;
3925 }
3926 TypeResolution::Value(crate::TypeInner::CooperativeMatrix {
3927 columns,
3928 rows,
3929 scalar,
3930 role: _,
3931 }) => {
3932 write!(
3933 self.out,
3934 "{}::simdgroup_{}{}x{}",
3935 NAMESPACE,
3936 scalar.to_msl_name(),
3937 columns as u32,
3938 rows as u32,
3939 )?;
3940 }
3941 TypeResolution::Value(ref other) => {
3942 log::warn!("Type {other:?} isn't a known local");
3943 return Err(Error::FeatureNotImplemented("weird local type".to_string()));
3944 }
3945 }
3946
3947 Ok(())
3948 }
3949
3950 pub(super) fn start_baking_expression(
3957 &mut self,
3958 handle: Handle<crate::Expression>,
3959 context: &ExpressionContext,
3960 name: &str,
3961 ) -> BackendResult {
3962 self.put_baked_type(handle, context)?;
3963 write!(self.out, " {name} = ")?;
3964
3965 Ok(())
3966 }
3967
3968 pub(super) fn write_result_assignment(
3979 &mut self,
3980 result: Handle<crate::Expression>,
3981 context: &ExpressionContext,
3982 ) -> Result<String, Error> {
3983 let name = match self.hoisted_loop_expressions.remove(&result) {
3984 Some(hoisted) => {
3985 write!(self.out, "{hoisted} = ")?;
3986 hoisted
3987 }
3988 None => {
3989 let name = Baked(result).to_string();
3990 self.start_baking_expression(result, context, &name)?;
3991 name
3992 }
3993 };
3994 self.named_expressions.insert(result, name.clone());
3995 Ok(name)
3996 }
3997
3998 fn put_cache_restricted_level(
4011 &mut self,
4012 load: Handle<crate::Expression>,
4013 image: Handle<crate::Expression>,
4014 mip_level: Option<Handle<crate::Expression>>,
4015 indent: back::Level,
4016 context: &StatementContext,
4017 ) -> BackendResult {
4018 let level_of_detail = match mip_level {
4021 Some(level) => level,
4022 None => return Ok(()),
4023 };
4024
4025 if context.expression.policies.image_load != index::BoundsCheckPolicy::Restrict
4026 || !context.expression.image_needs_lod(image)
4027 {
4028 return Ok(());
4029 }
4030
4031 write!(self.out, "{}uint {} = ", indent, ClampedLod(load),)?;
4032 self.put_restricted_scalar_image_index(
4033 image,
4034 level_of_detail,
4035 "get_num_mip_levels",
4036 &context.expression,
4037 )?;
4038 writeln!(self.out, ";")?;
4039
4040 Ok(())
4041 }
4042
4043 fn put_casting_to_packed_chars(
4049 &mut self,
4050 fun: crate::MathFunction,
4051 arg0: Handle<crate::Expression>,
4052 arg1: Handle<crate::Expression>,
4053 indent: back::Level,
4054 context: &StatementContext<'_>,
4055 ) -> Result<(), Error> {
4056 let packed_type = match fun {
4057 crate::MathFunction::Dot4I8Packed => "packed_char4",
4058 crate::MathFunction::Dot4U8Packed => "packed_uchar4",
4059 _ => unreachable!(),
4060 };
4061
4062 for arg in [arg0, arg1] {
4063 write!(
4064 self.out,
4065 "{indent}{packed_type} {0} = as_type<{packed_type}>(",
4066 Reinterpreted::new(packed_type, arg)
4067 )?;
4068 self.put_expression(arg, &context.expression, true)?;
4069 writeln!(self.out, ");")?;
4070 }
4071
4072 Ok(())
4073 }
4074
4075 pub(super) fn bakes_expression(
4077 &self,
4078 handle: Handle<crate::Expression>,
4079 context: &ExpressionContext,
4080 ) -> bool {
4081 if super::hoist::is_statement_result(&context.function.expressions[handle]) {
4084 return true;
4085 }
4086
4087 if context.guarded_indices.contains(handle) {
4091 return true;
4092 }
4093
4094 if self.need_bake_expressions.contains(&handle) {
4095 return true;
4096 }
4097
4098 if context.resolve_type(handle).pointer_space().is_some() {
4100 return false;
4101 }
4102
4103 context.function.named_expressions.contains_key(&handle)
4106 }
4107
4108 pub(super) fn baked_name(
4113 &mut self,
4114 handle: Handle<crate::Expression>,
4115 context: &ExpressionContext,
4116 ) -> Option<String> {
4117 if !self.bakes_expression(handle, context) {
4118 return None;
4119 }
4120 Some(match context.function.named_expressions.get(&handle) {
4121 Some(name) => self.namer.call(name),
4124 None => Baked(handle).to_string(),
4125 })
4126 }
4127
4128 fn put_block(
4129 &mut self,
4130 level: back::Level,
4131 statements: &[crate::Statement],
4132 context: &StatementContext,
4133 ) -> BackendResult {
4134 for statement in statements {
4135 log::trace!("statement[{}] {:?}", level.0, statement);
4136 match *statement {
4137 crate::Statement::Emit(ref range) => {
4138 for handle in range.clone() {
4139 use crate::MathFunction as Mf;
4140
4141 match context.expression.function.expressions[handle] {
4142 crate::Expression::ImageLoad {
4145 image,
4146 level: mip_level,
4147 ..
4148 } => {
4149 self.put_cache_restricted_level(
4150 handle, image, mip_level, level, context,
4151 )?;
4152 }
4153
4154 crate::Expression::Math {
4163 fun: fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed),
4164 arg,
4165 arg1,
4166 ..
4167 } if context.expression.lang_version >= (2, 1) => {
4168 self.put_casting_to_packed_chars(
4169 fun,
4170 arg,
4171 arg1.unwrap(),
4172 level,
4173 context,
4174 )?;
4175 }
4176
4177 _ => (),
4178 }
4179
4180 let hoisted = self.hoisted_loop_expressions.remove(&handle);
4184 let declare = hoisted.is_none();
4185 let expr_name =
4186 hoisted.or_else(|| self.baked_name(handle, &context.expression));
4187
4188 if let Some(name) = expr_name {
4189 write!(self.out, "{level}")?;
4190 if declare {
4191 self.start_baking_expression(handle, &context.expression, &name)?;
4192 } else {
4193 write!(self.out, "{name} = ")?;
4194 }
4195 self.put_expression(handle, &context.expression, true)?;
4196 self.named_expressions.insert(handle, name);
4199 writeln!(self.out, ";")?;
4200 }
4201 }
4202 }
4203 crate::Statement::Block(ref block) => {
4204 if !block.is_empty() {
4205 writeln!(self.out, "{level}{{")?;
4206 self.put_block(level.next(), block, context)?;
4207 writeln!(self.out, "{level}}}")?;
4208 }
4209 }
4210 crate::Statement::If {
4211 condition,
4212 ref accept,
4213 ref reject,
4214 } => {
4215 write!(self.out, "{level}if (")?;
4216 self.put_expression(condition, &context.expression, true)?;
4217 writeln!(self.out, ") {{")?;
4218 self.put_block(level.next(), accept, context)?;
4219 if !reject.is_empty() {
4220 writeln!(self.out, "{level}}} else {{")?;
4221 self.put_block(level.next(), reject, context)?;
4222 }
4223 writeln!(self.out, "{level}}}")?;
4224 }
4225 crate::Statement::Switch {
4226 selector,
4227 ref cases,
4228 } => {
4229 write!(self.out, "{level}switch(")?;
4230 self.put_expression(selector, &context.expression, true)?;
4231 writeln!(self.out, ") {{")?;
4232 let lcase = level.next();
4233 for case in cases.iter() {
4234 match case.value {
4235 crate::SwitchValue::I32(value) => {
4236 write!(self.out, "{lcase}case {value}:")?;
4237 }
4238 crate::SwitchValue::U32(value) => {
4239 write!(self.out, "{lcase}case {value}u:")?;
4240 }
4241 crate::SwitchValue::Default => {
4242 write!(self.out, "{lcase}default:")?;
4243 }
4244 }
4245
4246 let write_block_braces = !(case.fall_through && case.body.is_empty());
4247 if write_block_braces {
4248 writeln!(self.out, " {{")?;
4249 } else {
4250 writeln!(self.out)?;
4251 }
4252
4253 self.put_block(lcase.next(), &case.body, context)?;
4254 if !case.fall_through && case.body.last().is_none_or(|s| !s.is_terminator())
4255 {
4256 writeln!(self.out, "{}break;", lcase.next())?;
4257 }
4258
4259 if write_block_braces {
4260 writeln!(self.out, "{lcase}}}")?;
4261 }
4262 }
4263 writeln!(self.out, "{level}}}")?;
4264 }
4265 crate::Statement::Loop {
4266 ref body,
4267 ref continuing,
4268 break_if,
4269 } => {
4270 let force_loop_bound_statements =
4271 self.gen_force_bounded_loop_statements(level, context);
4272 let gate_name = (!continuing.is_empty() || break_if.is_some())
4273 .then(|| self.namer.call("loop_init"));
4274
4275 if gate_name.is_some() {
4278 self.hoist_continuing_dependencies(
4279 level, body, continuing, break_if, context,
4280 )?;
4281 }
4282
4283 if let Some((ref decl, _)) = force_loop_bound_statements {
4284 writeln!(self.out, "{decl}")?;
4285 }
4286 if let Some(ref gate_name) = gate_name {
4287 writeln!(self.out, "{level}bool {gate_name} = true;")?;
4288 }
4289
4290 writeln!(self.out, "{level}while(true) {{",)?;
4291 if let Some((_, ref break_and_inc)) = force_loop_bound_statements {
4292 writeln!(self.out, "{break_and_inc}")?;
4293 }
4294 if let Some(ref gate_name) = gate_name {
4295 let lif = level.next();
4296 let lcontinuing = lif.next();
4297 writeln!(self.out, "{lif}if (!{gate_name}) {{")?;
4298 self.put_block(lcontinuing, continuing, context)?;
4299 if let Some(condition) = break_if {
4300 write!(self.out, "{lcontinuing}if (")?;
4301 self.put_expression(condition, &context.expression, true)?;
4302 writeln!(self.out, ") {{")?;
4303 writeln!(self.out, "{}break;", lcontinuing.next())?;
4304 writeln!(self.out, "{lcontinuing}}}")?;
4305 }
4306 writeln!(self.out, "{lif}}}")?;
4307 writeln!(self.out, "{lif}{gate_name} = false;")?;
4308 }
4309 self.put_block(level.next(), body, context)?;
4310
4311 writeln!(self.out, "{level}}}")?;
4312 }
4313 crate::Statement::Break => {
4314 writeln!(self.out, "{level}break;")?;
4315 }
4316 crate::Statement::Continue => {
4317 writeln!(self.out, "{level}continue;")?;
4318 }
4319 crate::Statement::Return {
4320 value: Some(expr_handle),
4321 } => {
4322 self.put_return_value(
4323 level,
4324 expr_handle,
4325 context.result_struct,
4326 &context.expression,
4327 )?;
4328 }
4329 crate::Statement::Return { value: None } => {
4330 writeln!(self.out, "{level}return;")?;
4331 }
4332 crate::Statement::Kill => {
4333 writeln!(self.out, "{level}{NAMESPACE}::discard_fragment();")?;
4334 }
4335 crate::Statement::ControlBarrier(flags)
4336 | crate::Statement::MemoryBarrier(flags) => {
4337 self.write_barrier(flags, level)?;
4338 }
4339 crate::Statement::Store { pointer, value } => {
4340 self.put_store(pointer, value, level, context)?
4341 }
4342 crate::Statement::ImageStore {
4343 image,
4344 coordinate,
4345 array_index,
4346 value,
4347 } => {
4348 let address = TexelAddress {
4349 coordinate,
4350 array_index,
4351 sample: None,
4352 level: None,
4353 };
4354 self.put_image_store(level, image, &address, value, context)?
4355 }
4356 crate::Statement::Call {
4357 function,
4358 ref arguments,
4359 result,
4360 } => {
4361 write!(self.out, "{level}")?;
4362 if let Some(expr) = result {
4363 self.write_result_assignment(expr, &context.expression)?;
4364 }
4365 let fun_name = &self.names[&NameKey::Function(function)];
4366 write!(self.out, "{fun_name}(")?;
4367 for (i, &handle) in arguments.iter().enumerate() {
4369 if i != 0 {
4370 write!(self.out, ", ")?;
4371 }
4372 self.put_expression(handle, &context.expression, true)?;
4373 }
4374 let mut separate = !arguments.is_empty();
4376 let fun_info = &context.expression.mod_info[function];
4377 let mut needs_buffer_sizes = false;
4378 for (handle, var) in context.expression.module.global_variables.iter() {
4379 if fun_info[handle].is_empty() {
4380 continue;
4381 }
4382 if var.space.needs_pass_through() {
4383 let name = &self.names[&NameKey::GlobalVariable(handle)];
4384 if separate {
4385 write!(self.out, ", ")?;
4386 } else {
4387 separate = true;
4388 }
4389 write!(self.out, "{name}")?;
4390 }
4391 needs_buffer_sizes |= context.expression.module.types[var.ty]
4392 .inner
4393 .needs_host_buffer_byte_size(&context.expression.module.types);
4394 }
4395 if needs_buffer_sizes {
4396 if separate {
4397 write!(self.out, ", ")?;
4398 }
4399 write!(self.out, "_buffer_sizes")?;
4400 }
4401
4402 writeln!(self.out, ");")?;
4404 }
4405 crate::Statement::Atomic {
4406 pointer,
4407 ref fun,
4408 value,
4409 result,
4410 } => {
4411 let context = &context.expression;
4412
4413 write!(self.out, "{level}")?;
4418 let fun_key = if let Some(result) = result {
4419 self.write_result_assignment(result, context)?;
4420 fun.to_msl()
4421 } else if context.resolve_type(value).scalar_width() == Some(8) {
4422 fun.to_msl_64_bit()?
4423 } else {
4424 fun.to_msl()
4425 };
4426
4427 let policy = context.choose_bounds_check_policy(pointer);
4431 let checked = policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
4432 && self.put_bounds_checks(pointer, context, back::Level(0), "")?;
4433
4434 if checked {
4436 write!(self.out, " ? ")?;
4437 }
4438
4439 match *fun {
4441 crate::AtomicFunction::Exchange { compare: Some(cmp) } => {
4442 write!(self.out, "{ATOMIC_COMP_EXCH_FUNCTION}({ATOMIC_REFERENCE}")?;
4443 self.put_access_chain(pointer, policy, context)?;
4444 write!(self.out, ", ")?;
4445 self.put_expression(cmp, context, true)?;
4446 write!(self.out, ", ")?;
4447 self.put_expression(value, context, true)?;
4448 write!(self.out, ")")?;
4449 }
4450 _ => {
4451 write!(
4452 self.out,
4453 "{NAMESPACE}::atomic_{fun_key}_explicit({ATOMIC_REFERENCE}"
4454 )?;
4455 self.put_access_chain(pointer, policy, context)?;
4456 write!(self.out, ", ")?;
4457 self.put_expression(value, context, true)?;
4458 write!(self.out, ", {NAMESPACE}::memory_order_relaxed)")?;
4459 }
4460 }
4461
4462 if checked {
4464 write!(self.out, " : DefaultConstructible()")?;
4465 }
4466
4467 writeln!(self.out, ";")?;
4469 }
4470 crate::Statement::ImageAtomic {
4471 image,
4472 coordinate,
4473 array_index,
4474 fun,
4475 value,
4476 } => {
4477 let address = TexelAddress {
4478 coordinate,
4479 array_index,
4480 sample: None,
4481 level: None,
4482 };
4483 self.put_image_atomic(level, image, &address, fun, value, context)?
4484 }
4485 crate::Statement::WorkGroupUniformLoad { pointer, result } => {
4486 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
4487
4488 write!(self.out, "{level}")?;
4489 self.write_result_assignment(result, &context.expression)?;
4490 self.put_load(pointer, &context.expression, true)?;
4491
4492 writeln!(self.out, ";")?;
4493 self.write_barrier(crate::Barrier::WORK_GROUP, level)?;
4494 }
4495 crate::Statement::RayQuery { query, ref fun } => {
4496 self.write_ray_query_stmt(level, context, query, fun)?;
4497 }
4498 crate::Statement::SubgroupBallot { result, predicate } => {
4499 write!(self.out, "{level}")?;
4500 self.write_result_assignment(result, &context.expression)?;
4501 write!(
4502 self.out,
4503 "{NAMESPACE}::uint4((uint64_t){NAMESPACE}::simd_ballot("
4504 )?;
4505 if let Some(predicate) = predicate {
4506 self.put_expression(predicate, &context.expression, true)?;
4507 } else {
4508 write!(self.out, "true")?;
4509 }
4510 writeln!(self.out, "), 0, 0, 0);")?;
4511 }
4512 crate::Statement::SubgroupCollectiveOperation {
4513 op,
4514 collective_op,
4515 argument,
4516 result,
4517 } => {
4518 write!(self.out, "{level}")?;
4519 self.write_result_assignment(result, &context.expression)?;
4520 match (collective_op, op) {
4521 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::All) => {
4522 write!(self.out, "{NAMESPACE}::simd_all(")?
4523 }
4524 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Any) => {
4525 write!(self.out, "{NAMESPACE}::simd_any(")?
4526 }
4527 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Add) => {
4528 write!(self.out, "{NAMESPACE}::simd_sum(")?
4529 }
4530 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Mul) => {
4531 write!(self.out, "{NAMESPACE}::simd_product(")?
4532 }
4533 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Max) => {
4534 write!(self.out, "{NAMESPACE}::simd_max(")?
4535 }
4536 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Min) => {
4537 write!(self.out, "{NAMESPACE}::simd_min(")?
4538 }
4539 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::And) => {
4540 write!(self.out, "{NAMESPACE}::simd_and(")?
4541 }
4542 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Or) => {
4543 write!(self.out, "{NAMESPACE}::simd_or(")?
4544 }
4545 (crate::CollectiveOperation::Reduce, crate::SubgroupOperation::Xor) => {
4546 write!(self.out, "{NAMESPACE}::simd_xor(")?
4547 }
4548 (
4549 crate::CollectiveOperation::ExclusiveScan,
4550 crate::SubgroupOperation::Add,
4551 ) => write!(self.out, "{NAMESPACE}::simd_prefix_exclusive_sum(")?,
4552 (
4553 crate::CollectiveOperation::ExclusiveScan,
4554 crate::SubgroupOperation::Mul,
4555 ) => write!(self.out, "{NAMESPACE}::simd_prefix_exclusive_product(")?,
4556 (
4557 crate::CollectiveOperation::InclusiveScan,
4558 crate::SubgroupOperation::Add,
4559 ) => write!(self.out, "{NAMESPACE}::simd_prefix_inclusive_sum(")?,
4560 (
4561 crate::CollectiveOperation::InclusiveScan,
4562 crate::SubgroupOperation::Mul,
4563 ) => write!(self.out, "{NAMESPACE}::simd_prefix_inclusive_product(")?,
4564 _ => unimplemented!(),
4565 }
4566 self.put_expression(argument, &context.expression, true)?;
4567 writeln!(self.out, ");")?;
4568 }
4569 crate::Statement::SubgroupGather {
4570 mode,
4571 argument,
4572 result,
4573 } => {
4574 write!(self.out, "{level}")?;
4575 self.write_result_assignment(result, &context.expression)?;
4576 match mode {
4577 crate::GatherMode::BroadcastFirst => {
4578 write!(self.out, "{NAMESPACE}::simd_broadcast_first(")?;
4579 }
4580 crate::GatherMode::Broadcast(_) => {
4581 write!(self.out, "{NAMESPACE}::simd_broadcast(")?;
4582 }
4583 crate::GatherMode::Shuffle(_) => {
4584 write!(self.out, "{NAMESPACE}::simd_shuffle(")?;
4585 }
4586 crate::GatherMode::ShuffleDown(_) => {
4587 write!(self.out, "{NAMESPACE}::simd_shuffle_down(")?;
4588 }
4589 crate::GatherMode::ShuffleUp(_) => {
4590 write!(self.out, "{NAMESPACE}::simd_shuffle_up(")?;
4591 }
4592 crate::GatherMode::ShuffleXor(_) => {
4593 write!(self.out, "{NAMESPACE}::simd_shuffle_xor(")?;
4594 }
4595 crate::GatherMode::QuadBroadcast(_) => {
4596 write!(self.out, "{NAMESPACE}::quad_broadcast(")?;
4597 }
4598 crate::GatherMode::QuadSwap(_) => {
4599 write!(self.out, "{NAMESPACE}::quad_shuffle_xor(")?;
4600 }
4601 }
4602 self.put_expression(argument, &context.expression, true)?;
4603 match mode {
4604 crate::GatherMode::BroadcastFirst => {}
4605 crate::GatherMode::Broadcast(index)
4606 | crate::GatherMode::Shuffle(index)
4607 | crate::GatherMode::ShuffleDown(index)
4608 | crate::GatherMode::ShuffleUp(index)
4609 | crate::GatherMode::ShuffleXor(index)
4610 | crate::GatherMode::QuadBroadcast(index) => {
4611 write!(self.out, ", ")?;
4612 self.put_expression(index, &context.expression, true)?;
4613 }
4614 crate::GatherMode::QuadSwap(direction) => {
4615 write!(self.out, ", ")?;
4616 match direction {
4617 crate::Direction::X => {
4618 write!(self.out, "1u")?;
4619 }
4620 crate::Direction::Y => {
4621 write!(self.out, "2u")?;
4622 }
4623 crate::Direction::Diagonal => {
4624 write!(self.out, "3u")?;
4625 }
4626 }
4627 }
4628 }
4629 writeln!(self.out, ");")?;
4630 }
4631 crate::Statement::CooperativeStore { target, ref data } => {
4632 write!(self.out, "{level}simdgroup_store(")?;
4633 self.put_expression(target, &context.expression, true)?;
4634 write!(self.out, ", &")?;
4635 self.put_access_chain(
4636 data.pointer,
4637 context.expression.policies.index,
4638 &context.expression,
4639 )?;
4640 write!(self.out, ", ")?;
4641 self.put_expression(data.stride, &context.expression, true)?;
4642 if !data.row_major {
4647 let matrix_origin = "0";
4648 let transpose = true;
4649 write!(self.out, ", {matrix_origin}, {transpose}")?;
4650 }
4651 writeln!(self.out, ");")?;
4652 }
4653 crate::Statement::RayPipelineFunction(_) => unreachable!(),
4654 crate::Statement::DebugPrintf {
4655 ref format,
4656 ref arguments,
4657 } => {
4658 if context.expression.lang_version < (3, 2) {
4660 return Err(Error::UnsupportedDebugPrintf);
4661 }
4662
4663 write!(
4664 self.out,
4665 "{level}metal::os_log_default.log_info(\"{}\"",
4666 format
4667 )?;
4668 for &arg in arguments {
4669 write!(self.out, ", ")?;
4670 self.put_expression(arg, &context.expression, true)?;
4671 }
4672 writeln!(self.out, ");")?;
4673 }
4674 }
4675 }
4676
4677 Ok(())
4678 }
4679
4680 fn put_store(
4681 &mut self,
4682 pointer: Handle<crate::Expression>,
4683 value: Handle<crate::Expression>,
4684 level: back::Level,
4685 context: &StatementContext,
4686 ) -> BackendResult {
4687 let policy = context.expression.choose_bounds_check_policy(pointer);
4688 if policy == index::BoundsCheckPolicy::ReadZeroSkipWrite
4689 && self.put_bounds_checks(pointer, &context.expression, level, "if (")?
4690 {
4691 writeln!(self.out, ") {{")?;
4692 self.put_unchecked_store(pointer, value, policy, level.next(), context)?;
4693 writeln!(self.out, "{level}}}")?;
4694 } else {
4695 self.put_unchecked_store(pointer, value, policy, level, context)?;
4696 }
4697
4698 Ok(())
4699 }
4700
4701 fn put_unchecked_store(
4702 &mut self,
4703 pointer: Handle<crate::Expression>,
4704 value: Handle<crate::Expression>,
4705 policy: index::BoundsCheckPolicy,
4706 level: back::Level,
4707 context: &StatementContext,
4708 ) -> BackendResult {
4709 let is_atomic_pointer = context
4710 .expression
4711 .resolve_type(pointer)
4712 .is_atomic_pointer(&context.expression.module.types);
4713
4714 if is_atomic_pointer {
4715 write!(
4716 self.out,
4717 "{level}{NAMESPACE}::atomic_store_explicit({ATOMIC_REFERENCE}"
4718 )?;
4719 self.put_access_chain(pointer, policy, &context.expression)?;
4720 write!(self.out, ", ")?;
4721 self.put_expression(value, &context.expression, true)?;
4722 writeln!(self.out, ", {NAMESPACE}::memory_order_relaxed);")?;
4723 } else {
4724 write!(self.out, "{level}")?;
4725 self.put_access_chain(pointer, policy, &context.expression)?;
4726 write!(self.out, " = ")?;
4727 self.put_expression(value, &context.expression, true)?;
4728 writeln!(self.out, ";")?;
4729 }
4730
4731 Ok(())
4732 }
4733
4734 pub fn write(
4735 &mut self,
4736 module: &crate::Module,
4737 info: &valid::ModuleInfo,
4738 options: &Options,
4739 pipeline_options: &PipelineOptions,
4740 ) -> Result<TranslationInfo, Error> {
4741 self.emit_int_div_checks = options.emit_int_div_checks;
4742 self.names.clear();
4743 self.namer.reset(
4744 module,
4745 &super::keywords::RESERVED_SET,
4746 proc::KeywordSet::empty(),
4747 proc::CaseInsensitiveKeywordSet::empty(),
4748 &[
4749 CLAMPED_LOD_LOAD_PREFIX,
4750 super::ray::INTERSECTION_FUNCTION_NAME,
4751 super::ray::RAY_QUERY_TRACKER_VARIABLE_PREFIX,
4752 super::ray::RAY_QUERY_T_MAX_TRACKER_VARIABLE_PREFIX,
4753 ],
4754 &mut self.names,
4755 );
4756 self.wrapped_functions.clear();
4757 self.struct_member_pads.clear();
4758
4759 writeln!(
4760 self.out,
4761 "// language: metal{}.{}",
4762 options.lang_version.0, options.lang_version.1
4763 )?;
4764 writeln!(self.out, "#include <metal_stdlib>")?;
4765 writeln!(self.out, "#include <simd/simd.h>")?;
4766 writeln!(self.out)?;
4767 writeln!(self.out, "using {NAMESPACE}::uint;")?;
4769
4770 if module.uses_mesh_shaders() && options.lang_version < (3, 0) {
4771 return Err(Error::UnsupportedMeshShader);
4772 }
4773 self.needs_object_memory_barriers = module
4774 .entry_points
4775 .iter()
4776 .any(|e| e.stage == crate::ShaderStage::Task && e.task_payload.is_some());
4777
4778 if module.special_types.ray_desc.is_some()
4779 || module.special_types.ray_intersection.is_some()
4780 {
4781 if options.lang_version < (2, 4) {
4782 return Err(Error::UnsupportedRayTracing);
4783 }
4784 }
4785
4786 if options
4787 .bounds_check_policies
4788 .contains(index::BoundsCheckPolicy::ReadZeroSkipWrite)
4789 {
4790 self.put_default_constructible()?;
4791 }
4792 writeln!(self.out)?;
4793
4794 {
4795 let globals: Vec<Handle<crate::GlobalVariable>> = module
4798 .global_variables
4799 .iter()
4800 .filter(|&(_, var)| {
4801 module.types[var.ty]
4802 .inner
4803 .needs_host_buffer_byte_size(&module.types)
4804 })
4805 .map(|(handle, _)| handle)
4806 .collect();
4807
4808 let mut buffer_indices = vec![];
4809 for vbm in &pipeline_options.vertex_buffer_mappings {
4810 buffer_indices.push(vbm.id);
4811 }
4812
4813 if !globals.is_empty() || !buffer_indices.is_empty() {
4814 writeln!(self.out, "struct _mslBufferSizes {{")?;
4815
4816 for global in globals {
4817 let var = &module.global_variables[global];
4818 let var_ty = var.ty;
4819 match module.types[var_ty].inner {
4820 crate::TypeInner::BindingArray { .. } => {
4821 let n =
4822 Self::binding_array_layout_count(module, pipeline_options, global);
4823 writeln!(
4824 self.out,
4825 "{}uint {}[{n}];",
4826 back::INDENT,
4827 ArraySizeMember(global),
4828 )?;
4829 }
4830 _ => writeln!(
4831 self.out,
4832 "{}uint {};",
4833 back::INDENT,
4834 ArraySizeMember(global)
4835 )?,
4836 }
4837 }
4838
4839 for idx in buffer_indices {
4840 writeln!(self.out, "{}uint buffer_size{};", back::INDENT, idx)?;
4841 }
4842
4843 writeln!(self.out, "}};")?;
4844 writeln!(self.out)?;
4845 }
4846 };
4847
4848 self.write_type_defs(module)?;
4849 self.write_global_constants(module, info)?;
4850 self.write_functions(module, info, options, pipeline_options)
4851 }
4852
4853 fn put_default_constructible(&mut self) -> BackendResult {
4866 let tab = back::INDENT;
4867 writeln!(self.out, "struct DefaultConstructible {{")?;
4868 writeln!(self.out, "{tab}template<typename T>")?;
4869 writeln!(self.out, "{tab}operator T() && {{")?;
4870 writeln!(self.out, "{tab}{tab}return T {{}};")?;
4871 writeln!(self.out, "{tab}}}")?;
4872 writeln!(self.out, "}};")?;
4873 Ok(())
4874 }
4875
4876 fn write_type_defs(&mut self, module: &crate::Module) -> BackendResult {
4877 let mut generated_argument_buffer_wrapper = false;
4878 let mut generated_external_texture_wrapper = false;
4879 for (handle, ty) in module.types.iter() {
4880 match ty.inner {
4881 crate::TypeInner::BindingArray { .. } if !generated_argument_buffer_wrapper => {
4882 writeln!(self.out, "template <typename T>")?;
4883 writeln!(self.out, "struct {ARGUMENT_BUFFER_WRAPPER_STRUCT} {{")?;
4884 writeln!(self.out, "{}T {WRAPPED_ARRAY_FIELD};", back::INDENT)?;
4885 writeln!(self.out, "}};")?;
4886 generated_argument_buffer_wrapper = true;
4887 }
4888 crate::TypeInner::Image {
4889 class: crate::ImageClass::External,
4890 ..
4891 } if !generated_external_texture_wrapper => {
4892 let params_ty_name = &self.names
4893 [&NameKey::Type(module.special_types.external_texture_params.unwrap())];
4894 writeln!(self.out, "struct {EXTERNAL_TEXTURE_WRAPPER_STRUCT} {{")?;
4895 writeln!(
4896 self.out,
4897 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane0;",
4898 back::INDENT
4899 )?;
4900 writeln!(
4901 self.out,
4902 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane1;",
4903 back::INDENT
4904 )?;
4905 writeln!(
4906 self.out,
4907 "{}{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample> plane2;",
4908 back::INDENT
4909 )?;
4910 writeln!(self.out, "{}{params_ty_name} params;", back::INDENT)?;
4911 writeln!(self.out, "}};")?;
4912 generated_external_texture_wrapper = true;
4913 }
4914 _ => {}
4915 }
4916
4917 if !ty.needs_alias() {
4918 continue;
4919 }
4920 let name = &self.names[&NameKey::Type(handle)];
4921 match ty.inner {
4922 crate::TypeInner::Array {
4936 base,
4937 size,
4938 stride: _,
4939 } => {
4940 let base_name = TypeContext {
4941 handle: base,
4942 gctx: module.to_ctx(),
4943 names: &self.names,
4944 access: crate::StorageAccess::empty(),
4945 first_time: false,
4946 };
4947
4948 match size.resolve(module.to_ctx())? {
4949 proc::IndexableLength::Known(size) => {
4950 writeln!(self.out, "struct {name} {{")?;
4951 writeln!(
4952 self.out,
4953 "{}{} {}[{}];",
4954 back::INDENT,
4955 base_name,
4956 WRAPPED_ARRAY_FIELD,
4957 size
4958 )?;
4959 writeln!(self.out, "}};")?;
4960 }
4961 proc::IndexableLength::Dynamic => {
4962 writeln!(self.out, "typedef {base_name} {name}[1];")?;
4963 }
4964 }
4965 }
4966 crate::TypeInner::Struct {
4967 ref members, span, ..
4968 } => {
4969 writeln!(self.out, "struct {name} {{")?;
4970 let mut last_offset = 0;
4971 for (index, member) in members.iter().enumerate() {
4972 if member.offset > last_offset {
4973 self.struct_member_pads.insert((handle, index as u32));
4974 let pad = member.offset - last_offset;
4975 writeln!(self.out, "{}char _pad{}[{}];", back::INDENT, index, pad)?;
4976 }
4977 let ty_inner = &module.types[member.ty].inner;
4978 last_offset = member.offset + ty_inner.size(module.to_ctx());
4979
4980 let member_name = &self.names[&NameKey::StructMember(handle, index as u32)];
4981
4982 match should_pack_struct_member(members, span, index, module) {
4984 Some(scalar) => {
4985 writeln!(
4986 self.out,
4987 "{}{}::packed_{}3 {};",
4988 back::INDENT,
4989 NAMESPACE,
4990 scalar.to_msl_name(),
4991 member_name
4992 )?;
4993 }
4994 None => {
4995 let base_name = TypeContext {
4996 handle: member.ty,
4997 gctx: module.to_ctx(),
4998 names: &self.names,
4999 access: crate::StorageAccess::empty(),
5000 first_time: false,
5001 };
5002 writeln!(
5003 self.out,
5004 "{}{} {};",
5005 back::INDENT,
5006 base_name,
5007 member_name
5008 )?;
5009
5010 if let crate::TypeInner::Vector {
5012 size: crate::VectorSize::Tri,
5013 scalar,
5014 } = *ty_inner
5015 {
5016 last_offset += scalar.width as u32;
5017 }
5018 }
5019 }
5020 }
5021 if last_offset < span {
5022 let pad = span - last_offset;
5023 writeln!(
5024 self.out,
5025 "{}char _pad{}[{}];",
5026 back::INDENT,
5027 members.len(),
5028 pad
5029 )?;
5030 }
5031 writeln!(self.out, "}};")?;
5032 }
5033 _ => {
5034 let ty_name = TypeContext {
5035 handle,
5036 gctx: module.to_ctx(),
5037 names: &self.names,
5038 access: crate::StorageAccess::empty(),
5039 first_time: true,
5040 };
5041 writeln!(self.out, "typedef {ty_name} {name};")?;
5042 }
5043 }
5044 }
5045
5046 for (type_key, struct_ty) in module.special_types.predeclared_types.iter() {
5048 match type_key {
5049 &crate::PredeclaredType::ModfResult { size, scalar }
5050 | &crate::PredeclaredType::FrexpResult { size, scalar } => {
5051 let arg_type_name_owner;
5052 let arg_type_name = if let Some(size) = size {
5053 arg_type_name_owner = format!(
5054 "{NAMESPACE}::{}{}",
5055 if scalar.width == 8 { "double" } else { "float" },
5056 size as u8
5057 );
5058 &arg_type_name_owner
5059 } else if scalar.width == 8 {
5060 "double"
5061 } else {
5062 "float"
5063 };
5064
5065 let other_type_name_owner;
5066 let (defined_func_name, called_func_name, other_type_name) =
5067 if matches!(type_key, &crate::PredeclaredType::ModfResult { .. }) {
5068 (MODF_FUNCTION, "modf", arg_type_name)
5069 } else {
5070 let other_type_name = if let Some(size) = size {
5071 other_type_name_owner = format!("int{}", size as u8);
5072 &other_type_name_owner
5073 } else {
5074 "int"
5075 };
5076 (FREXP_FUNCTION, "frexp", other_type_name)
5077 };
5078
5079 let struct_name = &self.names[&NameKey::Type(*struct_ty)];
5080
5081 writeln!(self.out)?;
5082 writeln!(
5083 self.out,
5084 "{struct_name} {defined_func_name}({arg_type_name} arg) {{
5085 {other_type_name} other;
5086 {arg_type_name} fract = {NAMESPACE}::{called_func_name}(arg, other);
5087 return {struct_name}{{ fract, other }};
5088}}"
5089 )?;
5090 }
5091 &crate::PredeclaredType::AtomicCompareExchangeWeakResult(scalar) => {
5092 let arg_type_name = scalar.to_msl_name();
5093 let called_func_name = "atomic_compare_exchange_weak_explicit";
5094 let defined_func_name = ATOMIC_COMP_EXCH_FUNCTION;
5095 let struct_name = &self.names[&NameKey::Type(*struct_ty)];
5096
5097 writeln!(self.out)?;
5098
5099 for address_space_name in ["device", "threadgroup"] {
5100 writeln!(
5101 self.out,
5102 "\
5103template <typename A>
5104{struct_name} {defined_func_name}(
5105 {address_space_name} A *atomic_ptr,
5106 {arg_type_name} cmp,
5107 {arg_type_name} v
5108) {{
5109 bool swapped = {NAMESPACE}::{called_func_name}(
5110 atomic_ptr, &cmp, v,
5111 metal::memory_order_relaxed, metal::memory_order_relaxed
5112 );
5113 return {struct_name}{{cmp, swapped}};
5114}}"
5115 )?;
5116 }
5117 }
5118 }
5119 }
5120
5121 Ok(())
5122 }
5123
5124 fn write_global_constants(
5126 &mut self,
5127 module: &crate::Module,
5128 mod_info: &valid::ModuleInfo,
5129 ) -> BackendResult {
5130 let constants = module.constants.iter().filter(|&(_, c)| c.name.is_some());
5131
5132 for (handle, constant) in constants {
5133 let ty_name = TypeContext {
5134 handle: constant.ty,
5135 gctx: module.to_ctx(),
5136 names: &self.names,
5137 access: crate::StorageAccess::empty(),
5138 first_time: false,
5139 };
5140 let name = &self.names[&NameKey::Constant(handle)];
5141 write!(self.out, "constant {ty_name} {name} = ")?;
5142 self.put_const_expression(constant.init, module, mod_info, &module.global_expressions)?;
5143 writeln!(self.out, ";")?;
5144 }
5145
5146 Ok(())
5147 }
5148
5149 fn put_inline_sampler_properties(
5150 &mut self,
5151 level: back::Level,
5152 sampler: &sm::InlineSampler,
5153 options: &Options,
5154 ) -> BackendResult {
5155 for (&letter, address) in ['s', 't', 'r'].iter().zip(sampler.address.iter()) {
5156 writeln!(
5157 self.out,
5158 "{}{}::{}_address::{},",
5159 level,
5160 NAMESPACE,
5161 letter,
5162 address.as_str(),
5163 )?;
5164 }
5165 writeln!(
5166 self.out,
5167 "{}{}::mag_filter::{},",
5168 level,
5169 NAMESPACE,
5170 sampler.mag_filter.as_str(),
5171 )?;
5172 writeln!(
5173 self.out,
5174 "{}{}::min_filter::{},",
5175 level,
5176 NAMESPACE,
5177 sampler.min_filter.as_str(),
5178 )?;
5179 if let Some(filter) = sampler.mip_filter {
5180 writeln!(
5181 self.out,
5182 "{}{}::mip_filter::{},",
5183 level,
5184 NAMESPACE,
5185 filter.as_str(),
5186 )?;
5187 }
5188 if sampler.border_color != sm::BorderColor::TransparentBlack {
5190 writeln!(
5191 self.out,
5192 "{}{}::border_color::{},",
5193 level,
5194 NAMESPACE,
5195 sampler.border_color.as_str(),
5196 )?;
5197 }
5198
5199 if options.lang_version >= (1, 2) {
5200 if let Some(ref lod) = sampler.lod_clamp {
5201 writeln!(
5202 self.out,
5203 "{}{}::lod_clamp({},{}),",
5204 level, NAMESPACE, lod.start, lod.end,
5205 )?;
5206 }
5207 if let Some(aniso) = sampler.max_anisotropy {
5208 writeln!(
5209 self.out,
5210 "{}{}::max_anisotropy({}),",
5211 level,
5212 NAMESPACE,
5213 aniso.get(),
5214 )?;
5215 }
5216 }
5217
5218 if sampler.compare_func != sm::CompareFunc::Never {
5219 writeln!(
5220 self.out,
5221 "{}{}::compare_func::{},",
5222 level,
5223 NAMESPACE,
5224 sampler.compare_func.as_str(),
5225 )?;
5226 }
5227 writeln!(
5228 self.out,
5229 "{}{}::coord::{}",
5230 level,
5231 NAMESPACE,
5232 sampler.coord.as_str()
5233 )?;
5234 Ok(())
5235 }
5236
5237 fn write_unpacking_function(
5238 &mut self,
5239 format: nt::VertexFormat,
5240 ) -> Result<(String, u32, Option<crate::VectorSize>, crate::Scalar), Error> {
5241 use crate::{Scalar, VectorSize};
5242 use nt::VertexFormat::*;
5243 match format {
5244 Uint8 => {
5245 let name = self.namer.call("unpackUint8");
5246 writeln!(self.out, "uint {name}(metal::uchar b0) {{")?;
5247 writeln!(self.out, "{}return uint(b0);", back::INDENT)?;
5248 writeln!(self.out, "}}")?;
5249 Ok((name, 1, None, Scalar::U32))
5250 }
5251 Uint8x2 => {
5252 let name = self.namer.call("unpackUint8x2");
5253 writeln!(
5254 self.out,
5255 "metal::uint2 {name}(metal::uchar b0, \
5256 metal::uchar b1) {{"
5257 )?;
5258 writeln!(self.out, "{}return metal::uint2(b0, b1);", back::INDENT)?;
5259 writeln!(self.out, "}}")?;
5260 Ok((name, 2, Some(VectorSize::Bi), Scalar::U32))
5261 }
5262 Uint8x4 => {
5263 let name = self.namer.call("unpackUint8x4");
5264 writeln!(
5265 self.out,
5266 "metal::uint4 {name}(metal::uchar b0, \
5267 metal::uchar b1, \
5268 metal::uchar b2, \
5269 metal::uchar b3) {{"
5270 )?;
5271 writeln!(
5272 self.out,
5273 "{}return metal::uint4(b0, b1, b2, b3);",
5274 back::INDENT
5275 )?;
5276 writeln!(self.out, "}}")?;
5277 Ok((name, 4, Some(VectorSize::Quad), Scalar::U32))
5278 }
5279 Sint8 => {
5280 let name = self.namer.call("unpackSint8");
5281 writeln!(self.out, "int {name}(metal::uchar b0) {{")?;
5282 writeln!(self.out, "{}return int(as_type<char>(b0));", back::INDENT)?;
5283 writeln!(self.out, "}}")?;
5284 Ok((name, 1, None, Scalar::I32))
5285 }
5286 Sint8x2 => {
5287 let name = self.namer.call("unpackSint8x2");
5288 writeln!(
5289 self.out,
5290 "metal::int2 {name}(metal::uchar b0, \
5291 metal::uchar b1) {{"
5292 )?;
5293 writeln!(
5294 self.out,
5295 "{}return metal::int2(as_type<char>(b0), \
5296 as_type<char>(b1));",
5297 back::INDENT
5298 )?;
5299 writeln!(self.out, "}}")?;
5300 Ok((name, 2, Some(VectorSize::Bi), Scalar::I32))
5301 }
5302 Sint8x4 => {
5303 let name = self.namer.call("unpackSint8x4");
5304 writeln!(
5305 self.out,
5306 "metal::int4 {name}(metal::uchar b0, \
5307 metal::uchar b1, \
5308 metal::uchar b2, \
5309 metal::uchar b3) {{"
5310 )?;
5311 writeln!(
5312 self.out,
5313 "{}return metal::int4(as_type<char>(b0), \
5314 as_type<char>(b1), \
5315 as_type<char>(b2), \
5316 as_type<char>(b3));",
5317 back::INDENT
5318 )?;
5319 writeln!(self.out, "}}")?;
5320 Ok((name, 4, Some(VectorSize::Quad), Scalar::I32))
5321 }
5322 Unorm8 => {
5323 let name = self.namer.call("unpackUnorm8");
5324 writeln!(self.out, "float {name}(metal::uchar b0) {{")?;
5325 writeln!(
5326 self.out,
5327 "{}return float(float(b0) / 255.0f);",
5328 back::INDENT
5329 )?;
5330 writeln!(self.out, "}}")?;
5331 Ok((name, 1, None, Scalar::F32))
5332 }
5333 Unorm8x2 => {
5334 let name = self.namer.call("unpackUnorm8x2");
5335 writeln!(
5336 self.out,
5337 "metal::float2 {name}(metal::uchar b0, \
5338 metal::uchar b1) {{"
5339 )?;
5340 writeln!(
5341 self.out,
5342 "{}return metal::float2(float(b0) / 255.0f, \
5343 float(b1) / 255.0f);",
5344 back::INDENT
5345 )?;
5346 writeln!(self.out, "}}")?;
5347 Ok((name, 2, Some(VectorSize::Bi), Scalar::F32))
5348 }
5349 Unorm8x4 => {
5350 let name = self.namer.call("unpackUnorm8x4");
5351 writeln!(
5352 self.out,
5353 "metal::float4 {name}(metal::uchar b0, \
5354 metal::uchar b1, \
5355 metal::uchar b2, \
5356 metal::uchar b3) {{"
5357 )?;
5358 writeln!(
5359 self.out,
5360 "{}return metal::float4(float(b0) / 255.0f, \
5361 float(b1) / 255.0f, \
5362 float(b2) / 255.0f, \
5363 float(b3) / 255.0f);",
5364 back::INDENT
5365 )?;
5366 writeln!(self.out, "}}")?;
5367 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5368 }
5369 Snorm8 => {
5370 let name = self.namer.call("unpackSnorm8");
5371 writeln!(self.out, "float {name}(metal::uchar b0) {{")?;
5372 writeln!(
5373 self.out,
5374 "{}return float(metal::max(-1.0f, as_type<char>(b0) / 127.0f));",
5375 back::INDENT
5376 )?;
5377 writeln!(self.out, "}}")?;
5378 Ok((name, 1, None, Scalar::F32))
5379 }
5380 Snorm8x2 => {
5381 let name = self.namer.call("unpackSnorm8x2");
5382 writeln!(
5383 self.out,
5384 "metal::float2 {name}(metal::uchar b0, \
5385 metal::uchar b1) {{"
5386 )?;
5387 writeln!(
5388 self.out,
5389 "{}return metal::float2(metal::max(-1.0f, as_type<char>(b0) / 127.0f), \
5390 metal::max(-1.0f, as_type<char>(b1) / 127.0f));",
5391 back::INDENT
5392 )?;
5393 writeln!(self.out, "}}")?;
5394 Ok((name, 2, Some(VectorSize::Bi), Scalar::F32))
5395 }
5396 Snorm8x4 => {
5397 let name = self.namer.call("unpackSnorm8x4");
5398 writeln!(
5399 self.out,
5400 "metal::float4 {name}(metal::uchar b0, \
5401 metal::uchar b1, \
5402 metal::uchar b2, \
5403 metal::uchar b3) {{"
5404 )?;
5405 writeln!(
5406 self.out,
5407 "{}return metal::float4(metal::max(-1.0f, as_type<char>(b0) / 127.0f), \
5408 metal::max(-1.0f, as_type<char>(b1) / 127.0f), \
5409 metal::max(-1.0f, as_type<char>(b2) / 127.0f), \
5410 metal::max(-1.0f, as_type<char>(b3) / 127.0f));",
5411 back::INDENT
5412 )?;
5413 writeln!(self.out, "}}")?;
5414 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
5415 }
5416 Uint16 => {
5417 let name = self.namer.call("unpackUint16");
5418 writeln!(
5419 self.out,
5420 "metal::uint {name}(metal::uint b0, \
5421 metal::uint b1) {{"
5422 )?;
5423 writeln!(
5424 self.out,
5425 "{}return metal::uint(b1 << 8 | b0);",
5426 back::INDENT
5427 )?;
5428 writeln!(self.out, "}}")?;
5429 Ok((name, 2, None, Scalar::U32))
5430 }
5431 Uint16x2 => {
5432 let name = self.namer.call("unpackUint16x2");
5433 writeln!(
5434 self.out,
5435 "metal::uint2 {name}(metal::uint b0, \
5436 metal::uint b1, \
5437 metal::uint b2, \
5438 metal::uint b3) {{"
5439 )?;
5440 writeln!(
5441 self.out,
5442 "{}return metal::uint2(b1 << 8 | b0, \
5443 b3 << 8 | b2);",
5444 back::INDENT
5445 )?;
5446 writeln!(self.out, "}}")?;
5447 Ok((name, 4, Some(VectorSize::Bi), Scalar::U32))
5448 }
5449 Uint16x4 => {
5450 let name = self.namer.call("unpackUint16x4");
5451 writeln!(
5452 self.out,
5453 "metal::uint4 {name}(metal::uint b0, \
5454 metal::uint b1, \
5455 metal::uint b2, \
5456 metal::uint b3, \
5457 metal::uint b4, \
5458 metal::uint b5, \
5459 metal::uint b6, \
5460 metal::uint b7) {{"
5461 )?;
5462 writeln!(
5463 self.out,
5464 "{}return metal::uint4(b1 << 8 | b0, \
5465 b3 << 8 | b2, \
5466 b5 << 8 | b4, \
5467 b7 << 8 | b6);",
5468 back::INDENT
5469 )?;
5470 writeln!(self.out, "}}")?;
5471 Ok((name, 8, Some(VectorSize::Quad), Scalar::U32))
5472 }
5473 Sint16 => {
5474 let name = self.namer.call("unpackSint16");
5475 writeln!(
5476 self.out,
5477 "int {name}(metal::ushort b0, \
5478 metal::ushort b1) {{"
5479 )?;
5480 writeln!(
5481 self.out,
5482 "{}return int(as_type<short>(metal::ushort(b1 << 8 | b0)));",
5483 back::INDENT
5484 )?;
5485 writeln!(self.out, "}}")?;
5486 Ok((name, 2, None, Scalar::I32))
5487 }
5488 Sint16x2 => {
5489 let name = self.namer.call("unpackSint16x2");
5490 writeln!(
5491 self.out,
5492 "metal::int2 {name}(metal::ushort b0, \
5493 metal::ushort b1, \
5494 metal::ushort b2, \
5495 metal::ushort b3) {{"
5496 )?;
5497 writeln!(
5498 self.out,
5499 "{}return metal::int2(as_type<short>(metal::ushort(b1 << 8 | b0)), \
5500 as_type<short>(metal::ushort(b3 << 8 | b2)));",
5501 back::INDENT
5502 )?;
5503 writeln!(self.out, "}}")?;
5504 Ok((name, 4, Some(VectorSize::Bi), Scalar::I32))
5505 }
5506 Sint16x4 => {
5507 let name = self.namer.call("unpackSint16x4");
5508 writeln!(
5509 self.out,
5510 "metal::int4 {name}(metal::ushort b0, \
5511 metal::ushort b1, \
5512 metal::ushort b2, \
5513 metal::ushort b3, \
5514 metal::ushort b4, \
5515 metal::ushort b5, \
5516 metal::ushort b6, \
5517 metal::ushort b7) {{"
5518 )?;
5519 writeln!(
5520 self.out,
5521 "{}return metal::int4(as_type<short>(metal::ushort(b1 << 8 | b0)), \
5522 as_type<short>(metal::ushort(b3 << 8 | b2)), \
5523 as_type<short>(metal::ushort(b5 << 8 | b4)), \
5524 as_type<short>(metal::ushort(b7 << 8 | b6)));",
5525 back::INDENT
5526 )?;
5527 writeln!(self.out, "}}")?;
5528 Ok((name, 8, Some(VectorSize::Quad), Scalar::I32))
5529 }
5530 Unorm16 => {
5531 let name = self.namer.call("unpackUnorm16");
5532 writeln!(
5533 self.out,
5534 "float {name}(metal::ushort b0, \
5535 metal::ushort b1) {{"
5536 )?;
5537 writeln!(
5538 self.out,
5539 "{}return float(float(b1 << 8 | b0) / 65535.0f);",
5540 back::INDENT
5541 )?;
5542 writeln!(self.out, "}}")?;
5543 Ok((name, 2, None, Scalar::F32))
5544 }
5545 Unorm16x2 => {
5546 let name = self.namer.call("unpackUnorm16x2");
5547 writeln!(
5548 self.out,
5549 "metal::float2 {name}(metal::ushort b0, \
5550 metal::ushort b1, \
5551 metal::ushort b2, \
5552 metal::ushort b3) {{"
5553 )?;
5554 writeln!(
5555 self.out,
5556 "{}return metal::float2(float(b1 << 8 | b0) / 65535.0f, \
5557 float(b3 << 8 | b2) / 65535.0f);",
5558 back::INDENT
5559 )?;
5560 writeln!(self.out, "}}")?;
5561 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5562 }
5563 Unorm16x4 => {
5564 let name = self.namer.call("unpackUnorm16x4");
5565 writeln!(
5566 self.out,
5567 "metal::float4 {name}(metal::ushort b0, \
5568 metal::ushort b1, \
5569 metal::ushort b2, \
5570 metal::ushort b3, \
5571 metal::ushort b4, \
5572 metal::ushort b5, \
5573 metal::ushort b6, \
5574 metal::ushort b7) {{"
5575 )?;
5576 writeln!(
5577 self.out,
5578 "{}return metal::float4(float(b1 << 8 | b0) / 65535.0f, \
5579 float(b3 << 8 | b2) / 65535.0f, \
5580 float(b5 << 8 | b4) / 65535.0f, \
5581 float(b7 << 8 | b6) / 65535.0f);",
5582 back::INDENT
5583 )?;
5584 writeln!(self.out, "}}")?;
5585 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5586 }
5587 Snorm16 => {
5588 let name = self.namer.call("unpackSnorm16");
5589 writeln!(
5590 self.out,
5591 "float {name}(metal::ushort b0, \
5592 metal::ushort b1) {{"
5593 )?;
5594 writeln!(
5595 self.out,
5596 "{}return metal::unpack_snorm2x16_to_float(b1 << 8 | b0).x;",
5597 back::INDENT
5598 )?;
5599 writeln!(self.out, "}}")?;
5600 Ok((name, 2, None, Scalar::F32))
5601 }
5602 Snorm16x2 => {
5603 let name = self.namer.call("unpackSnorm16x2");
5604 writeln!(
5605 self.out,
5606 "metal::float2 {name}(uint b0, \
5607 uint b1, \
5608 uint b2, \
5609 uint b3) {{"
5610 )?;
5611 writeln!(
5612 self.out,
5613 "{}return metal::unpack_snorm2x16_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5614 back::INDENT
5615 )?;
5616 writeln!(self.out, "}}")?;
5617 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5618 }
5619 Snorm16x4 => {
5620 let name = self.namer.call("unpackSnorm16x4");
5621 writeln!(
5622 self.out,
5623 "metal::float4 {name}(uint b0, \
5624 uint b1, \
5625 uint b2, \
5626 uint b3, \
5627 uint b4, \
5628 uint b5, \
5629 uint b6, \
5630 uint b7) {{"
5631 )?;
5632 writeln!(
5633 self.out,
5634 "{}return metal::float4(metal::unpack_snorm2x16_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5635 metal::unpack_snorm2x16_to_float(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5636 back::INDENT
5637 )?;
5638 writeln!(self.out, "}}")?;
5639 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5640 }
5641 Float16 => {
5642 let name = self.namer.call("unpackFloat16");
5643 writeln!(
5644 self.out,
5645 "float {name}(metal::ushort b0, \
5646 metal::ushort b1) {{"
5647 )?;
5648 writeln!(
5649 self.out,
5650 "{}return float(as_type<half>(metal::ushort(b1 << 8 | b0)));",
5651 back::INDENT
5652 )?;
5653 writeln!(self.out, "}}")?;
5654 Ok((name, 2, None, Scalar::F32))
5655 }
5656 Float16x2 => {
5657 let name = self.namer.call("unpackFloat16x2");
5658 writeln!(
5659 self.out,
5660 "metal::float2 {name}(metal::ushort b0, \
5661 metal::ushort b1, \
5662 metal::ushort b2, \
5663 metal::ushort b3) {{"
5664 )?;
5665 writeln!(
5666 self.out,
5667 "{}return metal::float2(as_type<half>(metal::ushort(b1 << 8 | b0)), \
5668 as_type<half>(metal::ushort(b3 << 8 | b2)));",
5669 back::INDENT
5670 )?;
5671 writeln!(self.out, "}}")?;
5672 Ok((name, 4, Some(VectorSize::Bi), Scalar::F32))
5673 }
5674 Float16x4 => {
5675 let name = self.namer.call("unpackFloat16x4");
5676 writeln!(
5677 self.out,
5678 "metal::float4 {name}(metal::ushort b0, \
5679 metal::ushort b1, \
5680 metal::ushort b2, \
5681 metal::ushort b3, \
5682 metal::ushort b4, \
5683 metal::ushort b5, \
5684 metal::ushort b6, \
5685 metal::ushort b7) {{"
5686 )?;
5687 writeln!(
5688 self.out,
5689 "{}return metal::float4(as_type<half>(metal::ushort(b1 << 8 | b0)), \
5690 as_type<half>(metal::ushort(b3 << 8 | b2)), \
5691 as_type<half>(metal::ushort(b5 << 8 | b4)), \
5692 as_type<half>(metal::ushort(b7 << 8 | b6)));",
5693 back::INDENT
5694 )?;
5695 writeln!(self.out, "}}")?;
5696 Ok((name, 8, Some(VectorSize::Quad), Scalar::F32))
5697 }
5698 Float32 => {
5699 let name = self.namer.call("unpackFloat32");
5700 writeln!(
5701 self.out,
5702 "float {name}(uint b0, \
5703 uint b1, \
5704 uint b2, \
5705 uint b3) {{"
5706 )?;
5707 writeln!(
5708 self.out,
5709 "{}return as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5710 back::INDENT
5711 )?;
5712 writeln!(self.out, "}}")?;
5713 Ok((name, 4, None, Scalar::F32))
5714 }
5715 Float32x2 => {
5716 let name = self.namer.call("unpackFloat32x2");
5717 writeln!(
5718 self.out,
5719 "metal::float2 {name}(uint b0, \
5720 uint b1, \
5721 uint b2, \
5722 uint b3, \
5723 uint b4, \
5724 uint b5, \
5725 uint b6, \
5726 uint b7) {{"
5727 )?;
5728 writeln!(
5729 self.out,
5730 "{}return metal::float2(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5731 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5732 back::INDENT
5733 )?;
5734 writeln!(self.out, "}}")?;
5735 Ok((name, 8, Some(VectorSize::Bi), Scalar::F32))
5736 }
5737 Float32x3 => {
5738 let name = self.namer.call("unpackFloat32x3");
5739 writeln!(
5740 self.out,
5741 "metal::float3 {name}(uint b0, \
5742 uint b1, \
5743 uint b2, \
5744 uint b3, \
5745 uint b4, \
5746 uint b5, \
5747 uint b6, \
5748 uint b7, \
5749 uint b8, \
5750 uint b9, \
5751 uint b10, \
5752 uint b11) {{"
5753 )?;
5754 writeln!(
5755 self.out,
5756 "{}return metal::float3(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5757 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5758 as_type<float>(b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5759 back::INDENT
5760 )?;
5761 writeln!(self.out, "}}")?;
5762 Ok((name, 12, Some(VectorSize::Tri), Scalar::F32))
5763 }
5764 Float32x4 => {
5765 let name = self.namer.call("unpackFloat32x4");
5766 writeln!(
5767 self.out,
5768 "metal::float4 {name}(uint b0, \
5769 uint b1, \
5770 uint b2, \
5771 uint b3, \
5772 uint b4, \
5773 uint b5, \
5774 uint b6, \
5775 uint b7, \
5776 uint b8, \
5777 uint b9, \
5778 uint b10, \
5779 uint b11, \
5780 uint b12, \
5781 uint b13, \
5782 uint b14, \
5783 uint b15) {{"
5784 )?;
5785 writeln!(
5786 self.out,
5787 "{}return metal::float4(as_type<float>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5788 as_type<float>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5789 as_type<float>(b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5790 as_type<float>(b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5791 back::INDENT
5792 )?;
5793 writeln!(self.out, "}}")?;
5794 Ok((name, 16, Some(VectorSize::Quad), Scalar::F32))
5795 }
5796 Uint32 => {
5797 let name = self.namer.call("unpackUint32");
5798 writeln!(
5799 self.out,
5800 "uint {name}(uint b0, \
5801 uint b1, \
5802 uint b2, \
5803 uint b3) {{"
5804 )?;
5805 writeln!(
5806 self.out,
5807 "{}return (b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5808 back::INDENT
5809 )?;
5810 writeln!(self.out, "}}")?;
5811 Ok((name, 4, None, Scalar::U32))
5812 }
5813 Uint32x2 => {
5814 let name = self.namer.call("unpackUint32x2");
5815 writeln!(
5816 self.out,
5817 "uint2 {name}(uint b0, \
5818 uint b1, \
5819 uint b2, \
5820 uint b3, \
5821 uint b4, \
5822 uint b5, \
5823 uint b6, \
5824 uint b7) {{"
5825 )?;
5826 writeln!(
5827 self.out,
5828 "{}return uint2((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5829 (b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5830 back::INDENT
5831 )?;
5832 writeln!(self.out, "}}")?;
5833 Ok((name, 8, Some(VectorSize::Bi), Scalar::U32))
5834 }
5835 Uint32x3 => {
5836 let name = self.namer.call("unpackUint32x3");
5837 writeln!(
5838 self.out,
5839 "uint3 {name}(uint b0, \
5840 uint b1, \
5841 uint b2, \
5842 uint b3, \
5843 uint b4, \
5844 uint b5, \
5845 uint b6, \
5846 uint b7, \
5847 uint b8, \
5848 uint b9, \
5849 uint b10, \
5850 uint b11) {{"
5851 )?;
5852 writeln!(
5853 self.out,
5854 "{}return uint3((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5855 (b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5856 (b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5857 back::INDENT
5858 )?;
5859 writeln!(self.out, "}}")?;
5860 Ok((name, 12, Some(VectorSize::Tri), Scalar::U32))
5861 }
5862 Uint32x4 => {
5863 let name = self.namer.call("unpackUint32x4");
5864 writeln!(
5865 self.out,
5866 "{NAMESPACE}::uint4 {name}(uint b0, \
5867 uint b1, \
5868 uint b2, \
5869 uint b3, \
5870 uint b4, \
5871 uint b5, \
5872 uint b6, \
5873 uint b7, \
5874 uint b8, \
5875 uint b9, \
5876 uint b10, \
5877 uint b11, \
5878 uint b12, \
5879 uint b13, \
5880 uint b14, \
5881 uint b15) {{"
5882 )?;
5883 writeln!(
5884 self.out,
5885 "{}return {NAMESPACE}::uint4((b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5886 (b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5887 (b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5888 (b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5889 back::INDENT
5890 )?;
5891 writeln!(self.out, "}}")?;
5892 Ok((name, 16, Some(VectorSize::Quad), Scalar::U32))
5893 }
5894 Sint32 => {
5895 let name = self.namer.call("unpackSint32");
5896 writeln!(
5897 self.out,
5898 "int {name}(uint b0, \
5899 uint b1, \
5900 uint b2, \
5901 uint b3) {{"
5902 )?;
5903 writeln!(
5904 self.out,
5905 "{}return as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
5906 back::INDENT
5907 )?;
5908 writeln!(self.out, "}}")?;
5909 Ok((name, 4, None, Scalar::I32))
5910 }
5911 Sint32x2 => {
5912 let name = self.namer.call("unpackSint32x2");
5913 writeln!(
5914 self.out,
5915 "metal::int2 {name}(uint b0, \
5916 uint b1, \
5917 uint b2, \
5918 uint b3, \
5919 uint b4, \
5920 uint b5, \
5921 uint b6, \
5922 uint b7) {{"
5923 )?;
5924 writeln!(
5925 self.out,
5926 "{}return metal::int2(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5927 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4));",
5928 back::INDENT
5929 )?;
5930 writeln!(self.out, "}}")?;
5931 Ok((name, 8, Some(VectorSize::Bi), Scalar::I32))
5932 }
5933 Sint32x3 => {
5934 let name = self.namer.call("unpackSint32x3");
5935 writeln!(
5936 self.out,
5937 "metal::int3 {name}(uint b0, \
5938 uint b1, \
5939 uint b2, \
5940 uint b3, \
5941 uint b4, \
5942 uint b5, \
5943 uint b6, \
5944 uint b7, \
5945 uint b8, \
5946 uint b9, \
5947 uint b10, \
5948 uint b11) {{"
5949 )?;
5950 writeln!(
5951 self.out,
5952 "{}return metal::int3(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5953 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5954 as_type<int>(b11 << 24 | b10 << 16 | b9 << 8 | b8));",
5955 back::INDENT
5956 )?;
5957 writeln!(self.out, "}}")?;
5958 Ok((name, 12, Some(VectorSize::Tri), Scalar::I32))
5959 }
5960 Sint32x4 => {
5961 let name = self.namer.call("unpackSint32x4");
5962 writeln!(
5963 self.out,
5964 "metal::int4 {name}(uint b0, \
5965 uint b1, \
5966 uint b2, \
5967 uint b3, \
5968 uint b4, \
5969 uint b5, \
5970 uint b6, \
5971 uint b7, \
5972 uint b8, \
5973 uint b9, \
5974 uint b10, \
5975 uint b11, \
5976 uint b12, \
5977 uint b13, \
5978 uint b14, \
5979 uint b15) {{"
5980 )?;
5981 writeln!(
5982 self.out,
5983 "{}return metal::int4(as_type<int>(b3 << 24 | b2 << 16 | b1 << 8 | b0), \
5984 as_type<int>(b7 << 24 | b6 << 16 | b5 << 8 | b4), \
5985 as_type<int>(b11 << 24 | b10 << 16 | b9 << 8 | b8), \
5986 as_type<int>(b15 << 24 | b14 << 16 | b13 << 8 | b12));",
5987 back::INDENT
5988 )?;
5989 writeln!(self.out, "}}")?;
5990 Ok((name, 16, Some(VectorSize::Quad), Scalar::I32))
5991 }
5992 Unorm10_10_10_2 => {
5993 let name = self.namer.call("unpackUnorm10_10_10_2");
5994 writeln!(
5995 self.out,
5996 "metal::float4 {name}(uint b0, \
5997 uint b1, \
5998 uint b2, \
5999 uint b3) {{"
6000 )?;
6001 writeln!(
6002 self.out,
6003 "{}return metal::unpack_unorm10a2_to_float(b3 << 24 | b2 << 16 | b1 << 8 | b0);",
6015 back::INDENT
6016 )?;
6017 writeln!(self.out, "}}")?;
6018 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
6019 }
6020 Snorm10_10_10_2 => {
6021 let name = self.namer.call("unpackSnorm10_10_10_2");
6022 writeln!(
6023 self.out,
6024 "metal::float4 {name}(uint b0, \
6025 uint b1, \
6026 uint b2, \
6027 uint b3) {{"
6028 )?;
6029 writeln!(
6030 self.out,
6031 "{}uint v = (b3 << 24 | b2 << 16 | b1 << 8 | b0); \
6033 return metal::max(metal::float4(float(as_type<int>(v << 22) >> 22) / 511.0f, \
6034 float(as_type<int>(v << 12) >> 22) / 511.0f, \
6035 float(as_type<int>(v << 2) >> 22) / 511.0f, \
6036 float(as_type<int>(v) >> 30)), \
6037 metal::float4(-1.0f));",
6038 back::INDENT
6039 )?;
6040 writeln!(self.out, "}}")?;
6041 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
6042 }
6043 Unorm8x4Bgra => {
6044 let name = self.namer.call("unpackUnorm8x4Bgra");
6045 writeln!(
6046 self.out,
6047 "metal::float4 {name}(metal::uchar b0, \
6048 metal::uchar b1, \
6049 metal::uchar b2, \
6050 metal::uchar b3) {{"
6051 )?;
6052 writeln!(
6053 self.out,
6054 "{}return metal::float4(float(b2) / 255.0f, \
6055 float(b1) / 255.0f, \
6056 float(b0) / 255.0f, \
6057 float(b3) / 255.0f);",
6058 back::INDENT
6059 )?;
6060 writeln!(self.out, "}}")?;
6061 Ok((name, 4, Some(VectorSize::Quad), Scalar::F32))
6062 }
6063 Float64 | Float64x2 | Float64x3 | Float64x4 => unreachable!(),
6064 }
6065 }
6066
6067 fn write_wrapped_unary_op(
6068 &mut self,
6069 module: &crate::Module,
6070 func_ctx: &back::FunctionCtx,
6071 op: crate::UnaryOperator,
6072 operand: Handle<crate::Expression>,
6073 ) -> BackendResult {
6074 let operand_ty = func_ctx.resolve_type(operand, &module.types);
6075 match op {
6076 crate::UnaryOperator::Negate
6083 if operand_ty.scalar_kind() == Some(crate::ScalarKind::Sint) =>
6084 {
6085 let Some((vector_size, scalar)) = operand_ty.vector_size_and_scalar() else {
6086 return Ok(());
6087 };
6088 let wrapped = WrappedFunction::UnaryOp {
6089 op,
6090 ty: (vector_size, scalar),
6091 };
6092 if !self.wrapped_functions.insert(wrapped) {
6093 return Ok(());
6094 }
6095
6096 let unsigned_scalar = crate::Scalar {
6097 kind: crate::ScalarKind::Uint,
6098 ..scalar
6099 };
6100 let mut type_name = String::new();
6101 let mut unsigned_type_name = String::new();
6102 match vector_size {
6103 None => {
6104 put_numeric_type(&mut type_name, scalar, &[])?;
6105 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[])?
6106 }
6107 Some(size) => {
6108 put_numeric_type(&mut type_name, scalar, &[size])?;
6109 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[size])?;
6110 }
6111 };
6112
6113 writeln!(self.out, "{type_name} {NEG_FUNCTION}({type_name} val) {{")?;
6114 let level = back::Level(1);
6115 if scalar.width < 4 {
6119 writeln!(
6120 self.out,
6121 "{level}return as_type<{type_name}>(static_cast<{unsigned_type_name}>(-as_type<{unsigned_type_name}>(val)));"
6122 )?;
6123 } else {
6124 writeln!(
6125 self.out,
6126 "{level}return as_type<{type_name}>(-as_type<{unsigned_type_name}>(val));"
6127 )?;
6128 }
6129 writeln!(self.out, "}}")?;
6130 writeln!(self.out)?;
6131 }
6132 _ => {}
6133 }
6134 Ok(())
6135 }
6136
6137 fn write_wrapped_binary_op(
6138 &mut self,
6139 module: &crate::Module,
6140 func_ctx: &back::FunctionCtx,
6141 expr: Handle<crate::Expression>,
6142 op: crate::BinaryOperator,
6143 left: Handle<crate::Expression>,
6144 right: Handle<crate::Expression>,
6145 ) -> BackendResult {
6146 let expr_ty = func_ctx.resolve_type(expr, &module.types);
6147 let left_ty = func_ctx.resolve_type(left, &module.types);
6148 let right_ty = func_ctx.resolve_type(right, &module.types);
6149 match (op, expr_ty.scalar_kind()) {
6150 (
6157 crate::BinaryOperator::Divide,
6158 Some(crate::ScalarKind::Sint | crate::ScalarKind::Uint),
6159 ) if self.emit_int_div_checks => {
6160 let Some(left_wrapped_ty) = left_ty.vector_size_and_scalar() else {
6161 return Ok(());
6162 };
6163 let Some(right_wrapped_ty) = right_ty.vector_size_and_scalar() else {
6164 return Ok(());
6165 };
6166 let wrapped = WrappedFunction::BinaryOp {
6167 op,
6168 left_ty: left_wrapped_ty,
6169 right_ty: right_wrapped_ty,
6170 };
6171 if !self.wrapped_functions.insert(wrapped) {
6172 return Ok(());
6173 }
6174
6175 let Some((vector_size, scalar)) = expr_ty.vector_size_and_scalar() else {
6176 return Ok(());
6177 };
6178 let mut type_name = String::new();
6179 match vector_size {
6180 None => put_numeric_type(&mut type_name, scalar, &[])?,
6181 Some(size) => put_numeric_type(&mut type_name, scalar, &[size])?,
6182 };
6183 writeln!(
6184 self.out,
6185 "{type_name} {DIV_FUNCTION}({type_name} lhs, {type_name} rhs) {{"
6186 )?;
6187 let level = back::Level(1);
6188 let (lp, rp) = if scalar.width < 4 {
6192 (format!("{type_name}("), ")".to_string())
6193 } else {
6194 (String::new(), String::new())
6195 };
6196 match scalar.kind {
6197 crate::ScalarKind::Sint => {
6198 let min_val = match scalar.width {
6199 2 => crate::Literal::I16(i16::MIN),
6200 4 => crate::Literal::I32(i32::MIN),
6201 8 => crate::Literal::I64(i64::MIN),
6202 _ => {
6203 return Err(Error::GenericValidation(format!(
6204 "Unexpected width for scalar {scalar:?}"
6205 )));
6206 }
6207 };
6208 write!(
6209 self.out,
6210 "{level}return lhs / metal::select(rhs, {lp}1{rp}, (lhs == "
6211 )?;
6212 self.put_literal(min_val)?;
6213 writeln!(self.out, " & rhs == {lp}-1{rp}) | (rhs == {lp}0{rp}));")?
6214 }
6215 crate::ScalarKind::Uint => {
6216 let suffix = if scalar.width < 4 { "" } else { "u" };
6217 writeln!(
6218 self.out,
6219 "{level}return lhs / metal::select(rhs, {lp}1{suffix}{rp}, rhs == {lp}0{suffix}{rp});"
6220 )?
6221 }
6222 _ => unreachable!(),
6223 }
6224 writeln!(self.out, "}}")?;
6225 writeln!(self.out)?;
6226 }
6227 (
6240 crate::BinaryOperator::Modulo,
6241 Some(crate::ScalarKind::Sint | crate::ScalarKind::Uint),
6242 ) if self.emit_int_div_checks => {
6243 let Some(left_wrapped_ty) = left_ty.vector_size_and_scalar() else {
6244 return Ok(());
6245 };
6246 let Some((right_vector_size, right_scalar)) = right_ty.vector_size_and_scalar()
6247 else {
6248 return Ok(());
6249 };
6250 let wrapped = WrappedFunction::BinaryOp {
6251 op,
6252 left_ty: left_wrapped_ty,
6253 right_ty: (right_vector_size, right_scalar),
6254 };
6255 if !self.wrapped_functions.insert(wrapped) {
6256 return Ok(());
6257 }
6258
6259 let Some((vector_size, scalar)) = expr_ty.vector_size_and_scalar() else {
6260 return Ok(());
6261 };
6262 let mut type_name = String::new();
6263 match vector_size {
6264 None => put_numeric_type(&mut type_name, scalar, &[])?,
6265 Some(size) => put_numeric_type(&mut type_name, scalar, &[size])?,
6266 };
6267 let mut rhs_type_name = String::new();
6268 match right_vector_size {
6269 None => put_numeric_type(&mut rhs_type_name, right_scalar, &[])?,
6270 Some(size) => put_numeric_type(&mut rhs_type_name, right_scalar, &[size])?,
6271 };
6272
6273 writeln!(
6274 self.out,
6275 "{type_name} {MOD_FUNCTION}({type_name} lhs, {type_name} rhs) {{"
6276 )?;
6277 let level = back::Level(1);
6278 let (lp, rp) = if scalar.width < 4 {
6279 (format!("{type_name}("), ")".to_string())
6280 } else {
6281 (String::new(), String::new())
6282 };
6283 match scalar.kind {
6284 crate::ScalarKind::Sint => {
6285 let min_val = match scalar.width {
6286 2 => crate::Literal::I16(i16::MIN),
6287 4 => crate::Literal::I32(i32::MIN),
6288 8 => crate::Literal::I64(i64::MIN),
6289 _ => {
6290 return Err(Error::GenericValidation(format!(
6291 "Unexpected width for scalar {scalar:?}"
6292 )));
6293 }
6294 };
6295 write!(
6296 self.out,
6297 "{level}{rhs_type_name} divisor = metal::select(rhs, {lp}1{rp}, (lhs == "
6298 )?;
6299 self.put_literal(min_val)?;
6300 writeln!(self.out, " & rhs == {lp}-1{rp}) | (rhs == {lp}0{rp}));")?;
6301 writeln!(self.out, "{level}return lhs - (lhs / divisor) * divisor;")?
6302 }
6303 crate::ScalarKind::Uint => {
6304 let suffix = if scalar.width < 4 { "" } else { "u" };
6305 writeln!(
6306 self.out,
6307 "{level}return lhs % metal::select(rhs, {lp}1{suffix}{rp}, rhs == {lp}0{suffix}{rp});"
6308 )?
6309 }
6310 _ => unreachable!(),
6311 }
6312 writeln!(self.out, "}}")?;
6313 writeln!(self.out)?;
6314 }
6315 _ => {}
6316 }
6317 Ok(())
6318 }
6319
6320 fn get_dot_wrapper_function_helper_name(
6326 &self,
6327 scalar: crate::Scalar,
6328 size: crate::VectorSize,
6329 ) -> String {
6330 debug_assert!(concrete_int_scalars().any(|s| s == scalar));
6332
6333 let type_name = scalar.to_msl_name();
6334 let size_suffix = common::vector_size_str(size);
6335 format!("{DOT_FUNCTION_PREFIX}_{type_name}{size_suffix}")
6336 }
6337
6338 #[allow(clippy::too_many_arguments)]
6339 fn write_wrapped_math_function(
6340 &mut self,
6341 module: &crate::Module,
6342 func_ctx: &back::FunctionCtx,
6343 fun: crate::MathFunction,
6344 arg: Handle<crate::Expression>,
6345 _arg1: Option<Handle<crate::Expression>>,
6346 _arg2: Option<Handle<crate::Expression>>,
6347 _arg3: Option<Handle<crate::Expression>>,
6348 ) -> BackendResult {
6349 let arg_ty = func_ctx.resolve_type(arg, &module.types);
6350 match fun {
6351 crate::MathFunction::Abs if arg_ty.scalar_kind() == Some(crate::ScalarKind::Sint) => {
6359 let Some((vector_size, scalar)) = arg_ty.vector_size_and_scalar() else {
6360 return Ok(());
6361 };
6362 let wrapped = WrappedFunction::Math {
6363 fun,
6364 arg_ty: (vector_size, scalar),
6365 };
6366 if !self.wrapped_functions.insert(wrapped) {
6367 return Ok(());
6368 }
6369
6370 let unsigned_scalar = crate::Scalar {
6371 kind: crate::ScalarKind::Uint,
6372 ..scalar
6373 };
6374 let mut type_name = String::new();
6375 let mut unsigned_type_name = String::new();
6376 match vector_size {
6377 None => {
6378 put_numeric_type(&mut type_name, scalar, &[])?;
6379 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[])?
6380 }
6381 Some(size) => {
6382 put_numeric_type(&mut type_name, scalar, &[size])?;
6383 put_numeric_type(&mut unsigned_type_name, unsigned_scalar, &[size])?;
6384 }
6385 };
6386
6387 writeln!(self.out, "{type_name} {ABS_FUNCTION}({type_name} val) {{")?;
6388 let level = back::Level(1);
6389 let zero = if scalar.width < 4 {
6390 format!("{type_name}(0)")
6391 } else {
6392 "0".to_string()
6393 };
6394 let neg_expr = if scalar.width < 4 {
6395 format!(
6396 "static_cast<{unsigned_type_name}>(-as_type<{unsigned_type_name}>(val))"
6397 )
6398 } else {
6399 format!("-as_type<{unsigned_type_name}>(val)")
6400 };
6401 writeln!(self.out, "{level}return metal::select(as_type<{type_name}>({neg_expr}), val, val >= {zero});")?;
6402 writeln!(self.out, "}}")?;
6403 writeln!(self.out)?;
6404 }
6405
6406 crate::MathFunction::Dot => match *arg_ty {
6407 crate::TypeInner::Vector { size, scalar }
6408 if matches!(
6409 scalar.kind,
6410 crate::ScalarKind::Sint | crate::ScalarKind::Uint
6411 ) =>
6412 {
6413 let wrapped = WrappedFunction::Math {
6415 fun,
6416 arg_ty: (Some(size), scalar),
6417 };
6418 if !self.wrapped_functions.insert(wrapped) {
6419 return Ok(());
6420 }
6421
6422 let mut vec_ty = String::new();
6423 put_numeric_type(&mut vec_ty, scalar, &[size])?;
6424 let mut ret_ty = String::new();
6425 put_numeric_type(&mut ret_ty, scalar, &[])?;
6426
6427 let fun_name = self.get_dot_wrapper_function_helper_name(scalar, size);
6428
6429 writeln!(self.out, "{ret_ty} {fun_name}({vec_ty} a, {vec_ty} b) {{")?;
6431 let level = back::Level(1);
6432 write!(self.out, "{level}return ")?;
6433 self.put_dot_product("a", "b", size as usize, |writer, name, index| {
6434 write!(writer.out, "{name}.{}", back::COMPONENTS[index])?;
6435 Ok(())
6436 })?;
6437 writeln!(self.out, ";")?;
6438 writeln!(self.out, "}}")?;
6439 writeln!(self.out)?;
6440 }
6441 _ => {}
6442 },
6443
6444 _ => {}
6445 }
6446 Ok(())
6447 }
6448
6449 fn write_wrapped_cast(
6450 &mut self,
6451 module: &crate::Module,
6452 func_ctx: &back::FunctionCtx,
6453 expr: Handle<crate::Expression>,
6454 dst_scalar: crate::Scalar,
6455 ) -> BackendResult {
6456 let src_ty = func_ctx.resolve_type(expr, &module.types);
6467 let Some((vector_size, src_scalar)) = src_ty.vector_size_and_scalar() else {
6468 return Ok(());
6469 };
6470 if src_scalar.kind != crate::ScalarKind::Float
6471 || (dst_scalar.kind != crate::ScalarKind::Sint
6472 && dst_scalar.kind != crate::ScalarKind::Uint)
6473 {
6474 return Ok(());
6475 }
6476 let wrapped = WrappedFunction::Cast {
6477 src_scalar,
6478 vector_size,
6479 dst_scalar,
6480 };
6481 if !self.wrapped_functions.insert(wrapped) {
6482 return Ok(());
6483 }
6484 let (min, max) = proc::min_max_float_representable_by(src_scalar, dst_scalar);
6485
6486 let mut src_type_name = String::new();
6487 match vector_size {
6488 None => put_numeric_type(&mut src_type_name, src_scalar, &[])?,
6489 Some(size) => put_numeric_type(&mut src_type_name, src_scalar, &[size])?,
6490 };
6491 let mut dst_type_name = String::new();
6492 match vector_size {
6493 None => put_numeric_type(&mut dst_type_name, dst_scalar, &[])?,
6494 Some(size) => put_numeric_type(&mut dst_type_name, dst_scalar, &[size])?,
6495 };
6496 let fun_name = match dst_scalar {
6497 crate::Scalar::I32 => F2I32_FUNCTION,
6498 crate::Scalar::U32 => F2U32_FUNCTION,
6499 crate::Scalar::I64 => F2I64_FUNCTION,
6500 crate::Scalar::U64 => F2U64_FUNCTION,
6501 _ => unreachable!(),
6502 };
6503
6504 writeln!(
6505 self.out,
6506 "{dst_type_name} {fun_name}({src_type_name} value) {{"
6507 )?;
6508 let level = back::Level(1);
6509 write!(
6510 self.out,
6511 "{level}return static_cast<{dst_type_name}>({NAMESPACE}::clamp(value, "
6512 )?;
6513 self.put_literal(min)?;
6514 write!(self.out, ", ")?;
6515 self.put_literal(max)?;
6516 writeln!(self.out, "));")?;
6517 writeln!(self.out, "}}")?;
6518 writeln!(self.out)?;
6519 Ok(())
6520 }
6521
6522 fn write_convert_yuv_to_rgb_and_return(
6530 &mut self,
6531 level: back::Level,
6532 y: &str,
6533 uv: &str,
6534 params: &str,
6535 ) -> BackendResult {
6536 let l1 = level;
6537 let l2 = l1.next();
6538
6539 writeln!(
6541 self.out,
6542 "{l1}float3 srcGammaRgb = ({params}.yuv_conversion_matrix * float4({y}, {uv}, 1.0)).rgb;"
6543 )?;
6544
6545 writeln!(self.out, "{l1}float3 srcLinearRgb = {NAMESPACE}::select(")?;
6548 writeln!(self.out, "{l2}{NAMESPACE}::pow((srcGammaRgb + {params}.src_tf.a - 1.0) / {params}.src_tf.a, {params}.src_tf.g),")?;
6549 writeln!(self.out, "{l2}srcGammaRgb / {params}.src_tf.k,")?;
6550 writeln!(
6551 self.out,
6552 "{l2}srcGammaRgb < {params}.src_tf.k * {params}.src_tf.b);"
6553 )?;
6554
6555 writeln!(
6558 self.out,
6559 "{l1}float3 dstLinearRgb = {params}.gamut_conversion_matrix * srcLinearRgb;"
6560 )?;
6561
6562 writeln!(self.out, "{l1}float3 dstGammaRgb = {NAMESPACE}::select(")?;
6565 writeln!(self.out, "{l2}{params}.dst_tf.a * {NAMESPACE}::pow(dstLinearRgb, 1.0 / {params}.dst_tf.g) - ({params}.dst_tf.a - 1),")?;
6566 writeln!(self.out, "{l2}{params}.dst_tf.k * dstLinearRgb,")?;
6567 writeln!(self.out, "{l2}dstLinearRgb < {params}.dst_tf.b);")?;
6568
6569 writeln!(self.out, "{l1}return float4(dstGammaRgb, 1.0);")?;
6570 Ok(())
6571 }
6572
6573 #[allow(clippy::too_many_arguments)]
6574 fn write_wrapped_image_load(
6575 &mut self,
6576 module: &crate::Module,
6577 func_ctx: &back::FunctionCtx,
6578 image: Handle<crate::Expression>,
6579 _coordinate: Handle<crate::Expression>,
6580 _array_index: Option<Handle<crate::Expression>>,
6581 _sample: Option<Handle<crate::Expression>>,
6582 _level: Option<Handle<crate::Expression>>,
6583 ) -> BackendResult {
6584 let class = match *func_ctx.resolve_type(image, &module.types) {
6586 crate::TypeInner::Image { class, .. } => class,
6587 _ => unreachable!(),
6588 };
6589 if class != crate::ImageClass::External {
6590 return Ok(());
6591 }
6592 let wrapped = WrappedFunction::ImageLoad { class };
6593 if !self.wrapped_functions.insert(wrapped) {
6594 return Ok(());
6595 }
6596
6597 writeln!(self.out, "float4 {IMAGE_LOAD_EXTERNAL_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex, uint2 coords) {{")?;
6598 let l1 = back::Level(1);
6599 let l2 = l1.next();
6600 let l3 = l2.next();
6601 writeln!(
6602 self.out,
6603 "{l1}uint2 plane0_size = uint2(tex.plane0.get_width(), tex.plane0.get_height());"
6604 )?;
6605 writeln!(
6609 self.out,
6610 "{l1}uint2 cropped_size = {NAMESPACE}::any(tex.params.size != 0) ? tex.params.size : plane0_size;"
6611 )?;
6612 writeln!(
6613 self.out,
6614 "{l1}coords = {NAMESPACE}::min(coords, cropped_size - 1);"
6615 )?;
6616
6617 writeln!(self.out, "{l1}uint2 plane0_coords = uint2({NAMESPACE}::round(tex.params.load_transform * float3(float2(coords), 1.0)));")?;
6619 writeln!(self.out, "{l1}if (tex.params.num_planes == 1u) {{")?;
6620 writeln!(self.out, "{l2}return tex.plane0.read(plane0_coords);")?;
6622 writeln!(self.out, "{l1}}} else {{")?;
6623
6624 writeln!(
6626 self.out,
6627 "{l2}uint2 plane1_size = uint2(tex.plane1.get_width(), tex.plane1.get_height());"
6628 )?;
6629 writeln!(self.out, "{l2}uint2 plane1_coords = uint2({NAMESPACE}::floor(float2(plane0_coords) * float2(plane1_size) / float2(plane0_size)));")?;
6630
6631 writeln!(self.out, "{l2}float y = tex.plane0.read(plane0_coords).x;")?;
6633
6634 writeln!(self.out, "{l2}float2 uv;")?;
6635 writeln!(self.out, "{l2}if (tex.params.num_planes == 2u) {{")?;
6636 writeln!(self.out, "{l3}uv = tex.plane1.read(plane1_coords).xy;")?;
6638 writeln!(self.out, "{l2}}} else {{")?;
6639 writeln!(
6641 self.out,
6642 "{l2}uint2 plane2_size = uint2(tex.plane2.get_width(), tex.plane2.get_height());"
6643 )?;
6644 writeln!(self.out, "{l2}uint2 plane2_coords = uint2({NAMESPACE}::floor(float2(plane0_coords) * float2(plane2_size) / float2(plane0_size)));")?;
6645 writeln!(
6646 self.out,
6647 "{l3}uv = float2(tex.plane1.read(plane1_coords).x, tex.plane2.read(plane2_coords).x);"
6648 )?;
6649 writeln!(self.out, "{l2}}}")?;
6650
6651 self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "tex.params")?;
6652
6653 writeln!(self.out, "{l1}}}")?;
6654 writeln!(self.out, "}}")?;
6655 writeln!(self.out)?;
6656 Ok(())
6657 }
6658
6659 #[allow(clippy::too_many_arguments)]
6660 fn write_wrapped_image_sample(
6661 &mut self,
6662 module: &crate::Module,
6663 func_ctx: &back::FunctionCtx,
6664 image: Handle<crate::Expression>,
6665 _sampler: Handle<crate::Expression>,
6666 _gather: Option<crate::SwizzleComponent>,
6667 _coordinate: Handle<crate::Expression>,
6668 _array_index: Option<Handle<crate::Expression>>,
6669 _offset: Option<Handle<crate::Expression>>,
6670 _level: crate::SampleLevel,
6671 _depth_ref: Option<Handle<crate::Expression>>,
6672 clamp_to_edge: bool,
6673 ) -> BackendResult {
6674 if !clamp_to_edge {
6677 return Ok(());
6678 }
6679 let class = match *func_ctx.resolve_type(image, &module.types) {
6680 crate::TypeInner::Image { class, .. } => class,
6681 _ => unreachable!(),
6682 };
6683 let wrapped = WrappedFunction::ImageSample {
6684 class,
6685 clamp_to_edge: true,
6686 };
6687 if !self.wrapped_functions.insert(wrapped) {
6688 return Ok(());
6689 }
6690 match class {
6691 crate::ImageClass::External => {
6692 writeln!(self.out, "float4 {IMAGE_SAMPLE_BASE_CLAMP_TO_EDGE_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex, {NAMESPACE}::sampler samp, float2 coords) {{")?;
6693 let l1 = back::Level(1);
6694 let l2 = l1.next();
6695 let l3 = l2.next();
6696 writeln!(self.out, "{l1}uint2 plane0_size = uint2(tex.plane0.get_width(), tex.plane0.get_height());")?;
6697 writeln!(
6698 self.out,
6699 "{l1}coords = tex.params.sample_transform * float3(coords, 1.0);"
6700 )?;
6701
6702 writeln!(
6710 self.out,
6711 "{l1}float2 bounds_min = tex.params.sample_transform * float3(0.0, 0.0, 1.0);"
6712 )?;
6713 writeln!(
6714 self.out,
6715 "{l1}float2 bounds_max = tex.params.sample_transform * float3(1.0, 1.0, 1.0);"
6716 )?;
6717 writeln!(self.out, "{l1}float4 bounds = float4({NAMESPACE}::min(bounds_min, bounds_max), {NAMESPACE}::max(bounds_min, bounds_max));")?;
6718 writeln!(
6719 self.out,
6720 "{l1}float2 plane0_half_texel = float2(0.5, 0.5) / float2(plane0_size);"
6721 )?;
6722 writeln!(
6723 self.out,
6724 "{l1}float2 plane0_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane0_half_texel, bounds.zw - plane0_half_texel);"
6725 )?;
6726 writeln!(self.out, "{l1}if (tex.params.num_planes == 1u) {{")?;
6727 writeln!(
6729 self.out,
6730 "{l2}return tex.plane0.sample(samp, plane0_coords, {NAMESPACE}::level(0.0f));"
6731 )?;
6732 writeln!(self.out, "{l1}}} else {{")?;
6733 writeln!(self.out, "{l2}uint2 plane1_size = uint2(tex.plane1.get_width(), tex.plane1.get_height());")?;
6734 writeln!(
6735 self.out,
6736 "{l2}float2 plane1_half_texel = float2(0.5, 0.5) / float2(plane1_size);"
6737 )?;
6738 writeln!(
6739 self.out,
6740 "{l2}float2 plane1_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane1_half_texel, bounds.zw - plane1_half_texel);"
6741 )?;
6742
6743 writeln!(
6745 self.out,
6746 "{l2}float y = tex.plane0.sample(samp, plane0_coords, {NAMESPACE}::level(0.0f)).r;"
6747 )?;
6748 writeln!(self.out, "{l2}float2 uv = float2(0.0, 0.0);")?;
6749 writeln!(self.out, "{l2}if (tex.params.num_planes == 2u) {{")?;
6750 writeln!(
6752 self.out,
6753 "{l3}uv = tex.plane1.sample(samp, plane1_coords, {NAMESPACE}::level(0.0f)).xy;"
6754 )?;
6755 writeln!(self.out, "{l2}}} else {{")?;
6756 writeln!(self.out, "{l3}uint2 plane2_size = uint2(tex.plane2.get_width(), tex.plane2.get_height());")?;
6758 writeln!(
6759 self.out,
6760 "{l3}float2 plane2_half_texel = float2(0.5, 0.5) / float2(plane2_size);"
6761 )?;
6762 writeln!(
6763 self.out,
6764 "{l3}float2 plane2_coords = {NAMESPACE}::clamp(coords, bounds.xy + plane2_half_texel, bounds.zw - plane1_half_texel);"
6765 )?;
6766 writeln!(self.out, "{l3}uv.x = tex.plane1.sample(samp, plane1_coords, {NAMESPACE}::level(0.0f)).x;")?;
6767 writeln!(self.out, "{l3}uv.y = tex.plane2.sample(samp, plane2_coords, {NAMESPACE}::level(0.0f)).x;")?;
6768 writeln!(self.out, "{l2}}}")?;
6769
6770 self.write_convert_yuv_to_rgb_and_return(l2, "y", "uv", "tex.params")?;
6771
6772 writeln!(self.out, "{l1}}}")?;
6773 writeln!(self.out, "}}")?;
6774 writeln!(self.out)?;
6775 }
6776 _ => {
6777 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) {{")?;
6778 let l1 = back::Level(1);
6779 writeln!(self.out, "{l1}{NAMESPACE}::float2 half_texel = 0.5 / {NAMESPACE}::float2(tex.get_width(0u), tex.get_height(0u));")?;
6780 writeln!(
6781 self.out,
6782 "{l1}return tex.sample(samp, {NAMESPACE}::clamp(coords, half_texel, 1.0 - half_texel), {NAMESPACE}::level(0.0));"
6783 )?;
6784 writeln!(self.out, "}}")?;
6785 writeln!(self.out)?;
6786 }
6787 }
6788 Ok(())
6789 }
6790
6791 fn write_wrapped_image_query(
6792 &mut self,
6793 module: &crate::Module,
6794 func_ctx: &back::FunctionCtx,
6795 image: Handle<crate::Expression>,
6796 query: crate::ImageQuery,
6797 ) -> BackendResult {
6798 if !matches!(query, crate::ImageQuery::Size { .. }) {
6800 return Ok(());
6801 }
6802 let class = match *func_ctx.resolve_type(image, &module.types) {
6803 crate::TypeInner::Image { class, .. } => class,
6804 _ => unreachable!(),
6805 };
6806 if class != crate::ImageClass::External {
6807 return Ok(());
6808 }
6809 let wrapped = WrappedFunction::ImageQuerySize { class };
6810 if !self.wrapped_functions.insert(wrapped) {
6811 return Ok(());
6812 }
6813 writeln!(
6814 self.out,
6815 "uint2 {IMAGE_SIZE_EXTERNAL_FUNCTION}({EXTERNAL_TEXTURE_WRAPPER_STRUCT} tex) {{"
6816 )?;
6817 let l1 = back::Level(1);
6818 let l2 = l1.next();
6819 writeln!(
6820 self.out,
6821 "{l1}if ({NAMESPACE}::any(tex.params.size != uint2(0u))) {{"
6822 )?;
6823 writeln!(self.out, "{l2}return tex.params.size;")?;
6824 writeln!(self.out, "{l1}}} else {{")?;
6825 writeln!(
6827 self.out,
6828 "{l2}return uint2(tex.plane0.get_width(), tex.plane0.get_height());"
6829 )?;
6830 writeln!(self.out, "{l1}}}")?;
6831 writeln!(self.out, "}}")?;
6832 writeln!(self.out)?;
6833 Ok(())
6834 }
6835
6836 fn write_wrapped_cooperative_load(
6837 &mut self,
6838 module: &crate::Module,
6839 func_ctx: &back::FunctionCtx,
6840 columns: crate::CooperativeSize,
6841 rows: crate::CooperativeSize,
6842 pointer: Handle<crate::Expression>,
6843 ) -> BackendResult {
6844 let ptr_ty = func_ctx.resolve_type(pointer, &module.types);
6845 let space = ptr_ty.pointer_space().unwrap();
6846 let space_name = space.to_msl_name().unwrap_or_default();
6847 let scalar = ptr_ty
6848 .pointer_base_type()
6849 .unwrap()
6850 .inner_with(&module.types)
6851 .scalar()
6852 .unwrap();
6853 let wrapped = WrappedFunction::CooperativeLoad {
6854 space_name,
6855 columns,
6856 rows,
6857 scalar,
6858 };
6859 if !self.wrapped_functions.insert(wrapped) {
6860 return Ok(());
6861 }
6862 let scalar_name = scalar.to_msl_name();
6863 writeln!(
6864 self.out,
6865 "{NAMESPACE}::simdgroup_{scalar_name}{}x{} {COOPERATIVE_LOAD_FUNCTION}(const {space_name} {scalar_name}* ptr, int stride, bool is_row_major) {{",
6866 columns as u32, rows as u32,
6867 )?;
6868 let l1 = back::Level(1);
6869 writeln!(
6870 self.out,
6871 "{l1}{NAMESPACE}::simdgroup_{scalar_name}{}x{} m;",
6872 columns as u32, rows as u32
6873 )?;
6874 let matrix_origin = "0";
6875 writeln!(
6876 self.out,
6877 "{l1}simdgroup_load(m, ptr, stride, {matrix_origin}, is_row_major);"
6878 )?;
6879 writeln!(self.out, "{l1}return m;")?;
6880 writeln!(self.out, "}}")?;
6881 writeln!(self.out)?;
6882 Ok(())
6883 }
6884
6885 fn write_wrapped_cooperative_multiply_add(
6886 &mut self,
6887 module: &crate::Module,
6888 func_ctx: &back::FunctionCtx,
6889 space: crate::AddressSpace,
6890 a: Handle<crate::Expression>,
6891 b: Handle<crate::Expression>,
6892 c: Handle<crate::Expression>,
6893 ) -> BackendResult {
6894 let space_name = space.to_msl_name().unwrap_or_default();
6895 let (a_c, a_r, ab_scalar) = match *func_ctx.resolve_type(a, &module.types) {
6896 crate::TypeInner::CooperativeMatrix {
6897 columns,
6898 rows,
6899 scalar,
6900 ..
6901 } => (columns, rows, scalar),
6902 _ => unreachable!(),
6903 };
6904 let (b_c, b_r) = match *func_ctx.resolve_type(b, &module.types) {
6905 crate::TypeInner::CooperativeMatrix { columns, rows, .. } => (columns, rows),
6906 _ => unreachable!(),
6907 };
6908 let c_scalar = match *func_ctx.resolve_type(c, &module.types) {
6909 crate::TypeInner::CooperativeMatrix { scalar, .. } => scalar,
6910 _ => unreachable!(),
6911 };
6912 let wrapped = WrappedFunction::CooperativeMultiplyAdd {
6913 space_name,
6914 columns: b_c,
6915 rows: a_r,
6916 intermediate: a_c,
6917 ab_scalar,
6918 c_scalar,
6919 };
6920 if !self.wrapped_functions.insert(wrapped) {
6921 return Ok(());
6922 }
6923 let ab_scalar_name = ab_scalar.to_msl_name();
6924 let c_scalar_name = c_scalar.to_msl_name();
6925 writeln!(
6926 self.out,
6927 "{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) {{",
6928 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,
6929 )?;
6930 let l1 = back::Level(1);
6931 writeln!(
6932 self.out,
6933 "{l1}{NAMESPACE}::simdgroup_{c_scalar_name}{}x{} d;",
6934 b_c as u32, a_r as u32
6935 )?;
6936 writeln!(self.out, "{l1}simdgroup_multiply_accumulate(d,a,b,c);")?;
6937 writeln!(self.out, "{l1}return d;")?;
6938 writeln!(self.out, "}}")?;
6939 writeln!(self.out)?;
6940 Ok(())
6941 }
6942
6943 pub(super) fn write_wrapped_functions(
6944 &mut self,
6945 module: &crate::Module,
6946 func_ctx: &back::FunctionCtx,
6947 options: &Options,
6948 ) -> BackendResult {
6949 for (expr_handle, expr) in func_ctx.expressions.iter() {
6950 match *expr {
6951 crate::Expression::Unary { op, expr: operand } => {
6952 self.write_wrapped_unary_op(module, func_ctx, op, operand)?;
6953 }
6954 crate::Expression::Binary { op, left, right } => {
6955 self.write_wrapped_binary_op(module, func_ctx, expr_handle, op, left, right)?;
6956 }
6957 crate::Expression::Math {
6958 fun,
6959 arg,
6960 arg1,
6961 arg2,
6962 arg3,
6963 } => {
6964 self.write_wrapped_math_function(module, func_ctx, fun, arg, arg1, arg2, arg3)?;
6965 }
6966 crate::Expression::As { expr, target } => {
6967 self.write_wrapped_cast(module, func_ctx, expr, target)?;
6968 }
6969 crate::Expression::ImageLoad {
6970 image,
6971 coordinate,
6972 array_index,
6973 sample,
6974 level,
6975 } => {
6976 self.write_wrapped_image_load(
6977 module,
6978 func_ctx,
6979 image,
6980 coordinate,
6981 array_index,
6982 sample,
6983 level,
6984 )?;
6985 }
6986 crate::Expression::ImageSample {
6987 image,
6988 sampler,
6989 gather,
6990 coordinate,
6991 array_index,
6992 offset,
6993 level,
6994 depth_ref,
6995 clamp_to_edge,
6996 } => {
6997 self.write_wrapped_image_sample(
6998 module,
6999 func_ctx,
7000 image,
7001 sampler,
7002 gather,
7003 coordinate,
7004 array_index,
7005 offset,
7006 level,
7007 depth_ref,
7008 clamp_to_edge,
7009 )?;
7010 }
7011 crate::Expression::ImageQuery { image, query } => {
7012 self.write_wrapped_image_query(module, func_ctx, image, query)?;
7013 }
7014 crate::Expression::CooperativeLoad {
7015 columns,
7016 rows,
7017 role: _,
7018 ref data,
7019 } => {
7020 self.write_wrapped_cooperative_load(
7021 module,
7022 func_ctx,
7023 columns,
7024 rows,
7025 data.pointer,
7026 )?;
7027 }
7028 crate::Expression::CooperativeMultiplyAdd { a, b, c } => {
7029 let space = crate::AddressSpace::Private;
7030 self.write_wrapped_cooperative_multiply_add(module, func_ctx, space, a, b, c)?;
7031 }
7032 crate::Expression::RayQueryGetIntersection { committed, .. } => {
7033 self.write_rq_get_intersection_function(module, committed, options)?;
7034 }
7035 _ => {}
7036 }
7037 }
7038
7039 Ok(())
7040 }
7041
7042 fn write_functions(
7044 &mut self,
7045 module: &crate::Module,
7046 mod_info: &valid::ModuleInfo,
7047 options: &Options,
7048 pipeline_options: &PipelineOptions,
7049 ) -> Result<TranslationInfo, Error> {
7050 use nt::VertexFormat;
7051
7052 struct AttributeMappingResolved {
7055 ty_name: String,
7056 dimension: Option<crate::VectorSize>,
7057 scalar: crate::Scalar,
7058 name: String,
7059 }
7060 let mut am_resolved = FastHashMap::<u32, AttributeMappingResolved>::default();
7061
7062 struct VertexBufferMappingResolved<'a> {
7063 id: u32,
7064 stride: u32,
7065 step_mode: back::msl::VertexBufferStepMode,
7066 ty_name: String,
7067 param_name: String,
7068 elem_name: String,
7069 attributes: &'a Vec<back::msl::AttributeMapping>,
7070 }
7071 let mut vbm_resolved = Vec::<VertexBufferMappingResolved>::new();
7072
7073 struct UnpackingFunction {
7075 name: String,
7076 byte_count: u32,
7077 dimension: Option<crate::VectorSize>,
7078 scalar: crate::Scalar,
7079 }
7080 let mut unpacking_functions = FastHashMap::<VertexFormat, UnpackingFunction>::default();
7081
7082 let mut needs_vertex_id = false;
7088 let v_id = self.namer.call("v_id");
7089
7090 let mut needs_instance_id = false;
7091 let i_id = self.namer.call("i_id");
7092 if pipeline_options.vertex_pulling_transform {
7093 for vbm in &pipeline_options.vertex_buffer_mappings {
7094 let buffer_id = vbm.id;
7095 let buffer_stride = vbm.stride;
7096
7097 assert!(
7098 buffer_stride > 0,
7099 "Vertex pulling requires a non-zero buffer stride."
7100 );
7101
7102 match vbm.step_mode {
7103 back::msl::VertexBufferStepMode::Constant => {}
7104 back::msl::VertexBufferStepMode::ByVertex => {
7105 needs_vertex_id = true;
7106 }
7107 back::msl::VertexBufferStepMode::ByInstance => {
7108 needs_instance_id = true;
7109 }
7110 }
7111
7112 let buffer_ty = self.namer.call(format!("vb_{buffer_id}_type").as_str());
7113 let buffer_param = self.namer.call(format!("vb_{buffer_id}_in").as_str());
7114 let buffer_elem = self.namer.call(format!("vb_{buffer_id}_elem").as_str());
7115
7116 vbm_resolved.push(VertexBufferMappingResolved {
7117 id: buffer_id,
7118 stride: buffer_stride,
7119 step_mode: vbm.step_mode,
7120 ty_name: buffer_ty,
7121 param_name: buffer_param,
7122 elem_name: buffer_elem,
7123 attributes: &vbm.attributes,
7124 });
7125
7126 for attribute in &vbm.attributes {
7128 if unpacking_functions.contains_key(&attribute.format) {
7129 continue;
7130 }
7131 let (name, byte_count, dimension, scalar) =
7132 match self.write_unpacking_function(attribute.format) {
7133 Ok((name, byte_count, dimension, scalar)) => {
7134 (name, byte_count, dimension, scalar)
7135 }
7136 _ => {
7137 continue;
7138 }
7139 };
7140 unpacking_functions.insert(
7141 attribute.format,
7142 UnpackingFunction {
7143 name,
7144 byte_count,
7145 dimension,
7146 scalar,
7147 },
7148 );
7149 }
7150 }
7151 }
7152
7153 let mut pass_through_globals = Vec::new();
7154 for (fun_handle, fun) in module.functions.iter() {
7155 log::trace!(
7156 "function {:?}, handle {:?}",
7157 fun.name.as_deref().unwrap_or("(anonymous)"),
7158 fun_handle
7159 );
7160
7161 let ctx = back::FunctionCtx {
7162 ty: back::FunctionType::Function(fun_handle),
7163 info: &mod_info[fun_handle],
7164 expressions: &fun.expressions,
7165 named_expressions: &fun.named_expressions,
7166 };
7167
7168 writeln!(self.out)?;
7169 self.write_wrapped_functions(module, &ctx, options)?;
7170
7171 let fun_info = &mod_info[fun_handle];
7172 pass_through_globals.clear();
7173 let mut needs_buffer_sizes = false;
7174 for (handle, var) in module.global_variables.iter() {
7175 if !fun_info[handle].is_empty() {
7176 if var.space.needs_pass_through() {
7177 pass_through_globals.push(handle);
7178 }
7179 needs_buffer_sizes |= module.types[var.ty]
7180 .inner
7181 .needs_host_buffer_byte_size(&module.types);
7182 }
7183 }
7184
7185 let fun_name = &self.names[&NameKey::Function(fun_handle)];
7186 match fun.result {
7187 Some(ref result) => {
7188 let ty_name = TypeContext {
7189 handle: result.ty,
7190 gctx: module.to_ctx(),
7191 names: &self.names,
7192 access: crate::StorageAccess::empty(),
7193 first_time: false,
7194 };
7195 write!(self.out, "{ty_name}")?;
7196 }
7197 None => {
7198 write!(self.out, "void")?;
7199 }
7200 }
7201 writeln!(self.out, " {fun_name}(")?;
7202
7203 for (index, arg) in fun.arguments.iter().enumerate() {
7204 let name = &self.names[&NameKey::FunctionArgument(fun_handle, index as u32)];
7205 let param_type_name = TypeContext {
7206 handle: arg.ty,
7207 gctx: module.to_ctx(),
7208 names: &self.names,
7209 access: match module.types[arg.ty].inner {
7210 crate::TypeInner::Image {
7211 class: crate::ImageClass::Storage { access, .. },
7212 ..
7213 } => access,
7214 _ => crate::StorageAccess::empty(),
7215 },
7216 first_time: false,
7217 };
7218 let separator = separate(
7219 !pass_through_globals.is_empty()
7220 || index + 1 != fun.arguments.len()
7221 || needs_buffer_sizes,
7222 );
7223 writeln!(
7224 self.out,
7225 "{}{} {}{}",
7226 back::INDENT,
7227 param_type_name,
7228 name,
7229 separator
7230 )?;
7231 }
7232 for (index, &handle) in pass_through_globals.iter().enumerate() {
7233 let tyvar = TypedGlobalVariable {
7234 module,
7235 names: &self.names,
7236 handle,
7237 usage: fun_info[handle],
7238 reference: true,
7239 };
7240 let separator =
7241 separate(index + 1 != pass_through_globals.len() || needs_buffer_sizes);
7242 write!(self.out, "{}", back::INDENT)?;
7243 tyvar.try_fmt(&mut self.out)?;
7244 writeln!(self.out, "{separator}")?;
7245 }
7246
7247 if needs_buffer_sizes {
7248 writeln!(
7249 self.out,
7250 "{}constant _mslBufferSizes& _buffer_sizes",
7251 back::INDENT
7252 )?;
7253 }
7254
7255 writeln!(self.out, ") {{")?;
7256
7257 let guarded_indices =
7258 index::find_checked_indexes(module, fun, fun_info, options.bounds_check_policies);
7259
7260 let context = StatementContext {
7261 expression: ExpressionContext {
7262 function: fun,
7263 origin: FunctionOrigin::Handle(fun_handle),
7264 info: fun_info,
7265 lang_version: options.lang_version,
7266 policies: options.bounds_check_policies,
7267 guarded_indices,
7268 module,
7269 mod_info,
7270 pipeline_options,
7271 force_loop_bounding: options.force_loop_bounding,
7272 emit_int_div_checks: options.emit_int_div_checks,
7273 ray_query_initialization_tracking: options.ray_query_initialization_tracking,
7274 },
7275 result_struct: None,
7276 };
7277
7278 self.put_locals(&context.expression)?;
7279 self.update_expressions_to_bake(fun, fun_info, &context.expression);
7280 self.put_block(back::Level(1), &fun.body, &context)?;
7281 writeln!(self.out, "}}")?;
7282 self.named_expressions.clear();
7283 self.hoisted_loop_expressions.clear();
7284 }
7285
7286 let ep_range = get_entry_points(module, pipeline_options.entry_point.as_ref())
7287 .map_err(|(stage, name)| Error::EntryPointNotFound(stage, name))?;
7288
7289 let mut info = TranslationInfo {
7290 entry_point_names: Vec::with_capacity(ep_range.len()),
7291 };
7292
7293 for ep_index in ep_range {
7294 let ep = &module.entry_points[ep_index];
7295 let fun = &ep.function;
7296 let fun_info = mod_info.get_entry_point(ep_index);
7297 let mut ep_error = None;
7298
7299 let mut v_existing_id = None;
7303 let mut i_existing_id = None;
7304
7305 log::trace!(
7306 "entry point {:?}, index {:?}",
7307 fun.name.as_deref().unwrap_or("(anonymous)"),
7308 ep_index
7309 );
7310
7311 let ctx = back::FunctionCtx {
7312 ty: back::FunctionType::EntryPoint(ep_index as u16),
7313 info: fun_info,
7314 expressions: &fun.expressions,
7315 named_expressions: &fun.named_expressions,
7316 };
7317
7318 self.write_wrapped_functions(module, &ctx, options)?;
7319
7320 let (em_str, in_mode, out_mode, can_vertex_pull) = match ep.stage {
7321 crate::ShaderStage::Vertex => (
7322 Some("vertex"),
7323 LocationMode::VertexInput,
7324 LocationMode::VertexOutput,
7325 true,
7326 ),
7327 crate::ShaderStage::Fragment => (
7328 Some("fragment"),
7329 LocationMode::FragmentInput,
7330 LocationMode::FragmentOutput,
7331 false,
7332 ),
7333 crate::ShaderStage::Compute => (
7334 Some("kernel"),
7335 LocationMode::Uniform,
7336 LocationMode::Uniform,
7337 false,
7338 ),
7339 crate::ShaderStage::Task => {
7340 (None, LocationMode::Uniform, LocationMode::Uniform, false)
7341 }
7342 crate::ShaderStage::Mesh => {
7343 (None, LocationMode::Uniform, LocationMode::MeshOutput, false)
7344 }
7345 crate::ShaderStage::RayGeneration
7346 | crate::ShaderStage::AnyHit
7347 | crate::ShaderStage::ClosestHit
7348 | crate::ShaderStage::Miss => unimplemented!(),
7349 };
7350
7351 let do_vertex_pulling = can_vertex_pull
7353 && pipeline_options.vertex_pulling_transform
7354 && !pipeline_options.vertex_buffer_mappings.is_empty();
7355
7356 let needs_buffer_sizes = do_vertex_pulling
7358 || module
7359 .global_variables
7360 .iter()
7361 .filter(|&(handle, _)| !fun_info[handle].is_empty())
7362 .any(|(_, var)| {
7363 module.types[var.ty]
7364 .inner
7365 .needs_host_buffer_byte_size(&module.types)
7366 });
7367
7368 if !options.fake_missing_bindings {
7371 for (var_handle, var) in module.global_variables.iter() {
7372 if fun_info[var_handle].is_empty() {
7373 continue;
7374 }
7375 match var.space {
7376 crate::AddressSpace::Uniform
7377 | crate::AddressSpace::Storage { .. }
7378 | crate::AddressSpace::Handle => {
7379 let br = match var.binding {
7380 Some(ref br) => br,
7381 None => {
7382 let var_name = var.name.clone().unwrap_or_default();
7383 ep_error =
7384 Some(super::EntryPointError::MissingBinding(var_name));
7385 break;
7386 }
7387 };
7388 let target = options.get_resource_binding_target(ep, br);
7389 let good = match target {
7390 Some(target) => {
7391 match module.types[var.ty].inner {
7395 crate::TypeInner::Image {
7396 class: crate::ImageClass::External,
7397 ..
7398 } => target.external_texture.is_some(),
7399 crate::TypeInner::Image { .. } => target.texture.is_some(),
7400 crate::TypeInner::Sampler { .. } => {
7401 target.sampler.is_some()
7402 }
7403 _ => target.buffer.is_some(),
7404 }
7405 }
7406 None => false,
7407 };
7408 if !good {
7409 ep_error = Some(super::EntryPointError::MissingBindTarget(*br));
7410 break;
7411 }
7412 }
7413 crate::AddressSpace::Immediate => {
7414 if let Err(e) = options.resolve_immediates(ep) {
7415 ep_error = Some(e);
7416 break;
7417 }
7418 }
7419 crate::AddressSpace::Function
7420 | crate::AddressSpace::Private
7421 | crate::AddressSpace::WorkGroup
7422 | crate::AddressSpace::TaskPayload => {}
7423 crate::AddressSpace::RayPayload
7424 | crate::AddressSpace::IncomingRayPayload => unimplemented!(),
7425 }
7426 }
7427 if needs_buffer_sizes {
7428 if let Err(err) = options.resolve_sizes_buffer(ep) {
7429 ep_error = Some(err);
7430 }
7431 }
7432 }
7433
7434 if let Some(err) = ep_error {
7435 info.entry_point_names.push(Err(err));
7436 continue;
7437 }
7438 let fun_name = self.names[&NameKey::EntryPoint(ep_index as _)].clone();
7439 info.entry_point_names.push(Ok(fun_name.clone()));
7440
7441 writeln!(self.out)?;
7442
7443 let mut flattened_member_names = FastHashMap::default();
7449 let mut varyings_namer = proc::Namer::default();
7451
7452 let mut empty_names = FastHashMap::default(); varyings_namer.reset(
7454 module,
7455 &super::keywords::RESERVED_SET,
7456 proc::KeywordSet::empty(),
7457 proc::CaseInsensitiveKeywordSet::empty(),
7458 &[CLAMPED_LOD_LOAD_PREFIX],
7459 &mut empty_names,
7460 );
7461
7462 let mut flattened_arguments = Vec::new();
7467 for (arg_index, arg) in fun.arguments.iter().enumerate() {
7468 match module.types[arg.ty].inner {
7469 crate::TypeInner::Struct { ref members, .. } => {
7470 for (member_index, member) in members.iter().enumerate() {
7471 let member_index = member_index as u32;
7472 flattened_arguments.push((
7473 NameKey::StructMember(arg.ty, member_index),
7474 member.ty,
7475 member.binding.as_ref(),
7476 ));
7477 let name_key = NameKey::StructMember(arg.ty, member_index);
7478 let name = match member.binding {
7479 Some(crate::Binding::Location { .. }) => {
7480 if do_vertex_pulling {
7481 self.namer.call(&self.names[&name_key])
7482 } else {
7483 varyings_namer.call(&self.names[&name_key])
7484 }
7485 }
7486 _ => self.namer.call(&self.names[&name_key]),
7487 };
7488 flattened_member_names.insert(name_key, name);
7489 }
7490 }
7491 _ => flattened_arguments.push((
7492 NameKey::EntryPointArgument(ep_index as _, arg_index as u32),
7493 arg.ty,
7494 arg.binding.as_ref(),
7495 )),
7496 }
7497 }
7498
7499 let stage_in_name = self.namer.call(&format!("{fun_name}Input"));
7504 let varyings_member_name = self.namer.call("varyings");
7505 let mut has_varyings = false;
7506
7507 if !flattened_arguments.is_empty() {
7508 if !do_vertex_pulling {
7509 writeln!(self.out, "struct {stage_in_name} {{")?;
7510 }
7511 for &(ref name_key, ty, binding) in flattened_arguments.iter() {
7512 let Some(binding) = binding else {
7513 continue;
7514 };
7515 let name = match *name_key {
7516 NameKey::StructMember(..) => &flattened_member_names[name_key],
7517 _ => &self.names[name_key],
7518 };
7519 let ty_name = TypeContext {
7520 handle: ty,
7521 gctx: module.to_ctx(),
7522 names: &self.names,
7523 access: crate::StorageAccess::empty(),
7524 first_time: false,
7525 };
7526 let resolved = options.resolve_local_binding(binding, in_mode)?;
7527 let location = match *binding {
7528 crate::Binding::Location { location, .. } => Some(location),
7529 crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. }) => None,
7530 crate::Binding::BuiltIn(_) => continue,
7531 };
7532 if do_vertex_pulling {
7533 let Some(location) = location else {
7534 continue;
7535 };
7536 am_resolved.insert(
7538 location,
7539 AttributeMappingResolved {
7540 ty_name: ty_name.to_string(),
7541 dimension: ty_name.vector_size(),
7542 scalar: ty_name.scalar().unwrap(),
7543 name: name.to_string(),
7544 },
7545 );
7546 } else {
7547 has_varyings = true;
7548 if let super::ResolvedBinding::User {
7549 prefix,
7550 index,
7551 interpolation: Some(super::ResolvedInterpolation::PerVertex),
7552 } = resolved
7553 {
7554 if options.lang_version < (4, 0) {
7555 return Err(Error::PerVertexNotSupported);
7556 }
7557 write!(
7558 self.out,
7559 "{}{NAMESPACE}::vertex_value<{}> {name} [[user({prefix}{index})]]",
7560 back::INDENT,
7561 ty_name.unwrap_array()
7562 )?;
7563 } else {
7564 write!(self.out, "{}{} {}", back::INDENT, ty_name, name)?;
7565 resolved.try_fmt(&mut self.out)?;
7566 }
7567 writeln!(self.out, ";")?;
7568 }
7569 }
7570 if !do_vertex_pulling {
7571 writeln!(self.out, "}};")?;
7572 }
7573 }
7574
7575 let stage_out_name = self.namer.call(&format!("{fun_name}Output"));
7578 let result_member_name = self.namer.call("member");
7579 let result_type_name = match fun.result {
7580 Some(ref result)
7581 if ep.stage != crate::ShaderStage::Task
7582 && !only_writes_discarded_depth(module, ep) =>
7583 {
7584 let mut result_members = Vec::new();
7585 if let crate::TypeInner::Struct { ref members, .. } =
7586 module.types[result.ty].inner
7587 {
7588 for (member_index, member) in members.iter().enumerate() {
7589 result_members.push((
7590 &self.names[&NameKey::StructMember(result.ty, member_index as u32)],
7591 member.ty,
7592 member.binding.as_ref(),
7593 ));
7594 }
7595 } else {
7596 result_members.push((
7597 &result_member_name,
7598 result.ty,
7599 result.binding.as_ref(),
7600 ));
7601 }
7602
7603 writeln!(self.out, "struct {stage_out_name} {{")?;
7604 let mut has_point_size = false;
7605 for (name, ty, binding) in result_members {
7606 let ty_name = TypeContext {
7607 handle: ty,
7608 gctx: module.to_ctx(),
7609 names: &self.names,
7610 access: crate::StorageAccess::empty(),
7611 first_time: true,
7612 };
7613 let binding = binding.ok_or_else(|| {
7614 Error::GenericValidation("Expected binding, got None".into())
7615 })?;
7616
7617 if let crate::Binding::BuiltIn(crate::BuiltIn::PointSize) = *binding {
7618 has_point_size = true;
7619 if !pipeline_options.allow_and_force_point_size {
7620 continue;
7621 }
7622 }
7623
7624 if discards_frag_depth(ep) && is_frag_depth(Some(binding)) {
7625 continue;
7626 }
7627
7628 let array_len = match module.types[ty].inner {
7629 crate::TypeInner::Array {
7630 size: crate::ArraySize::Constant(size),
7631 ..
7632 } => Some(size),
7633 _ => None,
7634 };
7635 let resolved = match (
7636 options.resolve_local_binding(binding, out_mode)?,
7637 ep.early_depth_test,
7638 ) {
7639 (
7640 super::ResolvedBinding::BuiltIn(crate::BuiltIn::FragDepth),
7641 Some(crate::EarlyDepthTest::Allow { conservative }),
7642 ) => super::ResolvedBinding::ConservativeDepth(conservative),
7643 (resolved, _) => resolved,
7644 };
7645 write!(self.out, "{}{} {}", back::INDENT, ty_name, name)?;
7646 resolved.try_fmt(&mut self.out)?;
7647 if let Some(array_len) = array_len {
7648 write!(self.out, " [{array_len}]")?;
7649 }
7650 writeln!(self.out, ";")?;
7651 }
7652
7653 if pipeline_options.allow_and_force_point_size
7654 && ep.stage == crate::ShaderStage::Vertex
7655 && !has_point_size
7656 {
7657 writeln!(
7659 self.out,
7660 "{}float _point_size [[point_size]];",
7661 back::INDENT
7662 )?;
7663 }
7664 writeln!(self.out, "}};")?;
7665 &stage_out_name
7666 }
7667 Some(ref result) if ep.stage == crate::ShaderStage::Task => {
7668 assert_eq!(
7669 module.types[result.ty].inner,
7670 crate::TypeInner::Vector {
7671 size: crate::VectorSize::Tri,
7672 scalar: crate::Scalar::U32
7673 }
7674 );
7675
7676 "metal::uint3"
7677 }
7678 _ => "void",
7679 };
7680
7681 let out_mesh_info = if let Some(ref mesh_info) = ep.mesh_info {
7682 Some(self.write_mesh_output_types(
7683 mesh_info,
7684 &fun_name,
7685 module,
7686 pipeline_options.allow_and_force_point_size,
7687 options,
7688 )?)
7689 } else {
7690 None
7691 };
7692
7693 if do_vertex_pulling {
7696 for vbm in &vbm_resolved {
7697 let buffer_stride = vbm.stride;
7698 let buffer_ty = &vbm.ty_name;
7699
7700 writeln!(
7704 self.out,
7705 "struct {buffer_ty} {{ metal::uchar data[{buffer_stride}]; }};"
7706 )?;
7707 }
7708 }
7709
7710 let is_wrapped = matches!(
7711 ep.stage,
7712 crate::ShaderStage::Task | crate::ShaderStage::Mesh
7713 );
7714 let fun_name = fun_name.clone();
7715 let nested_fun_name = if is_wrapped {
7716 self.namer.call(&format!("_{fun_name}"))
7717 } else {
7718 fun_name.clone()
7719 };
7720
7721 if ep.stage == crate::ShaderStage::Compute && options.lang_version >= (2, 1) {
7723 let total_threads =
7724 ep.workgroup_size[0] * ep.workgroup_size[1] * ep.workgroup_size[2];
7725 write!(
7726 self.out,
7727 "[[max_total_threads_per_threadgroup({total_threads})]] "
7728 )?;
7729 }
7730
7731 if ep.stage == crate::ShaderStage::Fragment
7732 && ep.early_depth_test == Some(crate::EarlyDepthTest::Force)
7733 {
7734 if options.lang_version < (1, 2) {
7735 return Err(Error::UnsupportedAttribute(
7736 "early_fragment_tests".to_string(),
7737 ));
7738 }
7739 write!(self.out, "[[early_fragment_tests]] ")?;
7740 }
7741
7742 if let Some(em_str) = em_str {
7744 write!(self.out, "{em_str} ")?;
7745 }
7746 writeln!(self.out, "{result_type_name} {nested_fun_name}(")?;
7747
7748 let mut args = Vec::new();
7749
7750 if has_varyings {
7753 args.push(EntryPointArgument {
7754 ty_name: stage_in_name,
7755 name: varyings_member_name.clone(),
7756 binding: " [[stage_in]]".to_string(),
7757 init: None,
7758 });
7759 }
7760
7761 let mut local_invocation_index = None;
7762
7763 for &(ref name_key, ty, binding) in flattened_arguments.iter() {
7766 let binding = match binding {
7767 Some(&crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. })) => continue,
7768 Some(binding @ &crate::Binding::BuiltIn { .. }) => binding,
7769 _ => continue,
7770 };
7771 let name = match *name_key {
7772 NameKey::StructMember(..) => &flattened_member_names[name_key],
7773 _ => &self.names[name_key],
7774 };
7775
7776 if binding == &crate::Binding::BuiltIn(crate::BuiltIn::LocalInvocationIndex) {
7777 local_invocation_index = Some(name_key);
7778 }
7779
7780 let ty_name = TypeContext {
7781 handle: ty,
7782 gctx: module.to_ctx(),
7783 names: &self.names,
7784 access: crate::StorageAccess::empty(),
7785 first_time: false,
7786 };
7787
7788 match *binding {
7789 crate::Binding::BuiltIn(crate::BuiltIn::VertexIndex) => {
7790 v_existing_id = Some(name.clone());
7791 }
7792 crate::Binding::BuiltIn(crate::BuiltIn::InstanceIndex) => {
7793 i_existing_id = Some(name.clone());
7794 }
7795 _ => {}
7796 };
7797
7798 let resolved = options.resolve_local_binding(binding, in_mode)?;
7799 let mut binding = String::new();
7800 resolved.try_fmt(&mut binding)?;
7801
7802 args.push(EntryPointArgument {
7803 ty_name: format!("{ty_name}"),
7804 name: name.clone(),
7805 binding,
7806 init: None,
7807 });
7808 }
7809
7810 let need_workgroup_variables_initialization =
7811 self.need_workgroup_variables_initialization(options, ep, module, fun_info);
7812
7813 if local_invocation_index.is_none()
7814 && (need_workgroup_variables_initialization
7815 || ep.stage == crate::ShaderStage::Task
7816 || ep.stage == crate::ShaderStage::Mesh)
7817 {
7818 args.push(EntryPointArgument {
7819 ty_name: "uint".to_string(),
7820 name: "__local_invocation_index".to_string(),
7821 binding: " [[thread_index_in_threadgroup]]".to_string(),
7822 init: None,
7823 });
7824 }
7825
7826 for (handle, var) in module.global_variables.iter() {
7831 let usage = fun_info[handle];
7832 if usage.is_empty() || var.space == crate::AddressSpace::Private {
7833 continue;
7834 }
7835
7836 if options.lang_version < (1, 2) {
7837 match var.space {
7838 crate::AddressSpace::Storage { access }
7848 if access.contains(crate::StorageAccess::STORE)
7849 && ep.stage == crate::ShaderStage::Fragment =>
7850 {
7851 return Err(Error::UnsupportedWritableStorageBuffer)
7852 }
7853 crate::AddressSpace::Handle => {
7854 match module.types[var.ty].inner {
7855 crate::TypeInner::Image {
7856 class: crate::ImageClass::Storage { access, .. },
7857 ..
7858 } => {
7859 if access.contains(crate::StorageAccess::STORE)
7869 && (ep.stage == crate::ShaderStage::Vertex
7870 || ep.stage == crate::ShaderStage::Fragment)
7871 {
7872 return Err(Error::UnsupportedWritableStorageTexture(
7873 ep.stage,
7874 ));
7875 }
7876
7877 if access.contains(
7878 crate::StorageAccess::LOAD | crate::StorageAccess::STORE,
7879 ) {
7880 return Err(Error::UnsupportedRWStorageTexture);
7881 }
7882 }
7883 _ => {}
7884 }
7885 }
7886 _ => {}
7887 }
7888 }
7889
7890 match var.space {
7892 crate::AddressSpace::Handle => match module.types[var.ty].inner {
7893 crate::TypeInner::BindingArray { base, .. } => {
7894 match module.types[base].inner {
7895 crate::TypeInner::Sampler { .. } => {
7896 if options.lang_version < (2, 0) {
7897 return Err(Error::UnsupportedArrayOf(
7898 "samplers".to_string(),
7899 ));
7900 }
7901 }
7902 crate::TypeInner::Image { class, .. } => match class {
7903 crate::ImageClass::Sampled { .. }
7904 | crate::ImageClass::Depth { .. }
7905 | crate::ImageClass::Storage {
7906 access: crate::StorageAccess::LOAD,
7907 ..
7908 } => {
7909 if options.lang_version < (2, 0) {
7914 return Err(Error::UnsupportedArrayOf(
7915 "textures".to_string(),
7916 ));
7917 }
7918 }
7919 crate::ImageClass::Storage {
7920 access: crate::StorageAccess::STORE,
7921 ..
7922 } => {
7923 if options.lang_version < (2, 0) {
7928 return Err(Error::UnsupportedArrayOf(
7929 "write-only textures".to_string(),
7930 ));
7931 }
7932 }
7933 crate::ImageClass::Storage { .. } => {
7934 if options.lang_version < (3, 0) {
7935 return Err(Error::UnsupportedArrayOf(
7936 "read-write textures".to_string(),
7937 ));
7938 }
7939 }
7940 crate::ImageClass::External => {
7941 return Err(Error::UnsupportedArrayOf(
7942 "external textures".to_string(),
7943 ));
7944 }
7945 },
7946 _ => {
7947 return Err(Error::UnsupportedArrayOfType(base));
7948 }
7949 }
7950 }
7951 _ => {}
7952 },
7953 _ => {}
7954 }
7955
7956 let resolved = match var.space {
7958 crate::AddressSpace::Immediate => options.resolve_immediates(ep).ok(),
7959 crate::AddressSpace::WorkGroup => None,
7960 crate::AddressSpace::TaskPayload => Some(back::msl::ResolvedBinding::Payload),
7961 _ => options
7962 .resolve_resource_binding(ep, var.binding.as_ref().unwrap())
7963 .ok(),
7964 };
7965 if let Some(ref resolved) = resolved {
7966 if resolved.as_inline_sampler(options).is_some() {
7968 continue;
7969 }
7970 }
7971
7972 match module.types[var.ty].inner {
7973 crate::TypeInner::Image {
7974 class: crate::ImageClass::External,
7975 ..
7976 } => {
7977 let target = match resolved {
7981 Some(back::msl::ResolvedBinding::Resource(target)) => {
7982 target.external_texture
7983 }
7984 _ => None,
7985 };
7986
7987 for i in 0..3 {
7988 let plane_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
7989 handle,
7990 ExternalTextureNameKey::Plane(i),
7991 )];
7992 let ty_name = format!(
7993 "{NAMESPACE}::texture2d<float, {NAMESPACE}::access::sample>"
7994 );
7995 let name = plane_name.clone();
7996 let binding = if let Some(ref target) = target {
7997 format!(" [[texture({})]]", target.planes[i])
7998 } else {
7999 String::new()
8000 };
8001 args.push(EntryPointArgument {
8002 ty_name,
8003 name,
8004 binding,
8005 init: None,
8006 });
8007 }
8008 let params_ty_name = &self.names
8009 [&NameKey::Type(module.special_types.external_texture_params.unwrap())];
8010 let params_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
8011 handle,
8012 ExternalTextureNameKey::Params,
8013 )];
8014 let binding = if let Some(ref target) = target {
8015 format!(" [[buffer({})]]", target.params)
8016 } else {
8017 String::new()
8018 };
8019
8020 args.push(EntryPointArgument {
8021 ty_name: format!("constant {params_ty_name}&"),
8022 name: params_name.clone(),
8023 binding,
8024 init: None,
8025 });
8026 }
8027 _ => {
8028 if var.space == crate::AddressSpace::WorkGroup
8029 && ep.stage == crate::ShaderStage::Mesh
8030 {
8031 continue;
8032 }
8033 let tyvar = TypedGlobalVariable {
8034 module,
8035 names: &self.names,
8036 handle,
8037 usage,
8038 reference: true,
8039 };
8040 let parts = tyvar.to_parts()?;
8041 let mut binding = String::new();
8042 if let Some(resolved) = resolved {
8043 resolved.try_fmt(&mut binding)?;
8044 }
8045 args.push(EntryPointArgument {
8046 ty_name: parts.ty_name,
8047 name: parts.var_name,
8048 binding,
8049 init: var.init,
8050 });
8051 }
8052 }
8053 }
8054
8055 if do_vertex_pulling {
8056 if needs_vertex_id && v_existing_id.is_none() {
8057 args.push(EntryPointArgument {
8059 ty_name: "uint".to_string(),
8060 name: v_id.clone(),
8061 binding: " [[vertex_id]]".to_string(),
8062 init: None,
8063 });
8064 }
8065
8066 if needs_instance_id && i_existing_id.is_none() {
8067 args.push(EntryPointArgument {
8068 ty_name: "uint".to_string(),
8069 name: i_id.clone(),
8070 binding: " [[instance_id]]".to_string(),
8071 init: None,
8072 });
8073 }
8074
8075 for vbm in &vbm_resolved {
8078 let id = &vbm.id;
8079 let ty_name = &vbm.ty_name;
8080 let param_name = &vbm.param_name;
8081 args.push(EntryPointArgument {
8082 ty_name: format!("const device {ty_name}*"),
8083 name: param_name.clone(),
8084 binding: format!(" [[buffer({id})]]"),
8085 init: None,
8086 });
8087 }
8088 }
8089
8090 if needs_buffer_sizes {
8093 let resolved = options.resolve_sizes_buffer(ep).unwrap();
8095 let mut binding = String::new();
8096 resolved.try_fmt(&mut binding)?;
8097 args.push(EntryPointArgument {
8098 ty_name: "constant _mslBufferSizes&".to_string(),
8099 name: "_buffer_sizes".to_string(),
8100 binding,
8101 init: None,
8102 });
8103 }
8104
8105 let mut is_first_arg = true;
8106 for arg in &args {
8107 if is_first_arg {
8108 write!(self.out, " ")?;
8109 } else {
8110 write!(self.out, ", ")?;
8111 }
8112 is_first_arg = false;
8113 write!(self.out, "{} {}", arg.ty_name, arg.name)?;
8114 if !is_wrapped {
8115 write!(self.out, "{}", arg.binding)?;
8116 if let Some(init) = arg.init {
8117 write!(self.out, " = ")?;
8118 self.put_const_expression(
8119 init,
8120 module,
8121 mod_info,
8122 &module.global_expressions,
8123 )?;
8124 }
8125 }
8126 writeln!(self.out)?;
8127 }
8128 if ep.stage == crate::ShaderStage::Mesh {
8129 for (handle, var) in module.global_variables.iter() {
8130 if var.space != crate::AddressSpace::WorkGroup || fun_info[handle].is_empty() {
8131 continue;
8132 }
8133 if is_first_arg {
8134 write!(self.out, " ")?;
8135 } else {
8136 write!(self.out, ", ")?;
8137 }
8138 let ty_context = TypeContext {
8139 handle: module.global_variables[handle].ty,
8140 gctx: module.to_ctx(),
8141 names: &self.names,
8142 access: crate::StorageAccess::empty(),
8143 first_time: false,
8144 };
8145 writeln!(
8146 self.out,
8147 "threadgroup {ty_context}& {}",
8148 self.names[&NameKey::GlobalVariable(handle)]
8149 )?;
8150 }
8151 }
8152
8153 writeln!(self.out, ") {{")?;
8155
8156 if do_vertex_pulling {
8158 for vbm in &vbm_resolved {
8161 for attribute in vbm.attributes {
8162 let location = attribute.shader_location;
8163 let am_option = am_resolved.get(&location);
8164 if am_option.is_none() {
8165 continue;
8168 }
8169 let am = am_option.unwrap();
8170 let attribute_ty_name = &am.ty_name;
8171 let attribute_name = &am.name;
8172
8173 writeln!(
8174 self.out,
8175 "{}{attribute_ty_name} {attribute_name} = {{}};",
8176 back::Level(1)
8177 )?;
8178 }
8179
8180 write!(self.out, "{}if (", back::Level(1))?;
8183
8184 let idx = &vbm.id;
8185 let stride = &vbm.stride;
8186 let index_name = match vbm.step_mode {
8187 back::msl::VertexBufferStepMode::Constant => "0",
8188 back::msl::VertexBufferStepMode::ByVertex => {
8189 if let Some(ref name) = v_existing_id {
8190 name
8191 } else {
8192 &v_id
8193 }
8194 }
8195 back::msl::VertexBufferStepMode::ByInstance => {
8196 if let Some(ref name) = i_existing_id {
8197 name
8198 } else {
8199 &i_id
8200 }
8201 }
8202 };
8203 write!(
8204 self.out,
8205 "{index_name} < (_buffer_sizes.buffer_size{idx} / {stride})"
8206 )?;
8207
8208 writeln!(self.out, ") {{")?;
8209
8210 let ty_name = &vbm.ty_name;
8212 let elem_name = &vbm.elem_name;
8213 let param_name = &vbm.param_name;
8214
8215 writeln!(
8216 self.out,
8217 "{}const {ty_name} {elem_name} = {param_name}[{index_name}];",
8218 back::Level(2),
8219 )?;
8220
8221 for attribute in vbm.attributes {
8224 let location = attribute.shader_location;
8225 let Some(am) = am_resolved.get(&location) else {
8226 continue;
8230 };
8231 let attribute_name = &am.name;
8232 let attribute_ty_name = &am.ty_name;
8233
8234 let offset = attribute.offset;
8235 let func = unpacking_functions
8236 .get(&attribute.format)
8237 .expect("Should have generated this unpacking function earlier.");
8238 let func_name = &func.name;
8239
8240 let needs_padding_or_truncation = am.dimension.cmp(&func.dimension);
8246
8247 let needs_conversion = am.scalar != func.scalar;
8250
8251 if needs_padding_or_truncation != Ordering::Equal {
8252 writeln!(
8255 self.out,
8256 "{}// {attribute_ty_name} <- {:?}",
8257 back::Level(2),
8258 attribute.format
8259 )?;
8260 }
8261
8262 write!(self.out, "{}{attribute_name} = ", back::Level(2),)?;
8263
8264 if needs_padding_or_truncation == Ordering::Greater {
8265 write!(self.out, "{attribute_ty_name}(")?;
8267 }
8268
8269 if needs_conversion {
8271 put_numeric_type(&mut self.out, am.scalar, func.dimension.as_slice())?;
8272 write!(self.out, "(")?;
8273 }
8274 write!(self.out, "{func_name}({elem_name}.data[{offset}]")?;
8275 for i in (offset + 1)..(offset + func.byte_count) {
8276 write!(self.out, ", {elem_name}.data[{i}]")?;
8277 }
8278 write!(self.out, ")")?;
8279 if needs_conversion {
8280 write!(self.out, ")")?;
8281 }
8282
8283 match needs_padding_or_truncation {
8284 Ordering::Greater => {
8285 let ty_is_int = scalar_is_int(am.scalar);
8287 let zero_value = if ty_is_int { "0" } else { "0.0" };
8288 let one_value = if ty_is_int { "1" } else { "1.0" };
8289 for i in func.dimension.map_or(1, u8::from)
8290 ..am.dimension.map_or(1, u8::from)
8291 {
8292 write!(
8293 self.out,
8294 ", {}",
8295 if i == 3 { one_value } else { zero_value }
8296 )?;
8297 }
8298 }
8299 Ordering::Less => {
8300 write!(
8302 self.out,
8303 ".{}",
8304 &"xyzw"[0..usize::from(am.dimension.map_or(1, u8::from))]
8305 )?;
8306 }
8307 Ordering::Equal => {}
8308 }
8309
8310 if needs_padding_or_truncation == Ordering::Greater {
8311 write!(self.out, ")")?;
8312 }
8313
8314 writeln!(self.out, ";")?;
8315 }
8316
8317 writeln!(self.out, "{}}}", back::Level(1))?;
8319 }
8320 }
8321
8322 for (handle, var) in module.global_variables.iter() {
8325 let usage = fun_info[handle];
8326 if usage.is_empty() {
8327 continue;
8328 }
8329 if var.space == crate::AddressSpace::Private {
8330 let tyvar = TypedGlobalVariable {
8331 module,
8332 names: &self.names,
8333 handle,
8334 usage,
8335
8336 reference: false,
8337 };
8338 write!(self.out, "{}", back::INDENT)?;
8339 tyvar.try_fmt(&mut self.out)?;
8340 match var.init {
8341 Some(value) => {
8342 write!(self.out, " = ")?;
8343 self.put_const_expression(
8344 value,
8345 module,
8346 mod_info,
8347 &module.global_expressions,
8348 )?;
8349 writeln!(self.out, ";")?;
8350 }
8351 None => {
8352 writeln!(self.out, " = {{}};")?;
8353 }
8354 };
8355 } else if let Some(ref binding) = var.binding {
8356 let resolved = options.resolve_resource_binding(ep, binding).unwrap();
8357 if let Some(sampler) = resolved.as_inline_sampler(options) {
8358 let name = &self.names[&NameKey::GlobalVariable(handle)];
8360 writeln!(
8361 self.out,
8362 "{}constexpr {}::sampler {}(",
8363 back::INDENT,
8364 NAMESPACE,
8365 name
8366 )?;
8367 self.put_inline_sampler_properties(back::Level(2), sampler, options)?;
8368 writeln!(self.out, "{});", back::INDENT)?;
8369 } else if module.types[var.ty].inner.is_external_image() {
8370 let wrapper_name = &self.names[&NameKey::GlobalVariable(handle)];
8373 let l1 = back::Level(1);
8374 let l2 = l1.next();
8375 writeln!(
8376 self.out,
8377 "{l1}const {EXTERNAL_TEXTURE_WRAPPER_STRUCT} {wrapper_name} {{"
8378 )?;
8379 for i in 0..3 {
8380 let plane_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
8381 handle,
8382 ExternalTextureNameKey::Plane(i),
8383 )];
8384 writeln!(self.out, "{l2}.plane{i} = {plane_name},")?;
8385 }
8386 let params_name = &self.names[&NameKey::ExternalTextureGlobalVariable(
8387 handle,
8388 ExternalTextureNameKey::Params,
8389 )];
8390 writeln!(self.out, "{l2}.params = {params_name},")?;
8391 writeln!(self.out, "{l1}}};")?;
8392 }
8393 }
8394 }
8395
8396 if need_workgroup_variables_initialization {
8397 self.write_workgroup_variables_initialization(
8398 module,
8399 mod_info,
8400 fun_info,
8401 local_invocation_index,
8402 ep.stage,
8403 )?;
8404 }
8405
8406 for (arg_index, arg) in fun.arguments.iter().enumerate() {
8417 let arg_name =
8418 &self.names[&NameKey::EntryPointArgument(ep_index as _, arg_index as u32)];
8419 match module.types[arg.ty].inner {
8420 crate::TypeInner::Struct { ref members, .. } => {
8421 let struct_name = &self.names[&NameKey::Type(arg.ty)];
8422 write!(
8423 self.out,
8424 "{}const {} {} = {{ ",
8425 back::INDENT,
8426 struct_name,
8427 arg_name
8428 )?;
8429 for (member_index, member) in members.iter().enumerate() {
8430 let key = NameKey::StructMember(arg.ty, member_index as u32);
8431 let name = &flattened_member_names[&key];
8432 if member_index != 0 {
8433 write!(self.out, ", ")?;
8434 }
8435 if self
8437 .struct_member_pads
8438 .contains(&(arg.ty, member_index as u32))
8439 {
8440 write!(self.out, "{{}}, ")?;
8441 }
8442 match member.binding {
8443 Some(crate::Binding::Location {
8444 interpolation: Some(crate::Interpolation::PerVertex),
8445 ..
8446 }) => {
8447 writeln!(
8448 self.out,
8449 "{0}{{ {1}.{2}.get({NAMESPACE}::vertex_index::first), {1}.{2}.get({NAMESPACE}::vertex_index::second), {1}.{2}.get({NAMESPACE}::vertex_index::third) }}",
8450 back::INDENT,
8451 varyings_member_name,
8452 arg_name,
8453 )?;
8454 continue;
8455 }
8456 Some(crate::Binding::Location { .. }) if has_varyings => {
8457 write!(self.out, "{varyings_member_name}.")?;
8458 }
8459 _ => (),
8460 }
8461 write!(self.out, "{name}")?;
8462 }
8463 writeln!(self.out, " }};")?;
8464 }
8465 _ => match arg.binding {
8466 Some(crate::Binding::Location {
8467 interpolation: Some(crate::Interpolation::PerVertex),
8468 ..
8469 }) => {
8470 let ty_name = TypeContext {
8471 handle: arg.ty,
8472 gctx: module.to_ctx(),
8473 names: &self.names,
8474 access: crate::StorageAccess::empty(),
8475 first_time: false,
8476 };
8477 writeln!(
8478 self.out,
8479 "{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) }};",
8480 back::INDENT,
8481 varyings_member_name,
8482 arg_name,
8483 )?;
8484 }
8485 Some(crate::Binding::Location { .. })
8486 | Some(crate::Binding::BuiltIn(crate::BuiltIn::Barycentric { .. }))
8487 if has_varyings =>
8488 {
8489 writeln!(
8490 self.out,
8491 "{}const auto {} = {}.{};",
8492 back::INDENT,
8493 arg_name,
8494 varyings_member_name,
8495 arg_name
8496 )?;
8497 }
8498 _ => {}
8499 },
8500 }
8501 }
8502
8503 let guarded_indices =
8504 index::find_checked_indexes(module, fun, fun_info, options.bounds_check_policies);
8505
8506 let context = StatementContext {
8507 expression: ExpressionContext {
8508 function: fun,
8509 origin: FunctionOrigin::EntryPoint(ep_index as _),
8510 info: fun_info,
8511 lang_version: options.lang_version,
8512 policies: options.bounds_check_policies,
8513 guarded_indices,
8514 module,
8515 mod_info,
8516 pipeline_options,
8517 force_loop_bounding: options.force_loop_bounding,
8518 emit_int_div_checks: options.emit_int_div_checks,
8519 ray_query_initialization_tracking: options.ray_query_initialization_tracking,
8520 },
8521 result_struct: if ep.stage == crate::ShaderStage::Task {
8522 None
8523 } else {
8524 Some(&stage_out_name)
8525 },
8526 };
8527
8528 self.put_locals(&context.expression)?;
8531 self.update_expressions_to_bake(fun, fun_info, &context.expression);
8532 self.put_block(back::Level(1), &fun.body, &context)?;
8533 writeln!(self.out, "}}")?;
8534 if ep_index + 1 != module.entry_points.len() {
8535 writeln!(self.out)?;
8536 }
8537 self.named_expressions.clear();
8538 self.hoisted_loop_expressions.clear();
8539
8540 if is_wrapped {
8541 self.write_wrapper_function(NestedFunctionInfo {
8542 options,
8543 ep,
8544 module,
8545 mod_info,
8546 fun_info,
8547 args,
8548 local_invocation_index,
8549 nested_name: &nested_fun_name,
8550 outer_name: &fun_name,
8551 out_mesh_info,
8552 })?;
8553 }
8554 }
8555
8556 Ok(info)
8557 }
8558
8559 pub(super) fn write_barrier(
8560 &mut self,
8561 flags: crate::Barrier,
8562 level: back::Level,
8563 ) -> BackendResult {
8564 if flags.is_empty() {
8567 writeln!(
8568 self.out,
8569 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_none);",
8570 )?;
8571 }
8572 if flags.contains(crate::Barrier::STORAGE) {
8573 writeln!(
8574 self.out,
8575 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_device);",
8576 )?;
8577 }
8578 if flags.contains(crate::Barrier::WORK_GROUP) {
8579 writeln!(
8580 self.out,
8581 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_threadgroup);",
8582 )?;
8583 if self.needs_object_memory_barriers {
8584 writeln!(
8585 self.out,
8586 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_object_data);",
8587 )?;
8588 }
8589 }
8590 if flags.contains(crate::Barrier::SUB_GROUP) {
8591 writeln!(
8592 self.out,
8593 "{level}{NAMESPACE}::simdgroup_barrier({NAMESPACE}::mem_flags::mem_threadgroup);",
8594 )?;
8595 }
8596 if flags.contains(crate::Barrier::TEXTURE) {
8597 writeln!(
8598 self.out,
8599 "{level}{NAMESPACE}::threadgroup_barrier({NAMESPACE}::mem_flags::mem_texture);",
8600 )?;
8601 }
8602 Ok(())
8603 }
8604}
8605
8606mod workgroup_mem_init {
8609 use crate::EntryPoint;
8610
8611 use super::*;
8612
8613 enum Access {
8614 GlobalVariable(Handle<crate::GlobalVariable>),
8615 StructMember(Handle<crate::Type>, u32),
8616 Array(usize),
8617 }
8618
8619 impl Access {
8620 fn write<W: Write>(
8621 &self,
8622 writer: &mut W,
8623 names: &FastHashMap<NameKey, String>,
8624 ) -> Result<(), core::fmt::Error> {
8625 match *self {
8626 Access::GlobalVariable(handle) => {
8627 write!(writer, "{}", names[&NameKey::GlobalVariable(handle)])
8628 }
8629 Access::StructMember(handle, index) => {
8630 write!(writer, ".{}", names[&NameKey::StructMember(handle, index)])
8631 }
8632 Access::Array(depth) => write!(writer, ".{WRAPPED_ARRAY_FIELD}[__i{depth}]"),
8633 }
8634 }
8635 }
8636
8637 struct AccessStack {
8638 stack: Vec<Access>,
8639 array_depth: usize,
8640 }
8641
8642 impl AccessStack {
8643 const fn new() -> Self {
8644 Self {
8645 stack: Vec::new(),
8646 array_depth: 0,
8647 }
8648 }
8649
8650 fn enter_array<R>(&mut self, cb: impl FnOnce(&mut Self, usize) -> R) -> R {
8651 let array_depth = self.array_depth;
8652 self.stack.push(Access::Array(array_depth));
8653 self.array_depth += 1;
8654 let res = cb(self, array_depth);
8655 self.stack.pop();
8656 self.array_depth -= 1;
8657 res
8658 }
8659
8660 fn enter<R>(&mut self, new: Access, cb: impl FnOnce(&mut Self) -> R) -> R {
8661 self.stack.push(new);
8662 let res = cb(self);
8663 self.stack.pop();
8664 res
8665 }
8666
8667 fn write<W: Write>(
8668 &self,
8669 writer: &mut W,
8670 names: &FastHashMap<NameKey, String>,
8671 ) -> Result<(), core::fmt::Error> {
8672 for next in self.stack.iter() {
8673 next.write(writer, names)?;
8674 }
8675 Ok(())
8676 }
8677 }
8678
8679 impl<W: Write> Writer<W> {
8680 pub(super) fn need_workgroup_variables_initialization(
8681 &mut self,
8682 options: &Options,
8683 ep: &EntryPoint,
8684 module: &crate::Module,
8685 fun_info: &valid::FunctionInfo,
8686 ) -> bool {
8687 let is_task = ep.stage == crate::ShaderStage::Task;
8688 options.zero_initialize_workgroup_memory
8689 && ep.stage.compute_like()
8690 && module.global_variables.iter().any(|(handle, var)| {
8691 let is_right_address_space = var.space == crate::AddressSpace::WorkGroup
8692 || (var.space == crate::AddressSpace::TaskPayload && is_task);
8693 !fun_info[handle].is_empty() && is_right_address_space
8694 })
8695 }
8696
8697 pub fn write_workgroup_variables_initialization(
8698 &mut self,
8699 module: &crate::Module,
8700 module_info: &valid::ModuleInfo,
8701 fun_info: &valid::FunctionInfo,
8702 local_invocation_index: Option<&NameKey>,
8703 stage: crate::ShaderStage,
8704 ) -> BackendResult {
8705 let level = back::Level(1);
8706
8707 writeln!(
8708 self.out,
8709 "{}if ({} == 0u) {{",
8710 level,
8711 local_invocation_index
8712 .map(|name_key| self.names[name_key].as_str())
8713 .unwrap_or("__local_invocation_index"),
8714 )?;
8715
8716 let mut access_stack = AccessStack::new();
8717
8718 let is_task = stage == crate::ShaderStage::Task;
8719 let vars = module.global_variables.iter().filter(|&(handle, var)| {
8720 let is_right_address_space = var.space == crate::AddressSpace::WorkGroup
8721 || (var.space == crate::AddressSpace::TaskPayload && is_task);
8722 !fun_info[handle].is_empty() && is_right_address_space
8723 });
8724
8725 for (handle, var) in vars {
8726 access_stack.enter(Access::GlobalVariable(handle), |access_stack| {
8727 self.write_workgroup_variable_initialization(
8728 module,
8729 module_info,
8730 var.ty,
8731 access_stack,
8732 level.next(),
8733 )
8734 })?;
8735 }
8736
8737 writeln!(self.out, "{level}}}")?;
8738 self.write_barrier(crate::Barrier::WORK_GROUP, level)
8739 }
8740
8741 fn write_workgroup_variable_initialization(
8742 &mut self,
8743 module: &crate::Module,
8744 module_info: &valid::ModuleInfo,
8745 ty: Handle<crate::Type>,
8746 access_stack: &mut AccessStack,
8747 level: back::Level,
8748 ) -> BackendResult {
8749 if module_info[ty].contains(valid::TypeFlags::CONSTRUCTIBLE) {
8750 write!(self.out, "{level}")?;
8751 access_stack.write(&mut self.out, &self.names)?;
8752 writeln!(self.out, " = {{}};")?;
8753 } else {
8754 match module.types[ty].inner {
8755 crate::TypeInner::Atomic { .. } => {
8756 write!(
8757 self.out,
8758 "{level}{NAMESPACE}::atomic_store_explicit({ATOMIC_REFERENCE}"
8759 )?;
8760 access_stack.write(&mut self.out, &self.names)?;
8761 writeln!(self.out, ", 0, {NAMESPACE}::memory_order_relaxed);")?;
8762 }
8763 crate::TypeInner::Array { base, size, .. } => {
8764 let count = match size.resolve(module.to_ctx())? {
8765 proc::IndexableLength::Known(count) => count,
8766 proc::IndexableLength::Dynamic => unreachable!(),
8767 };
8768
8769 access_stack.enter_array(|access_stack, array_depth| {
8770 writeln!(
8771 self.out,
8772 "{level}for (int __i{array_depth} = 0; __i{array_depth} < {count}; __i{array_depth}++) {{"
8773 )?;
8774 self.write_workgroup_variable_initialization(
8775 module,
8776 module_info,
8777 base,
8778 access_stack,
8779 level.next(),
8780 )?;
8781 writeln!(self.out, "{level}}}")?;
8782 BackendResult::Ok(())
8783 })?;
8784 }
8785 crate::TypeInner::Struct { ref members, .. } => {
8786 for (index, member) in members.iter().enumerate() {
8787 access_stack.enter(
8788 Access::StructMember(ty, index as u32),
8789 |access_stack| {
8790 self.write_workgroup_variable_initialization(
8791 module,
8792 module_info,
8793 member.ty,
8794 access_stack,
8795 level,
8796 )
8797 },
8798 )?;
8799 }
8800 }
8801 _ => unreachable!(),
8802 }
8803 }
8804
8805 Ok(())
8806 }
8807 }
8808}
8809
8810impl crate::AtomicFunction {
8811 const fn to_msl(self) -> &'static str {
8812 match self {
8813 Self::Add => "fetch_add",
8814 Self::Subtract => "fetch_sub",
8815 Self::And => "fetch_and",
8816 Self::InclusiveOr => "fetch_or",
8817 Self::ExclusiveOr => "fetch_xor",
8818 Self::Min => "fetch_min",
8819 Self::Max => "fetch_max",
8820 Self::Exchange { compare: None } => "exchange",
8821 Self::Exchange { compare: Some(_) } => ATOMIC_COMP_EXCH_FUNCTION,
8822 }
8823 }
8824
8825 fn to_msl_64_bit(self) -> Result<&'static str, Error> {
8826 Ok(match self {
8827 Self::Min => "min",
8828 Self::Max => "max",
8829 _ => Err(Error::FeatureNotImplemented(
8830 "64-bit atomic operation other than min/max".to_string(),
8831 ))?,
8832 })
8833 }
8834}