cubecl_wgpu/compiler/wgsl/
compiler.rs

1use super::Item;
2use super::Subgroup;
3use super::shader::ComputeShader;
4use crate::compiler::wgsl::{self, SharedValue};
5
6use cubecl_common::backtrace::BackTrace;
7use cubecl_core::ir::{Processor, UIntKind};
8use cubecl_core::{
9    Info,
10    post_processing::{
11        checked_io::CheckedIoVisitor, optimize_scope, saturating::SaturatingArithmeticProcessor,
12        unroll::UnrollVisitor,
13    },
14};
15use cubecl_core::{
16    Metadata, WgpuCompilationOptions,
17    ir::{self as cube, Scope},
18    prelude::expand_erf,
19};
20use cubecl_core::{post_processing::disaggregate::DisaggregateVisitor, prelude::*};
21use cubecl_ir::AddressSpace;
22use cubecl_runtime::compiler::CompilationError;
23use cubecl_runtime::kernel;
24use hashbrown::HashMap;
25
26pub const MAX_VECTOR_SIZE: usize = 4;
27
28/// Wgsl Compiler.
29#[derive(Clone, Default)]
30pub struct WgslCompiler {
31    kernel_name: String,
32    info: Info,
33    ext_meta_pos: HashMap<cube::Value, u32>,
34    buffer_vis: Vec<Visibility>,
35    local_invocation_index: bool,
36    local_invocation_id: bool,
37    global_invocation_id: bool,
38    workgroup_id: bool,
39    subgroup_size: bool,
40    subgroup_id: bool,
41    subgroup_invocation_id: bool,
42    id: bool,
43    num_workgroups: bool,
44    workgroup_id_no_axis: bool,
45    workgroup_size_no_axis: bool,
46    num_workgroup_no_axis: bool,
47    shared_values: Vec<SharedValue>,
48    #[allow(dead_code)]
49    compilation_options: WgpuCompilationOptions,
50    strategy: ExecutionMode,
51    subgroup_instructions_used: bool,
52    f16_used: bool,
53}
54
55impl core::fmt::Debug for WgslCompiler {
56    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        f.write_str("WgslCompiler")
58    }
59}
60
61impl cubecl_core::Compiler for WgslCompiler {
62    type Representation = ComputeShader;
63    type CompilationOptions = WgpuCompilationOptions;
64
65    fn compile(
66        &mut self,
67        shader: kernel::KernelDefinition,
68        compilation_options: &Self::CompilationOptions,
69        mode: ExecutionMode,
70        address_type: StorageType,
71    ) -> Result<Self::Representation, CompilationError> {
72        self.compilation_options = *compilation_options;
73        self.compile_shader(shader, mode, address_type)
74    }
75
76    fn elem_size(&self, elem: cube::ElemType) -> usize {
77        elem.size()
78    }
79
80    fn extension(&self) -> &'static str {
81        "wgsl"
82    }
83}
84
85impl WgslCompiler {
86    fn compile_shader(
87        &mut self,
88        value: kernel::KernelDefinition,
89        mode: ExecutionMode,
90        address_type: StorageType,
91    ) -> Result<wgsl::ComputeShader, CompilationError> {
92        let errors = value.body.pop_errors();
93        if !errors.is_empty() {
94            let mut reason = "Can't compile wgsl kernel".to_string();
95            for error in errors {
96                reason += error.as_str();
97                reason += "\n";
98            }
99
100            return Err(CompilationError::Validation {
101                reason,
102                backtrace: BackTrace::capture(),
103            });
104        }
105
106        self.strategy = mode;
107        self.kernel_name = value.options.kernel_name.clone();
108
109        let num_meta = value.buffers.len();
110
111        self.ext_meta_pos = HashMap::new();
112        let mut num_ext = 0;
113
114        for binding in value.buffers.iter() {
115            self.ext_meta_pos.insert(binding.value, num_ext);
116            if binding.has_extended_meta {
117                num_ext += 1;
118            }
119        }
120
121        let metadata = Metadata::new(num_meta as u32, num_ext);
122        self.info = Info::new(&value.scalars, metadata, address_type);
123
124        CheckedIoVisitor::new(self.strategy, self.kernel_name.clone()).apply(&value.body);
125        DisaggregateVisitor::apply(&value.body);
126        UnrollVisitor::new(MAX_VECTOR_SIZE).apply(&value.body);
127
128        self.buffer_vis = optimize_scope(&value.body).into();
129        self.buffer_vis
130            .resize(value.num_global_buffers(), Visibility::Read);
131
132        let address_type = self.compile_storage_type(address_type);
133        let instructions = self.compile_scope(&value.body);
134        let extensions = register_extensions(&instructions);
135        let body = wgsl::Body {
136            instructions,
137            id: self.id,
138            address_type,
139        };
140
141        Ok(wgsl::ComputeShader {
142            address_type,
143            buffers: value
144                .buffers
145                .into_iter()
146                .map(|mut it| {
147                    // This is safe when combined with the unroll transform that adjusts all indices.
148                    // Must not be used alone
149                    if it.value.ty.vector_size() > MAX_VECTOR_SIZE {
150                        it.value.ty = it.value.ty.with_vector_size(MAX_VECTOR_SIZE);
151                    }
152                    self.compile_binding(it)
153                })
154                .collect(),
155            scalars: value
156                .scalars
157                .into_iter()
158                .map(|binding| (self.compile_storage_type(binding.ty), binding.count))
159                .collect(),
160            shared_values: self.shared_values.clone(),
161            static_meta_len: self.info.metadata.static_len() as usize,
162            info: self.info.clone(),
163            workgroup_size: value.cube_dim,
164            global_invocation_id: self.global_invocation_id || self.id,
165            local_invocation_index: self.local_invocation_index,
166            local_invocation_id: self.local_invocation_id,
167            num_workgroups: self.id
168                || self.num_workgroups
169                || self.num_workgroup_no_axis
170                || self.workgroup_id_no_axis,
171            workgroup_id: self.workgroup_id || self.workgroup_id_no_axis,
172            subgroup_size: self.subgroup_size,
173            subgroup_id: self.subgroup_id,
174            subgroup_invocation_id: self.subgroup_invocation_id,
175            body,
176            extensions,
177            num_workgroups_no_axis: self.num_workgroup_no_axis,
178            workgroup_id_no_axis: self.workgroup_id_no_axis,
179            workgroup_size_no_axis: self.workgroup_size_no_axis,
180            subgroup_instructions_used: self.subgroup_instructions_used,
181            f16_used: self.f16_used,
182            kernel_name: value.options.kernel_name,
183        })
184    }
185
186    fn compile_type(&mut self, item: cube::Type) -> Item {
187        match item {
188            cube::Type::Scalar(ty) => wgsl::Item::Scalar(self.compile_storage_type(ty)),
189            cube::Type::Vector(ty, size) => {
190                let elem = self.compile_storage_type(ty.storage_type());
191                wgsl::Item::Vector(elem, size)
192            }
193            cube::Type::Atomic(ty) => {
194                let inner = self.compile_type(*ty);
195                wgsl::Item::Atomic(inner.intern())
196            }
197            cube::Type::Pointer(ty, class) => {
198                let inner = self.compile_type(*ty);
199                let class = self.compile_pointer_class(class);
200                wgsl::Item::Pointer(inner.intern(), class)
201            }
202            cube::Type::Array(ty, size) => {
203                let inner = self.compile_type(*ty);
204                wgsl::Item::Array(inner.intern(), size)
205            }
206            cube::Type::DynamicArray(ty) => {
207                let inner = self.compile_type(*ty);
208                wgsl::Item::DynamicArray(inner.intern())
209            }
210            cube::Type::Opaque(_) => unimplemented!("Can't compile opaque type"),
211            cube::Type::Semantic(_) => unimplemented!("Can't compile semantic type"),
212            cube::Type::Matrix(_) => unimplemented!("Matrices not yet supported in WGSL"),
213            cube::Type::Aggregate(_) => {
214                unreachable!("Should be disaggregated at this point")
215            }
216        }
217    }
218
219    fn compile_storage_type(&mut self, ty: cube::StorageType) -> wgsl::Elem {
220        match ty {
221            cube::StorageType::Scalar(ty) => self.compile_elem(ty),
222            cube::StorageType::Packed(_, _) => {
223                unimplemented!("Packed types not yet supported in WGSL")
224            }
225        }
226    }
227
228    fn compile_elem(&mut self, value: cube::ElemType) -> wgsl::Elem {
229        match value {
230            cube::ElemType::Float(f) => match f {
231                cube::FloatKind::E2M1
232                | cube::FloatKind::E2M3
233                | cube::FloatKind::E3M2
234                | cube::FloatKind::E4M3
235                | cube::FloatKind::E5M2
236                | cube::FloatKind::UE8M0 => panic!("Minifloat is not a valid WgpuElement"),
237                cube::FloatKind::F16 => {
238                    self.f16_used = true;
239                    wgsl::Elem::F16
240                }
241                cube::FloatKind::BF16 => panic!("bf16 is not a valid WgpuElement"),
242                cube::FloatKind::TF32 => panic!("tf32 is not a valid WgpuElement"),
243                cube::FloatKind::Flex32 => wgsl::Elem::F32,
244                cube::FloatKind::F32 => wgsl::Elem::F32,
245                cube::FloatKind::F64 => wgsl::Elem::F64,
246            },
247            cube::ElemType::Int(i) => match i {
248                cube::IntKind::I32 => wgsl::Elem::I32,
249                cube::IntKind::I64 => wgsl::Elem::I64,
250                kind => panic!("{kind:?} is not a valid WgpuElement"),
251            },
252            cube::ElemType::UInt(kind) => match kind {
253                cube::UIntKind::U32 => wgsl::Elem::U32,
254                cube::UIntKind::U64 => wgsl::Elem::U64,
255                kind => panic!("{kind:?} is not a valid WgpuElement"),
256            },
257            cube::ElemType::Bool => wgsl::Elem::Bool,
258        }
259    }
260
261    fn compile_pointer_class(&self, class: cube::AddressSpace) -> wgsl::PointerClass {
262        match class {
263            cubecl_ir::AddressSpace::Global(id) => {
264                wgsl::PointerClass::Global(self.buffer_vis[id as usize])
265            }
266            cubecl_ir::AddressSpace::Shared => wgsl::PointerClass::Shared,
267            cubecl_ir::AddressSpace::Local => wgsl::PointerClass::Local,
268        }
269    }
270
271    fn ext_meta_pos(&self, val: &cube::Value) -> u32 {
272        self.ext_meta_pos[val]
273    }
274
275    pub(crate) fn compile_value(&mut self, value: cube::Value) -> wgsl::Value {
276        let item = value.ty;
277        match value.kind {
278            cube::ValueKind::Value { id } => wgsl::Value::Value {
279                id,
280                item: self.compile_type(item),
281            },
282            cube::ValueKind::Constant(value) => {
283                wgsl::Value::Constant(value, self.compile_type(item))
284            }
285        }
286    }
287
288    fn constant_var(&mut self, value: u32) -> wgsl::Value {
289        let val = cube::Value::constant(value.into(), UIntKind::U32);
290        self.compile_value(val)
291    }
292
293    fn compile_scope(&mut self, scope: &cube::Scope) -> Vec<wgsl::Instruction> {
294        let mut instructions = Vec::new();
295
296        let saturating: Box<dyn Processor> = Box::new(SaturatingArithmeticProcessor::new(true));
297        let processing = scope.process([&*saturating]);
298
299        processing
300            .instructions
301            .into_iter()
302            .for_each(|op| self.compile_operation(&mut instructions, op.operation, op.out, scope));
303
304        instructions
305    }
306
307    fn compile_operation(
308        &mut self,
309        instructions: &mut Vec<wgsl::Instruction>,
310        operation: cube::Operation,
311        out: Option<cube::Value>,
312        scope: &cube::Scope,
313    ) {
314        match operation {
315            cube::Operation::Copy(value) => instructions.push(wgsl::Instruction::Assign {
316                input: self.compile_value(value),
317                out: self.compile_value(out.unwrap()),
318            }),
319            cube::Operation::DeclareVariable {
320                value_ty,
321                addr_space: AddressSpace::Local,
322                ..
323            } => instructions.push(wgsl::Instruction::DeclareVariable {
324                val: self.compile_value(out.unwrap()),
325                value_ty: self.compile_type(value_ty),
326            }),
327            cube::Operation::DeclareVariable {
328                value_ty,
329                addr_space: AddressSpace::Shared,
330                alignment,
331            } => {
332                let ty = self.compile_type(value_ty);
333                let value = self.compile_value(out.unwrap());
334                self.shared_values
335                    .push(SharedValue::new(ty, value, alignment as u32));
336            }
337            cube::Operation::DeclareVariable { addr_space, .. } => {
338                unimplemented!("Unsupported declare address space {addr_space}")
339            }
340            cube::Operation::Memory(memory) => self.compile_memory(memory, out, instructions),
341            cube::Operation::Arithmetic(op) => {
342                self.compile_arithmetic(op, out, instructions, scope)
343            }
344            cube::Operation::Comparison(op) => self.compile_cmp(op, out, instructions),
345            cube::Operation::Bitwise(op) => self.compile_bitwise(op, out, instructions),
346            cube::Operation::Operator(op) => self.compile_operator(op, out, instructions),
347            cube::Operation::Atomic(op) => instructions.push(self.compile_atomic(op, out)),
348            cube::Operation::Metadata(op) => instructions.push(self.compile_metadata(op, out)),
349            cube::Operation::Branch(val) => self.compile_branch(instructions, val),
350            cube::Operation::Synchronization(val) => {
351                self.compile_synchronization(instructions, val)
352            }
353            cube::Operation::WorkgroupUniformLoad(op) => {
354                instructions.push(wgsl::Instruction::WorkgroupUniformLoad {
355                    input: self.compile_value(op),
356                    out: self.compile_value(out.unwrap()),
357                });
358            }
359            cube::Operation::Plane(op) => self.compile_subgroup(instructions, op, out),
360            cube::Operation::CoopMma(_) => {
361                panic!("Cooperative matrix-multiply and accumulate isn't supported on wgpu.")
362            }
363            cube::Operation::NonSemantic(cube::NonSemantic::Comment { content }) => {
364                self.compile_comment(instructions, content)
365            }
366            cube::Operation::NonSemantic(_) => {}
367            cube::Operation::Barrier(_) => {
368                panic!("Barrier isn't supported on wgpu.")
369            }
370            cube::Operation::Tma(_) => panic!("TMA isn't supported on wgpu."),
371            cube::Operation::TensorIndexing(_) => panic!("TMA isn't supported on wgpu."),
372            cube::Operation::Marker(_) => {}
373            cube::Operation::ConstructAggregate(..)
374            | cube::Operation::ExtractAggregateField(..) => {
375                unreachable!("Should be disaggregated at this point")
376            }
377        }
378    }
379
380    fn compile_subgroup(
381        &mut self,
382        instructions: &mut Vec<wgsl::Instruction>,
383        subgroup: cube::Plane,
384        out: Option<cube::Value>,
385    ) {
386        self.subgroup_instructions_used = true;
387
388        let out = out.unwrap();
389        let op = match subgroup {
390            cube::Plane::Elect => Subgroup::Elect {
391                out: self.compile_value(out),
392            },
393            cube::Plane::All(op) => Subgroup::All {
394                input: self.compile_value(op.input),
395                out: self.compile_value(out),
396            },
397            cube::Plane::Any(op) => Subgroup::Any {
398                input: self.compile_value(op.input),
399                out: self.compile_value(out),
400            },
401            cube::Plane::Ballot(op) => Subgroup::Ballot {
402                input: self.compile_value(op.input),
403                out: self.compile_value(out),
404            },
405
406            cube::Plane::Broadcast(op) => Subgroup::Broadcast {
407                lhs: self.compile_value(op.lhs),
408                rhs: self.compile_value(op.rhs),
409                out: self.compile_value(out),
410            },
411
412            cube::Plane::Sum(op) => Subgroup::Sum {
413                input: self.compile_value(op.input),
414                out: self.compile_value(out),
415            },
416
417            cube::Plane::ExclusiveSum(op) => Subgroup::ExclusiveSum {
418                input: self.compile_value(op.input),
419                out: self.compile_value(out),
420            },
421            cube::Plane::InclusiveSum(op) => Subgroup::InclusiveSum {
422                input: self.compile_value(op.input),
423                out: self.compile_value(out),
424            },
425            cube::Plane::Prod(op) => Subgroup::Prod {
426                input: self.compile_value(op.input),
427                out: self.compile_value(out),
428            },
429            cube::Plane::ExclusiveProd(op) => Subgroup::ExclusiveProd {
430                input: self.compile_value(op.input),
431                out: self.compile_value(out),
432            },
433            cube::Plane::InclusiveProd(op) => Subgroup::InclusiveProd {
434                input: self.compile_value(op.input),
435                out: self.compile_value(out),
436            },
437            cube::Plane::Min(op) => Subgroup::Min {
438                input: self.compile_value(op.input),
439                out: self.compile_value(out),
440            },
441            cube::Plane::Max(op) => Subgroup::Max {
442                input: self.compile_value(op.input),
443                out: self.compile_value(out),
444            },
445            cube::Plane::Shuffle(op) => Subgroup::Shuffle {
446                lhs: self.compile_value(op.lhs),
447                rhs: self.compile_value(op.rhs),
448                out: self.compile_value(out),
449            },
450            cube::Plane::ShuffleXor(op) => Subgroup::ShuffleXor {
451                lhs: self.compile_value(op.lhs),
452                rhs: self.compile_value(op.rhs),
453                out: self.compile_value(out),
454            },
455            cube::Plane::ShuffleUp(op) => Subgroup::ShuffleUp {
456                lhs: self.compile_value(op.lhs),
457                rhs: self.compile_value(op.rhs),
458                out: self.compile_value(out),
459            },
460            cube::Plane::ShuffleDown(op) => Subgroup::ShuffleDown {
461                lhs: self.compile_value(op.lhs),
462                rhs: self.compile_value(op.rhs),
463                out: self.compile_value(out),
464            },
465        };
466
467        instructions.push(wgsl::Instruction::Subgroup(op));
468    }
469
470    fn compile_branch(&mut self, instructions: &mut Vec<wgsl::Instruction>, branch: cube::Branch) {
471        match branch {
472            cube::Branch::If(op) => instructions.push(wgsl::Instruction::If {
473                cond: self.compile_value(op.cond),
474                instructions: self.compile_scope(&op.scope),
475            }),
476            cube::Branch::IfElse(op) => instructions.push(wgsl::Instruction::IfElse {
477                cond: self.compile_value(op.cond),
478                instructions_if: self.compile_scope(&op.scope_if),
479                instructions_else: self.compile_scope(&op.scope_else),
480            }),
481            cube::Branch::Switch(op) => instructions.push(wgsl::Instruction::Switch {
482                value: self.compile_value(op.value),
483                instructions_default: self.compile_scope(&op.scope_default),
484                cases: op
485                    .cases
486                    .into_iter()
487                    .map(|(val, scope)| (self.compile_value(val), self.compile_scope(&scope)))
488                    .collect(),
489            }),
490            cube::Branch::Return => instructions.push(wgsl::Instruction::Return),
491            // No unreachable hint in WGSL
492            cube::Branch::Unreachable => instructions.push(wgsl::Instruction::Return),
493            cube::Branch::Break => instructions.push(wgsl::Instruction::Break),
494            cube::Branch::RangeLoop(range_loop) => {
495                instructions.push(wgsl::Instruction::RangeLoop {
496                    i: self.compile_value(range_loop.i),
497                    start: self.compile_value(range_loop.start),
498                    end: self.compile_value(range_loop.end),
499                    step: range_loop.step.map(|it| self.compile_value(it)),
500                    inclusive: range_loop.inclusive,
501                    instructions: self.compile_scope(&range_loop.scope),
502                })
503            }
504            cube::Branch::Loop(op) => instructions.push(wgsl::Instruction::Loop {
505                instructions: self.compile_scope(&op.scope),
506            }),
507        };
508    }
509
510    fn compile_synchronization(
511        &mut self,
512        instructions: &mut Vec<wgsl::Instruction>,
513        synchronization: cube::Synchronization,
514    ) {
515        match synchronization {
516            cube::Synchronization::SyncCube => {
517                instructions.push(wgsl::Instruction::WorkgroupBarrier)
518            }
519            cube::Synchronization::SyncPlane => {
520                panic!("Synchronization within a plane is not supported in WGSL")
521            }
522            cube::Synchronization::SyncStorage => {
523                instructions.push(wgsl::Instruction::StorageBarrier)
524            }
525            cube::Synchronization::SyncAsyncProxyShared => panic!("TMA is not supported in WGSL"),
526        };
527    }
528
529    fn compile_comment(&mut self, instructions: &mut Vec<wgsl::Instruction>, content: String) {
530        instructions.push(wgsl::Instruction::Comment { content })
531    }
532
533    fn compile_metadata(
534        &mut self,
535        metadata: cube::Metadata,
536        out: Option<cube::Value>,
537    ) -> wgsl::Instruction {
538        let out = out.unwrap();
539        match metadata {
540            cube::Metadata::Stride { dim, list } => {
541                let position = self.ext_meta_pos(&list);
542                let offset = self.info.metadata.stride_offset_index(position);
543                wgsl::Instruction::ExtendedMeta {
544                    info_offset: self.compile_value(offset.into()),
545                    dim: self.compile_value(dim),
546                    out: self.compile_value(out),
547                }
548            }
549            cube::Metadata::Shape { dim, list } => {
550                let position = self.ext_meta_pos(&list);
551                let offset = self.info.metadata.shape_offset_index(position);
552                wgsl::Instruction::ExtendedMeta {
553                    info_offset: self.compile_value(offset.into()),
554                    dim: self.compile_value(dim),
555                    out: self.compile_value(out),
556                }
557            }
558            cube::Metadata::BufferLength { list } => match list.address_space() {
559                cube::AddressSpace::Global(id) => {
560                    let offset = self.info.metadata.buffer_len_index(id);
561                    wgsl::Instruction::Metadata {
562                        out: self.compile_value(out),
563                        info_offset: self.compile_value(offset.into()),
564                    }
565                }
566                _ => wgsl::Instruction::Length {
567                    list: self.compile_value(list),
568                    out: self.compile_value(out),
569                },
570            },
571        }
572    }
573
574    fn compile_memory(
575        &mut self,
576        value: cube::Memory,
577        out: Option<cube::Value>,
578        instructions: &mut Vec<wgsl::Instruction>,
579    ) {
580        match value {
581            cube::Memory::Index(op) => {
582                instructions.push(wgsl::Instruction::Index {
583                    lhs: self.compile_value(op.list),
584                    rhs: self.compile_value(op.index),
585                    out: self.compile_value(out.unwrap()),
586                });
587            }
588            cube::Memory::Load(value) => instructions.push(wgsl::Instruction::Load {
589                input: self.compile_value(value),
590                out: self.compile_value(out.unwrap()),
591            }),
592            cube::Memory::Store(op) => instructions.push(wgsl::Instruction::Store {
593                input: self.compile_value(op.value),
594                out: self.compile_value(op.ptr),
595            }),
596            cube::Memory::CopyMemory(op) => instructions.push(wgsl::Instruction::CopyBulk {
597                source: self.compile_value(op.source),
598                target: self.compile_value(op.target),
599                len: op.len as u32,
600            }),
601        }
602    }
603
604    fn compile_arithmetic(
605        &mut self,
606        value: cube::Arithmetic,
607        out: Option<cube::Value>,
608        instructions: &mut Vec<wgsl::Instruction>,
609        scope: &Scope,
610    ) {
611        let out = out.unwrap();
612        match value {
613            cube::Arithmetic::Max(op) => instructions.push(wgsl::Instruction::Max {
614                lhs: self.compile_value(op.lhs),
615                rhs: self.compile_value(op.rhs),
616                out: self.compile_value(out),
617            }),
618            cube::Arithmetic::Min(op) => instructions.push(wgsl::Instruction::Min {
619                lhs: self.compile_value(op.lhs),
620                rhs: self.compile_value(op.rhs),
621                out: self.compile_value(out),
622            }),
623            cube::Arithmetic::Add(op) => instructions.push(wgsl::Instruction::Add {
624                lhs: self.compile_value(op.lhs),
625                rhs: self.compile_value(op.rhs),
626                out: self.compile_value(out),
627            }),
628            cube::Arithmetic::SaturatingAdd(_) => {
629                unreachable!("Saturating add should be removed by processor");
630            }
631            cube::Arithmetic::Fma(op) => instructions.push(wgsl::Instruction::Fma {
632                a: self.compile_value(op.a),
633                b: self.compile_value(op.b),
634                c: self.compile_value(op.c),
635                out: self.compile_value(out),
636            }),
637            cube::Arithmetic::ModFloor(op) => instructions.push(wgsl::Instruction::ModFloor {
638                lhs: self.compile_value(op.lhs),
639                rhs: self.compile_value(op.rhs),
640                out: self.compile_value(out),
641            }),
642            cube::Arithmetic::Sub(op) => instructions.push(wgsl::Instruction::Sub {
643                lhs: self.compile_value(op.lhs),
644                rhs: self.compile_value(op.rhs),
645                out: self.compile_value(out),
646            }),
647            cube::Arithmetic::SaturatingSub(_) => {
648                unreachable!("Saturating sub should be removed by processor");
649            }
650            cube::Arithmetic::Mul(op) => instructions.push(wgsl::Instruction::Mul {
651                lhs: self.compile_value(op.lhs),
652                rhs: self.compile_value(op.rhs),
653                out: self.compile_value(out),
654            }),
655            cube::Arithmetic::Div(op) => instructions.push(wgsl::Instruction::Div {
656                lhs: self.compile_value(op.lhs),
657                rhs: self.compile_value(op.rhs),
658                out: self.compile_value(out),
659            }),
660            cube::Arithmetic::Abs(op) => instructions.push(wgsl::Instruction::Abs {
661                input: self.compile_value(op.input),
662                out: self.compile_value(out),
663            }),
664            cube::Arithmetic::Exp(op) => instructions.push(wgsl::Instruction::Exp {
665                input: self.compile_value(op.input),
666                out: self.compile_value(out),
667            }),
668            cube::Arithmetic::Log(op) => instructions.push(wgsl::Instruction::Log {
669                input: self.compile_value(op.input),
670                out: self.compile_value(out),
671            }),
672            cube::Arithmetic::Log1p(op) => instructions.push(wgsl::Instruction::Log1p {
673                input: self.compile_value(op.input),
674                out: self.compile_value(out),
675            }),
676            cube::Arithmetic::Expm1(op) => instructions.push(wgsl::Instruction::Expm1 {
677                input: self.compile_value(op.input),
678                out: self.compile_value(out),
679            }),
680            cube::Arithmetic::Cos(op) => instructions.push(wgsl::Instruction::Cos {
681                input: self.compile_value(op.input),
682                out: self.compile_value(out),
683            }),
684            cube::Arithmetic::Sin(op) => instructions.push(wgsl::Instruction::Sin {
685                input: self.compile_value(op.input),
686                out: self.compile_value(out),
687            }),
688            cube::Arithmetic::Tan(op) => instructions.push(wgsl::Instruction::Tan {
689                input: self.compile_value(op.input),
690                out: self.compile_value(out),
691            }),
692            cube::Arithmetic::Tanh(op) => instructions.push(wgsl::Instruction::Tanh {
693                input: self.compile_value(op.input),
694                out: self.compile_value(out),
695            }),
696            cube::Arithmetic::Sinh(op) => instructions.push(wgsl::Instruction::Sinh {
697                input: self.compile_value(op.input),
698                out: self.compile_value(out),
699            }),
700            cube::Arithmetic::Cosh(op) => instructions.push(wgsl::Instruction::Cosh {
701                input: self.compile_value(op.input),
702                out: self.compile_value(out),
703            }),
704            cube::Arithmetic::ArcCos(op) => instructions.push(wgsl::Instruction::ArcCos {
705                input: self.compile_value(op.input),
706                out: self.compile_value(out),
707            }),
708            cube::Arithmetic::ArcSin(op) => instructions.push(wgsl::Instruction::ArcSin {
709                input: self.compile_value(op.input),
710                out: self.compile_value(out),
711            }),
712            cube::Arithmetic::ArcTan(op) => instructions.push(wgsl::Instruction::ArcTan {
713                input: self.compile_value(op.input),
714                out: self.compile_value(out),
715            }),
716            cube::Arithmetic::ArcSinh(op) => instructions.push(wgsl::Instruction::ArcSinh {
717                input: self.compile_value(op.input),
718                out: self.compile_value(out),
719            }),
720            cube::Arithmetic::ArcCosh(op) => instructions.push(wgsl::Instruction::ArcCosh {
721                input: self.compile_value(op.input),
722                out: self.compile_value(out),
723            }),
724            cube::Arithmetic::ArcTanh(op) => instructions.push(wgsl::Instruction::ArcTanh {
725                input: self.compile_value(op.input),
726                out: self.compile_value(out),
727            }),
728            cube::Arithmetic::Degrees(op) => instructions.push(wgsl::Instruction::Degrees {
729                input: self.compile_value(op.input),
730                out: self.compile_value(out),
731            }),
732            cube::Arithmetic::Radians(op) => instructions.push(wgsl::Instruction::Radians {
733                input: self.compile_value(op.input),
734                out: self.compile_value(out),
735            }),
736            cube::Arithmetic::ArcTan2(op) => instructions.push(wgsl::Instruction::ArcTan2 {
737                lhs: self.compile_value(op.lhs),
738                rhs: self.compile_value(op.rhs),
739                out: self.compile_value(out),
740            }),
741            // No powi in WGSL
742            cube::Arithmetic::Powf(op) | cube::Arithmetic::Powi(op) => {
743                instructions.push(wgsl::Instruction::Powf {
744                    lhs: self.compile_value(op.lhs),
745                    rhs: self.compile_value(op.rhs),
746                    out: self.compile_value(out),
747                })
748            }
749            cube::Arithmetic::Hypot(op) => {
750                let scope = scope.child();
751                expand_hypot(&scope, op.lhs, op.rhs, out);
752                instructions.extend(self.compile_scope(&scope));
753            }
754            cube::Arithmetic::Rhypot(op) => {
755                let scope = scope.child();
756                expand_rhypot(&scope, op.lhs, op.rhs, out);
757                instructions.extend(self.compile_scope(&scope));
758            }
759
760            cube::Arithmetic::Sqrt(op) => instructions.push(wgsl::Instruction::Sqrt {
761                input: self.compile_value(op.input),
762                out: self.compile_value(out),
763            }),
764            cube::Arithmetic::InverseSqrt(op) => {
765                instructions.push(wgsl::Instruction::InverseSqrt {
766                    input: self.compile_value(op.input),
767                    out: self.compile_value(out),
768                })
769            }
770            cube::Arithmetic::Round(op) => instructions.push(wgsl::Instruction::Round {
771                input: self.compile_value(op.input),
772                out: self.compile_value(out),
773            }),
774            cube::Arithmetic::Floor(op) => instructions.push(wgsl::Instruction::Floor {
775                input: self.compile_value(op.input),
776                out: self.compile_value(out),
777            }),
778            cube::Arithmetic::Ceil(op) => instructions.push(wgsl::Instruction::Ceil {
779                input: self.compile_value(op.input),
780                out: self.compile_value(out),
781            }),
782            cube::Arithmetic::Trunc(op) => instructions.push(wgsl::Instruction::Trunc {
783                input: self.compile_value(op.input),
784                out: self.compile_value(out),
785            }),
786            cube::Arithmetic::Erf(op) => {
787                let scope = scope.child();
788                expand_erf(&scope, op.input, out);
789                instructions.extend(self.compile_scope(&scope));
790            }
791            cube::Arithmetic::MulHi(op) => {
792                let scope = scope.child();
793                match self.compilation_options.supports_u64 {
794                    true => expand_himul_64(&scope, op.lhs, op.rhs, out),
795                    false => expand_himul_sim(&scope, op.lhs, op.rhs, out),
796                }
797                instructions.extend(self.compile_scope(&scope));
798            }
799            cube::Arithmetic::Recip(op) => instructions.push(wgsl::Instruction::Recip {
800                input: self.compile_value(op.input),
801                out: self.compile_value(out),
802            }),
803            cube::Arithmetic::Clamp(op) => instructions.push(wgsl::Instruction::Clamp {
804                input: self.compile_value(op.input),
805                min_value: self.compile_value(op.min_value),
806                max_value: self.compile_value(op.max_value),
807                out: self.compile_value(out),
808            }),
809            cube::Arithmetic::Rem(op) => instructions.push(wgsl::Instruction::Remainder {
810                lhs: self.compile_value(op.lhs),
811                rhs: self.compile_value(op.rhs),
812                out: self.compile_value(out),
813            }),
814            cube::Arithmetic::Neg(op) => instructions.push(wgsl::Instruction::Negate {
815                input: self.compile_value(op.input),
816                out: self.compile_value(out),
817            }),
818            cube::Arithmetic::Magnitude(op) => instructions.push(wgsl::Instruction::Magnitude {
819                input: self.compile_value(op.input),
820                out: self.compile_value(out),
821            }),
822            cube::Arithmetic::Normalize(op) => instructions.push(wgsl::Instruction::Normalize {
823                input: self.compile_value(op.input),
824                out: self.compile_value(out),
825            }),
826            cube::Arithmetic::Dot(op) => instructions.push(wgsl::Instruction::Dot {
827                lhs: self.compile_value(op.lhs),
828                rhs: self.compile_value(op.rhs),
829                out: self.compile_value(out),
830            }),
831            cube::Arithmetic::VectorSum(op) => instructions.push(wgsl::Instruction::VectorSum {
832                input: self.compile_value(op.input),
833                out: self.compile_value(out),
834            }),
835        }
836    }
837
838    fn compile_cmp(
839        &mut self,
840        value: cube::Comparison,
841        out: Option<cube::Value>,
842        instructions: &mut Vec<wgsl::Instruction>,
843    ) {
844        let out = out.unwrap();
845        match value {
846            cube::Comparison::Equal(op) => instructions.push(wgsl::Instruction::Equal {
847                lhs: self.compile_value(op.lhs),
848                rhs: self.compile_value(op.rhs),
849                out: self.compile_value(out),
850            }),
851            cube::Comparison::Lower(op) => instructions.push(wgsl::Instruction::Lower {
852                lhs: self.compile_value(op.lhs),
853                rhs: self.compile_value(op.rhs),
854                out: self.compile_value(out),
855            }),
856            cube::Comparison::Greater(op) => instructions.push(wgsl::Instruction::Greater {
857                lhs: self.compile_value(op.lhs),
858                rhs: self.compile_value(op.rhs),
859                out: self.compile_value(out),
860            }),
861            cube::Comparison::LowerEqual(op) => instructions.push(wgsl::Instruction::LowerEqual {
862                lhs: self.compile_value(op.lhs),
863                rhs: self.compile_value(op.rhs),
864                out: self.compile_value(out),
865            }),
866            cube::Comparison::GreaterEqual(op) => {
867                instructions.push(wgsl::Instruction::GreaterEqual {
868                    lhs: self.compile_value(op.lhs),
869                    rhs: self.compile_value(op.rhs),
870                    out: self.compile_value(out),
871                })
872            }
873            cube::Comparison::NotEqual(op) => instructions.push(wgsl::Instruction::NotEqual {
874                lhs: self.compile_value(op.lhs),
875                rhs: self.compile_value(op.rhs),
876                out: self.compile_value(out),
877            }),
878            cube::Comparison::IsNan(op) => instructions.push(wgsl::Instruction::IsNan {
879                input: self.compile_value(op.input),
880                out: self.compile_value(out),
881            }),
882            cube::Comparison::IsInf(op) => instructions.push(wgsl::Instruction::IsInf {
883                input: self.compile_value(op.input),
884                out: self.compile_value(out),
885            }),
886        }
887    }
888
889    fn compile_bitwise(
890        &mut self,
891        value: cube::Bitwise,
892        out: Option<cube::Value>,
893        instructions: &mut Vec<wgsl::Instruction>,
894    ) {
895        let out = out.unwrap();
896        match value {
897            cube::Bitwise::BitwiseOr(op) => instructions.push(wgsl::Instruction::BitwiseOr {
898                lhs: self.compile_value(op.lhs),
899                rhs: self.compile_value(op.rhs),
900                out: self.compile_value(out),
901            }),
902            cube::Bitwise::BitwiseAnd(op) => instructions.push(wgsl::Instruction::BitwiseAnd {
903                lhs: self.compile_value(op.lhs),
904                rhs: self.compile_value(op.rhs),
905                out: self.compile_value(out),
906            }),
907            cube::Bitwise::BitwiseXor(op) => instructions.push(wgsl::Instruction::BitwiseXor {
908                lhs: self.compile_value(op.lhs),
909                rhs: self.compile_value(op.rhs),
910                out: self.compile_value(out),
911            }),
912            cube::Bitwise::CountOnes(op) => instructions.push(wgsl::Instruction::CountBits {
913                input: self.compile_value(op.input),
914                out: self.compile_value(out),
915            }),
916            cube::Bitwise::ReverseBits(op) => instructions.push(wgsl::Instruction::ReverseBits {
917                input: self.compile_value(op.input),
918                out: self.compile_value(out),
919            }),
920            cube::Bitwise::ShiftLeft(op) => instructions.push(wgsl::Instruction::ShiftLeft {
921                lhs: self.compile_value(op.lhs),
922                rhs: self.compile_value(op.rhs),
923                out: self.compile_value(out),
924            }),
925            cube::Bitwise::ShiftRight(op) => instructions.push(wgsl::Instruction::ShiftRight {
926                lhs: self.compile_value(op.lhs),
927                rhs: self.compile_value(op.rhs),
928                out: self.compile_value(out),
929            }),
930            cube::Bitwise::BitwiseNot(op) => instructions.push(wgsl::Instruction::BitwiseNot {
931                input: self.compile_value(op.input),
932                out: self.compile_value(out),
933            }),
934            cube::Bitwise::LeadingZeros(op) => instructions.push(wgsl::Instruction::LeadingZeros {
935                input: self.compile_value(op.input),
936                out: self.compile_value(out),
937            }),
938            cube::Bitwise::TrailingZeros(op) => {
939                instructions.push(wgsl::Instruction::TrailingZeros {
940                    input: self.compile_value(op.input),
941                    out: self.compile_value(out),
942                })
943            }
944            cube::Bitwise::FindFirstSet(op) => instructions.push(wgsl::Instruction::FindFirstSet {
945                input: self.compile_value(op.input),
946                out: self.compile_value(out),
947            }),
948        }
949    }
950
951    fn compile_operator(
952        &mut self,
953        value: cube::Operator,
954        out: Option<cube::Value>,
955        instructions: &mut Vec<wgsl::Instruction>,
956    ) {
957        let out = out.unwrap();
958        match value {
959            cube::Operator::Cast(op) => instructions.push(wgsl::Instruction::Assign {
960                input: self.compile_value(op.input),
961                out: self.compile_value(out),
962            }),
963
964            cube::Operator::And(op) => instructions.push(wgsl::Instruction::And {
965                lhs: self.compile_value(op.lhs),
966                rhs: self.compile_value(op.rhs),
967                out: self.compile_value(out),
968            }),
969            cube::Operator::Or(op) => instructions.push(wgsl::Instruction::Or {
970                lhs: self.compile_value(op.lhs),
971                rhs: self.compile_value(op.rhs),
972                out: self.compile_value(out),
973            }),
974            cube::Operator::Not(op) => instructions.push(wgsl::Instruction::Not {
975                input: self.compile_value(op.input),
976                out: self.compile_value(out),
977            }),
978            cube::Operator::Reinterpret(op) => instructions.push(wgsl::Instruction::Bitcast {
979                input: self.compile_value(op.input),
980                out: self.compile_value(out),
981            }),
982            cube::Operator::InitVector(op) => instructions.push(wgsl::Instruction::VecInit {
983                inputs: op
984                    .inputs
985                    .into_iter()
986                    .map(|val| self.compile_value(val))
987                    .collect(),
988                out: self.compile_value(out),
989            }),
990            cube::Operator::ExtractComponent(op) => instructions.push(wgsl::Instruction::Extract {
991                vector: self.compile_value(op.lhs),
992                index: self.compile_value(op.rhs),
993                out: self.compile_value(out),
994            }),
995            cube::Operator::InsertComponent(op) => instructions.push(wgsl::Instruction::Insert {
996                vector: self.compile_value(op.vector),
997                index: self.compile_value(op.index),
998                value: self.compile_value(op.value),
999                out: self.compile_value(out),
1000            }),
1001            cube::Operator::Select(op) => instructions.push(wgsl::Instruction::Select {
1002                cond: self.compile_value(op.cond),
1003                then: self.compile_value(op.then),
1004                or_else: self.compile_value(op.or_else),
1005                out: self.compile_value(out),
1006            }),
1007            cube::Operator::ReadBuiltin(builtin) => {
1008                let out = self.compile_value(out);
1009                let constant = {
1010                    let out = out.clone();
1011                    |value| {
1012                        instructions.push(wgsl::Instruction::Assign { input: value, out });
1013                    }
1014                };
1015                let builtin = match builtin {
1016                    cube::Builtin::AbsolutePos => {
1017                        self.id = true;
1018                        wgsl::Builtin::Id
1019                    }
1020                    cube::Builtin::UnitPos => {
1021                        self.local_invocation_index = true;
1022                        wgsl::Builtin::LocalInvocationIndex
1023                    }
1024                    cube::Builtin::UnitPosX => {
1025                        self.local_invocation_id = true;
1026                        wgsl::Builtin::LocalInvocationIdX
1027                    }
1028                    cube::Builtin::UnitPosY => {
1029                        self.local_invocation_id = true;
1030                        wgsl::Builtin::LocalInvocationIdY
1031                    }
1032                    cube::Builtin::UnitPosZ => {
1033                        self.local_invocation_id = true;
1034                        wgsl::Builtin::LocalInvocationIdZ
1035                    }
1036                    cube::Builtin::CubePosX => {
1037                        self.workgroup_id = true;
1038                        wgsl::Builtin::WorkgroupIdX
1039                    }
1040                    cube::Builtin::CubePosY => {
1041                        self.workgroup_id = true;
1042                        wgsl::Builtin::WorkgroupIdY
1043                    }
1044                    cube::Builtin::CubePosZ => {
1045                        self.workgroup_id = true;
1046                        wgsl::Builtin::WorkgroupIdZ
1047                    }
1048                    cube::Builtin::CubePosCluster
1049                    | cube::Builtin::CubePosClusterX
1050                    | cube::Builtin::CubePosClusterY
1051                    | cube::Builtin::CubePosClusterZ => {
1052                        constant(self.constant_var(1));
1053                        return;
1054                    }
1055                    cube::Builtin::AbsolutePosX => {
1056                        self.global_invocation_id = true;
1057                        wgsl::Builtin::GlobalInvocationIdX
1058                    }
1059                    cube::Builtin::AbsolutePosY => {
1060                        self.global_invocation_id = true;
1061                        wgsl::Builtin::GlobalInvocationIdY
1062                    }
1063                    cube::Builtin::AbsolutePosZ => {
1064                        self.global_invocation_id = true;
1065                        wgsl::Builtin::GlobalInvocationIdZ
1066                    }
1067                    cube::Builtin::CubeDimX => wgsl::Builtin::WorkgroupSizeX,
1068                    cube::Builtin::CubeDimY => wgsl::Builtin::WorkgroupSizeY,
1069                    cube::Builtin::CubeDimZ => wgsl::Builtin::WorkgroupSizeZ,
1070                    cube::Builtin::CubeClusterDim
1071                    | cube::Builtin::CubeClusterDimX
1072                    | cube::Builtin::CubeClusterDimY
1073                    | cube::Builtin::CubeClusterDimZ => {
1074                        constant(self.constant_var(1));
1075                        return;
1076                    }
1077                    cube::Builtin::CubeCountX => {
1078                        self.num_workgroups = true;
1079                        wgsl::Builtin::NumWorkgroupsX
1080                    }
1081                    cube::Builtin::CubeCountY => {
1082                        self.num_workgroups = true;
1083                        wgsl::Builtin::NumWorkgroupsY
1084                    }
1085                    cube::Builtin::CubeCountZ => {
1086                        self.num_workgroups = true;
1087                        wgsl::Builtin::NumWorkgroupsZ
1088                    }
1089                    cube::Builtin::CubePos => {
1090                        self.workgroup_id_no_axis = true;
1091                        wgsl::Builtin::WorkgroupId
1092                    }
1093                    cube::Builtin::CubeDim => {
1094                        self.workgroup_size_no_axis = true;
1095                        wgsl::Builtin::WorkgroupSize
1096                    }
1097                    cube::Builtin::CubeCount => {
1098                        self.num_workgroup_no_axis = true;
1099                        wgsl::Builtin::NumWorkgroups
1100                    }
1101                    cube::Builtin::PlaneDim => {
1102                        self.subgroup_size = true;
1103                        wgsl::Builtin::SubgroupSize
1104                    }
1105                    cube::Builtin::PlanePos => {
1106                        self.subgroup_id = true;
1107                        wgsl::Builtin::SubgroupId
1108                    }
1109                    cube::Builtin::UnitPosPlane => {
1110                        self.subgroup_invocation_id = true;
1111                        wgsl::Builtin::SubgroupInvocationId
1112                    }
1113                };
1114                instructions.push(wgsl::Instruction::ReadBuiltin { builtin, out });
1115            }
1116            cube::Operator::ReadScalar(id) => instructions.push(wgsl::Instruction::ReadScalar {
1117                id,
1118                out: self.compile_value(out),
1119            }),
1120        }
1121    }
1122
1123    fn compile_atomic(
1124        &mut self,
1125        atomic: cube::AtomicOp,
1126        out: Option<cube::Value>,
1127    ) -> wgsl::Instruction {
1128        match atomic {
1129            cube::AtomicOp::Add(op) => wgsl::Instruction::AtomicAdd {
1130                ptr: self.compile_value(op.ptr),
1131                value: self.compile_value(op.value),
1132                out: self.compile_value(out.unwrap()),
1133            },
1134            cube::AtomicOp::Sub(op) => wgsl::Instruction::AtomicSub {
1135                ptr: self.compile_value(op.ptr),
1136                value: self.compile_value(op.value),
1137                out: self.compile_value(out.unwrap()),
1138            },
1139            cube::AtomicOp::Max(op) => wgsl::Instruction::AtomicMax {
1140                ptr: self.compile_value(op.ptr),
1141                value: self.compile_value(op.value),
1142                out: self.compile_value(out.unwrap()),
1143            },
1144            cube::AtomicOp::Min(op) => wgsl::Instruction::AtomicMin {
1145                ptr: self.compile_value(op.ptr),
1146                value: self.compile_value(op.value),
1147                out: self.compile_value(out.unwrap()),
1148            },
1149            cube::AtomicOp::And(op) => wgsl::Instruction::AtomicAnd {
1150                ptr: self.compile_value(op.ptr),
1151                value: self.compile_value(op.value),
1152                out: self.compile_value(out.unwrap()),
1153            },
1154            cube::AtomicOp::Or(op) => wgsl::Instruction::AtomicOr {
1155                ptr: self.compile_value(op.ptr),
1156                value: self.compile_value(op.value),
1157                out: self.compile_value(out.unwrap()),
1158            },
1159            cube::AtomicOp::Xor(op) => wgsl::Instruction::AtomicXor {
1160                ptr: self.compile_value(op.ptr),
1161                value: self.compile_value(op.value),
1162                out: self.compile_value(out.unwrap()),
1163            },
1164            cube::AtomicOp::Load(ptr) => wgsl::Instruction::AtomicLoad {
1165                input: self.compile_value(ptr),
1166                out: self.compile_value(out.unwrap()),
1167            },
1168            cube::AtomicOp::Store(op) => wgsl::Instruction::AtomicStore {
1169                input: self.compile_value(op.value),
1170                out: self.compile_value(op.ptr),
1171            },
1172            cube::AtomicOp::Swap(op) => wgsl::Instruction::AtomicSwap {
1173                lhs: self.compile_value(op.ptr),
1174                rhs: self.compile_value(op.value),
1175                out: self.compile_value(out.unwrap()),
1176            },
1177            cube::AtomicOp::CompareAndSwap(op) => wgsl::Instruction::AtomicCompareExchangeWeak {
1178                ptr: self.compile_value(op.ptr),
1179                cmp: self.compile_value(op.cmp),
1180                value: self.compile_value(op.val),
1181                out: self.compile_value(out.unwrap()),
1182            },
1183        }
1184    }
1185
1186    fn compile_binding(&mut self, arg: kernel::KernelArg) -> wgsl::KernelArg {
1187        wgsl::KernelArg {
1188            id: arg.id,
1189            visibility: self.buffer_vis[arg.id as usize],
1190            value: self.compile_value(arg.value),
1191        }
1192    }
1193}
1194
1195fn register_extensions(instructions: &[wgsl::Instruction]) -> Vec<wgsl::Extension> {
1196    let mut extensions = Vec::new();
1197
1198    let mut register_extension = |extension: wgsl::Extension| {
1199        if !extensions.contains(&extension) {
1200            extensions.push(extension);
1201        }
1202    };
1203
1204    // Since not all instructions are native to WGSL, we need to add the custom ones.
1205    for instruction in instructions {
1206        match instruction {
1207            wgsl::Instruction::Powf { lhs: _, rhs, out } => {
1208                register_extension(wgsl::Extension::PowfPrimitive(out.elem()));
1209                register_extension(wgsl::powf_extension(rhs, out));
1210            }
1211            #[cfg(target_os = "macos")]
1212            wgsl::Instruction::Tanh { input, out: _ } => {
1213                register_extension(wgsl::Extension::SafeTanhPrimitive(input.elem()));
1214                register_extension(wgsl::Extension::SafeTanh(input.item()));
1215            }
1216            wgsl::Instruction::IsNan { input, out } => {
1217                register_extension(wgsl::Extension::IsNanPrimitive(input.elem()));
1218                register_extension(wgsl::Extension::IsNan(input.item(), out.item()));
1219            }
1220            wgsl::Instruction::IsInf { input, out } => {
1221                register_extension(wgsl::Extension::IsInfPrimitive(input.elem()));
1222                register_extension(wgsl::Extension::IsInf(input.item(), out.item()));
1223            }
1224            wgsl::Instruction::If { instructions, .. } => {
1225                for extension in register_extensions(instructions) {
1226                    register_extension(extension);
1227                }
1228            }
1229            wgsl::Instruction::IfElse {
1230                instructions_if,
1231                instructions_else,
1232                ..
1233            } => {
1234                for extension in register_extensions(instructions_if) {
1235                    register_extension(extension);
1236                }
1237                for extension in register_extensions(instructions_else) {
1238                    register_extension(extension);
1239                }
1240            }
1241            wgsl::Instruction::Loop { instructions } => {
1242                for extension in register_extensions(instructions) {
1243                    register_extension(extension);
1244                }
1245            }
1246            wgsl::Instruction::RangeLoop { instructions, .. } => {
1247                for extension in register_extensions(instructions) {
1248                    register_extension(extension);
1249                }
1250            }
1251            _ => {}
1252        }
1253    }
1254
1255    extensions
1256}
cubecl_wgpu/compiler/wgsl/compiler.rs

cubecl_wgpu/compiler/wgsl/
compiler.rs
