cubecl_spirv/
item.rs

1use cubecl_core::ir::{self as core, AddressSpace, ClampMode, FloatKind, IntKind, UIntKind};
2use rspirv::spirv::{
3    Capability, CooperativeMatrixLayout, CooperativeMatrixUse, FPEncoding, Scope, StorageClass,
4    TensorClampMode, Word,
5};
6use serde::{Deserialize, Serialize};
7
8use crate::{compiler::SpirvCompiler, target::SpirvTarget, value::ConstVal};
9
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub enum Item {
12    Scalar(Elem),
13    // Vector of scalars. Must be 2, 3, or 4, unless long vectors extension is enabled
14    Vector(Elem, u32),
15    Pointer(StorageClass, Box<Item>),
16    Array(Box<Item>, u32),
17    DynamicArray(Box<Item>),
18    CoopMatrix {
19        ty: Elem,
20        rows: u32,
21        columns: u32,
22        ident: CooperativeMatrixUse,
23        layout: Option<CooperativeMatrixLayout>,
24        scope: Scope,
25    },
26    TensorLayout {
27        dims: usize,
28        clamp_mode: TensorClampMode,
29    },
30    TensorView {
31        dims: usize,
32        has_dims: bool,
33        permutation: Vec<u32>,
34    },
35}
36
37impl Item {
38    pub fn id<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>) -> Word {
39        let id = match self {
40            Item::Scalar(elem) => elem.id(b),
41            Item::Vector(elem, vec) => {
42                let elem = elem.id(b);
43                if b.compilation_options.vulkan.supports_long_vectors {
44                    let len = b.const_u32(*vec);
45                    b.type_vector_id_ext(elem, len)
46                } else {
47                    b.type_vector(elem, *vec)
48                }
49            }
50            Item::Pointer(storage_class, item) => {
51                let item = item.id(b);
52                b.type_pointer(None, *storage_class, item)
53            }
54            Item::Array(item, size) => {
55                let item = item.id(b);
56                let id = b.id();
57                let size = b.const_u32(*size);
58                b.type_array_id(Some(id), item, size)
59            }
60            Item::DynamicArray(item) => {
61                let item = item.id(b);
62                let id = b.id();
63                b.type_runtime_array_id(Some(id), item)
64            }
65            Item::CoopMatrix {
66                ty,
67                rows,
68                columns,
69                ident,
70                scope,
71                ..
72            } => {
73                let ty = ty.id(b);
74                let scope = b.const_u32(*scope as u32);
75                let usage = b.const_u32(*ident as u32);
76                b.type_cooperative_matrix_khr(ty, scope, *rows, *columns, usage)
77            }
78            Item::TensorLayout { dims, clamp_mode } => {
79                let dim = b.const_u32(*dims as u32);
80                let clamp_mode = b.const_u32(*clamp_mode as u32);
81                b.type_tensor_layout_nv(dim, clamp_mode)
82            }
83            Item::TensorView {
84                dims,
85                has_dims,
86                permutation,
87            } => {
88                let bool = b.type_bool();
89                let dim = b.const_u32(*dims as u32);
90                let has_dims = if *has_dims {
91                    b.constant_true(bool)
92                } else {
93                    b.constant_false(bool)
94                };
95                let permutation = permutation
96                    .iter()
97                    .map(|it| b.const_u32(*it))
98                    .collect::<Vec<_>>();
99                b.type_tensor_view_nv(dim, has_dims, permutation)
100            }
101        };
102        if b.debug_symbols && !b.state.debug_types.contains(&id) {
103            b.debug_name(id, format!("{self}"));
104            b.state.debug_types.insert(id);
105        }
106        id
107    }
108
109    pub fn builtin_u32() -> Self {
110        Item::Scalar(Elem::Int(32, false))
111    }
112
113    pub fn value_type(&self) -> Item {
114        match self {
115            Item::Pointer(_, item) => item.value_type(),
116            Item::Array(item, _) => item.value_type(),
117            Item::DynamicArray(item) => item.value_type(),
118            other => other.clone(),
119        }
120    }
121
122    pub fn unwrap_ptr(&self) -> Item {
123        match self {
124            Item::Pointer(_, item) => item.as_ref().clone(),
125            other => other.clone(),
126        }
127    }
128
129    pub fn size(&self) -> u32 {
130        match self {
131            Item::Scalar(elem) => elem.size(),
132            Item::Vector(elem, factor) => elem.size() * *factor,
133            Item::Pointer(_, item) => item.size(),
134            Item::Array(item, size) => item.size() * *size,
135            Item::DynamicArray(item) => item.size(),
136            Item::CoopMatrix { ty, .. } => ty.size(),
137            Item::TensorLayout { .. } => 1,
138            Item::TensorView { .. } => 1,
139        }
140    }
141
142    pub fn elem(&self) -> Elem {
143        match self {
144            Item::Scalar(elem) => *elem,
145            Item::Vector(elem, _) => *elem,
146            Item::Pointer(_, item) => item.elem(),
147            Item::Array(item, _) => item.elem(),
148            Item::DynamicArray(item) => item.elem(),
149            Item::CoopMatrix { ty, .. } => *ty,
150            Item::TensorLayout { .. } => Elem::Void,
151            Item::TensorView { .. } => Elem::Void,
152        }
153    }
154
155    pub fn same_vectorization(&self, elem: Elem) -> Item {
156        match self {
157            Item::Scalar(_) => Item::Scalar(elem),
158            Item::Vector(_, factor) => Item::Vector(elem, *factor),
159            _ => unreachable!(),
160        }
161    }
162
163    pub fn vectorization(&self) -> u32 {
164        match self {
165            Item::Vector(_, factor) => *factor,
166            _ => 1,
167        }
168    }
169
170    pub fn constant<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>, value: ConstVal) -> Word {
171        let scalar = self.elem().constant(b, value);
172        let ty = self.id(b);
173        match self {
174            Item::Scalar(_) => scalar,
175            Item::Vector(_, vec) => b.constant_composite(ty, (0..*vec).map(|_| scalar)),
176            Item::Pointer(_, _) => unimplemented!("Can't create constant pointer"),
177            Item::Array(_, _) => unimplemented!("Can't create constant pointer"),
178            Item::DynamicArray(..) => unimplemented!("Can't create constant pointer"),
179            Item::CoopMatrix { .. } => unimplemented!("Can't create constant cmma matrix"),
180            Item::TensorLayout { .. } => unimplemented!("Can't create constant cmma matrix"),
181            Item::TensorView { .. } => unimplemented!("Can't create constant cmma matrix"),
182        }
183    }
184
185    pub fn const_u32<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>, value: u32) -> Word {
186        b.static_cast(ConstVal::Bit32(value), &Elem::Int(32, false), self)
187            .0
188    }
189
190    /// Broadcast a scalar to a vector if needed, ex: f32 -> vec2<f32>, vec2<f32> -> vec2<f32>
191    pub fn broadcast<T: SpirvTarget>(
192        &self,
193        b: &mut SpirvCompiler<T>,
194        obj: Word,
195        out_id: Option<Word>,
196        other: &Item,
197    ) -> Word {
198        match (self, other) {
199            (Item::Scalar(elem), Item::Vector(_, factor)) => {
200                let item = Item::Vector(*elem, *factor);
201                let ty = item.id(b);
202                b.composite_construct(ty, out_id, (0..*factor).map(|_| obj).collect::<Vec<_>>())
203                    .unwrap()
204            }
205            _ => obj,
206        }
207    }
208
209    pub fn cast_to<T: SpirvTarget>(
210        &self,
211        b: &mut SpirvCompiler<T>,
212        out_id: Option<Word>,
213        obj: Word,
214        other: &Item,
215    ) -> Word {
216        let ty = other.id(b);
217
218        let matching_vec = match (self, other) {
219            (Item::Scalar(_), Item::Scalar(_)) => true,
220            (Item::Scalar(_), Item::Vector(..)) => false,
221            (Item::Vector(_, factor_from), Item::Vector(_, factor_to)) => factor_from == factor_to,
222            _ => true,
223        };
224        let matching_elem = self.elem() == other.elem();
225
226        let convert_i_width =
227            |b: &mut SpirvCompiler<T>, obj: Word, out_id: Option<Word>, signed: bool| {
228                if signed {
229                    b.s_convert(ty, out_id, obj).unwrap()
230                } else {
231                    b.u_convert(ty, out_id, obj).unwrap()
232                }
233            };
234
235        let convert_int = |b: &mut SpirvCompiler<T>,
236                           obj: Word,
237                           out_id: Option<Word>,
238                           (width_self, signed_self),
239                           (width_other, signed_other)| {
240            let width_differs = width_self != width_other;
241            let sign_extend = signed_self && signed_other;
242            match width_differs {
243                true => convert_i_width(b, obj, out_id, sign_extend),
244                false => b.copy_object(ty, out_id, obj).unwrap(),
245            }
246        };
247
248        let cast_elem = |b: &mut SpirvCompiler<T>, obj: Word, out_id: Option<Word>| -> Word {
249            match (self.elem(), other.elem()) {
250                (Elem::Bool, Elem::Int(_, _)) => {
251                    let one = other.const_u32(b, 1);
252                    let zero = other.const_u32(b, 0);
253                    b.select(ty, out_id, obj, one, zero).unwrap()
254                }
255                (Elem::Bool, Elem::Float(_, _)) | (Elem::Bool, Elem::Relaxed) => {
256                    let one = other.const_u32(b, 1);
257                    let zero = other.const_u32(b, 0);
258                    b.select(ty, out_id, obj, one, zero).unwrap()
259                }
260                (Elem::Int(_, _), Elem::Bool) => {
261                    let zero = self.const_u32(b, 0);
262                    b.i_not_equal(ty, out_id, obj, zero).unwrap()
263                }
264                (Elem::Int(width_self, signed_self), Elem::Int(width_other, signed_other)) => {
265                    convert_int(
266                        b,
267                        obj,
268                        out_id,
269                        (width_self, signed_self),
270                        (width_other, signed_other),
271                    )
272                }
273                (Elem::Int(_, false), Elem::Float(_, _)) | (Elem::Int(_, false), Elem::Relaxed) => {
274                    b.convert_u_to_f(ty, out_id, obj).unwrap()
275                }
276                (Elem::Int(_, true), Elem::Float(_, _)) | (Elem::Int(_, true), Elem::Relaxed) => {
277                    b.convert_s_to_f(ty, out_id, obj).unwrap()
278                }
279                (Elem::Float(_, _), Elem::Bool) | (Elem::Relaxed, Elem::Bool) => {
280                    let zero = self.const_u32(b, 0);
281                    b.f_unord_not_equal(ty, out_id, obj, zero).unwrap()
282                }
283                (Elem::Float(_, _), Elem::Int(_, false)) | (Elem::Relaxed, Elem::Int(_, false)) => {
284                    b.convert_f_to_u(ty, out_id, obj).unwrap()
285                }
286                (Elem::Float(_, _), Elem::Int(_, true)) | (Elem::Relaxed, Elem::Int(_, true)) => {
287                    b.convert_f_to_s(ty, out_id, obj).unwrap()
288                }
289                (Elem::Float(32, _), Elem::Relaxed) | (Elem::Relaxed, Elem::Float(32, _)) => {
290                    if out_id.is_some() {
291                        b.copy_object(ty, out_id, obj).unwrap()
292                    } else {
293                        obj
294                    }
295                }
296                (Elem::Float(_, _), Elem::Float(_, _))
297                | (Elem::Float(_, _), Elem::Relaxed)
298                | (Elem::Relaxed, Elem::Float(_, _)) => b.f_convert(ty, out_id, obj).unwrap(),
299                (Elem::Bool, Elem::Bool) => b.copy_object(ty, out_id, obj).unwrap(),
300                (Elem::Relaxed, Elem::Relaxed) => b.copy_object(ty, out_id, obj).unwrap(),
301                (from, to) => panic!("Invalid cast from {from:?} to {to:?}"),
302            }
303        };
304
305        match (matching_vec, matching_elem) {
306            (true, true) if out_id.is_some() => b.copy_object(ty, out_id, obj).unwrap(),
307            (true, true) => obj,
308            (true, false) => cast_elem(b, obj, out_id),
309            (false, true) => self.broadcast(b, obj, out_id, other),
310            (false, false) => {
311                let broadcast = self.broadcast(b, obj, None, other);
312                cast_elem(b, broadcast, out_id)
313            }
314        }
315    }
316
317    pub fn is_ptr(&self) -> bool {
318        matches!(self, Item::Pointer(..))
319    }
320
321    pub fn is_array(&self) -> bool {
322        matches!(self, Item::Array(..) | Item::DynamicArray(..))
323    }
324}
325
326#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
327pub enum Elem {
328    Void,
329    Bool,
330    Int(u32, bool),
331    Float(u32, Option<FPEncoding>),
332    Relaxed,
333}
334
335impl Elem {
336    pub fn id<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>) -> Word {
337        let id = match self {
338            Elem::Void => b.type_void(),
339            Elem::Bool => b.type_bool(),
340            Elem::Int(width, _) => b.type_int(*width, 0),
341            Elem::Float(width, encoding) => b.type_float(*width, *encoding),
342            Elem::Relaxed => b.type_float(32, None),
343        };
344        if b.debug_symbols && !b.state.debug_types.contains(&id) {
345            b.debug_name(id, format!("{self}"));
346            b.state.debug_types.insert(id);
347        }
348        id
349    }
350
351    pub fn size(&self) -> u32 {
352        match self {
353            Elem::Void => 0,
354            Elem::Bool => 1,
355            Elem::Int(size, _) => *size / 8,
356            Elem::Float(size, _) => *size / 8,
357            Elem::Relaxed => 4,
358        }
359    }
360
361    pub fn constant<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>, value: ConstVal) -> Word {
362        let ty = self.id(b);
363        match self {
364            Elem::Void => unreachable!(),
365            Elem::Bool if value.as_u64() != 0 => b.constant_true(ty),
366            Elem::Bool => b.constant_false(ty),
367            _ => match value {
368                ConstVal::Bit32(val) => b.dedup_constant_bit32(ty, val),
369                ConstVal::Bit64(val) => b.dedup_constant_bit64(ty, val),
370            },
371        }
372    }
373
374    pub fn float_encoding(&self) -> Option<FPEncoding> {
375        match self {
376            Elem::Float(_, encoding) => *encoding,
377            _ => None,
378        }
379    }
380
381    pub fn width(&self) -> u32 {
382        self.size() * 8
383    }
384}
385
386impl<T: SpirvTarget> SpirvCompiler<T> {
387    pub fn compile_type(&mut self, item: core::Type) -> Item {
388        match item {
389            core::Type::Scalar(storage) => Item::Scalar(self.compile_storage_type(storage)),
390            core::Type::Vector(inner, size) => {
391                Item::Vector(self.compile_storage_type(inner.storage_type()), size as u32)
392            }
393            core::Type::Atomic(inner) => {
394                if let core::Type::Vector(ty, _) = inner.value_type()
395                    && ty.is_float()
396                    && ty.size() == 2
397                {
398                    self.capabilities.insert(Capability::AtomicFloat16VectorNV);
399                }
400                self.compile_type(*inner)
401            }
402            core::Type::Pointer(inner, class) => {
403                let storage_class = compile_pointer_class(class);
404                let item = self.compile_type(*inner);
405                Item::Pointer(storage_class, Box::new(item))
406            }
407            core::Type::Semantic(semantic) => match semantic {
408                core::SemanticType::TensorLayout(dims, clamp_mode) => Item::TensorLayout {
409                    dims,
410                    clamp_mode: compile_clamp_mode(clamp_mode),
411                },
412                core::SemanticType::TensorView(dims, has_dims, permutation) => Item::TensorView {
413                    dims,
414                    has_dims,
415                    permutation: permutation[..dims].to_vec(),
416                },
417            },
418            core::Type::Array(inner, size) => {
419                let item = self.compile_type(*inner);
420                Item::Array(Box::new(item), size as u32)
421            }
422            core::Type::DynamicArray(inner) => {
423                let item = self.compile_type(*inner);
424                Item::DynamicArray(Box::new(item))
425            }
426            core::Type::Matrix(ty) => {
427                let mat = self.compile_matrix(&ty);
428                Item::CoopMatrix {
429                    ty: mat.elem,
430                    rows: self.matrix_rows(&mat),
431                    columns: self.matrix_columns(&mat),
432                    ident: mat.ident,
433                    layout: mat.layout,
434                    scope: mat.scope,
435                }
436            }
437            core::Type::Opaque(opaque_type) => match opaque_type {
438                core::OpaqueType::Barrier(..) | core::OpaqueType::BarrierToken(..) => {
439                    panic!("Barrier not supported in SPIR-V")
440                }
441                core::OpaqueType::TensorMap => panic!("Tensor map not supported in SPIR-V"),
442            },
443            core::Type::Aggregate(_) => {
444                unreachable!("Should be disaggregated at this point")
445            }
446        }
447    }
448
449    pub fn compile_storage_type(&mut self, ty: core::StorageType) -> Elem {
450        match ty {
451            core::StorageType::Scalar(ty) => self.compile_elem(ty),
452            core::StorageType::Packed(_, _) => {
453                unimplemented!("Packed types not yet supported in SPIR-V")
454            }
455        }
456    }
457
458    pub fn compile_elem(&mut self, elem: core::ElemType) -> Elem {
459        match elem {
460            core::ElemType::Float(
461                core::FloatKind::E2M1
462                | core::FloatKind::E2M3
463                | core::FloatKind::E3M2
464                | core::FloatKind::UE8M0,
465            ) => panic!("Minifloat not supported in SPIR-V"),
466            core::ElemType::Float(core::FloatKind::E4M3) => {
467                self.capabilities.insert(Capability::Float8EXT);
468                Elem::Float(8, Some(FPEncoding::Float8E4M3EXT))
469            }
470            core::ElemType::Float(core::FloatKind::E5M2) => {
471                self.capabilities.insert(Capability::Float8EXT);
472                Elem::Float(8, Some(FPEncoding::Float8E5M2EXT))
473            }
474            core::ElemType::Float(core::FloatKind::BF16) => {
475                self.capabilities.insert(Capability::BFloat16TypeKHR);
476                Elem::Float(16, Some(FPEncoding::BFloat16KHR))
477            }
478            core::ElemType::Float(FloatKind::F16) => {
479                self.capabilities.insert(Capability::Float16);
480                Elem::Float(16, None)
481            }
482            core::ElemType::Float(FloatKind::TF32) => panic!("TF32 not supported in SPIR-V"),
483            core::ElemType::Float(FloatKind::Flex32) => Elem::Relaxed,
484            core::ElemType::Float(FloatKind::F32) => Elem::Float(32, None),
485            core::ElemType::Float(FloatKind::F64) => {
486                self.capabilities.insert(Capability::Float64);
487                Elem::Float(64, None)
488            }
489            core::ElemType::Int(IntKind::I8) => {
490                self.capabilities.insert(Capability::Int8);
491                Elem::Int(8, true)
492            }
493            core::ElemType::Int(IntKind::I16) => {
494                self.capabilities.insert(Capability::Int16);
495                Elem::Int(16, true)
496            }
497            core::ElemType::Int(IntKind::I32) => Elem::Int(32, true),
498            core::ElemType::Int(IntKind::I64) => {
499                self.capabilities.insert(Capability::Int64);
500                Elem::Int(64, true)
501            }
502            core::ElemType::UInt(UIntKind::U64) => {
503                self.capabilities.insert(Capability::Int64);
504                Elem::Int(64, false)
505            }
506            core::ElemType::UInt(UIntKind::U32) => Elem::Int(32, false),
507            core::ElemType::UInt(UIntKind::U16) => {
508                self.capabilities.insert(Capability::Int16);
509                Elem::Int(16, false)
510            }
511            core::ElemType::UInt(UIntKind::U8) => {
512                self.capabilities.insert(Capability::Int8);
513                Elem::Int(8, false)
514            }
515            core::ElemType::Bool => Elem::Bool,
516        }
517    }
518
519    pub fn compile_function_param_type(&mut self, val: core::Value) -> Word {
520        match val.kind {
521            core::ValueKind::Value { .. } | core::ValueKind::Constant(..) => {
522                self.compile_type(val.ty).id(self)
523            }
524        }
525    }
526
527    pub fn static_cast(&mut self, val: ConstVal, from: &Elem, item: &Item) -> (Word, ConstVal) {
528        let elem_cast = match (*from, item.elem()) {
529            (Elem::Bool, Elem::Int(width, _)) => ConstVal::from_uint(val.as_u32() as u64, width),
530            (Elem::Bool, Elem::Float(width, encoding)) => {
531                ConstVal::from_float(val.as_u32() as f64, width, encoding)
532            }
533            (Elem::Bool, Elem::Relaxed) => ConstVal::from_float(val.as_u32() as f64, 32, None),
534            (Elem::Int(_, _), Elem::Bool) => ConstVal::from_bool(val.as_u64() != 0),
535            (Elem::Int(_, false), Elem::Int(width, _)) => ConstVal::from_uint(val.as_u64(), width),
536            (Elem::Int(w_in, true), Elem::Int(width, _)) => {
537                ConstVal::from_uint(val.as_int(w_in) as u64, width)
538            }
539            (Elem::Int(_, false), Elem::Float(width, encoding)) => {
540                ConstVal::from_float(val.as_u64() as f64, width, encoding)
541            }
542            (Elem::Int(_, false), Elem::Relaxed) => {
543                ConstVal::from_float(val.as_u64() as f64, 32, None)
544            }
545            (Elem::Int(in_w, true), Elem::Float(width, encoding)) => {
546                ConstVal::from_float(val.as_int(in_w) as f64, width, encoding)
547            }
548            (Elem::Int(in_w, true), Elem::Relaxed) => {
549                ConstVal::from_float(val.as_int(in_w) as f64, 32, None)
550            }
551            (Elem::Float(in_w, encoding), Elem::Bool) => {
552                ConstVal::from_bool(val.as_float(in_w, encoding) != 0.0)
553            }
554            (Elem::Relaxed, Elem::Bool) => ConstVal::from_bool(val.as_float(32, None) != 0.0),
555            (Elem::Float(in_w, encoding), Elem::Int(out_w, false)) => {
556                ConstVal::from_uint(val.as_float(in_w, encoding) as u64, out_w)
557            }
558            (Elem::Relaxed, Elem::Int(out_w, false)) => {
559                ConstVal::from_uint(val.as_float(32, None) as u64, out_w)
560            }
561            (Elem::Float(in_w, encoding), Elem::Int(out_w, true)) => {
562                ConstVal::from_int(val.as_float(in_w, encoding) as i64, out_w)
563            }
564            (Elem::Relaxed, Elem::Int(out_w, true)) => {
565                ConstVal::from_int(val.as_float(32, None) as i64, out_w)
566            }
567            (Elem::Float(in_w, encoding), Elem::Float(out_w, encoding_out)) => {
568                ConstVal::from_float(val.as_float(in_w, encoding), out_w, encoding_out)
569            }
570            (Elem::Relaxed, Elem::Float(out_w, encoding)) => {
571                ConstVal::from_float(val.as_float(32, None), out_w, encoding)
572            }
573            (Elem::Float(in_w, encoding), Elem::Relaxed) => {
574                ConstVal::from_float(val.as_float(in_w, encoding), 32, None)
575            }
576            (Elem::Bool, Elem::Bool) => val,
577            (Elem::Relaxed, Elem::Relaxed) => val,
578            (_, Elem::Void) | (Elem::Void, _) => unreachable!(),
579        };
580        let id = item.constant(self, elem_cast);
581        (id, elem_cast)
582    }
583}
584
585pub fn compile_pointer_class(class: AddressSpace) -> StorageClass {
586    match class {
587        AddressSpace::Global(_) => StorageClass::PhysicalStorageBuffer,
588        AddressSpace::Shared => StorageClass::Workgroup,
589        AddressSpace::Local => StorageClass::Function,
590    }
591}
592
593impl std::fmt::Display for Item {
594    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
595        match self {
596            Item::Scalar(elem) => write!(f, "{elem}"),
597            Item::Vector(elem, factor) => write!(f, "vec{factor}<{elem}>"),
598            Item::Pointer(class, item) => write!(f, "ptr<{class:?}, {item}>"),
599            Item::Array(item, size) => write!(f, "array<{item}, {size}>"),
600            Item::DynamicArray(item) => write!(f, "array<{item}>"),
601            Item::CoopMatrix { ty, ident, .. } => write!(f, "matrix<{ty}, {ident:?}>"),
602            Item::TensorLayout { dims, clamp_mode } => {
603                write!(f, "tensor_layout<{dims}, {clamp_mode:?}>")
604            }
605            Item::TensorView {
606                dims,
607                has_dims,
608                permutation,
609            } => {
610                write!(
611                    f,
612                    "tensor_view<{:?}, has_dims: {has_dims}>",
613                    &permutation[..*dims]
614                )
615            }
616        }
617    }
618}
619
620impl std::fmt::Display for Elem {
621    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
622        match self {
623            Elem::Void => write!(f, "void"),
624            Elem::Bool => write!(f, "bool"),
625            Elem::Int(width, false) => write!(f, "u{width}"),
626            Elem::Int(width, true) => write!(f, "i{width}"),
627            Elem::Float(width, None) => write!(f, "f{width}"),
628            Elem::Float(_, Some(FPEncoding::BFloat16KHR)) => write!(f, "bf16"),
629            Elem::Float(_, Some(FPEncoding::Float8E4M3EXT)) => write!(f, "e4m3"),
630            Elem::Float(_, Some(FPEncoding::Float8E5M2EXT)) => write!(f, "e5m2"),
631            Elem::Relaxed => write!(f, "flex32"),
632        }
633    }
634}
635
636fn compile_clamp_mode(clamp_mode: ClampMode) -> TensorClampMode {
637    match clamp_mode {
638        ClampMode::Undefined => TensorClampMode::Undefined,
639        ClampMode::Constant(_) => TensorClampMode::Constant,
640        ClampMode::ClampToEdge => TensorClampMode::ClampToEdge,
641        ClampMode::Repeat => TensorClampMode::Repeat,
642        ClampMode::RepeatMirrored => TensorClampMode::RepeatMirrored,
643    }
644}
cubecl_spirv/item.rs

cubecl_spirv/
item.rs
