cubecl_spirv/
value.rs

1#![allow(unknown_lints, unnecessary_transmutes)]
2
3use std::mem::transmute;
4
5use crate::{
6    SpirvCompiler, SpirvTarget,
7    item::{Elem, Item},
8};
9use cubecl_core::ir::{self, ConstantValue, Id};
10use rspirv::spirv::{self, FPEncoding, MemoryAccess, StorageClass, Word};
11
12#[allow(clippy::enum_variant_names)]
13#[derive(Debug, Clone, PartialEq)]
14pub enum Value {
15    Constant(Word, ConstVal, Item),
16    Value { id: Id, item: Item },
17}
18
19impl Value {
20    pub fn scope(&self) -> spirv::Scope {
21        match self.item() {
22            Item::Pointer(class, _) => match class {
23                StorageClass::StorageBuffer
24                | StorageClass::PhysicalStorageBuffer
25                | StorageClass::Uniform => spirv::Scope::Device,
26                StorageClass::Workgroup => spirv::Scope::Workgroup,
27                _ => spirv::Scope::Invocation,
28            },
29            Item::CoopMatrix { scope, .. } => scope,
30            _ => spirv::Scope::Invocation,
31        }
32    }
33}
34
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
36pub enum ConstVal {
37    Bit32(u32),
38    Bit64(u64),
39}
40
41impl ConstVal {
42    pub fn as_u64(&self) -> u64 {
43        match self {
44            ConstVal::Bit32(val) => *val as u64,
45            ConstVal::Bit64(val) => *val,
46        }
47    }
48
49    pub fn as_u32(&self) -> u32 {
50        match self {
51            ConstVal::Bit32(val) => *val,
52            ConstVal::Bit64(_) => panic!("Truncating 64 bit value to 32 bit"),
53        }
54    }
55
56    pub fn as_float(&self, width: u32, encoding: Option<FPEncoding>) -> f64 {
57        match (width, encoding) {
58            (64, _) => f64::from_bits(self.as_u64()),
59            (32, _) => f32::from_bits(self.as_u32()) as f64,
60            (16, None) => half::f16::from_bits(self.as_u32() as u16).to_f64(),
61            (_, Some(FPEncoding::BFloat16KHR)) => {
62                half::bf16::from_bits(self.as_u32() as u16).to_f64()
63            }
64            (_, Some(FPEncoding::Float8E4M3EXT)) => {
65                cubecl_common::e4m3::from_bits(self.as_u32() as u8).to_f64()
66            }
67            (_, Some(FPEncoding::Float8E5M2EXT)) => {
68                cubecl_common::e5m2::from_bits(self.as_u32() as u8).to_f64()
69            }
70            _ => unreachable!(),
71        }
72    }
73
74    pub fn as_int(&self, width: u32) -> i64 {
75        unsafe {
76            match width {
77                64 => transmute::<u64, i64>(self.as_u64()),
78                32 => transmute::<u32, i32>(self.as_u32()) as i64,
79                16 => transmute::<u16, i16>(self.as_u32() as u16) as i64,
80                8 => transmute::<u8, i8>(self.as_u32() as u8) as i64,
81                _ => unreachable!(),
82            }
83        }
84    }
85
86    pub fn from_float(value: f64, width: u32, encoding: Option<FPEncoding>) -> Self {
87        match (width, encoding) {
88            (64, _) => ConstVal::Bit64(value.to_bits()),
89            (32, _) => ConstVal::Bit32((value as f32).to_bits()),
90            (16, None) => ConstVal::Bit32(half::f16::from_f64(value).to_bits() as u32),
91            (_, Some(FPEncoding::BFloat16KHR)) => {
92                ConstVal::Bit32(half::bf16::from_f64(value).to_bits() as u32)
93            }
94            (_, Some(FPEncoding::Float8E4M3EXT)) => {
95                ConstVal::Bit32(cubecl_common::e4m3::from_f64(value).to_bits() as u32)
96            }
97            (_, Some(FPEncoding::Float8E5M2EXT)) => {
98                ConstVal::Bit32(cubecl_common::e5m2::from_f64(value).to_bits() as u32)
99            }
100            _ => unreachable!(),
101        }
102    }
103
104    pub fn from_int(value: i64, width: u32) -> Self {
105        match width {
106            64 => ConstVal::Bit64(unsafe { transmute::<i64, u64>(value) }),
107            32 => ConstVal::Bit32(unsafe { transmute::<i32, u32>(value as i32) }),
108            16 => ConstVal::Bit32(unsafe { transmute::<i16, u16>(value as i16) } as u32),
109            8 => ConstVal::Bit32(unsafe { transmute::<i8, u8>(value as i8) } as u32),
110            _ => unreachable!(),
111        }
112    }
113
114    pub fn from_uint(value: u64, width: u32) -> Self {
115        match width {
116            64 => ConstVal::Bit64(value),
117            32 => ConstVal::Bit32(value as u32),
118            16 => ConstVal::Bit32(value as u16 as u32),
119            8 => ConstVal::Bit32(value as u8 as u32),
120            _ => unreachable!(),
121        }
122    }
123
124    pub fn from_bool(value: bool) -> Self {
125        ConstVal::Bit32(value as u32)
126    }
127}
128
129impl From<(ConstantValue, Item)> for ConstVal {
130    fn from((value, ty): (ConstantValue, Item)) -> Self {
131        let elem = ty.elem();
132        let width = elem.size() * 8;
133        match value {
134            ConstantValue::Int(val) => ConstVal::from_int(val, width),
135            ConstantValue::Float(val) => ConstVal::from_float(val, width, elem.float_encoding()),
136            ConstantValue::UInt(val) => ConstVal::from_uint(val, width),
137            ConstantValue::Bool(val) => ConstVal::from_bool(val),
138        }
139    }
140}
141
142impl From<u32> for ConstVal {
143    fn from(value: u32) -> Self {
144        ConstVal::Bit32(value)
145    }
146}
147
148impl From<f32> for ConstVal {
149    fn from(value: f32) -> Self {
150        ConstVal::Bit32(value.to_bits())
151    }
152}
153
154impl Value {
155    pub fn id<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>) -> Word {
156        match self {
157            Value::Constant(id, _, _) => *id,
158            Value::Value { id, .. } => b.get_value(*id),
159        }
160    }
161
162    pub fn item(&self) -> Item {
163        match self {
164            Value::Constant(_, _, item) => item.clone(),
165            Value::Value { item, .. } => item.clone(),
166        }
167    }
168
169    pub fn elem(&self) -> Elem {
170        self.item().elem()
171    }
172
173    pub fn as_const(&self) -> Option<ConstVal> {
174        match self {
175            Self::Constant(_, val, _) => Some(*val),
176            _ => None,
177        }
178    }
179}
180
181impl<T: SpirvTarget> SpirvCompiler<T> {
182    pub fn compile_value(&mut self, value: ir::Value) -> Value {
183        let item = value.ty;
184        match value.kind {
185            ir::ValueKind::Constant(value) => {
186                let item = self.compile_type(item);
187                let const_val = (value, item.clone()).into();
188
189                if let Some(existing) = self.state.constants.get(&(const_val, item.clone())) {
190                    Value::Constant(*existing, const_val, item)
191                } else {
192                    let id = item.constant(self, const_val);
193                    self.state.constants.insert((const_val, item.clone()), id);
194                    Value::Constant(id, const_val, item)
195                }
196            }
197            ir::ValueKind::Value { id } => {
198                let item = self.compile_type(item);
199                Value::Value { id, item }
200            }
201        }
202    }
203
204    pub fn read(&mut self, value: &Value) -> Word {
205        value.id(self)
206    }
207
208    pub fn read_as(&mut self, value: &Value, item: &Item) -> Word {
209        if let Some(as_const) = value.as_const() {
210            self.static_cast(as_const, &value.elem(), item).0
211        } else {
212            let id = self.read(value);
213            value.item().cast_to(self, None, id, item)
214        }
215    }
216
217    pub fn index(&mut self, list: &Value, index: &Value, out: &Value) -> Word {
218        let list = self.read(list);
219        let index_id = self.read(index);
220        let write_id = self.write_id(out);
221        let ptr_ty = out.item().id(self);
222        self.in_bounds_access_chain(ptr_ty, Some(write_id), list, [index_id])
223            .unwrap()
224    }
225
226    pub fn write_id(&mut self, value: &Value) -> Word {
227        match value {
228            Value::Value { id, .. } => self.get_value(*id),
229            Value::Constant(_, _, _) => panic!("Can't write to constant scalar"),
230        }
231    }
232
233    /// Like [`write_id`], but specialized for the hacky semantics of CMMA ops. Unlike normal ops, they
234    /// can have mutable variables as direct outputs.
235    pub fn write_id_cmma(&mut self, value: &Value) -> Word {
236        match value {
237            Value::Value { item, .. } if item.is_ptr() => self.id(),
238            Value::Value { id, .. } => self.get_value(*id),
239            Value::Constant(_, _, _) => panic!("Can't write to constant scalar"),
240        }
241    }
242
243    pub fn write(&mut self, variable: &Value, value: Word) {
244        if let Value::Value { id, .. } = variable {
245            self.state.values.insert(*id, value);
246        }
247    }
248
249    /// Like [`write`], but specialized for the hacky semantics of CMMA ops. Unlike normal ops, they
250    /// can have mutable variables as direct outputs.
251    pub fn write_cmma(&mut self, variable: &Value, value: Word) {
252        match variable {
253            ptr @ Value::Value { item, .. } if item.is_ptr() => {
254                let ptr = self.read(ptr);
255
256                self.store(ptr, value, None, []).unwrap()
257            }
258            Value::Value { id, .. } => {
259                self.state.values.insert(*id, value);
260            }
261            _ => {}
262        }
263    }
264
265    pub fn load_aligned(&mut self, ptr: &Value, out: &Value) -> Word {
266        let out_ty = out.item().id(self);
267        let write_id = self.write_id(out);
268        let align = ptr.item().size();
269
270        let ptr = self.read(ptr);
271
272        self.load(
273            out_ty,
274            Some(write_id),
275            ptr,
276            Some(MemoryAccess::ALIGNED),
277            [align.into()],
278        )
279        .unwrap()
280    }
281
282    pub fn store_aligned(&mut self, ptr: &Value, value: &Value) {
283        let align = ptr.item().size();
284
285        let ptr = self.read(ptr);
286        let value = self.read(value);
287
288        self.store(ptr, value, Some(MemoryAccess::ALIGNED), [align.into()])
289            .unwrap()
290    }
291}
cubecl_spirv/value.rs

cubecl_spirv/
value.rs
