Skip to main content

cubecl_cpp/shared/
base.rs

1use super::{
2    BinaryInstruction, Body, Component, ComputeKernel, Dialect, Elem, FP4Kind, FP6Kind, FP8Kind,
3    FragmentIdent, FragmentLayout, FragmentType, IndexInstruction, Instruction, Item, KernelArg,
4    SharedMemory, UnaryInstruction, Value, WarpInstruction, WmmaInstruction, barrier::BarrierOps,
5};
6use crate::shared::{Builtin, MmaShape, PointerClass};
7use cubecl_common::backtrace::BackTrace;
8use cubecl_core::{
9    CubeDim,
10    ir::{
11        self as ir, AddressSpace, DeviceProperties, ElemType, FloatKind, InstructionModes,
12        OpaqueType, Operation, Processor, SourceLoc, StorageType, Type,
13        features::{AtomicUsage, EnumSet, TypeUsage},
14    },
15    post_processing::{self, checked_io::CheckedIoVisitor, disaggregate::DisaggregateVisitor},
16    prelude::{FastMath, KernelDefinition, Visibility},
17    server::ExecutionMode,
18};
19use cubecl_opt::{Optimizer, SharedLiveness};
20use cubecl_runtime::compiler::{CompilationError, Compiler};
21use std::{
22    collections::{HashMap, HashSet},
23    fmt::Debug,
24};
25
26pub(super) static COUNTER_TMP_VAR: std::sync::atomic::AtomicU32 =
27    std::sync::atomic::AtomicU32::new(0);
28
29#[derive(Clone, Debug)]
30pub struct CompilationOptions {
31    pub warp_size: u32,
32    pub supports_features: CppSupportedFeatures,
33}
34
35#[derive(Clone, Debug, Default)]
36pub struct CppSupportedFeatures {
37    pub grid_constants: bool,
38    pub clusters: bool,
39    pub fast_math: bool,
40    pub fast_tanh: bool,
41    pub elect_sync: bool,
42}
43
44impl Default for CompilationOptions {
45    fn default() -> Self {
46        Self {
47            warp_size: 32,
48            supports_features: Default::default(),
49        }
50    }
51}
52
53/// Cube indexes flags.
54/// When true the corresponding index is declared and computed as needed in the kernel.
55#[derive(Debug, Clone, Default)]
56pub struct CubeIndexFlags {
57    pub absolute_pos: bool,
58    pub absolute_pos_tuple: bool,
59    pub cube_count: bool,
60    pub cube_count_tuple: bool,
61    pub cube_dim: bool,
62    pub cube_dim_tuple: bool,
63    pub cube_pos: bool,
64    pub cube_pos_tuple: bool,
65    pub plane_dim: bool,
66    pub plane_dim_checked: bool,
67    pub plane_pos: bool,
68    pub unit_pos: bool,
69    pub unit_pos_tuple: bool,
70    pub unit_pos_plane: bool,
71    pub cluster_pos: bool,
72}
73
74/// Flags gathered during Cube IR translation for the kernel compilation.
75#[derive(Debug, Clone)]
76pub struct Flags<D: Dialect> {
77    pub elem_fp4: bool,
78    pub elem_fp6: bool,
79    pub elem_fp8: bool,
80    pub elem_bf16: bool,
81    pub elem_f16: bool,
82    pub elem_tf32: bool,
83    pub indexes: CubeIndexFlags,
84    pub op_barrier: bool,
85    pub thread_block: bool,
86    pub inst_tma: bool,
87    pub inst_tma_im2col: bool,
88    pub inst_wmma: bool,
89    pub inst_ptx_wrappers: bool,
90    pub inst_async_copy: bool,
91    pub use_grid_constants: bool,
92    pub static_meta_length: usize,
93    pub has_dynamic_meta: bool,
94    pub has_info: bool,
95    pub cube_dim: CubeDim,
96    pub cluster_dim: Option<CubeDim>,
97    pub address_type: Item<D>,
98}
99
100#[allow(clippy::too_many_arguments)]
101#[derive(Clone, Debug)]
102pub struct CppCompiler<D: Dialect> {
103    kernel_name: String,
104    buffer_vis: Vec<Visibility>,
105    barriers: Vec<BarrierOps<D>>,
106    compilation_options: CompilationOptions,
107    ext_meta_positions: HashMap<ir::Value, u32>,
108    cluster_dim: CubeDim,
109    extensions: Vec<D::Extension>,
110    flags: Flags<D>,
111    items: HashSet<Item<D>>,
112    info: cubecl_core::Info,
113    source_loc: Option<SourceLoc>,
114    strategy: ExecutionMode,
115    addr_type: Item<D>,
116}
117
118impl<D: Dialect> Default for Flags<D> {
119    fn default() -> Self {
120        Self {
121            elem_fp4: Default::default(),
122            elem_fp6: Default::default(),
123            elem_fp8: Default::default(),
124            elem_bf16: Default::default(),
125            elem_f16: Default::default(),
126            elem_tf32: Default::default(),
127            indexes: Default::default(),
128            op_barrier: Default::default(),
129            thread_block: Default::default(),
130            inst_tma: Default::default(),
131            inst_tma_im2col: Default::default(),
132            inst_wmma: Default::default(),
133            inst_ptx_wrappers: Default::default(),
134            inst_async_copy: Default::default(),
135            use_grid_constants: Default::default(),
136            static_meta_length: Default::default(),
137            has_info: Default::default(),
138            has_dynamic_meta: Default::default(),
139            cube_dim: CubeDim::new_single(),
140            cluster_dim: Default::default(),
141            address_type: Item::Scalar(Elem::U32),
142        }
143    }
144}
145
146impl<D: Dialect> Default for CppCompiler<D> {
147    fn default() -> Self {
148        Self {
149            kernel_name: Default::default(),
150            buffer_vis: Default::default(),
151            barriers: Default::default(),
152            compilation_options: Default::default(),
153            ext_meta_positions: Default::default(),
154            cluster_dim: CubeDim::new_single(),
155            extensions: Default::default(),
156            flags: Flags::default(),
157            items: Default::default(),
158            info: Default::default(),
159            source_loc: Default::default(),
160            strategy: Default::default(),
161            addr_type: Item::Scalar(Elem::U32),
162        }
163    }
164}
165
166impl<D: Dialect> Compiler for CppCompiler<D> {
167    type Representation = ComputeKernel<D>;
168    type CompilationOptions = CompilationOptions;
169
170    fn compile(
171        &mut self,
172        mut kernel: KernelDefinition,
173        compilation_options: &Self::CompilationOptions,
174        strategy: ExecutionMode,
175        addr_type: StorageType,
176    ) -> Result<Self::Representation, CompilationError> {
177        let errors = kernel.body.pop_errors();
178        if !errors.is_empty() {
179            let mut reason = "Can't compile cpp kernel\nCaused by:\n  ".to_string();
180            for error in errors {
181                reason += error.as_str();
182                reason += "\n";
183            }
184
185            return Err(CompilationError::Validation {
186                reason,
187                backtrace: BackTrace::capture(),
188            });
189        }
190
191        self.addr_type = self.compile_type(addr_type.into());
192        self.compilation_options = compilation_options.clone();
193        self.strategy = strategy;
194        self.kernel_name = kernel.options.kernel_name.clone();
195
196        if !self.compilation_options.supports_features.clusters {
197            kernel.options.cluster_dim = None;
198        }
199        self.cluster_dim = kernel.options.cluster_dim.unwrap_or(CubeDim::new_single());
200
201        let ir = self.clone().compile_ir(kernel, addr_type);
202        COUNTER_TMP_VAR.store(0, std::sync::atomic::Ordering::Relaxed);
203        Ok(ir)
204    }
205
206    fn elem_size(&self, elem: ir::ElemType) -> usize {
207        elem.size()
208    }
209
210    fn extension(&self) -> &'static str {
211        "cpp"
212    }
213}
214
215impl<D: Dialect> CppCompiler<D> {
216    fn compile_ir(
217        mut self,
218        value: KernelDefinition,
219        address_type: StorageType,
220    ) -> ComputeKernel<D> {
221        let metadata = self.build_metadata(&value);
222        self.info = cubecl_core::Info::new(&value.scalars, metadata, address_type);
223
224        let scope_state = value.body.state().clone_deep();
225
226        let mut opt = Optimizer::shared_only(value.body.clone(), value.cube_dim);
227        let shared_allocs = opt.main.analysis::<SharedLiveness>(&opt.global_state);
228
229        *value.body.state_mut() = scope_state;
230
231        CheckedIoVisitor::new(self.strategy, self.kernel_name.clone()).apply(&value.body);
232        DisaggregateVisitor::apply(&value.body);
233
234        self.buffer_vis = post_processing::optimize_scope(&value.body).into();
235        self.buffer_vis
236            .resize(value.num_global_buffers(), Visibility::Read);
237
238        let address_type = self.compile_type(address_type.into());
239        let instructions = self.compile_scope(&value.body);
240
241        let tensor_maps = value
242            .tensor_maps
243            .into_iter()
244            .map(|b| self.compile_binding(b))
245            .collect();
246        let buffers = value
247            .buffers
248            .into_iter()
249            .map(|b| self.compile_binding(b))
250            .collect();
251        let scalars = value
252            .scalars
253            .into_iter()
254            .map(|binding| (self.compile_storage_type(binding.ty), binding.count))
255            .collect::<Vec<_>>();
256
257        let shared_memories = shared_allocs
258            .allocations
259            .values()
260            .map(|alloc| SharedMemory {
261                ptr: self.compile_value(ir::Value::new(alloc.id, alloc.smem.root_ptr.ty)),
262                value_ty: self.compile_type(alloc.smem.value_ty),
263                align: alloc.smem.alignment,
264                offset: alloc.offset,
265            })
266            .collect();
267
268        let body = Body {
269            instructions,
270            shared_memories,
271            barriers: self.barriers,
272            info_by_ptr: !self.compilation_options.supports_features.grid_constants,
273            has_dynamic_meta: self.info.has_dynamic_meta,
274            address_type: self.addr_type,
275        };
276
277        // translation flags
278        let flags = Flags {
279            indexes: D::builtin_rules(&self.flags.indexes),
280            inst_wmma: self.flags.inst_wmma,
281            thread_block: self.flags.thread_block,
282            op_barrier: self.flags.op_barrier,
283            elem_fp4: self.flags.elem_fp4,
284            elem_fp6: self.flags.elem_fp6,
285            elem_fp8: self.flags.elem_fp8,
286            elem_bf16: self.flags.elem_bf16,
287            elem_f16: self.flags.elem_f16,
288            elem_tf32: self.flags.elem_tf32,
289            inst_tma: self.flags.inst_tma,
290            inst_tma_im2col: self.flags.inst_tma_im2col,
291            inst_async_copy: self.flags.inst_async_copy,
292            inst_ptx_wrappers: self.flags.inst_ptx_wrappers,
293            use_grid_constants: self.compilation_options.supports_features.grid_constants,
294            has_info: self.info.has_info(),
295            has_dynamic_meta: self.info.has_dynamic_meta,
296            static_meta_length: self.info.metadata.static_len() as usize,
297            cube_dim: value.cube_dim,
298            cluster_dim: value.options.cluster_dim,
299            address_type,
300        };
301
302        let mut cluster_dim = value.options.cluster_dim;
303        if !self.compilation_options.supports_features.clusters {
304            cluster_dim = None;
305        }
306
307        ComputeKernel {
308            tensor_maps,
309            buffers,
310            scalars,
311            meta_static_len: self.info.metadata.static_len() as usize,
312            cube_dim: value.cube_dim,
313            body,
314            extensions: self.extensions,
315            flags,
316            items: self.items,
317            kernel_name: value.options.kernel_name,
318            cluster_dim,
319            info: self.info.clone(),
320        }
321    }
322
323    fn build_metadata(&mut self, value: &KernelDefinition) -> cubecl_core::Metadata {
324        let mut num_ext = 0;
325
326        let mut all_meta: Vec<_> = value
327            .buffers
328            .iter()
329            .chain(value.tensor_maps.iter())
330            .map(|buf| (buf.id, buf.value, buf.has_extended_meta))
331            .collect();
332
333        all_meta.sort_by_key(|(id, _, _)| *id);
334
335        for (_, value, has_extended_meta) in &all_meta {
336            self.ext_meta_positions.insert(*value, num_ext);
337            if *has_extended_meta {
338                num_ext += 1;
339            }
340        }
341
342        let num_meta = all_meta.len();
343
344        cubecl_core::Metadata::new(num_meta as u32, num_ext)
345    }
346
347    pub(crate) fn ext_meta_position(&self, val: &ir::Value) -> u32 {
348        self.ext_meta_positions[val]
349    }
350
351    fn compile_scope(&mut self, scope: &ir::Scope) -> Vec<Instruction<D>> {
352        let mut instructions = Vec::new();
353
354        let dialect_processors = D::processors();
355        let mut processors: Vec<&dyn Processor> = vec![];
356        processors.extend(dialect_processors.iter().map(|it| &**it));
357
358        let processing = scope.process(processors);
359
360        processing
361            .instructions
362            .into_iter()
363            .for_each(|op| self.compile_instruction(&mut instructions, op));
364
365        instructions
366    }
367
368    fn compile_instruction(
369        &mut self,
370        instructions: &mut Vec<Instruction<D>>,
371        instruction: ir::Instruction,
372    ) {
373        self.update_debug_loc(instructions, &instruction);
374        let out = instruction.out;
375
376        match instruction.operation {
377            ir::Operation::Copy(value) => {
378                instructions.push(Instruction::Assign(UnaryInstruction {
379                    input: self.compile_value(value),
380                    out: self.compile_value(out.unwrap()),
381                }));
382            }
383            ir::Operation::DeclareVariable {
384                addr_space: AddressSpace::Local,
385                value_ty,
386                ..
387            } => instructions.push(Instruction::DeclareVariable {
388                val: self.compile_value(out.unwrap()),
389                value_ty: self.compile_type(value_ty),
390            }),
391            ir::Operation::DeclareVariable {
392                addr_space: AddressSpace::Shared,
393                ..
394            } => {
395                // Optimizer handles parsing and allocating these
396            }
397            ir::Operation::DeclareVariable { addr_space, .. } => {
398                unimplemented!("Unsupported declaration address space {addr_space}")
399            }
400            ir::Operation::Arithmetic(op) => {
401                self.compile_arithmetic(op, out, instruction.modes, instructions)
402            }
403            ir::Operation::Memory(op) => self.compile_memory(op, out, instructions),
404            ir::Operation::Comparison(op) => self.compile_comparison(op, out, instructions),
405            ir::Operation::Bitwise(op) => self.compile_bitwise(op, out, instructions),
406            ir::Operation::Operator(op) => self.compile_operator(op, out, instructions),
407            ir::Operation::Atomic(op) => self.compile_atomic(op, out, instructions),
408            ir::Operation::Metadata(op) => instructions.push(self.compile_metadata(op, out)),
409            ir::Operation::Branch(val) => self.compile_branch(instructions, val),
410            ir::Operation::Synchronization(val) => match val {
411                ir::Synchronization::SyncCube => instructions.push(Instruction::SyncThreads),
412                ir::Synchronization::SyncPlane => instructions.push(Instruction::SyncWarp),
413                ir::Synchronization::SyncStorage => instructions.push(Instruction::SyncThreads),
414                ir::Synchronization::SyncAsyncProxyShared => {
415                    self.flags.inst_tma = true;
416                    instructions.push(Instruction::ProxyAsyncToSharedFence)
417                }
418            },
419            ir::Operation::WorkgroupUniformLoad(input) => {
420                let is_atomic = input.ty.is_atomic();
421                instructions.push(Instruction::SyncThreads);
422                let load = UnaryInstruction {
423                    input: self.compile_value(input),
424                    out: self.compile_value(out.unwrap()),
425                };
426                if is_atomic {
427                    instructions.push(Instruction::AtomicLoad(load));
428                } else {
429                    instructions.push(Instruction::Load(load));
430                }
431            }
432            ir::Operation::Plane(op) => {
433                self.flags.indexes.plane_dim_checked = true;
434                let out = self.compile_value(out.unwrap());
435                match op {
436                    ir::Plane::Sum(op) => {
437                        let instruction = WarpInstruction::ReduceSum {
438                            input: self.compile_value(op.input),
439                            out,
440                        };
441                        D::register_warp_instruction_extension(&mut self.extensions, &instruction);
442                        instructions.push(Instruction::Warp(instruction));
443                    }
444                    ir::Plane::InclusiveSum(op) => {
445                        self.flags.indexes.unit_pos_plane = true;
446                        instructions.push(Instruction::Warp(WarpInstruction::InclusiveSum {
447                            input: self.compile_value(op.input),
448                            out,
449                        }))
450                    }
451                    ir::Plane::InclusiveProd(op) => {
452                        self.flags.indexes.unit_pos_plane = true;
453                        instructions.push(Instruction::Warp(WarpInstruction::InclusiveProd {
454                            input: self.compile_value(op.input),
455                            out,
456                        }))
457                    }
458                    ir::Plane::ExclusiveSum(op) => {
459                        self.flags.indexes.unit_pos_plane = true;
460                        instructions.push(Instruction::Warp(WarpInstruction::ExclusiveSum {
461                            input: self.compile_value(op.input),
462                            out,
463                        }))
464                    }
465                    ir::Plane::ExclusiveProd(op) => {
466                        self.flags.indexes.unit_pos_plane = true;
467                        instructions.push(Instruction::Warp(WarpInstruction::ExclusiveProd {
468                            input: self.compile_value(op.input),
469                            out,
470                        }))
471                    }
472                    ir::Plane::Prod(op) => {
473                        let instruction = WarpInstruction::ReduceProd {
474                            input: self.compile_value(op.input),
475                            out,
476                        };
477                        D::register_warp_instruction_extension(&mut self.extensions, &instruction);
478                        instructions.push(Instruction::Warp(instruction))
479                    }
480                    ir::Plane::Max(op) => {
481                        let instruction = WarpInstruction::ReduceMax {
482                            input: self.compile_value(op.input),
483                            out,
484                        };
485                        D::register_warp_instruction_extension(&mut self.extensions, &instruction);
486                        instructions.push(Instruction::Warp(instruction))
487                    }
488                    ir::Plane::Min(op) => {
489                        let instruction = WarpInstruction::ReduceMin {
490                            input: self.compile_value(op.input),
491                            out,
492                        };
493                        D::register_warp_instruction_extension(&mut self.extensions, &instruction);
494                        instructions.push(Instruction::Warp(instruction))
495                    }
496                    ir::Plane::Elect => {
497                        if self.compilation_options.supports_features.elect_sync {
498                            self.flags.inst_ptx_wrappers = true;
499                            instructions.push(Instruction::Warp(WarpInstruction::Elect { out }))
500                        } else {
501                            instructions
502                                .push(Instruction::Warp(WarpInstruction::ElectFallback { out }))
503                        }
504                    }
505                    ir::Plane::All(op) => {
506                        instructions.push(Instruction::Warp(WarpInstruction::All {
507                            input: self.compile_value(op.input),
508                            out,
509                        }))
510                    }
511                    ir::Plane::Any(op) => {
512                        instructions.push(Instruction::Warp(WarpInstruction::Any {
513                            input: self.compile_value(op.input),
514                            out,
515                        }))
516                    }
517                    ir::Plane::Ballot(op) => {
518                        instructions.push(Instruction::Warp(WarpInstruction::Ballot {
519                            input: self.compile_value(op.input),
520                            out,
521                        }))
522                    }
523                    ir::Plane::Broadcast(op) => {
524                        instructions.push(Instruction::Warp(WarpInstruction::Broadcast {
525                            input: self.compile_value(op.lhs),
526                            id: self.compile_value(op.rhs),
527                            out,
528                        }))
529                    }
530                    ir::Plane::Shuffle(op) => {
531                        instructions.push(Instruction::Warp(WarpInstruction::Shuffle {
532                            input: self.compile_value(op.lhs),
533                            src_lane: self.compile_value(op.rhs),
534                            out,
535                        }))
536                    }
537                    ir::Plane::ShuffleXor(op) => {
538                        instructions.push(Instruction::Warp(WarpInstruction::ShuffleXor {
539                            input: self.compile_value(op.lhs),
540                            mask: self.compile_value(op.rhs),
541                            out,
542                        }))
543                    }
544                    ir::Plane::ShuffleUp(op) => {
545                        instructions.push(Instruction::Warp(WarpInstruction::ShuffleUp {
546                            input: self.compile_value(op.lhs),
547                            delta: self.compile_value(op.rhs),
548                            out,
549                        }))
550                    }
551                    ir::Plane::ShuffleDown(op) => {
552                        instructions.push(Instruction::Warp(WarpInstruction::ShuffleDown {
553                            input: self.compile_value(op.lhs),
554                            delta: self.compile_value(op.rhs),
555                            out,
556                        }))
557                    }
558                }
559            }
560            ir::Operation::CoopMma(cmma) => instructions.push(self.compile_cmma(cmma, out)),
561            ir::Operation::NonSemantic(debug) => match debug {
562                ir::NonSemantic::Print {
563                    format_string,
564                    args,
565                } => instructions.push(Instruction::Printf {
566                    format_string,
567                    args: args
568                        .into_iter()
569                        .map(|arg| self.compile_value(arg))
570                        .collect(),
571                }),
572                ir::NonSemantic::Comment { content } => {
573                    instructions.push(Instruction::Comment { content })
574                }
575                // Don't need to handle scopes
576                _ => {}
577            },
578            ir::Operation::TensorIndexing(_) => panic!("Tensor indexing only supported in Vulkan"),
579            ir::Operation::Barrier(barrier_ops) => match barrier_ops {
580                ir::BarrierOps::Init {
581                    barrier,
582                    is_elected,
583                    arrival_count,
584                } => {
585                    let Type::Opaque(OpaqueType::Barrier(level)) = barrier.ty else {
586                        unreachable!()
587                    };
588                    let barrier = self.compile_value(barrier);
589                    let arrival_count = self.compile_value(arrival_count);
590                    instructions.push(Instruction::Barrier(super::barrier::BarrierOps::Init {
591                        barrier,
592                        is_elected: self.compile_value(is_elected),
593                        arrival_count,
594                        level,
595                    }));
596                }
597                ir::BarrierOps::InitManual {
598                    barrier,
599                    arrival_count,
600                } => {
601                    let barrier = self.compile_value(barrier);
602                    let arrival_count = self.compile_value(arrival_count);
603                    instructions.push(Instruction::Barrier(
604                        super::barrier::BarrierOps::InitManual {
605                            barrier,
606                            arrival_count,
607                        },
608                    ));
609                }
610                ir::BarrierOps::MemCopyAsync {
611                    barrier,
612                    source,
613                    destination,
614                    source_length,
615                } => {
616                    instructions.push(Instruction::Barrier(
617                        super::barrier::BarrierOps::MemCopyAsync {
618                            barrier: self.compile_value(barrier),
619                            source: self.compile_value(source),
620                            destination: self.compile_value(destination),
621                            source_length: self.compile_value(source_length),
622                            cooperative: false,
623                        },
624                    ));
625                }
626                ir::BarrierOps::MemCopyAsyncCooperative {
627                    barrier,
628                    source,
629                    destination,
630                    source_length,
631                } => {
632                    self.flags.thread_block = true;
633                    instructions.push(Instruction::Barrier(
634                        super::barrier::BarrierOps::MemCopyAsync {
635                            barrier: self.compile_value(barrier),
636                            source: self.compile_value(source),
637                            destination: self.compile_value(destination),
638                            source_length: self.compile_value(source_length),
639                            cooperative: true,
640                        },
641                    ));
642                }
643                ir::BarrierOps::MemCopyAsyncTx {
644                    barrier,
645                    source,
646                    destination,
647                    source_length,
648                } => {
649                    instructions.push(Instruction::Barrier(
650                        super::barrier::BarrierOps::MemCopyAsyncTx {
651                            barrier: self.compile_value(barrier),
652                            source: self.compile_value(source),
653                            destination: self.compile_value(destination),
654                            source_length: self.compile_value(source_length),
655                        },
656                    ));
657                }
658                ir::BarrierOps::CopyAsync {
659                    source,
660                    destination,
661                    source_length,
662                    copy_length,
663                    checked,
664                } => {
665                    self.flags.inst_async_copy = true;
666                    instructions.push(Instruction::Barrier(
667                        super::barrier::BarrierOps::CopyAsync {
668                            source: self.compile_value(source),
669                            destination: self.compile_value(destination),
670                            source_length: self.compile_value(source_length),
671                            copy_size: copy_length,
672                            checked,
673                        },
674                    ));
675                }
676                ir::BarrierOps::TmaLoad {
677                    barrier,
678                    tensor_map,
679                    destination,
680                    indices,
681                } => {
682                    instructions.push(Instruction::Barrier(
683                        super::barrier::BarrierOps::MemCopyAsyncTensorGlobalToShared {
684                            barrier: self.compile_value(barrier),
685                            smem_buffer: self.compile_value(destination),
686                            tensor_map: self.compile_value(tensor_map),
687                            indices: indices
688                                .into_iter()
689                                .map(|it| self.compile_value(it))
690                                .collect(),
691                        },
692                    ));
693                }
694                ir::BarrierOps::TmaLoadIm2col {
695                    barrier,
696                    tensor_map,
697                    destination,
698                    indices,
699                    offsets,
700                } => {
701                    self.flags.inst_tma_im2col = true;
702                    instructions.push(Instruction::Barrier(
703                        super::barrier::BarrierOps::TmaLoadIm2col {
704                            barrier: self.compile_value(barrier),
705                            smem_buffer: self.compile_value(destination),
706                            tensor_map: self.compile_value(tensor_map),
707                            indices: indices
708                                .into_iter()
709                                .map(|it| self.compile_value(it))
710                                .collect(),
711                            offsets: offsets
712                                .into_iter()
713                                .map(|it| self.compile_value(it))
714                                .collect(),
715                        },
716                    ));
717                }
718                ir::BarrierOps::Arrive { barrier } => {
719                    instructions.push(Instruction::Barrier(super::barrier::BarrierOps::Arrive {
720                        barrier: self.compile_value(barrier),
721                        token: self.compile_value(out.unwrap()),
722                    }))
723                }
724                ir::BarrierOps::ArriveTx {
725                    barrier,
726                    arrive_count_update,
727                    transaction_count_update,
728                } => {
729                    instructions.push(Instruction::Barrier(super::barrier::BarrierOps::ArriveTx {
730                        barrier: self.compile_value(barrier),
731                        token: self.compile_value(out.unwrap()),
732                        arrive_count_update: self.compile_value(arrive_count_update),
733                        transaction_count_update: self.compile_value(transaction_count_update),
734                    }))
735                }
736                ir::BarrierOps::CommitCopyAsync { barrier } => {
737                    self.flags.inst_async_copy = true;
738                    instructions.push(Instruction::Barrier(
739                        super::barrier::BarrierOps::ArriveCopyAsync {
740                            barrier: self.compile_value(barrier),
741                        },
742                    ))
743                }
744                ir::BarrierOps::ExpectTx {
745                    barrier,
746                    transaction_count_update,
747                } => {
748                    instructions.push(Instruction::Barrier(super::barrier::BarrierOps::ExpectTx {
749                        barrier: self.compile_value(barrier),
750                        transaction_count_update: self.compile_value(transaction_count_update),
751                    }))
752                }
753                ir::BarrierOps::Wait { barrier, token } => {
754                    instructions.push(Instruction::Barrier(super::barrier::BarrierOps::Wait {
755                        barrier: self.compile_value(barrier),
756                        token: self.compile_value(token),
757                    }))
758                }
759                ir::BarrierOps::WaitParity { barrier, phase } => instructions.push(
760                    Instruction::Barrier(super::barrier::BarrierOps::WaitParity {
761                        barrier: self.compile_value(barrier),
762                        phase: self.compile_value(phase),
763                    }),
764                ),
765                ir::BarrierOps::ArriveAndWait { barrier } => {
766                    let Type::Opaque(OpaqueType::Barrier(level)) = barrier.ty else {
767                        unreachable!()
768                    };
769                    instructions.push(Instruction::Barrier(
770                        super::barrier::BarrierOps::ArriveAndWait {
771                            barrier: self.compile_value(barrier),
772                            level,
773                        },
774                    ))
775                }
776            },
777            ir::Operation::Tma(tma_ops) => {
778                self.flags.inst_tma = true;
779                match tma_ops {
780                    ir::TmaOps::TmaStore {
781                        source,
782                        coordinates,
783                    } => {
784                        instructions.push(Instruction::MemCopyAsyncTensorSharedToGlobal {
785                            smem_buffer: self.compile_value(source),
786                            tensor_map: self.compile_value(out.unwrap()),
787                            indices: coordinates
788                                .into_iter()
789                                .map(|it| self.compile_value(it))
790                                .collect(),
791                        });
792                    }
793                    ir::TmaOps::CommitGroup => {
794                        instructions.push(Instruction::BulkCommitGroup);
795                    }
796                    ir::TmaOps::WaitGroup { max_pending } => {
797                        instructions.push(Instruction::BulkWaitGroup { max_pending });
798                    }
799                    ir::TmaOps::WaitGroupRead { max_pending } => {
800                        instructions.push(Instruction::BulkWaitGroupRead { max_pending });
801                    }
802                }
803            }
804            ir::Operation::Marker(_) => {}
805            ir::Operation::ConstructAggregate(..) | ir::Operation::ExtractAggregateField(..) => {
806                unreachable!("Should be disaggregated at this point")
807            }
808        }
809    }
810
811    fn update_debug_loc(&mut self, instructions: &mut Vec<Instruction<D>>, inst: &ir::Instruction) {
812        if !matches!(inst.operation, Operation::NonSemantic(_)) {
813            match &inst.source_loc {
814                Some(loc) if Some(loc) != self.source_loc.as_ref() => {
815                    self.source_loc = Some(loc.clone());
816                    instructions.push(Instruction::Line {
817                        file: loc.source.file.clone(),
818                        line: loc.line,
819                    });
820                }
821                _ => {}
822            }
823        }
824    }
825
826    fn compile_cmma(&mut self, cmma: ir::CoopMma, out: Option<ir::Value>) -> Instruction<D> {
827        self.flags.inst_wmma = true;
828
829        let inst = match cmma {
830            ir::CoopMma::Fill { value } => WmmaInstruction::Fill {
831                frag: self.compile_value(out.unwrap()),
832                value: self.compile_value(value),
833            },
834            ir::CoopMma::Load {
835                ptr,
836                stride,
837                layout,
838            } => WmmaInstruction::Load {
839                frag: self.compile_value(out.unwrap()),
840                ptr: self.compile_value(ptr),
841                stride: self.compile_value(stride),
842                layout: layout.and_then(|l| self.compile_matrix_layout(l)),
843            },
844            ir::CoopMma::Execute {
845                mat_a,
846                mat_b,
847                mat_c,
848            } => WmmaInstruction::Execute {
849                frag_a: self.compile_value(mat_a),
850                frag_b: self.compile_value(mat_b),
851                frag_c: self.compile_value(mat_c),
852                frag_d: self.compile_value(out.unwrap()),
853                warp_size: self.compilation_options.warp_size,
854            },
855            ir::CoopMma::ExecuteManual {
856                matrix,
857                registers_a,
858                registers_b,
859                registers_c,
860            } => WmmaInstruction::ExecuteManual {
861                shape: MmaShape::new(matrix.m as u32, matrix.n as u32, matrix.k as u32),
862                frag_a: self.compile_value(registers_a),
863                frag_b: self.compile_value(registers_b),
864                frag_c: self.compile_value(registers_c),
865                frag_d: self.compile_value(out.unwrap()),
866            },
867            ir::CoopMma::ExecuteScaled {
868                matrix,
869                registers_a,
870                registers_b,
871                registers_c,
872                scales_a,
873                scales_b,
874                scales_factor,
875            } => WmmaInstruction::ExecuteScaled {
876                shape: MmaShape::new(matrix.m as u32, matrix.n as u32, matrix.k as u32),
877                frag_a: self.compile_value(registers_a),
878                frag_b: self.compile_value(registers_b),
879                frag_c: self.compile_value(registers_c),
880                frag_d: self.compile_value(out.unwrap()),
881
882                scales_a: self.compile_value(scales_a),
883                scales_b: self.compile_value(scales_b),
884                scales_factor: scales_factor as u32,
885            },
886            ir::CoopMma::ExecuteElementwise { .. } => {
887                panic!("Elementwise only supported in Vulkan")
888            }
889            ir::CoopMma::Store {
890                mat,
891                stride,
892                destination,
893                layout,
894            } => {
895                self.flags.indexes.unit_pos = true;
896                self.flags.indexes.plane_pos = true;
897                WmmaInstruction::Store {
898                    destination: self.compile_value(destination),
899                    frag: self.compile_value(mat),
900                    stride: self.compile_value(stride),
901                    layout: self
902                        .compile_matrix_layout(layout)
903                        .expect("Layout required for store instruction"),
904                }
905            }
906            ir::CoopMma::LoadMatrix {
907                ptr,
908                factor,
909                transpose,
910            } => WmmaInstruction::LdMatrix {
911                output: self.compile_value(out.unwrap()),
912                ptr: self.compile_value(ptr),
913                factor: factor as u32,
914                transpose,
915            },
916            ir::CoopMma::StoreMatrix {
917                registers,
918                factor,
919                transpose,
920                destination,
921            } => WmmaInstruction::StMatrix {
922                registers: self.compile_value(registers),
923                ptr: self.compile_value(destination),
924                factor: factor as u32,
925                transpose,
926            },
927            ir::CoopMma::Cast { input } => WmmaInstruction::Cast {
928                input: self.compile_value(input),
929                output: self.compile_value(out.unwrap()),
930            },
931            ir::CoopMma::RowIndex { .. } | ir::CoopMma::ColIndex { .. } => {
932                panic!("Row/Col index should be handled by processors")
933            }
934            ir::CoopMma::LoadTensor { .. } | ir::CoopMma::StoreTensor { .. } => {
935                panic!("Load/store tensor is only supported in Vulkan")
936            }
937        };
938
939        D::register_wmma_instruction_extension(&mut self.extensions, &inst);
940
941        Instruction::Wmma(inst)
942    }
943
944    fn compile_metadata(
945        &mut self,
946        metadata: ir::Metadata,
947        out: Option<ir::Value>,
948    ) -> Instruction<D> {
949        let out = out.unwrap();
950        match metadata {
951            ir::Metadata::Stride { dim, list } => {
952                let position = self.ext_meta_position(&list);
953                let offset = self.info.metadata.stride_offset_index(position);
954                Instruction::ExtendedMetadata {
955                    info_offset: self.compile_value(offset.into()),
956                    dim: self.compile_value(dim),
957                    out: self.compile_value(out),
958                }
959            }
960            ir::Metadata::Shape { dim, list } => {
961                let position = self.ext_meta_position(&list);
962                let offset = self.info.metadata.shape_offset_index(position);
963                Instruction::ExtendedMetadata {
964                    info_offset: self.compile_value(offset.into()),
965                    dim: self.compile_value(dim),
966                    out: self.compile_value(out),
967                }
968            }
969            ir::Metadata::BufferLength { list } => {
970                let out = self.compile_value(out);
971
972                let AddressSpace::Global(id) = list.address_space() else {
973                    unreachable!("Value should have id")
974                };
975                let offset = self.info.metadata.buffer_len_index(id);
976                Instruction::Metadata {
977                    info_offset: self.compile_value(offset.into()),
978                    out,
979                }
980            }
981        }
982    }
983
984    fn compile_branch(&mut self, instructions: &mut Vec<Instruction<D>>, branch: ir::Branch) {
985        match branch {
986            ir::Branch::If(op) => instructions.push(Instruction::If {
987                cond: self.compile_value(op.cond),
988                instructions: self.compile_scope(&op.scope),
989            }),
990            ir::Branch::IfElse(op) => instructions.push(Instruction::IfElse {
991                cond: self.compile_value(op.cond),
992                instructions_if: self.compile_scope(&op.scope_if),
993                instructions_else: self.compile_scope(&op.scope_else),
994            }),
995            ir::Branch::Switch(op) => instructions.push(Instruction::Switch {
996                value: self.compile_value(op.value),
997                instructions_default: self.compile_scope(&op.scope_default),
998                instructions_cases: op
999                    .cases
1000                    .into_iter()
1001                    .map(|(val, block)| (self.compile_value(val), self.compile_scope(&block)))
1002                    .collect(),
1003            }),
1004            ir::Branch::Return => instructions.push(Instruction::Return),
1005            ir::Branch::Break => instructions.push(Instruction::Break),
1006            ir::Branch::Unreachable => instructions.push(Instruction::Unreachable),
1007            ir::Branch::RangeLoop(range_loop) => instructions.push(Instruction::RangeLoop {
1008                i: self.compile_value(range_loop.i),
1009                start: self.compile_value(range_loop.start),
1010                end: self.compile_value(range_loop.end),
1011                step: range_loop.step.map(|it| self.compile_value(it)),
1012                inclusive: range_loop.inclusive,
1013                instructions: self.compile_scope(&range_loop.scope),
1014            }),
1015            ir::Branch::Loop(op) => instructions.push(Instruction::Loop {
1016                instructions: self.compile_scope(&op.scope),
1017            }),
1018        };
1019    }
1020
1021    fn compile_atomic(
1022        &mut self,
1023        value: ir::AtomicOp,
1024        out: Option<ir::Value>,
1025        instructions: &mut Vec<Instruction<D>>,
1026    ) {
1027        match value {
1028            ir::AtomicOp::Load(ptr) => {
1029                instructions.push(Instruction::AtomicLoad(UnaryInstruction {
1030                    input: self.compile_value(ptr),
1031                    out: self.compile_value(out.unwrap()),
1032                }))
1033            }
1034            ir::AtomicOp::Store(op) => {
1035                instructions.push(Instruction::AtomicStore(UnaryInstruction {
1036                    input: self.compile_value(op.value),
1037                    out: self.compile_value(op.ptr),
1038                }))
1039            }
1040            ir::AtomicOp::Swap(op) => instructions.push(Instruction::AtomicSwap(
1041                self.compile_atomic_binary(op, out.unwrap()),
1042            )),
1043            ir::AtomicOp::Add(op) => instructions.push(Instruction::AtomicAdd(
1044                self.compile_atomic_binary(op, out.unwrap()),
1045            )),
1046            ir::AtomicOp::Sub(op) => instructions.push(Instruction::AtomicSub(
1047                self.compile_atomic_binary(op, out.unwrap()),
1048            )),
1049            ir::AtomicOp::Max(op) => instructions.push(Instruction::AtomicMax(
1050                self.compile_atomic_binary(op, out.unwrap()),
1051            )),
1052            ir::AtomicOp::Min(op) => instructions.push(Instruction::AtomicMin(
1053                self.compile_atomic_binary(op, out.unwrap()),
1054            )),
1055            ir::AtomicOp::And(op) => instructions.push(Instruction::AtomicAnd(
1056                self.compile_atomic_binary(op, out.unwrap()),
1057            )),
1058            ir::AtomicOp::Or(op) => instructions.push(Instruction::AtomicOr(
1059                self.compile_atomic_binary(op, out.unwrap()),
1060            )),
1061            ir::AtomicOp::Xor(op) => instructions.push(Instruction::AtomicXor(
1062                self.compile_atomic_binary(op, out.unwrap()),
1063            )),
1064            ir::AtomicOp::CompareAndSwap(op) => instructions.push(Instruction::AtomicCAS {
1065                input: self.compile_value(op.ptr),
1066                cmp: self.compile_value(op.cmp),
1067                val: self.compile_value(op.val),
1068                out: self.compile_value(out.unwrap()),
1069            }),
1070        }
1071    }
1072
1073    fn compile_arithmetic(
1074        &mut self,
1075        value: ir::Arithmetic,
1076        out: Option<ir::Value>,
1077        modes: InstructionModes,
1078        instructions: &mut Vec<Instruction<D>>,
1079    ) {
1080        let out = out.unwrap();
1081        match value {
1082            ir::Arithmetic::Add(op) => {
1083                instructions.push(Instruction::Add(self.compile_binary(op, out)))
1084            }
1085            ir::Arithmetic::SaturatingAdd(op) => {
1086                instructions.push(Instruction::SaturatingAdd(self.compile_binary(op, out)))
1087            }
1088            ir::Arithmetic::Mul(op) => {
1089                instructions.push(Instruction::Mul(self.compile_binary(op, out)))
1090            }
1091            ir::Arithmetic::Div(op) => {
1092                let op = self.compile_binary(op, out);
1093                instructions.push(self.select_fast_float(
1094                    out.ty,
1095                    modes,
1096                    FastMath::AllowReciprocal
1097                        | FastMath::ReducedPrecision
1098                        | FastMath::UnsignedZero
1099                        | FastMath::NotInf,
1100                    Instruction::Div(op),
1101                    Instruction::FastDiv(op),
1102                ))
1103            }
1104            ir::Arithmetic::Sub(op) => {
1105                instructions.push(Instruction::Sub(self.compile_binary(op, out)))
1106            }
1107            ir::Arithmetic::SaturatingSub(op) => {
1108                instructions.push(Instruction::SaturatingSub(self.compile_binary(op, out)))
1109            }
1110            ir::Arithmetic::MulHi(op) => {
1111                let instruction = Instruction::HiMul(self.compile_binary(op, out));
1112                D::register_instruction_extension(&mut self.extensions, &instruction);
1113                instructions.push(instruction)
1114            }
1115            ir::Arithmetic::Abs(op) => {
1116                instructions.push(Instruction::Abs(self.compile_unary(op, out)))
1117            }
1118            ir::Arithmetic::Exp(op) => {
1119                let op = self.compile_unary(op, out);
1120                instructions.push(self.select_fast_float(
1121                    out.ty,
1122                    modes,
1123                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1124                    Instruction::Exp(op),
1125                    Instruction::FastExp(op),
1126                ));
1127            }
1128            ir::Arithmetic::Log(op) => {
1129                let op = self.compile_unary(op, out);
1130                instructions.push(self.select_fast_float(
1131                    out.ty,
1132                    modes,
1133                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1134                    Instruction::Log(op),
1135                    Instruction::FastLog(op),
1136                ));
1137            }
1138            ir::Arithmetic::Log1p(op) => {
1139                instructions.push(Instruction::Log1p(self.compile_unary(op, out)))
1140            }
1141            ir::Arithmetic::Expm1(op) => {
1142                instructions.push(Instruction::Expm1(self.compile_unary(op, out)))
1143            }
1144            ir::Arithmetic::Cos(op) => {
1145                let op = self.compile_unary(op, out);
1146                instructions.push(self.select_fast_float(
1147                    out.ty,
1148                    modes,
1149                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1150                    Instruction::Cos(op),
1151                    Instruction::FastCos(op),
1152                ));
1153            }
1154            ir::Arithmetic::Sin(op) => {
1155                let op = self.compile_unary(op, out);
1156                instructions.push(self.select_fast_float(
1157                    out.ty,
1158                    modes,
1159                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1160                    Instruction::Sin(op),
1161                    Instruction::FastSin(op),
1162                ));
1163            }
1164            ir::Arithmetic::Tan(op) => {
1165                instructions.push(Instruction::Tan(self.compile_unary(op, out)))
1166            }
1167            ir::Arithmetic::Tanh(op) => {
1168                let op = self.compile_unary(op, out);
1169                let instruction = Instruction::Tanh(op);
1170                D::register_instruction_extension(&mut self.extensions, &instruction);
1171                if self.compilation_options.supports_features.fast_tanh {
1172                    instructions.push(self.select_fast_float(
1173                        out.ty,
1174                        modes,
1175                        FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1176                        instruction,
1177                        Instruction::FastTanh(op),
1178                    ))
1179                } else {
1180                    instructions.push(instruction);
1181                }
1182            }
1183            ir::Arithmetic::Sinh(op) => {
1184                let instruction = Instruction::Sinh(self.compile_unary(op, out));
1185                D::register_instruction_extension(&mut self.extensions, &instruction);
1186                instructions.push(instruction)
1187            }
1188            ir::Arithmetic::Cosh(op) => {
1189                let instruction = Instruction::Cosh(self.compile_unary(op, out));
1190                D::register_instruction_extension(&mut self.extensions, &instruction);
1191                instructions.push(instruction)
1192            }
1193            ir::Arithmetic::ArcCos(op) => {
1194                let instruction = Instruction::ArcCos(self.compile_unary(op, out));
1195                D::register_instruction_extension(&mut self.extensions, &instruction);
1196                instructions.push(instruction)
1197            }
1198            ir::Arithmetic::ArcSin(op) => {
1199                let instruction = Instruction::ArcSin(self.compile_unary(op, out));
1200                D::register_instruction_extension(&mut self.extensions, &instruction);
1201                instructions.push(instruction)
1202            }
1203            ir::Arithmetic::ArcTan(op) => {
1204                let instruction = Instruction::ArcTan(self.compile_unary(op, out));
1205                D::register_instruction_extension(&mut self.extensions, &instruction);
1206                instructions.push(instruction)
1207            }
1208            ir::Arithmetic::ArcSinh(op) => {
1209                let instruction = Instruction::ArcSinh(self.compile_unary(op, out));
1210                D::register_instruction_extension(&mut self.extensions, &instruction);
1211                instructions.push(instruction)
1212            }
1213            ir::Arithmetic::ArcCosh(op) => {
1214                let instruction = Instruction::ArcCosh(self.compile_unary(op, out));
1215                D::register_instruction_extension(&mut self.extensions, &instruction);
1216                instructions.push(instruction)
1217            }
1218            ir::Arithmetic::ArcTanh(op) => {
1219                let instruction = Instruction::ArcTanh(self.compile_unary(op, out));
1220                D::register_instruction_extension(&mut self.extensions, &instruction);
1221                instructions.push(instruction)
1222            }
1223            ir::Arithmetic::Degrees(op) => {
1224                let instruction = Instruction::Degrees(self.compile_unary(op, out));
1225                D::register_instruction_extension(&mut self.extensions, &instruction);
1226                instructions.push(instruction)
1227            }
1228            ir::Arithmetic::Radians(op) => {
1229                let instruction = Instruction::Radians(self.compile_unary(op, out));
1230                D::register_instruction_extension(&mut self.extensions, &instruction);
1231                instructions.push(instruction)
1232            }
1233            ir::Arithmetic::ArcTan2(op) => {
1234                let instruction = Instruction::ArcTan2(self.compile_binary(op, out));
1235                D::register_instruction_extension(&mut self.extensions, &instruction);
1236                instructions.push(instruction)
1237            }
1238            ir::Arithmetic::Powf(op) => {
1239                let op = self.compile_binary(op, out);
1240                instructions.push(self.select_fast_float(
1241                    out.ty,
1242                    modes,
1243                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1244                    Instruction::Powf(op),
1245                    Instruction::FastPowf(op),
1246                ))
1247            }
1248            ir::Arithmetic::Powi(op) => {
1249                instructions.push(Instruction::Powi(self.compile_binary(op, out)))
1250            }
1251            ir::Arithmetic::Hypot(op) => {
1252                let instruction = Instruction::Hypot(self.compile_binary(op, out));
1253                D::register_instruction_extension(&mut self.extensions, &instruction);
1254                instructions.push(instruction)
1255            }
1256            ir::Arithmetic::Rhypot(op) => {
1257                let instruction = Instruction::Rhypot(self.compile_binary(op, out));
1258                D::register_instruction_extension(&mut self.extensions, &instruction);
1259                instructions.push(instruction)
1260            }
1261            ir::Arithmetic::Sqrt(op) => {
1262                let op = self.compile_unary(op, out);
1263                instructions.push(self.select_fast_float(
1264                    out.ty,
1265                    modes,
1266                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1267                    Instruction::Sqrt(op),
1268                    Instruction::FastSqrt(op),
1269                ))
1270            }
1271            ir::Arithmetic::InverseSqrt(op) => {
1272                let op = self.compile_unary(op, out);
1273                instructions.push(self.select_fast_float(
1274                    out.ty,
1275                    modes,
1276                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1277                    Instruction::InverseSqrt(op),
1278                    Instruction::FastInverseSqrt(op),
1279                ))
1280            }
1281            ir::Arithmetic::Erf(op) => {
1282                let instruction = Instruction::Erf(self.compile_unary(op, out));
1283                D::register_instruction_extension(&mut self.extensions, &instruction);
1284                instructions.push(instruction)
1285            }
1286            ir::Arithmetic::Max(op) => {
1287                let instruction = Instruction::Max(self.compile_binary(op, out));
1288                D::register_instruction_extension(&mut self.extensions, &instruction);
1289                instructions.push(instruction)
1290            }
1291            ir::Arithmetic::Min(op) => {
1292                let instruction = Instruction::Min(self.compile_binary(op, out));
1293                D::register_instruction_extension(&mut self.extensions, &instruction);
1294                instructions.push(instruction)
1295            }
1296            ir::Arithmetic::Clamp(op) => instructions.push(Instruction::Clamp {
1297                input: self.compile_value(op.input),
1298                min_value: self.compile_value(op.min_value),
1299                max_value: self.compile_value(op.max_value),
1300                out: self.compile_value(out),
1301            }),
1302            ir::Arithmetic::Recip(op) => {
1303                let elem = op.input.ty.elem_type();
1304                let input = self.compile_value(op.input);
1305                let out = self.compile_value(out);
1306                let lhs = match elem {
1307                    ir::ElemType::Float(_) => ir::ConstantValue::Float(1.0),
1308                    ir::ElemType::Int(_) => ir::ConstantValue::Int(1),
1309                    ir::ElemType::UInt(_) => ir::ConstantValue::UInt(1),
1310                    ir::ElemType::Bool => ir::ConstantValue::Bool(true),
1311                };
1312                let div = Instruction::Div(BinaryInstruction {
1313                    lhs: Value::Constant(lhs, self.compile_type(op.input.ty)),
1314                    rhs: input,
1315                    out,
1316                });
1317                let recip = Instruction::FastRecip(UnaryInstruction { input, out });
1318
1319                let instruction = self.select_fast_float(
1320                    elem.into(),
1321                    modes,
1322                    FastMath::AllowReciprocal
1323                        | FastMath::ReducedPrecision
1324                        | FastMath::UnsignedZero
1325                        | FastMath::NotInf,
1326                    div,
1327                    recip,
1328                );
1329                D::register_instruction_extension(&mut self.extensions, &instruction);
1330                instructions.push(instruction);
1331            }
1332            ir::Arithmetic::Round(op) => {
1333                instructions.push(Instruction::Round(self.compile_unary(op, out)))
1334            }
1335            ir::Arithmetic::Floor(op) => {
1336                instructions.push(Instruction::Floor(self.compile_unary(op, out)))
1337            }
1338            ir::Arithmetic::Ceil(op) => {
1339                instructions.push(Instruction::Ceil(self.compile_unary(op, out)))
1340            }
1341            ir::Arithmetic::Trunc(op) => {
1342                instructions.push(Instruction::Trunc(self.compile_unary(op, out)))
1343            }
1344            ir::Arithmetic::Rem(op) => {
1345                instructions.push(Instruction::Rem(self.compile_binary(op, out)))
1346            }
1347            ir::Arithmetic::ModFloor(op) => {
1348                instructions.push(Instruction::ModFloor(self.compile_binary(op, out)));
1349            }
1350            ir::Arithmetic::Fma(op) => instructions.push(Instruction::Fma {
1351                a: self.compile_value(op.a),
1352                b: self.compile_value(op.b),
1353                c: self.compile_value(op.c),
1354                out: self.compile_value(out),
1355            }),
1356            ir::Arithmetic::Neg(op) => {
1357                instructions.push(Instruction::Neg(self.compile_unary(op, out)))
1358            }
1359            ir::Arithmetic::Normalize(op) => {
1360                let op = self.compile_unary(op, out);
1361                instructions.push(self.select_fast_float(
1362                    out.ty,
1363                    modes,
1364                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1365                    Instruction::Normalize(op),
1366                    Instruction::FastNormalize(op),
1367                ))
1368            }
1369            ir::Arithmetic::Magnitude(op) => {
1370                let op = self.compile_unary(op, out);
1371                instructions.push(self.select_fast_float(
1372                    out.ty,
1373                    modes,
1374                    FastMath::ReducedPrecision | FastMath::NotNaN | FastMath::NotInf,
1375                    Instruction::Magnitude(op),
1376                    Instruction::FastMagnitude(op),
1377                ))
1378            }
1379            ir::Arithmetic::Dot(op) => {
1380                instructions.push(Instruction::Dot(self.compile_binary(op, out)))
1381            }
1382            ir::Arithmetic::VectorSum(op) => {
1383                instructions.push(Instruction::VectorSum(self.compile_unary(op, out)))
1384            }
1385        };
1386    }
1387
1388    fn select_fast_float(
1389        &self,
1390        ty: ir::Type,
1391        modes: InstructionModes,
1392        required_flags: EnumSet<FastMath>,
1393        default: Instruction<D>,
1394        fast: Instruction<D>,
1395    ) -> Instruction<D> {
1396        if !self.compilation_options.supports_features.fast_math
1397            || !matches!(ty.elem_type(), ElemType::Float(FloatKind::F32))
1398        {
1399            return default;
1400        }
1401
1402        if modes.fp_math_mode.is_superset(required_flags) {
1403            fast
1404        } else {
1405            default
1406        }
1407    }
1408
1409    fn compile_comparison(
1410        &mut self,
1411        value: ir::Comparison,
1412        out: Option<ir::Value>,
1413        instructions: &mut Vec<Instruction<D>>,
1414    ) {
1415        let out = out.unwrap();
1416        match value {
1417            ir::Comparison::Equal(op) => {
1418                instructions.push(Instruction::Equal(self.compile_binary(op, out)))
1419            }
1420            ir::Comparison::Lower(op) => {
1421                instructions.push(Instruction::Lower(self.compile_binary(op, out)))
1422            }
1423            ir::Comparison::Greater(op) => {
1424                instructions.push(Instruction::Greater(self.compile_binary(op, out)))
1425            }
1426            ir::Comparison::LowerEqual(op) => {
1427                instructions.push(Instruction::LowerEqual(self.compile_binary(op, out)))
1428            }
1429            ir::Comparison::GreaterEqual(op) => {
1430                instructions.push(Instruction::GreaterEqual(self.compile_binary(op, out)))
1431            }
1432            ir::Comparison::NotEqual(op) => {
1433                instructions.push(Instruction::NotEqual(self.compile_binary(op, out)))
1434            }
1435            ir::Comparison::IsNan(op) => {
1436                instructions.push(Instruction::IsNan(self.compile_unary(op, out)))
1437            }
1438            ir::Comparison::IsInf(op) => {
1439                instructions.push(Instruction::IsInf(self.compile_unary(op, out)))
1440            }
1441        };
1442    }
1443
1444    fn compile_bitwise(
1445        &mut self,
1446        value: ir::Bitwise,
1447        out: Option<ir::Value>,
1448        instructions: &mut Vec<Instruction<D>>,
1449    ) {
1450        let out = out.unwrap();
1451        match value {
1452            ir::Bitwise::BitwiseOr(op) => {
1453                instructions.push(Instruction::BitwiseOr(self.compile_binary(op, out)))
1454            }
1455            ir::Bitwise::BitwiseAnd(op) => {
1456                instructions.push(Instruction::BitwiseAnd(self.compile_binary(op, out)))
1457            }
1458            ir::Bitwise::BitwiseXor(op) => {
1459                instructions.push(Instruction::BitwiseXor(self.compile_binary(op, out)))
1460            }
1461            ir::Bitwise::CountOnes(op) => {
1462                instructions.push(Instruction::CountBits(self.compile_unary(op, out)))
1463            }
1464            ir::Bitwise::ReverseBits(op) => {
1465                instructions.push(Instruction::ReverseBits(self.compile_unary(op, out)))
1466            }
1467            ir::Bitwise::ShiftLeft(op) => {
1468                instructions.push(Instruction::ShiftLeft(self.compile_binary(op, out)))
1469            }
1470            ir::Bitwise::ShiftRight(op) => {
1471                instructions.push(Instruction::ShiftRight(self.compile_binary(op, out)))
1472            }
1473            ir::Bitwise::BitwiseNot(op) => {
1474                instructions.push(Instruction::BitwiseNot(self.compile_unary(op, out)))
1475            }
1476            ir::Bitwise::LeadingZeros(op) => {
1477                instructions.push(Instruction::LeadingZeros(self.compile_unary(op, out)))
1478            }
1479            ir::Bitwise::TrailingZeros(op) => {
1480                instructions.push(Instruction::TrailingZeros(self.compile_unary(op, out)))
1481            }
1482            ir::Bitwise::FindFirstSet(op) => {
1483                let instruction = Instruction::FindFirstSet(self.compile_unary(op, out));
1484                D::register_instruction_extension(&mut self.extensions, &instruction);
1485                instructions.push(instruction)
1486            }
1487        };
1488    }
1489
1490    fn compile_memory(
1491        &mut self,
1492        value: ir::Memory,
1493        out: Option<ir::Value>,
1494        instructions: &mut Vec<Instruction<D>>,
1495    ) {
1496        match value {
1497            ir::Memory::Index(op) => {
1498                instructions.push(Instruction::Index(self.compile_index(op, out.unwrap())))
1499            }
1500            ir::Memory::Load(value) => instructions.push(Instruction::Load(UnaryInstruction {
1501                input: self.compile_value(value),
1502                out: self.compile_value(out.unwrap()),
1503            })),
1504            ir::Memory::Store(op) => instructions.push(Instruction::Store(UnaryInstruction {
1505                input: self.compile_value(op.value),
1506                out: self.compile_value(op.ptr),
1507            })),
1508            ir::Memory::CopyMemory(op) => instructions.push(Instruction::Copy {
1509                source: self.compile_value(op.source),
1510                dest: self.compile_value(op.target),
1511                len: op.len as u32,
1512            }),
1513        };
1514    }
1515
1516    fn compile_operator(
1517        &mut self,
1518        value: ir::Operator,
1519        out: Option<ir::Value>,
1520        instructions: &mut Vec<Instruction<D>>,
1521    ) {
1522        let out = out.unwrap();
1523        match value {
1524            ir::Operator::And(op) => {
1525                instructions.push(Instruction::And(self.compile_binary(op, out)))
1526            }
1527            ir::Operator::Or(op) => {
1528                instructions.push(Instruction::Or(self.compile_binary(op, out)))
1529            }
1530            ir::Operator::Not(op) => {
1531                instructions.push(Instruction::Not(self.compile_unary(op, out)))
1532            }
1533            ir::Operator::InitVector(op) => instructions.push(Instruction::VecInit {
1534                inputs: op
1535                    .inputs
1536                    .into_iter()
1537                    .map(|it| self.compile_value(it))
1538                    .collect(),
1539                out: self.compile_value(out),
1540            }),
1541            ir::Operator::InsertComponent(op) => instructions.push(Instruction::InsertComponent {
1542                vector: self.compile_value(op.vector),
1543                index: self.compile_value(op.index),
1544                value: self.compile_value(op.value),
1545                out: self.compile_value(out),
1546            }),
1547            ir::Operator::ExtractComponent(op) => {
1548                instructions.push(Instruction::ExtractComponent(self.compile_binary(op, out)))
1549            }
1550            ir::Operator::Select(op) => instructions.push(Instruction::Select {
1551                cond: self.compile_value(op.cond),
1552                then: self.compile_value(op.then),
1553                or_else: self.compile_value(op.or_else),
1554                out: self.compile_value(out),
1555            }),
1556            // Needs special conversion semantics
1557            ir::Operator::Cast(op)
1558                if (is_fp4_fp6_fp8(op.input.elem_type()) || is_fp4_fp6_fp8(out.elem_type()))
1559                // Trivial broadcast shouldn't use special cast logic
1560                    && op.input.elem_type() != out.elem_type() =>
1561            {
1562                // We may need these for intermediates
1563                self.flags.elem_f16 = true;
1564                self.flags.elem_bf16 = true;
1565                let vec_in = op.input.ty.vector_size();
1566                let packing = out.storage_type().packing_factor();
1567                self.compile_type(op.input.ty.with_vector_size(packing));
1568                self.compile_type(
1569                    ir::Type::scalar(ir::ElemType::Float(FloatKind::F16)).with_vector_size(vec_in),
1570                );
1571                self.compile_type(
1572                    ir::Type::scalar(ir::ElemType::Float(FloatKind::BF16)).with_vector_size(vec_in),
1573                );
1574                self.compile_type(
1575                    ir::Type::scalar(ir::ElemType::Float(FloatKind::F16)).with_vector_size(packing),
1576                );
1577                self.compile_type(
1578                    ir::Type::scalar(ir::ElemType::Float(FloatKind::BF16))
1579                        .with_vector_size(packing),
1580                );
1581
1582                let inst = self.compile_unary(op, out);
1583
1584                instructions.push(Instruction::SpecialCast(inst));
1585            }
1586            ir::Operator::Cast(op) => {
1587                let op = self.compile_unary(op, out);
1588
1589                if op.input.elem() == Elem::TF32 || op.out.elem() == Elem::TF32 {
1590                    self.flags.elem_tf32 = true;
1591                }
1592
1593                instructions.push(Instruction::Assign(op))
1594            }
1595            ir::Operator::Reinterpret(op) => {
1596                instructions.push(Instruction::Bitcast(self.compile_unary(op, out)))
1597            }
1598            ir::Operator::ReadBuiltin(builtin) => {
1599                let out = self.compile_value(out);
1600                let mut assign = |input| {
1601                    instructions.push(Instruction::Assign(UnaryInstruction { input, out }));
1602                };
1603                let builtin = match builtin {
1604                    ir::Builtin::AbsolutePos => {
1605                        self.flags.indexes.absolute_pos = true;
1606                        Builtin::AbsolutePos(*self.addr_type.elem())
1607                    }
1608                    ir::Builtin::CubePosCluster
1609                        if self.compilation_options.supports_features.clusters =>
1610                    {
1611                        self.flags.indexes.cluster_pos = true;
1612                        Builtin::ClusterRank
1613                    }
1614                    ir::Builtin::CubePosClusterX
1615                        if self.compilation_options.supports_features.clusters =>
1616                    {
1617                        self.flags.indexes.cluster_pos = true;
1618                        Builtin::ClusterIndexX
1619                    }
1620                    ir::Builtin::CubePosClusterY
1621                        if self.compilation_options.supports_features.clusters =>
1622                    {
1623                        self.flags.indexes.cluster_pos = true;
1624                        Builtin::ClusterIndexY
1625                    }
1626                    ir::Builtin::CubePosClusterZ
1627                        if self.compilation_options.supports_features.clusters =>
1628                    {
1629                        self.flags.indexes.cluster_pos = true;
1630                        Builtin::ClusterIndexZ
1631                    }
1632                    // Fallback if clusters aren't supported, ID is always 0 since clusters are always
1633                    // (1, 1, 1) if unsupported
1634                    ir::Builtin::CubePosCluster
1635                    | ir::Builtin::CubePosClusterX
1636                    | ir::Builtin::CubePosClusterY
1637                    | ir::Builtin::CubePosClusterZ => {
1638                        assign(const_u32(0));
1639                        return;
1640                    }
1641                    ir::Builtin::AbsolutePosX => {
1642                        self.flags.indexes.absolute_pos_tuple = true;
1643                        Builtin::AbsolutePosX
1644                    }
1645                    ir::Builtin::AbsolutePosY => {
1646                        self.flags.indexes.absolute_pos_tuple = true;
1647                        Builtin::AbsolutePosY
1648                    }
1649                    ir::Builtin::AbsolutePosZ => {
1650                        self.flags.indexes.absolute_pos_tuple = true;
1651                        Builtin::AbsolutePosZ
1652                    }
1653                    ir::Builtin::CubeDim => {
1654                        self.flags.indexes.cube_dim = true;
1655                        Builtin::CubeDim
1656                    }
1657                    ir::Builtin::CubeDimX => {
1658                        self.flags.indexes.cube_dim_tuple = true;
1659                        Builtin::CubeDimX
1660                    }
1661                    ir::Builtin::CubeDimY => {
1662                        self.flags.indexes.cube_dim_tuple = true;
1663                        Builtin::CubeDimY
1664                    }
1665                    ir::Builtin::CubeDimZ => {
1666                        self.flags.indexes.cube_dim_tuple = true;
1667                        Builtin::CubeDimZ
1668                    }
1669                    ir::Builtin::CubeClusterDim => {
1670                        assign(const_u32(self.cluster_dim.num_elems()));
1671                        return;
1672                    }
1673                    ir::Builtin::CubeClusterDimX => {
1674                        assign(const_u32(self.cluster_dim.x));
1675                        return;
1676                    }
1677                    ir::Builtin::CubeClusterDimY => {
1678                        assign(const_u32(self.cluster_dim.y));
1679                        return;
1680                    }
1681                    ir::Builtin::CubeClusterDimZ => {
1682                        assign(const_u32(self.cluster_dim.z));
1683                        return;
1684                    }
1685                    ir::Builtin::CubePos => {
1686                        self.flags.indexes.cube_pos = true;
1687                        Builtin::CubePos(*self.addr_type.elem())
1688                    }
1689                    ir::Builtin::CubePosX => {
1690                        self.flags.indexes.cube_pos_tuple = true;
1691                        Builtin::CubePosX
1692                    }
1693                    ir::Builtin::CubePosY => {
1694                        self.flags.indexes.cube_pos_tuple = true;
1695                        Builtin::CubePosY
1696                    }
1697                    ir::Builtin::CubePosZ => {
1698                        self.flags.indexes.cube_pos_tuple = true;
1699                        Builtin::CubePosZ
1700                    }
1701                    ir::Builtin::CubeCount => {
1702                        self.flags.indexes.cube_count = true;
1703                        Builtin::CubeCount(*self.addr_type.elem())
1704                    }
1705                    ir::Builtin::CubeCountX => {
1706                        self.flags.indexes.cube_count_tuple = true;
1707                        Builtin::CubeCountX
1708                    }
1709                    ir::Builtin::CubeCountY => {
1710                        self.flags.indexes.cube_count_tuple = true;
1711                        Builtin::CubeCountY
1712                    }
1713                    ir::Builtin::CubeCountZ => {
1714                        self.flags.indexes.cube_count_tuple = true;
1715                        Builtin::CubeCountZ
1716                    }
1717                    ir::Builtin::UnitPos => {
1718                        self.flags.indexes.unit_pos = true;
1719                        Builtin::UnitPos
1720                    }
1721                    ir::Builtin::UnitPosX => {
1722                        self.flags.indexes.unit_pos_tuple = true;
1723                        Builtin::UnitPosX
1724                    }
1725                    ir::Builtin::UnitPosY => {
1726                        self.flags.indexes.unit_pos_tuple = true;
1727                        Builtin::UnitPosY
1728                    }
1729                    ir::Builtin::UnitPosZ => {
1730                        self.flags.indexes.unit_pos_tuple = true;
1731                        Builtin::UnitPosZ
1732                    }
1733                    ir::Builtin::PlaneDim => {
1734                        self.flags.indexes.plane_dim = true;
1735                        Builtin::PlaneDim
1736                    }
1737                    ir::Builtin::PlanePos => {
1738                        self.flags.indexes.plane_pos = true;
1739                        Builtin::PlanePos
1740                    }
1741                    ir::Builtin::UnitPosPlane => {
1742                        self.flags.indexes.unit_pos_plane = true;
1743                        Builtin::UnitPosPlane
1744                    }
1745                };
1746                instructions.push(Instruction::ReadBuiltin { builtin, out })
1747            }
1748            ir::Operator::ReadScalar(id) => instructions.push(Instruction::ReadScalar {
1749                id,
1750                out: self.compile_value(out),
1751            }),
1752        };
1753    }
1754
1755    fn compile_binary(
1756        &mut self,
1757        value: ir::BinaryOperands,
1758        out: ir::Value,
1759    ) -> BinaryInstruction<D> {
1760        BinaryInstruction {
1761            lhs: self.compile_value(value.lhs),
1762            rhs: self.compile_value(value.rhs),
1763            out: self.compile_value(out),
1764        }
1765    }
1766
1767    fn compile_atomic_binary(
1768        &mut self,
1769        value: ir::AtomicBinaryOperands,
1770        out: ir::Value,
1771    ) -> BinaryInstruction<D> {
1772        BinaryInstruction {
1773            lhs: self.compile_value(value.ptr),
1774            rhs: self.compile_value(value.value),
1775            out: self.compile_value(out),
1776        }
1777    }
1778
1779    fn compile_index(&mut self, value: ir::IndexOperands, out: ir::Value) -> IndexInstruction<D> {
1780        IndexInstruction {
1781            list: self.compile_value(value.list),
1782            index: self.compile_value(value.index),
1783            out: self.compile_value(out),
1784        }
1785    }
1786
1787    fn compile_unary(&mut self, value: ir::UnaryOperands, out: ir::Value) -> UnaryInstruction<D> {
1788        UnaryInstruction {
1789            input: self.compile_value(value.input),
1790            out: self.compile_value(out),
1791        }
1792    }
1793
1794    fn compile_value(&mut self, value: ir::Value) -> Value<D> {
1795        let item = value.ty;
1796        match value.kind {
1797            ir::ValueKind::Value { id } => Value::Value {
1798                id,
1799                item: self.compile_type(item),
1800            },
1801            ir::ValueKind::Constant(value) => Value::Constant(value, self.compile_type(item)),
1802        }
1803    }
1804
1805    fn compile_matrix(&mut self, matrix: ir::MatrixType) -> FragmentType<D> {
1806        FragmentType {
1807            ident: self.compile_matrix_ident(matrix.ident),
1808            m: matrix.m as u32,
1809            n: matrix.n as u32,
1810            k: matrix.k as u32,
1811            elem: self.compile_storage_type(matrix.storage),
1812            layout: self.compile_matrix_layout(matrix.layout),
1813        }
1814    }
1815
1816    fn compile_matrix_ident(&mut self, ident: ir::MatrixIdent) -> FragmentIdent<D> {
1817        match ident {
1818            ir::MatrixIdent::A => FragmentIdent::A,
1819            ir::MatrixIdent::B => FragmentIdent::B,
1820            ir::MatrixIdent::Accumulator => FragmentIdent::Accumulator,
1821        }
1822    }
1823
1824    fn compile_matrix_layout(&mut self, layout: ir::MatrixLayout) -> Option<FragmentLayout<D>> {
1825        match layout {
1826            ir::MatrixLayout::ColMajor => Some(FragmentLayout::ColMajor),
1827            ir::MatrixLayout::RowMajor => Some(FragmentLayout::RowMajor),
1828            ir::MatrixLayout::Undefined => None,
1829        }
1830    }
1831
1832    fn compile_binding(&mut self, binding: cubecl_runtime::kernel::KernelArg) -> KernelArg<D> {
1833        KernelArg {
1834            id: binding.id,
1835            value: self.compile_value(binding.value),
1836            vis: self.buffer_vis[binding.id as usize],
1837        }
1838    }
1839
1840    fn compile_type(&mut self, ty: ir::Type) -> Item<D> {
1841        let item = match ty {
1842            ir::Type::Scalar(ty) => Item::Scalar(self.compile_storage_type(ty)),
1843            ir::Type::Vector(ty, vector_size) => {
1844                Item::Vector(self.compile_type(*ty).intern(), vector_size)
1845            }
1846            ir::Type::Atomic(ty) => {
1847                let item = self.compile_type(*ty);
1848                Item::Atomic(item.intern())
1849            }
1850            ir::Type::Pointer(ty, class) => {
1851                let item = self.compile_type(*ty);
1852                let class = match class {
1853                    ir::AddressSpace::Global(id) => {
1854                        PointerClass::Global(self.buffer_vis[id as usize])
1855                    }
1856                    ir::AddressSpace::Shared => PointerClass::Shared,
1857                    ir::AddressSpace::Local => PointerClass::Local,
1858                };
1859                Item::Pointer(item.intern(), class)
1860            }
1861            ir::Type::Array(ty, size) => {
1862                let ty = self.compile_type(*ty);
1863                Item::Array(ty.intern(), size)
1864            }
1865            ir::Type::DynamicArray(ty) => {
1866                let ty = self.compile_type(*ty);
1867                Item::DynamicArray(ty.intern())
1868            }
1869            ir::Type::Matrix(ty) => {
1870                let ty = self.compile_matrix(ty);
1871                Item::Fragment(ty)
1872            }
1873            ir::Type::Semantic(ty) => self.compile_semantic_type(ty),
1874            ir::Type::Opaque(ty) => self.compile_opaque_type(ty),
1875            ir::Type::Aggregate(_) => {
1876                unreachable!("Should be disaggregated at this point")
1877            }
1878        };
1879        if *item.elem() != super::Elem::TF32 {
1880            self.items.insert(item);
1881            self.items.insert(item.optimized());
1882        } else {
1883            // TF32 is represented as `float` in C++
1884            let item = item.with_elem(super::Elem::F32);
1885            self.items.insert(item);
1886        }
1887
1888        item
1889    }
1890
1891    fn compile_storage_type(&mut self, value: ir::StorageType) -> Elem<D> {
1892        match value {
1893            ir::StorageType::Scalar(ty) => self.compile_elem(ty),
1894            ir::StorageType::Packed(ir::ElemType::Float(kind), 2) => match kind {
1895                FloatKind::E2M1 => {
1896                    self.flags.elem_fp4 = true;
1897                    Elem::FP4x2(FP4Kind::E2M1)
1898                }
1899                FloatKind::E2M3 => {
1900                    self.flags.elem_fp6 = true;
1901                    Elem::FP6x2(FP6Kind::E2M3)
1902                }
1903                FloatKind::E3M2 => {
1904                    self.flags.elem_fp6 = true;
1905                    Elem::FP6(FP6Kind::E3M2)
1906                }
1907                FloatKind::E4M3 => {
1908                    self.flags.elem_fp8 = true;
1909                    Elem::FP8x2(FP8Kind::E4M3)
1910                }
1911                FloatKind::E5M2 => {
1912                    self.flags.elem_fp8 = true;
1913                    Elem::FP8x2(FP8Kind::E5M2)
1914                }
1915                FloatKind::UE8M0 => {
1916                    self.flags.elem_fp8 = true;
1917                    Elem::FP8x2(FP8Kind::UE8M0)
1918                }
1919                FloatKind::F16 => {
1920                    self.flags.elem_f16 = true;
1921                    Elem::F16x2
1922                }
1923                FloatKind::BF16 => {
1924                    self.flags.elem_bf16 = true;
1925                    Elem::BF16x2
1926                }
1927                other => unimplemented!("Unsupported storage type: packed<{other:?}, 2>"),
1928            },
1929            ir::StorageType::Packed(other, factor) => {
1930                unimplemented!("Unsupported storage type: packed<{other}, {factor}>")
1931            }
1932        }
1933    }
1934
1935    fn compile_elem(&mut self, value: ir::ElemType) -> Elem<D> {
1936        match value {
1937            ir::ElemType::Float(kind) => match kind {
1938                ir::FloatKind::E2M1 => {
1939                    self.flags.elem_fp4 = true;
1940                    Elem::FP4(FP4Kind::E2M1)
1941                }
1942                ir::FloatKind::E2M3 => {
1943                    self.flags.elem_fp6 = true;
1944                    Elem::FP6(FP6Kind::E2M3)
1945                }
1946                ir::FloatKind::E3M2 => {
1947                    self.flags.elem_fp6 = true;
1948                    Elem::FP6(FP6Kind::E3M2)
1949                }
1950                ir::FloatKind::E4M3 => {
1951                    self.flags.elem_fp8 = true;
1952                    Elem::FP8(FP8Kind::E4M3)
1953                }
1954                ir::FloatKind::E5M2 => {
1955                    self.flags.elem_fp8 = true;
1956                    Elem::FP8(FP8Kind::E5M2)
1957                }
1958                ir::FloatKind::UE8M0 => {
1959                    self.flags.elem_fp8 = true;
1960                    Elem::FP8(FP8Kind::UE8M0)
1961                }
1962                ir::FloatKind::F16 => {
1963                    self.flags.elem_f16 = true;
1964                    Elem::F16
1965                }
1966                ir::FloatKind::BF16 => {
1967                    self.flags.elem_bf16 = true;
1968                    Elem::BF16
1969                }
1970                ir::FloatKind::TF32 => Elem::TF32,
1971                ir::FloatKind::Flex32 => Elem::F32,
1972                ir::FloatKind::F32 => Elem::F32,
1973                ir::FloatKind::F64 => Elem::F64,
1974            },
1975            ir::ElemType::Int(kind) => match kind {
1976                ir::IntKind::I8 => Elem::I8,
1977                ir::IntKind::I16 => Elem::I16,
1978                ir::IntKind::I32 => Elem::I32,
1979                ir::IntKind::I64 => Elem::I64,
1980            },
1981            ir::ElemType::UInt(kind) => match kind {
1982                ir::UIntKind::U8 => Elem::U8,
1983                ir::UIntKind::U16 => Elem::U16,
1984                ir::UIntKind::U32 => Elem::U32,
1985                ir::UIntKind::U64 => Elem::U64,
1986            },
1987            ir::ElemType::Bool => Elem::Bool,
1988        }
1989    }
1990
1991    fn compile_semantic_type(&mut self, value: ir::SemanticType) -> Item<D> {
1992        match value {
1993            ir::SemanticType::TensorLayout(..) | ir::SemanticType::TensorView(..) => {
1994                panic!("Tensor addressing is only supported on Vulkan")
1995            }
1996        }
1997    }
1998
1999    fn compile_opaque_type(&mut self, value: ir::OpaqueType) -> Item<D> {
2000        match value {
2001            ir::OpaqueType::Barrier(barrier_level) => {
2002                self.flags.op_barrier = true;
2003                Item::Barrier(barrier_level)
2004            }
2005            ir::OpaqueType::BarrierToken(barrier_level) => {
2006                self.flags.op_barrier = true;
2007                Item::BarrierToken(barrier_level)
2008            }
2009            ir::OpaqueType::TensorMap => Item::TensorMap,
2010        }
2011    }
2012}
2013
2014fn is_fp4_fp6_fp8(elem: ir::ElemType) -> bool {
2015    match elem {
2016        ir::ElemType::Float(kind) => matches!(
2017            kind,
2018            FloatKind::E2M1
2019                | FloatKind::E2M3
2020                | FloatKind::E3M2
2021                | FloatKind::E4M3
2022                | FloatKind::E5M2
2023                | FloatKind::UE8M0
2024        ),
2025        _ => false,
2026    }
2027}
2028
2029fn const_u32<D: Dialect>(value: u32) -> Value<D> {
2030    Value::Constant(
2031        ir::ConstantValue::UInt(value as u64),
2032        Item::Scalar(Elem::U32),
2033    )
2034}
2035
2036pub fn register_supported_types(props: &mut DeviceProperties) {
2037    props.register_address_type(ir::AddressType::U32);
2038    props.register_address_type(ir::AddressType::U64);
2039
2040    let supported_types = [
2041        ir::ElemType::UInt(ir::UIntKind::U8),
2042        ir::ElemType::UInt(ir::UIntKind::U16),
2043        ir::ElemType::UInt(ir::UIntKind::U32),
2044        ir::ElemType::UInt(ir::UIntKind::U64),
2045        ir::ElemType::Int(ir::IntKind::I8),
2046        ir::ElemType::Int(ir::IntKind::I16),
2047        ir::ElemType::Int(ir::IntKind::I32),
2048        ir::ElemType::Int(ir::IntKind::I64),
2049        ir::ElemType::Float(ir::FloatKind::BF16),
2050        ir::ElemType::Float(ir::FloatKind::F16),
2051        ir::ElemType::Float(ir::FloatKind::F32),
2052        ir::ElemType::Float(ir::FloatKind::Flex32),
2053        ir::ElemType::Float(ir::FloatKind::F64),
2054        ir::ElemType::Bool,
2055    ];
2056
2057    let supported_atomic_types = [
2058        ir::ElemType::Int(ir::IntKind::I32),
2059        ir::ElemType::Int(ir::IntKind::I64),
2060        ir::ElemType::UInt(ir::UIntKind::U32),
2061        ir::ElemType::UInt(ir::UIntKind::U64),
2062        ir::ElemType::Float(ir::FloatKind::F32),
2063    ];
2064
2065    for ty in supported_types {
2066        props.register_type_usage(ty, TypeUsage::all());
2067    }
2068
2069    for ty in supported_atomic_types {
2070        // Restricted to 32-bit integers because not every min/max/bitwise/CAS overload
2071        // exists for 64-bit and float atomics across the C++ dialects (CUDA, HIP, Metal).
2072        let usage = match ty {
2073            ir::ElemType::Int(ir::IntKind::I32) | ir::ElemType::UInt(ir::UIntKind::U32) => {
2074                AtomicUsage::all()
2075            }
2076            _ => AtomicUsage::Add | AtomicUsage::LoadStore,
2077        };
2078        props.register_atomic_type_usage(Type::atomic(ty), usage);
2079    }
2080}