cubecl_spirv/
arithmetic.rs

1use crate::{
2    SpirvCompiler, SpirvTarget,
3    item::{Elem, Item},
4    value::ConstVal,
5};
6use cubecl_core::ir::{self as core, Arithmetic, InstructionModes};
7use rspirv::spirv::{Capability, Decoration, FPEncoding};
8
9impl<T: SpirvTarget> SpirvCompiler<T> {
10    pub fn compile_arithmetic(
11        &mut self,
12        op: Arithmetic,
13        out: Option<core::Value>,
14        modes: InstructionModes,
15        uniform: bool,
16    ) {
17        let out = out.unwrap();
18        match op {
19            Arithmetic::Add(op) => {
20                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
21                    match out_ty.elem() {
22                        Elem::Int(_, _) => b.i_add(ty, Some(out), lhs, rhs).unwrap(),
23                        Elem::Float(..) => {
24                            b.declare_math_mode(modes, out);
25                            b.f_add(ty, Some(out), lhs, rhs).unwrap()
26                        }
27                        Elem::Relaxed => {
28                            b.decorate(out, Decoration::RelaxedPrecision, []);
29                            b.declare_math_mode(modes, out);
30                            b.f_add(ty, Some(out), lhs, rhs).unwrap()
31                        }
32                        _ => unreachable!(),
33                    };
34                });
35            }
36            Arithmetic::SaturatingAdd(_) => {
37                unimplemented!("Should be replaced by polyfill");
38            }
39            Arithmetic::Sub(op) => {
40                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
41                    match out_ty.elem() {
42                        Elem::Int(_, _) => b.i_sub(ty, Some(out), lhs, rhs).unwrap(),
43                        Elem::Float(..) => {
44                            b.declare_math_mode(modes, out);
45                            b.f_sub(ty, Some(out), lhs, rhs).unwrap()
46                        }
47                        Elem::Relaxed => {
48                            b.decorate(out, Decoration::RelaxedPrecision, []);
49                            b.declare_math_mode(modes, out);
50                            b.f_sub(ty, Some(out), lhs, rhs).unwrap()
51                        }
52                        _ => unreachable!(),
53                    };
54                });
55            }
56            Arithmetic::SaturatingSub(_) => {
57                unimplemented!("Should be replaced by polyfill");
58            }
59            Arithmetic::Mul(op) => {
60                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
61                    match out_ty.elem() {
62                        Elem::Int(_, _) => b.i_mul(ty, Some(out), lhs, rhs).unwrap(),
63                        Elem::Float(..) => {
64                            b.declare_math_mode(modes, out);
65                            b.f_mul(ty, Some(out), lhs, rhs).unwrap()
66                        }
67                        Elem::Relaxed => {
68                            b.decorate(out, Decoration::RelaxedPrecision, []);
69                            b.declare_math_mode(modes, out);
70                            b.f_mul(ty, Some(out), lhs, rhs).unwrap()
71                        }
72                        _ => unreachable!(),
73                    };
74                });
75            }
76            Arithmetic::MulHi(op) => {
77                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
78                    let out_st = b.type_struct([ty, ty]);
79                    let extended = match out_ty.elem() {
80                        Elem::Int(_, false) => b.u_mul_extended(out_st, None, lhs, rhs).unwrap(),
81                        Elem::Int(_, true) => b.s_mul_extended(out_st, None, lhs, rhs).unwrap(),
82                        _ => unreachable!(),
83                    };
84                    b.composite_extract(ty, Some(out), extended, [1]).unwrap();
85                });
86            }
87            Arithmetic::Div(op) => {
88                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
89                    match out_ty.elem() {
90                        Elem::Int(_, false) => b.u_div(ty, Some(out), lhs, rhs).unwrap(),
91                        Elem::Int(_, true) => b.s_div(ty, Some(out), lhs, rhs).unwrap(),
92                        Elem::Float(..) => {
93                            b.declare_math_mode(modes, out);
94                            b.f_div(ty, Some(out), lhs, rhs).unwrap()
95                        }
96                        Elem::Relaxed => {
97                            b.decorate(out, Decoration::RelaxedPrecision, []);
98                            b.declare_math_mode(modes, out);
99                            b.f_div(ty, Some(out), lhs, rhs).unwrap()
100                        }
101                        _ => unreachable!(),
102                    };
103                });
104            }
105            Arithmetic::ModFloor(op) => {
106                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
107                    match out_ty.elem() {
108                        Elem::Int(_, false) => b.u_mod(ty, Some(out), lhs, rhs).unwrap(),
109                        Elem::Int(_, true) => {
110                            // Convert to float and use `f_mod` (floored division) instead of `s_mod`
111                            // (truncated division) to match remainder semantics across dtypes
112                            // e.g. remainder(-2, 3) = 1, not 2
113                            let f_ty = match out_ty {
114                                Item::Scalar(_elem) => Item::Scalar(Elem::Relaxed),
115                                Item::Vector(_elem, factor) => Item::Vector(Elem::Relaxed, factor),
116                                _ => unreachable!(),
117                            };
118                            let f_ty = f_ty.id(b);
119                            let lhs_f = b.convert_s_to_f(f_ty, None, lhs).unwrap();
120                            let rhs_f = b.convert_s_to_f(f_ty, None, rhs).unwrap();
121                            let rem = b.f_mod(f_ty, None, lhs_f, rhs_f).unwrap();
122                            b.convert_f_to_s(ty, Some(out), rem).unwrap()
123                        }
124                        Elem::Float(..) => {
125                            b.declare_math_mode(modes, out);
126                            b.f_mod(ty, Some(out), lhs, rhs).unwrap()
127                        }
128                        Elem::Relaxed => {
129                            b.decorate(out, Decoration::RelaxedPrecision, []);
130                            b.declare_math_mode(modes, out);
131                            b.f_mod(ty, Some(out), lhs, rhs).unwrap()
132                        }
133                        _ => unreachable!(),
134                    };
135                });
136            }
137            Arithmetic::Rem(op) => {
138                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
139                    match out_ty.elem() {
140                        Elem::Int(_, false) => b.u_mod(ty, Some(out), lhs, rhs).unwrap(),
141                        Elem::Int(_, true) => b.s_rem(ty, Some(out), lhs, rhs).unwrap(),
142                        Elem::Float(..) => {
143                            b.declare_math_mode(modes, out);
144                            b.f_rem(ty, Some(out), lhs, rhs).unwrap()
145                        }
146                        Elem::Relaxed => {
147                            b.decorate(out, Decoration::RelaxedPrecision, []);
148                            b.declare_math_mode(modes, out);
149                            b.f_rem(ty, Some(out), lhs, rhs).unwrap()
150                        }
151                        _ => unreachable!(),
152                    };
153                });
154            }
155            Arithmetic::Dot(op) => {
156                if op.lhs.ty.vector_size() == 1 {
157                    self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
158                        match out_ty.elem() {
159                            Elem::Int(_, _) => b.i_mul(ty, Some(out), lhs, rhs).unwrap(),
160                            Elem::Float(..) => {
161                                b.declare_math_mode(modes, out);
162                                b.f_mul(ty, Some(out), lhs, rhs).unwrap()
163                            }
164                            Elem::Relaxed => {
165                                b.decorate(out, Decoration::RelaxedPrecision, []);
166                                b.declare_math_mode(modes, out);
167                                b.f_mul(ty, Some(out), lhs, rhs).unwrap()
168                            }
169                            _ => unreachable!(),
170                        };
171                    });
172                } else {
173                    let lhs = self.compile_value(op.lhs);
174                    let rhs = self.compile_value(op.rhs);
175                    let out = self.compile_value(out);
176                    let ty = out.item().id(self);
177
178                    let lhs_id = self.read(&lhs);
179                    let rhs_id = self.read(&rhs);
180                    let out_id = self.write_id(&out);
181                    self.mark_uniformity(out_id, uniform);
182
183                    if matches!(lhs.elem(), Elem::Int(_, _)) {
184                        self.capabilities.insert(Capability::DotProduct);
185                    }
186                    if matches!(lhs.elem(), Elem::Float(16, Some(FPEncoding::BFloat16KHR))) {
187                        self.capabilities.insert(Capability::BFloat16DotProductKHR);
188                    }
189
190                    match (lhs.elem(), rhs.elem()) {
191                        (Elem::Int(_, false), Elem::Int(_, false)) => {
192                            self.u_dot(ty, Some(out_id), lhs_id, rhs_id, None)
193                        }
194                        (Elem::Int(_, true), Elem::Int(_, false)) => {
195                            self.su_dot(ty, Some(out_id), lhs_id, rhs_id, None)
196                        }
197                        (Elem::Int(_, false), Elem::Int(_, true)) => {
198                            self.su_dot(ty, Some(out_id), rhs_id, lhs_id, None)
199                        }
200                        (Elem::Int(_, true), Elem::Int(_, true)) => {
201                            self.s_dot(ty, Some(out_id), lhs_id, rhs_id, None)
202                        }
203                        (Elem::Float(..), Elem::Float(..))
204                        | (Elem::Relaxed, Elem::Float(..))
205                        | (Elem::Float(..), Elem::Relaxed) => {
206                            self.dot(ty, Some(out_id), lhs_id, rhs_id)
207                        }
208                        (Elem::Relaxed, Elem::Relaxed) => {
209                            self.decorate(out_id, Decoration::RelaxedPrecision, []);
210                            self.dot(ty, Some(out_id), lhs_id, rhs_id)
211                        }
212                        _ => unreachable!(),
213                    }
214                    .unwrap();
215                    self.write(&out, out_id);
216                }
217            }
218            Arithmetic::Fma(op) => {
219                let a = self.compile_value(op.a);
220                let b = self.compile_value(op.b);
221                let c = self.compile_value(op.c);
222                let out = self.compile_value(out);
223                let out_ty = out.item();
224                let relaxed = matches!(
225                    (a.item().elem(), b.item().elem(), c.item().elem()),
226                    (Elem::Relaxed, Elem::Relaxed, Elem::Relaxed)
227                );
228
229                let a_id = self.read_as(&a, &out_ty);
230                let b_id = self.read_as(&b, &out_ty);
231                let c_id = self.read_as(&c, &out_ty);
232                let out_id = self.write_id(&out);
233                self.mark_uniformity(out_id, uniform);
234
235                let ty = out_ty.id(self);
236
237                let mul = self.f_mul(ty, None, a_id, b_id).unwrap();
238                self.mark_uniformity(mul, uniform);
239                self.declare_math_mode(modes, mul);
240                self.f_add(ty, Some(out_id), mul, c_id).unwrap();
241                self.declare_math_mode(modes, out_id);
242                if relaxed {
243                    self.decorate(mul, Decoration::RelaxedPrecision, []);
244                    self.decorate(out_id, Decoration::RelaxedPrecision, []);
245                }
246                self.write(&out, out_id);
247            }
248            Arithmetic::Recip(op) => {
249                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
250                    let one = b
251                        .static_cast(ConstVal::Bit32(1), &Elem::Int(32, false), &out_ty)
252                        .0;
253                    b.declare_math_mode(modes, out);
254                    b.f_div(ty, Some(out), one, input).unwrap();
255                });
256            }
257            Arithmetic::Neg(op) => {
258                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
259                    match out_ty.elem() {
260                        Elem::Int(_, true) => b.s_negate(ty, Some(out), input).unwrap(),
261                        Elem::Float(..) => {
262                            b.declare_math_mode(modes, out);
263                            b.f_negate(ty, Some(out), input).unwrap()
264                        }
265                        Elem::Relaxed => {
266                            b.decorate(out, Decoration::RelaxedPrecision, []);
267                            b.declare_math_mode(modes, out);
268                            b.f_negate(ty, Some(out), input).unwrap()
269                        }
270                        _ => unreachable!(),
271                    };
272                });
273            }
274            Arithmetic::Erf(_) => {
275                unreachable!("Replaced by transformer")
276            }
277
278            // Extension functions
279            Arithmetic::Normalize(op) => {
280                self.compile_unary_op(op, out, uniform, |b, out_ty, ty, input, out| {
281                    b.declare_math_mode(modes, out);
282                    T::normalize(b, ty, input, out);
283                    if matches!(out_ty.elem(), Elem::Relaxed) {
284                        b.decorate(out, Decoration::RelaxedPrecision, []);
285                    }
286                });
287            }
288            Arithmetic::VectorSum(op) => {
289                let input_ir = op.input;
290                let input = self.compile_value(input_ir);
291                let out = self.compile_value(out);
292                let in_item = input.item();
293                let out_ty = out.item();
294                let vec_size = in_item.vectorization();
295                let scalar_ty = out_ty.id(self);
296                let input_id = self.read(&input);
297                let out_id = self.write_id(&out);
298                self.mark_uniformity(out_id, uniform);
299
300                if vec_size <= 1 {
301                    // Scalar: identity
302                    self.copy_object(scalar_ty, Some(out_id), input_id).unwrap();
303                } else if matches!(out_ty.elem(), Elem::Float(..) | Elem::Relaxed) {
304                    // Float vector: use OpDot with ones vector for optimal single instruction
305                    self.declare_math_mode(modes, out_id);
306                    let ones = in_item.constant(
307                        self,
308                        ConstVal::from_float(
309                            1.0,
310                            in_item.elem().width(),
311                            in_item.elem().float_encoding(),
312                        ),
313                    );
314                    self.dot(scalar_ty, Some(out_id), input_id, ones).unwrap();
315                    if matches!(out_ty.elem(), Elem::Relaxed) {
316                        self.decorate(out_id, Decoration::RelaxedPrecision, []);
317                    }
318                } else {
319                    // Integer vector: extract and add
320                    let elem_ty = out_ty.id(self);
321                    let mut acc = self
322                        .composite_extract(elem_ty, None, input_id, vec![0])
323                        .unwrap();
324                    for i in 1..vec_size {
325                        let elem = self
326                            .composite_extract(elem_ty, None, input_id, vec![i])
327                            .unwrap();
328                        acc = self.i_add(elem_ty, None, acc, elem).unwrap();
329                    }
330                    self.copy_object(scalar_ty, Some(out_id), acc).unwrap();
331                }
332                self.write(&out, out_id);
333            }
334            Arithmetic::Magnitude(op) => {
335                self.compile_unary_op(op, out, uniform, |b, out_ty, ty, input, out| {
336                    b.declare_math_mode(modes, out);
337                    T::magnitude(b, ty, input, out);
338                    if matches!(out_ty.elem(), Elem::Relaxed) {
339                        b.decorate(out, Decoration::RelaxedPrecision, []);
340                    }
341                });
342            }
343            Arithmetic::Abs(op) => {
344                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
345                    match out_ty.elem() {
346                        Elem::Int(_, _) => T::s_abs(b, ty, input, out),
347                        Elem::Float(..) => {
348                            b.declare_math_mode(modes, out);
349                            T::f_abs(b, ty, input, out)
350                        }
351                        Elem::Relaxed => {
352                            b.decorate(out, Decoration::RelaxedPrecision, []);
353                            b.declare_math_mode(modes, out);
354                            T::f_abs(b, ty, input, out)
355                        }
356                        _ => unreachable!(),
357                    }
358                });
359            }
360            Arithmetic::Exp(op) => {
361                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
362                    b.declare_math_mode(modes, out);
363                    T::exp(b, ty, input, out);
364                    if matches!(out_ty.elem(), Elem::Relaxed) {
365                        b.decorate(out, Decoration::RelaxedPrecision, []);
366                    }
367                });
368            }
369            Arithmetic::Log(op) => {
370                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
371                    b.declare_math_mode(modes, out);
372                    T::log(b, ty, input, out);
373                    if matches!(out_ty.elem(), Elem::Relaxed) {
374                        b.decorate(out, Decoration::RelaxedPrecision, []);
375                    }
376                })
377            }
378            Arithmetic::Log1p(op) => {
379                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
380                    let one = b
381                        .static_cast(ConstVal::Bit32(1), &Elem::Int(32, false), &out_ty)
382                        .0;
383                    let relaxed = matches!(out_ty.elem(), Elem::Relaxed);
384                    let add = match out_ty.elem() {
385                        Elem::Int(_, _) => b.i_add(ty, None, input, one).unwrap(),
386                        Elem::Float(..) | Elem::Relaxed => {
387                            b.declare_math_mode(modes, out);
388                            b.f_add(ty, None, input, one).unwrap()
389                        }
390                        _ => unreachable!(),
391                    };
392                    b.mark_uniformity(add, uniform);
393                    if relaxed {
394                        b.decorate(add, Decoration::RelaxedPrecision, []);
395                        b.decorate(out, Decoration::RelaxedPrecision, []);
396                    }
397                    b.declare_math_mode(modes, out);
398                    T::log(b, ty, add, out)
399                });
400            }
401            Arithmetic::Expm1(op) => {
402                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
403                    let relaxed = matches!(out_ty.elem(), Elem::Relaxed);
404                    let bool = out_ty.same_vectorization(Elem::Bool).id(b);
405                    let one = b
406                        .static_cast(
407                            ConstVal::from_float(1.0, 32, None),
408                            &Elem::Float(32, None),
409                            &out_ty,
410                        )
411                        .0;
412                    let half = b
413                        .static_cast(
414                            ConstVal::from_float(0.5, 32, None),
415                            &Elem::Float(32, None),
416                            &out_ty,
417                        )
418                        .0;
419                    let sixth = b
420                        .static_cast(
421                            ConstVal::from_float(1.0 / 6.0, 32, None),
422                            &Elem::Float(32, None),
423                            &out_ty,
424                        )
425                        .0;
426                    let threshold = b
427                        .static_cast(
428                            ConstVal::from_float(1.0e-5, 32, None),
429                            &Elem::Float(32, None),
430                            &out_ty,
431                        )
432                        .0;
433                    let abs = b.id();
434                    b.declare_math_mode(modes, abs);
435                    T::f_abs(b, ty, input, abs);
436                    let is_small = b.f_ord_less_than(bool, None, abs, threshold).unwrap();
437                    b.declare_math_mode(modes, is_small);
438                    let squared = b.f_mul(ty, None, input, input).unwrap();
439                    b.declare_math_mode(modes, squared);
440                    let cubed = b.f_mul(ty, None, squared, input).unwrap();
441                    b.declare_math_mode(modes, cubed);
442                    let half_squared = b.f_mul(ty, None, squared, half).unwrap();
443                    b.declare_math_mode(modes, half_squared);
444                    let sixth_cubed = b.f_mul(ty, None, cubed, sixth).unwrap();
445                    b.declare_math_mode(modes, sixth_cubed);
446                    let linear_plus_quad = b.f_add(ty, None, input, half_squared).unwrap();
447                    b.declare_math_mode(modes, linear_plus_quad);
448                    let taylor = b.f_add(ty, None, linear_plus_quad, sixth_cubed).unwrap();
449                    b.declare_math_mode(modes, taylor);
450                    let exp = b.id();
451                    b.declare_math_mode(modes, exp);
452                    T::exp(b, ty, input, exp);
453                    let native = b.f_sub(ty, None, exp, one).unwrap();
454                    b.declare_math_mode(modes, native);
455                    for id in [
456                        abs,
457                        squared,
458                        cubed,
459                        half_squared,
460                        sixth_cubed,
461                        linear_plus_quad,
462                        taylor,
463                        exp,
464                        native,
465                    ] {
466                        b.mark_uniformity(id, uniform);
467                        if relaxed {
468                            b.decorate(id, Decoration::RelaxedPrecision, []);
469                        }
470                    }
471                    if relaxed {
472                        b.decorate(out, Decoration::RelaxedPrecision, []);
473                    }
474                    b.select(ty, Some(out), is_small, taylor, native).unwrap();
475                });
476            }
477            Arithmetic::Cos(op) => {
478                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
479                    b.declare_math_mode(modes, out);
480                    T::cos(b, ty, input, out);
481                    if matches!(out_ty.elem(), Elem::Relaxed) {
482                        b.decorate(out, Decoration::RelaxedPrecision, []);
483                    }
484                })
485            }
486            Arithmetic::Sin(op) => {
487                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
488                    b.declare_math_mode(modes, out);
489                    T::sin(b, ty, input, out);
490                    if matches!(out_ty.elem(), Elem::Relaxed) {
491                        b.decorate(out, Decoration::RelaxedPrecision, []);
492                    }
493                })
494            }
495            Arithmetic::Tan(op) => {
496                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
497                    b.declare_math_mode(modes, out);
498                    T::tan(b, ty, input, out);
499                    if matches!(out_ty.elem(), Elem::Relaxed) {
500                        b.decorate(out, Decoration::RelaxedPrecision, []);
501                    }
502                })
503            }
504            Arithmetic::Tanh(op) => {
505                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
506                    b.declare_math_mode(modes, out);
507                    T::tanh(b, ty, input, out);
508                    if matches!(out_ty.elem(), Elem::Relaxed) {
509                        b.decorate(out, Decoration::RelaxedPrecision, []);
510                    }
511                })
512            }
513            Arithmetic::Sinh(op) => {
514                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
515                    b.declare_math_mode(modes, out);
516                    T::sinh(b, ty, input, out);
517                    if matches!(out_ty.elem(), Elem::Relaxed) {
518                        b.decorate(out, Decoration::RelaxedPrecision, []);
519                    }
520                })
521            }
522            Arithmetic::Cosh(op) => {
523                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
524                    b.declare_math_mode(modes, out);
525                    T::cosh(b, ty, input, out);
526                    if matches!(out_ty.elem(), Elem::Relaxed) {
527                        b.decorate(out, Decoration::RelaxedPrecision, []);
528                    }
529                })
530            }
531            Arithmetic::ArcCos(op) => {
532                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
533                    b.declare_math_mode(modes, out);
534                    T::acos(b, ty, input, out);
535                    if matches!(out_ty.elem(), Elem::Relaxed) {
536                        b.decorate(out, Decoration::RelaxedPrecision, []);
537                    }
538                })
539            }
540            Arithmetic::ArcSin(op) => {
541                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
542                    b.declare_math_mode(modes, out);
543                    T::asin(b, ty, input, out);
544                    if matches!(out_ty.elem(), Elem::Relaxed) {
545                        b.decorate(out, Decoration::RelaxedPrecision, []);
546                    }
547                })
548            }
549            Arithmetic::ArcTan(op) => {
550                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
551                    b.declare_math_mode(modes, out);
552                    T::atan(b, ty, input, out);
553                    if matches!(out_ty.elem(), Elem::Relaxed) {
554                        b.decorate(out, Decoration::RelaxedPrecision, []);
555                    }
556                })
557            }
558            Arithmetic::ArcSinh(op) => {
559                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
560                    b.declare_math_mode(modes, out);
561                    T::asinh(b, ty, input, out);
562                    if matches!(out_ty.elem(), Elem::Relaxed) {
563                        b.decorate(out, Decoration::RelaxedPrecision, []);
564                    }
565                })
566            }
567            Arithmetic::ArcCosh(op) => {
568                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
569                    b.declare_math_mode(modes, out);
570                    T::acosh(b, ty, input, out);
571                    if matches!(out_ty.elem(), Elem::Relaxed) {
572                        b.decorate(out, Decoration::RelaxedPrecision, []);
573                    }
574                })
575            }
576            Arithmetic::ArcTanh(op) => {
577                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
578                    b.declare_math_mode(modes, out);
579                    T::atanh(b, ty, input, out);
580                    if matches!(out_ty.elem(), Elem::Relaxed) {
581                        b.decorate(out, Decoration::RelaxedPrecision, []);
582                    }
583                })
584            }
585            Arithmetic::Degrees(op) => {
586                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
587                    b.declare_math_mode(modes, out);
588                    T::degrees(b, ty, input, out);
589                    if matches!(out_ty.elem(), Elem::Relaxed) {
590                        b.decorate(out, Decoration::RelaxedPrecision, []);
591                    }
592                })
593            }
594            Arithmetic::Radians(op) => {
595                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
596                    b.declare_math_mode(modes, out);
597                    T::radians(b, ty, input, out);
598                    if matches!(out_ty.elem(), Elem::Relaxed) {
599                        b.decorate(out, Decoration::RelaxedPrecision, []);
600                    }
601                })
602            }
603            Arithmetic::ArcTan2(op) => {
604                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
605                    b.declare_math_mode(modes, out);
606                    T::atan2(b, ty, lhs, rhs, out);
607                    if matches!(out_ty.elem(), Elem::Relaxed) {
608                        b.decorate(out, Decoration::RelaxedPrecision, []);
609                    }
610                })
611            }
612            // No powi for Vulkan, just auto-cast to float
613            Arithmetic::Powf(op) | Arithmetic::Powi(op) => {
614                self.compile_binary_op(op, out, uniform, |b, out_ty, ty, lhs, rhs, out| {
615                    let bool = match out_ty {
616                        Item::Scalar(_) => Elem::Bool.id(b),
617                        Item::Vector(_, factor) => Item::Vector(Elem::Bool, factor).id(b),
618                        _ => unreachable!(),
619                    };
620                    let relaxed = matches!(out_ty.elem(), Elem::Relaxed);
621                    let zero = out_ty.const_u32(b, 0);
622                    let one = out_ty.const_u32(b, 1);
623                    let two = out_ty.const_u32(b, 2);
624                    let modulo = b.f_rem(ty, None, rhs, two).unwrap();
625                    b.declare_math_mode(modes, modulo);
626                    let is_zero = b.f_ord_equal(bool, None, modulo, zero).unwrap();
627                    b.declare_math_mode(modes, is_zero);
628                    let abs = b.id();
629                    b.declare_math_mode(modes, abs);
630                    T::f_abs(b, ty, lhs, abs);
631                    let even = b.id();
632                    b.declare_math_mode(modes, even);
633                    T::pow(b, ty, abs, rhs, even);
634                    let cond2_0 = b.f_ord_equal(bool, None, modulo, one).unwrap();
635                    b.declare_math_mode(modes, cond2_0);
636                    let cond2_1 = b.f_ord_less_than(bool, None, lhs, zero).unwrap();
637                    b.declare_math_mode(modes, cond2_1);
638                    let cond2 = b.logical_and(bool, None, cond2_0, cond2_1).unwrap();
639                    let neg_lhs = b.f_negate(ty, None, lhs).unwrap();
640                    b.declare_math_mode(modes, neg_lhs);
641                    let pow2 = b.id();
642                    b.declare_math_mode(modes, pow2);
643                    T::pow(b, ty, neg_lhs, rhs, pow2);
644                    let pow2_neg = b.f_negate(ty, None, pow2).unwrap();
645                    b.declare_math_mode(modes, pow2_neg);
646                    let default = b.id();
647                    b.declare_math_mode(modes, default);
648                    T::pow(b, ty, lhs, rhs, default);
649                    let ids = [
650                        modulo, is_zero, abs, even, cond2_0, cond2_1, neg_lhs, pow2, pow2_neg,
651                        default,
652                    ];
653                    for id in ids {
654                        b.mark_uniformity(id, uniform);
655                        if relaxed {
656                            b.decorate(id, Decoration::RelaxedPrecision, []);
657                        }
658                    }
659                    let sel1 = b.select(ty, None, cond2, pow2_neg, default).unwrap();
660                    b.mark_uniformity(sel1, uniform);
661                    b.select(ty, Some(out), is_zero, even, sel1).unwrap();
662                })
663            }
664            Arithmetic::Hypot(_op) => {
665                unreachable!("Replaced by transformer");
666            }
667            Arithmetic::Rhypot(_op) => {
668                unreachable!("Replaced by transformer");
669            }
670            Arithmetic::Sqrt(op) => {
671                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
672                    b.declare_math_mode(modes, out);
673                    T::sqrt(b, ty, input, out);
674                    if matches!(out_ty.elem(), Elem::Relaxed) {
675                        b.decorate(out, Decoration::RelaxedPrecision, []);
676                    }
677                })
678            }
679            Arithmetic::InverseSqrt(op) => {
680                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
681                    b.declare_math_mode(modes, out);
682                    T::inverse_sqrt(b, ty, input, out);
683                    if matches!(out_ty.elem(), Elem::Relaxed) {
684                        b.decorate(out, Decoration::RelaxedPrecision, []);
685                    }
686                })
687            }
688            Arithmetic::Round(op) => {
689                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
690                    T::round(b, ty, input, out);
691                    if matches!(out_ty.elem(), Elem::Relaxed) {
692                        b.decorate(out, Decoration::RelaxedPrecision, []);
693                    }
694                })
695            }
696            Arithmetic::Floor(op) => {
697                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
698                    b.declare_math_mode(modes, out);
699                    T::floor(b, ty, input, out);
700                    if matches!(out_ty.elem(), Elem::Relaxed) {
701                        b.decorate(out, Decoration::RelaxedPrecision, []);
702                    }
703                })
704            }
705            Arithmetic::Ceil(op) => {
706                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
707                    b.declare_math_mode(modes, out);
708                    T::ceil(b, ty, input, out);
709                    if matches!(out_ty.elem(), Elem::Relaxed) {
710                        b.decorate(out, Decoration::RelaxedPrecision, []);
711                    }
712                })
713            }
714            Arithmetic::Trunc(op) => {
715                self.compile_unary_op_cast(op, out, uniform, |b, out_ty, ty, input, out| {
716                    b.declare_math_mode(modes, out);
717                    T::trunc(b, ty, input, out);
718                    if matches!(out_ty.elem(), Elem::Relaxed) {
719                        b.decorate(out, Decoration::RelaxedPrecision, []);
720                    }
721                })
722            }
723            Arithmetic::Clamp(op) => {
724                let input = self.compile_value(op.input);
725                let min = self.compile_value(op.min_value);
726                let max = self.compile_value(op.max_value);
727                let out = self.compile_value(out);
728                let out_ty = out.item();
729
730                let input = self.read_as(&input, &out_ty);
731                let min = self.read_as(&min, &out_ty);
732                let max = self.read_as(&max, &out_ty);
733                let out_id = self.write_id(&out);
734                self.mark_uniformity(out_id, uniform);
735
736                let ty = out_ty.id(self);
737
738                match out_ty.elem() {
739                    Elem::Int(_, false) => T::u_clamp(self, ty, input, min, max, out_id),
740                    Elem::Int(_, true) => T::s_clamp(self, ty, input, min, max, out_id),
741                    Elem::Float(..) => {
742                        self.declare_math_mode(modes, out_id);
743                        T::f_clamp(self, ty, input, min, max, out_id)
744                    }
745                    Elem::Relaxed => {
746                        self.decorate(out_id, Decoration::RelaxedPrecision, []);
747                        self.declare_math_mode(modes, out_id);
748                        T::f_clamp(self, ty, input, min, max, out_id)
749                    }
750                    _ => unreachable!(),
751                }
752                self.write(&out, out_id);
753            }
754
755            Arithmetic::Max(op) => self.compile_binary_op(
756                op,
757                out,
758                uniform,
759                |b, out_ty, ty, lhs, rhs, out| match out_ty.elem() {
760                    Elem::Int(_, false) => T::u_max(b, ty, lhs, rhs, out),
761                    Elem::Int(_, true) => T::s_max(b, ty, lhs, rhs, out),
762                    Elem::Float(..) => {
763                        b.declare_math_mode(modes, out);
764                        T::f_max(b, ty, lhs, rhs, out)
765                    }
766                    Elem::Relaxed => {
767                        b.decorate(out, Decoration::RelaxedPrecision, []);
768                        b.declare_math_mode(modes, out);
769                        T::f_max(b, ty, lhs, rhs, out)
770                    }
771                    _ => unreachable!(),
772                },
773            ),
774            Arithmetic::Min(op) => self.compile_binary_op(
775                op,
776                out,
777                uniform,
778                |b, out_ty, ty, lhs, rhs, out| match out_ty.elem() {
779                    Elem::Int(_, false) => T::u_min(b, ty, lhs, rhs, out),
780                    Elem::Int(_, true) => T::s_min(b, ty, lhs, rhs, out),
781                    Elem::Float(..) => {
782                        b.declare_math_mode(modes, out);
783                        T::f_min(b, ty, lhs, rhs, out)
784                    }
785                    Elem::Relaxed => {
786                        b.decorate(out, Decoration::RelaxedPrecision, []);
787                        b.declare_math_mode(modes, out);
788                        T::f_min(b, ty, lhs, rhs, out)
789                    }
790                    _ => unreachable!(),
791                },
792            ),
793        }
794    }
795}
cubecl_spirv/arithmetic.rs

cubecl_spirv/
arithmetic.rs
