1use cubecl_core::ir::{
2 self as core, AddressSpace, BinaryOperands, Comparison, Instruction, InstructionModes, Memory,
3 Operation, Operator, UnaryOperands,
4};
5use rspirv::spirv::{Decoration, MemoryAccess, Word};
6
7use crate::{
8 SpirvCompiler, SpirvTarget,
9 item::{Elem, Item},
10};
11
12impl<T: SpirvTarget> SpirvCompiler<T> {
13 pub fn compile_operation(&mut self, inst: Instruction) {
14 if !matches!(inst.operation, Operation::NonSemantic(_)) {
16 self.set_source_loc(&inst.source_loc);
17 }
18 let uniform = inst
19 .out
20 .is_some_and(|out| self.uniformity.is_val_uniform(out));
21 match inst.operation {
22 Operation::Copy(val) => {
23 let input = self.compile_value(val);
24 let out = self.compile_value(inst.out());
25 let ty = out.item().id(self);
26 let in_id = self.read(&input);
27 let in_id = input.item().broadcast(self, in_id, None, &out.item());
28 let out_id = self.write_id(&out);
29
30 self.copy_object(ty, Some(out_id), in_id).unwrap();
31 self.mark_uniformity(out_id, uniform);
32 self.write(&out, out_id);
33 }
34 Operation::DeclareVariable {
35 addr_space: AddressSpace::Local,
36 ..
37 } => {
38 let out = self.compile_value(inst.out());
39 let ty = out.item().id(self);
40 let id = self.declare_function_variable(ty, None);
41 self.write(&out, id);
42 }
43 Operation::DeclareVariable {
44 addr_space: AddressSpace::Shared,
45 ..
46 } => {
47 let out = inst.out().id();
49 let id = self.state.lookups.shared[&out].id;
50 self.insert_value(out, id);
51 }
52 Operation::DeclareVariable { addr_space, .. } => {
53 unimplemented!("Unsupported declare address space {addr_space}")
54 }
55 Operation::Memory(mem) => self.compile_memory(mem, inst.out),
56 Operation::Arithmetic(operator) => {
57 self.compile_arithmetic(operator, inst.out, inst.modes, uniform)
58 }
59 Operation::Comparison(operator) => {
60 self.compile_cmp(operator, inst.out, inst.modes, uniform)
61 }
62 Operation::Bitwise(operator) => self.compile_bitwise(operator, inst.out, uniform),
63 Operation::Operator(operator) => self.compile_operator(operator, inst.out, uniform),
64 Operation::Atomic(atomic) => self.compile_atomic(atomic, inst.out, inst.modes),
65 Operation::Branch(_) => unreachable!("Branches shouldn't exist in optimized IR"),
66 Operation::Metadata(meta) => self.compile_meta(meta, inst.out, uniform),
67 Operation::Plane(plane) => self.compile_plane(plane, inst.out, uniform),
68 Operation::Synchronization(sync) => self.compile_sync(sync),
69 Operation::WorkgroupUniformLoad(op) => {
70 self.compile_sync(core::Synchronization::SyncCube);
71 if op.ty.is_atomic() {
72 self.compile_atomic(core::AtomicOp::Load(op), inst.out, inst.modes);
73 } else {
74 self.compile_memory(core::Memory::Load(op), inst.out);
75 }
76 }
77 Operation::CoopMma(cmma) => self.compile_cmma(cmma, inst.out),
78 Operation::TensorIndexing(tensor) => self.compile_tensor_indexing(tensor, inst.out),
79 Operation::NonSemantic(debug) => self.compile_debug(debug),
80 Operation::Barrier(_) => panic!("Barrier not supported in SPIR-V"),
81 Operation::Tma(_) => panic!("TMA not supported in SPIR-V"),
82 Operation::Marker(_) => {}
83 Operation::ConstructAggregate(..) | Operation::ExtractAggregateField(..) => {
84 unreachable!("Should be disaggregated at this point")
85 }
86 }
87 }
88
89 pub fn compile_cmp(
90 &mut self,
91 op: Comparison,
92 out: Option<core::Value>,
93 modes: InstructionModes,
94 uniform: bool,
95 ) {
96 let out = out.unwrap();
97 match op {
98 Comparison::Equal(op) => {
99 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
100 match lhs_ty.elem() {
101 Elem::Bool => b.logical_equal(ty, Some(out), lhs, rhs),
102 Elem::Int(_, _) => b.i_equal(ty, Some(out), lhs, rhs),
103 Elem::Float(..) => {
104 b.declare_math_mode(modes, out);
105 b.f_ord_equal(ty, Some(out), lhs, rhs)
106 }
107 Elem::Relaxed => {
108 b.decorate(out, Decoration::RelaxedPrecision, []);
109 b.declare_math_mode(modes, out);
110 b.f_ord_equal(ty, Some(out), lhs, rhs)
111 }
112 Elem::Void => unreachable!(),
113 }
114 .unwrap();
115 });
116 }
117 Comparison::NotEqual(op) => {
118 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
119 match lhs_ty.elem() {
120 Elem::Bool => b.logical_not_equal(ty, Some(out), lhs, rhs),
121 Elem::Int(_, _) => b.i_not_equal(ty, Some(out), lhs, rhs),
122 Elem::Float(..) => {
123 b.declare_math_mode(modes, out);
124 b.f_ord_not_equal(ty, Some(out), lhs, rhs)
125 }
126 Elem::Relaxed => {
127 b.decorate(out, Decoration::RelaxedPrecision, []);
128 b.declare_math_mode(modes, out);
129 b.f_ord_not_equal(ty, Some(out), lhs, rhs)
130 }
131 Elem::Void => unreachable!(),
132 }
133 .unwrap();
134 });
135 }
136 Comparison::Lower(op) => {
137 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
138 match lhs_ty.elem() {
139 Elem::Int(_, false) => b.u_less_than(ty, Some(out), lhs, rhs),
140 Elem::Int(_, true) => b.s_less_than(ty, Some(out), lhs, rhs),
141 Elem::Float(..) => {
142 b.declare_math_mode(modes, out);
143 b.f_ord_less_than(ty, Some(out), lhs, rhs)
144 }
145 Elem::Relaxed => {
146 b.decorate(out, Decoration::RelaxedPrecision, []);
147 b.declare_math_mode(modes, out);
148 b.f_ord_less_than(ty, Some(out), lhs, rhs)
149 }
150 _ => unreachable!(),
151 }
152 .unwrap();
153 });
154 }
155 Comparison::LowerEqual(op) => {
156 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
157 match lhs_ty.elem() {
158 Elem::Int(_, false) => b.u_less_than_equal(ty, Some(out), lhs, rhs),
159 Elem::Int(_, true) => b.s_less_than_equal(ty, Some(out), lhs, rhs),
160 Elem::Float(..) => {
161 b.declare_math_mode(modes, out);
162 b.f_ord_less_than_equal(ty, Some(out), lhs, rhs)
163 }
164 Elem::Relaxed => {
165 b.decorate(out, Decoration::RelaxedPrecision, []);
166 b.declare_math_mode(modes, out);
167 b.f_ord_less_than_equal(ty, Some(out), lhs, rhs)
168 }
169 _ => unreachable!(),
170 }
171 .unwrap();
172 });
173 }
174 Comparison::Greater(op) => {
175 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
176 match lhs_ty.elem() {
177 Elem::Int(_, false) => b.u_greater_than(ty, Some(out), lhs, rhs),
178 Elem::Int(_, true) => b.s_greater_than(ty, Some(out), lhs, rhs),
179 Elem::Float(..) => {
180 b.declare_math_mode(modes, out);
181 b.f_ord_greater_than(ty, Some(out), lhs, rhs)
182 }
183 Elem::Relaxed => {
184 b.decorate(out, Decoration::RelaxedPrecision, []);
185 b.declare_math_mode(modes, out);
186 b.f_ord_greater_than(ty, Some(out), lhs, rhs)
187 }
188 _ => unreachable!(),
189 }
190 .unwrap();
191 });
192 }
193 Comparison::GreaterEqual(op) => {
194 self.compile_binary_op_bool(op, out, uniform, |b, lhs_ty, ty, lhs, rhs, out| {
195 match lhs_ty.elem() {
196 Elem::Int(_, false) => b.u_greater_than_equal(ty, Some(out), lhs, rhs),
197 Elem::Int(_, true) => b.s_greater_than_equal(ty, Some(out), lhs, rhs),
198 Elem::Float(..) => {
199 b.declare_math_mode(modes, out);
200 b.f_ord_greater_than_equal(ty, Some(out), lhs, rhs)
201 }
202 Elem::Relaxed => {
203 b.decorate(out, Decoration::RelaxedPrecision, []);
204 b.declare_math_mode(modes, out);
205 b.f_ord_greater_than_equal(ty, Some(out), lhs, rhs)
206 }
207 _ => unreachable!(),
208 }
209 .unwrap();
210 });
211 }
212 Comparison::IsNan(op) => {
213 self.compile_unary_op(op, out, uniform, |b, _, ty, input, out| {
214 b.is_nan(ty, Some(out), input).unwrap();
215 });
216 }
217 Comparison::IsInf(op) => {
218 self.compile_unary_op(op, out, uniform, |b, _, ty, input, out| {
219 b.is_inf(ty, Some(out), input).unwrap();
220 });
221 }
222 }
223 }
224
225 pub fn compile_memory(&mut self, mem: Memory, out: Option<core::Value>) {
226 match mem {
227 Memory::Index(op) => {
228 let list = self.compile_value(op.list);
229 let index = self.compile_value(op.index);
230 let out = self.compile_value(out.unwrap());
231
232 let ptr = self.index(&list, &index, &out);
233
234 self.write(&out, ptr);
235 }
236 Memory::Load(value) => {
237 let ptr = self.compile_value(value);
238 let out = self.compile_value(out.unwrap());
239
240 let id = self.load_aligned(&ptr, &out);
241 self.write(&out, id);
242 }
243 Memory::Store(op) => {
244 let ptr = self.compile_value(op.ptr);
245 let value = self.compile_value(op.value);
246
247 self.store_aligned(&ptr, &value);
248 }
249 Memory::CopyMemory(op) => {
250 let source = self.compile_value(op.source);
251 let target = self.compile_value(op.target);
252
253 let out_ty = target.item();
254 let align = source.item().size().max(target.item().size());
255
256 let source = self.read(&source);
257 let target = self.read(&target);
258
259 if op.len == 1 {
260 self.copy_memory(
261 target,
262 source,
263 Some(MemoryAccess::ALIGNED),
264 [align.into()],
265 None,
266 [],
267 )
268 .unwrap();
269 } else {
270 let size = op.len as u32 * out_ty.size();
271 let size_id = self.const_u32(size);
272
273 self.copy_memory_sized(
274 target,
275 source,
276 size_id,
277 Some(MemoryAccess::ALIGNED),
278 [size.into()],
279 None,
280 [],
281 )
282 .unwrap();
283 }
284 }
285 }
286 }
287
288 pub fn compile_operator(&mut self, op: Operator, out: Option<core::Value>, uniform: bool) {
289 let out = out.unwrap();
290 match op {
291 Operator::Cast(op) => {
292 let input = self.compile_value(op.input);
293 let out = self.compile_value(out);
294 let ty = out.item().id(self);
295 let in_id = self.read(&input);
296 let out_id = self.write_id(&out);
297 self.mark_uniformity(out_id, uniform);
298
299 if let Some(as_const) = input.as_const() {
300 let cast = self.static_cast(as_const, &input.elem(), &out.item()).0;
301 self.copy_object(ty, Some(out_id), cast).unwrap();
302 } else {
303 input.item().cast_to(self, Some(out_id), in_id, &out.item());
304 }
305
306 self.write(&out, out_id);
307 }
308 Operator::And(op) => {
309 self.compile_binary_op(op, out, uniform, |b, _, ty, lhs, rhs, out| {
310 b.logical_and(ty, Some(out), lhs, rhs).unwrap();
311 });
312 }
313 Operator::Or(op) => {
314 self.compile_binary_op(op, out, uniform, |b, _, ty, lhs, rhs, out| {
315 b.logical_or(ty, Some(out), lhs, rhs).unwrap();
316 });
317 }
318 Operator::Not(op) => {
319 self.compile_unary_op_cast(op, out, uniform, |b, _, ty, input, out| {
320 b.logical_not(ty, Some(out), input).unwrap();
321 });
322 }
323 Operator::Reinterpret(op) => {
324 self.compile_unary_op(op, out, uniform, |b, _, ty, input, out| {
325 b.bitcast(ty, Some(out), input).unwrap();
326 })
327 }
328 Operator::InitVector(op) => {
329 let values = op
330 .inputs
331 .into_iter()
332 .map(|input| self.compile_value(input))
333 .collect::<Vec<_>>()
334 .into_iter()
335 .map(|it| self.read(&it))
336 .collect::<Vec<_>>();
337 let item = self.compile_type(out.ty);
338 let out = self.compile_value(out);
339 let out_id = self.write_id(&out);
340 self.mark_uniformity(out_id, uniform);
341 let ty = item.id(self);
342 self.composite_construct(ty, Some(out_id), values).unwrap();
343 self.write(&out, out_id);
344 }
345 Operator::InsertComponent(op) => {
346 let vector = self.compile_value(op.vector);
347 let value = self.compile_value(op.value);
348 let output = self.compile_value(out);
349
350 let vector = self.read(&vector);
351 let value = self.read(&value);
352 let out_ty = output.item().id(self);
353 let write_id = self.write_id(&output);
354
355 if let Some(index) = op.index.as_const() {
356 let index = index.as_u32();
357 self.composite_insert(out_ty, Some(write_id), value, vector, [index])
358 .unwrap();
359 } else {
360 let index = self.compile_value(op.index);
361 let index = self.read(&index);
362
363 self.vector_insert_dynamic(out_ty, Some(write_id), vector, value, index)
364 .unwrap();
365 }
366
367 self.write(&output, write_id);
368 }
369 Operator::ExtractComponent(op) => {
370 let vector = self.compile_value(op.lhs);
371 let output = self.compile_value(out);
372
373 let vector = self.read(&vector);
374 let out_ty = output.item().id(self);
375 let write_id = self.write_id(&output);
376
377 if let Some(index) = op.rhs.as_const() {
378 let index = index.as_u32();
379 self.composite_extract(out_ty, Some(write_id), vector, [index])
380 .unwrap();
381 } else {
382 let index = self.compile_value(op.rhs);
383 let index = self.read(&index);
384
385 self.vector_extract_dynamic(out_ty, Some(write_id), vector, index)
386 .unwrap();
387 }
388
389 self.write(&output, write_id);
390 }
391 Operator::Select(op) => self.compile_select(op.cond, op.then, op.or_else, out, uniform),
392 Operator::ReadBuiltin(builtin) => {
393 let out = self.compile_value(out);
394 let value = self.compile_builtin(builtin, &out.item());
395 self.write(&out, value);
396 }
397 Operator::ReadScalar(id) => {
398 let value = self.global_scalar(id, out.storage_type());
399 let out = self.compile_value(out);
400 self.write(&out, value);
401 }
402 }
403 }
404
405 pub fn compile_unary_op_cast(
406 &mut self,
407 op: UnaryOperands,
408 out: core::Value,
409 uniform: bool,
410 exec: impl FnOnce(&mut Self, Item, Word, Word, Word),
411 ) {
412 let input = self.compile_value(op.input);
413 let out = self.compile_value(out);
414 let out_ty = out.item();
415
416 let input_id = self.read_as(&input, &out_ty);
417 let out_id = self.write_id(&out);
418 self.mark_uniformity(out_id, uniform);
419
420 let ty = out_ty.id(self);
421
422 exec(self, out_ty, ty, input_id, out_id);
423 self.write(&out, out_id);
424 }
425
426 pub fn compile_unary_op(
427 &mut self,
428 op: UnaryOperands,
429 out: core::Value,
430 uniform: bool,
431 exec: impl FnOnce(&mut Self, Item, Word, Word, Word),
432 ) {
433 let input = self.compile_value(op.input);
434 let out = self.compile_value(out);
435 let out_ty = out.item();
436
437 let input_id = self.read(&input);
438 let out_id = self.write_id(&out);
439 self.mark_uniformity(out_id, uniform);
440
441 let ty = out_ty.id(self);
442
443 exec(self, out_ty, ty, input_id, out_id);
444 self.write(&out, out_id);
445 }
446
447 pub fn compile_unary_op_bool(
448 &mut self,
449 op: UnaryOperands,
450 out: core::Value,
451 uniform: bool,
452 exec: impl FnOnce(&mut Self, Item, Word, Word, Word),
453 ) {
454 let input = self.compile_value(op.input);
455 let out = self.compile_value(out);
456 let in_ty = input.item();
457
458 let input_id = self.read(&input);
459 let out_id = self.write_id(&out);
460 self.mark_uniformity(out_id, uniform);
461
462 let ty = out.item().id(self);
463
464 exec(self, in_ty, ty, input_id, out_id);
465 self.write(&out, out_id);
466 }
467
468 pub fn compile_binary_op(
469 &mut self,
470 op: BinaryOperands,
471 out: core::Value,
472 uniform: bool,
473 exec: impl FnOnce(&mut Self, Item, Word, Word, Word, Word),
474 ) {
475 let lhs = self.compile_value(op.lhs);
476 let rhs = self.compile_value(op.rhs);
477 let out = self.compile_value(out);
478 let out_ty = out.item();
479
480 let lhs_id = self.read_as(&lhs, &out_ty);
481 let rhs_id = self.read_as(&rhs, &out_ty);
482 let out_id = self.write_id(&out);
483 self.mark_uniformity(out_id, uniform);
484
485 let ty = out_ty.id(self);
486
487 exec(self, out_ty, ty, lhs_id, rhs_id, out_id);
488 self.write(&out, out_id);
489 }
490
491 pub fn compile_binary_op_no_cast(
492 &mut self,
493 op: BinaryOperands,
494 out: core::Value,
495 uniform: bool,
496 exec: impl FnOnce(&mut Self, Item, Word, Word, Word, Word),
497 ) {
498 let lhs = self.compile_value(op.lhs);
499 let rhs = self.compile_value(op.rhs);
500 let out = self.compile_value(out);
501 let out_ty = out.item();
502
503 let lhs_id = self.read(&lhs);
504 let rhs_id = self.read(&rhs);
505 let out_id = self.write_id(&out);
506 self.mark_uniformity(out_id, uniform);
507
508 let ty = out_ty.id(self);
509
510 exec(self, out_ty, ty, lhs_id, rhs_id, out_id);
511 self.write(&out, out_id);
512 }
513
514 pub fn compile_binary_op_bool(
515 &mut self,
516 op: BinaryOperands,
517 out: core::Value,
518 uniform: bool,
519 exec: impl FnOnce(&mut Self, Item, Word, Word, Word, Word),
520 ) {
521 let lhs = self.compile_value(op.lhs);
522 let rhs = self.compile_value(op.rhs);
523 let out = self.compile_value(out);
524
525 let in_ty = out.item().same_vectorization(lhs.elem());
526
527 let lhs_id = self.read_as(&lhs, &in_ty);
528 let rhs_id = self.read_as(&rhs, &in_ty);
529 let out_id = self.write_id(&out);
530 self.mark_uniformity(out_id, uniform);
531
532 let ty = out.item().id(self);
533
534 exec(self, in_ty, ty, lhs_id, rhs_id, out_id);
535 self.write(&out, out_id);
536 }
537
538 pub fn compile_select(
539 &mut self,
540 cond: core::Value,
541 then: core::Value,
542 or_else: core::Value,
543 out: core::Value,
544 uniform: bool,
545 ) {
546 let cond = self.compile_value(cond);
547 let then = self.compile_value(then);
548 let or_else = self.compile_value(or_else);
549 let out = self.compile_value(out);
550
551 let out_ty = out.item();
552 let ty = out_ty.id(self);
553
554 let cond_id = self.read(&cond);
555 let then = self.read_as(&then, &out_ty);
556 let or_else = self.read_as(&or_else, &out_ty);
557 let out_id = self.write_id(&out);
558 self.mark_uniformity(out_id, uniform);
559
560 self.select(ty, Some(out_id), cond_id, then, or_else)
561 .unwrap();
562 self.write(&out, out_id);
563 }
564
565 pub fn mark_uniformity(&mut self, id: Word, uniform: bool) {
566 if uniform {
567 self.decorate(id, Decoration::Uniform, []);
568 }
569 }
570}