1use cubecl_macros_internal::{const_eval, cube_op, simplify};
2use half::{bf16, f16};
3use pliron::{
4 builtin::{attributes::IntegerAttr, types::IntegerType},
5 r#type::TypeHandle,
6};
7
8use crate::{
9 CanMaterialize, ConstantValue, PropagatesUniformity, Pure,
10 attributes::{BoolAttr, FloatAttr, IndexAttr, IntAttrExt},
11 dialect::{base::pure_binop, math::int_attr},
12 interfaces::{TriviallyUnrollable, TypedExt},
13 prelude::*,
14 types::{VectorType, scalar::BoolType},
15};
16
17pure_binop!("cmp.s_min", SMinOp);
18const_eval!(SMinOp, {
19 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| lhs.min(rhs),
20 custom: |lhs, rhs| {
22 let const_val = lhs.or(rhs)?;
23 let ty = const_val.get_type(ctx);
24 match const_val.as_const_val(ctx) {
25 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
26 _ => None
27 }
28 }
29});
30simplify!(SMinOp, {
31 |lhs, _| {
33 let ty = lhs?.get_type(ctx);
34 match lhs?.as_const_val(ctx) {
35 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => Some(self.rhs(ctx)),
36 _ => None,
37 }
38 },
39 |_, rhs| {
41 let ty = rhs?.get_type(ctx);
42 match rhs?.as_const_val(ctx) {
43 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => Some(self.lhs(ctx)),
44 _ => None,
45 }
46 },
47 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
49 true => Some(self.lhs(ctx)),
50 false => None
51 }
52});
53
54pure_binop!("cmp.u_min", UMinOp);
55const_eval!(UMinOp, {
56 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs.min(rhs),
57 custom: |lhs, rhs| match lhs.or(rhs)?.as_int(ctx)?.is_zero() {
59 true => Some(int_attr(ctx, self.result_type(ctx), 0)),
60 false => None,
61 }
62});
63simplify!(UMinOp, {
64 |lhs, _| {
66 let ty = lhs?.get_type(ctx);
67 match lhs?.as_const_val(ctx) {
68 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => Some(self.rhs(ctx)),
69 _ => None,
70 }
71 },
72 |_, rhs| {
74 let ty = rhs?.get_type(ctx);
75 match rhs?.as_const_val(ctx) {
76 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => Some(self.lhs(ctx)),
77 _ => None,
78 }
79 },
80 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
82 true => Some(self.lhs(ctx)),
83 false => None
84 }
85});
86
87pure_binop!("cmp.f_min", FMinOp);
88const_eval!(FMinOp, { [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| lhs.min(rhs) });
89simplify!(FMinOp, {
90 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
92 true => Some(self.lhs(ctx)),
93 false => None,
94 }
95});
96
97pure_binop!("cmp.s_max", SMaxOp);
98const_eval!(SMaxOp, {
99 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| lhs.max(rhs),
100 custom: |lhs, rhs| {
102 let const_val = lhs.or(rhs)?;
103 let ty = const_val.get_type(ctx);
104 match const_val.as_const_val(ctx) {
105 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
106 _ => None
107 }
108 }
109});
110simplify!(SMaxOp, {
111 |lhs, _| {
113 let ty = lhs?.get_type(ctx);
114 match lhs?.as_const_val(ctx) {
115 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => Some(self.rhs(ctx)),
116 _ => None,
117 }
118 },
119 |_, rhs| {
121 let ty = rhs?.get_type(ctx);
122 match rhs?.as_const_val(ctx) {
123 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => Some(self.lhs(ctx)),
124 _ => None,
125 }
126 },
127 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
129 true => Some(self.lhs(ctx)),
130 false => None,
131 }
132});
133
134pure_binop!("cmp.u_max", UMaxOp);
135const_eval!(UMaxOp, {
136 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| lhs.max(rhs),
137 custom: |lhs, rhs| {
139 let const_val = lhs.or(rhs)?;
140 let ty = const_val.get_type(ctx);
141 match const_val.as_const_val(ctx) {
142 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
143 _ => None
144 }
145 }
146});
147simplify!(UMaxOp, {
148 |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
150 true => Some(self.rhs(ctx)),
151 false => None,
152 },
153 |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
155 true => Some(self.lhs(ctx)),
156 false => None,
157 },
158 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
160 true => Some(self.lhs(ctx)),
161 false => None,
162 }
163});
164
165pure_binop!("cmp.f_max", FMaxOp);
166const_eval!(FMaxOp, {
167 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| lhs.max(rhs),
168});
169simplify!(FMaxOp, {
170 |_, _| match self.lhs(ctx) == self.rhs(ctx) {
172 true => Some(self.lhs(ctx)),
173 false => None,
174 }
175});
176
177#[cube_op(name = "cmp.s_clamp")]
178#[result_ty(same_as = input)]
179#[op_interfaces(SameOperandsType, SameOperandsAndResultType, TriviallyUnrollable)]
180#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
181pub struct SClampOp {
182 pub input: Value,
183 pub min: Value,
184 pub max: Value,
185}
186const_eval!(SClampOp, {
187 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |inp, min, max| inp.clamp(min, max),
188 custom: |_, min, _| {
190 let ty = min?.get_type(ctx);
191 match min?.as_const_val(ctx) {
192 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
193 _ => None
194 }
195 },
196 custom: |_, _, max| {
198 let ty = max?.get_type(ctx);
199 match max?.as_const_val(ctx) {
200 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
201 _ => None
202 }
203 }
204});
205
206#[cube_op(name = "cmp.u_clamp")]
207#[result_ty(same_as = input)]
208#[op_interfaces(SameOperandsType, SameOperandsAndResultType, TriviallyUnrollable)]
209#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
210pub struct UClampOp {
211 pub input: Value,
212 pub min: Value,
213 pub max: Value,
214}
215const_eval!(UClampOp, {
216 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |inp, min, max| inp.clamp(min, max),
217 custom: |_, min, _| {
219 let ty = min?.get_type(ctx);
220 match min?.as_const_val(ctx) {
221 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => Some(int_attr(ctx, ty, val as i128)),
222 _ => None
223 }
224 },
225 custom: |_, _, max| match max?.as_int(ctx)?.is_zero() {
227 true => Some(int_attr(ctx, self.result_type(ctx), 0)),
228 false => None
229 }
230});
231
232#[cube_op(name = "cmp.f_clamp")]
233#[result_ty(same_as = input)]
234#[op_interfaces(SameOperandsType, SameOperandsAndResultType, TriviallyUnrollable)]
235#[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
236pub struct FClampOp {
237 pub input: Value,
238 pub min: Value,
239 pub max: Value,
240}
241const_eval!(FClampOp, {
242 [FloatAttr(f16, bf16, f32, f64)]: |inp, min, max| inp.clamp(min, max),
243});
244
245macro_rules! cmp_binop {
246 ($name: literal, $ty: ident) => {
247 #[cubecl_macros_internal::cube_op(name = $name)]
248 #[result_ty(from_inputs = cmp_result_ty)]
249 #[$crate::prelude::op_interfaces(SameOperandsType, TriviallyUnrollable)]
250 #[op_traits(Pure, CanMaterialize, PropagatesUniformity)]
251 pub struct $ty {
252 pub lhs: Value,
253 pub rhs: Value,
254 }
255 };
256}
257
258cmp_binop!("cmp.s_less_than", SLessThanOp);
259const_eval!(SLessThanOp, {
260 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| -> BoolAttr {
261 (lhs < rhs).into()
262 },
263 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
265 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
266 false => None
267 },
268 custom: |_, rhs| {
270 let ty = rhs?.get_type(ctx);
271 match rhs?.as_const_val(ctx) {
272 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
273 _ => None
274 }
275 },
276 custom: |lhs, _| {
278 let ty = lhs?.get_type(ctx);
279 match lhs?.as_const_val(ctx) {
280 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
281 _ => None
282 }
283 },
284});
285
286cmp_binop!("cmp.u_less_than", ULessThanOp);
287const_eval!(ULessThanOp, {
288 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
289 (lhs < rhs).into()
290 },
291 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
293 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
294 false => None
295 },
296 custom: |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
298 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
299 false => None
300 },
301 custom: |lhs, _| {
303 let ty = lhs?.get_type(ctx);
304 match lhs?.as_const_val(ctx) {
305 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
306 _ => None
307 }
308 },
309});
310
311cmp_binop!("cmp.f_less_than", FLessThanOp);
312const_eval!(FLessThanOp, {
313 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
314 (lhs < rhs).into()
315 },
316});
317
318cmp_binop!("cmp.s_greater_than", SGreaterThanOp);
319const_eval!(SGreaterThanOp, {
320 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| -> BoolAttr {
321 (lhs > rhs).into()
322 },
323 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
325 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
326 false => None
327 },
328 custom: |lhs, _| {
330 let ty = lhs?.get_type(ctx);
331 match lhs?.as_const_val(ctx) {
332 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
333 _ => None
334 }
335 },
336 custom: |_, rhs| {
338 let ty = rhs?.get_type(ctx);
339 match rhs?.as_const_val(ctx) {
340 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
341 _ => None
342 }
343 }
344});
345
346cmp_binop!("cmp.u_greater_than", UGreaterThanOp);
347const_eval!(UGreaterThanOp, {
348 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
349 (lhs > rhs).into()
350 },
351 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
353 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
354 false => None
355 },
356 custom: |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
358 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
359 false => None
360 },
361 custom: |_, rhs| {
363 let ty = rhs?.get_type(ctx);
364 match rhs?.as_const_val(ctx) {
365 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
366 _ => None
367 }
368 }
369});
370
371cmp_binop!("cmp.f_greater_than", FGreaterThanOp);
372const_eval!(FGreaterThanOp, {
373 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
374 (lhs > rhs).into()
375 },
376});
377
378cmp_binop!("cmp.s_less_than_or_equal", SLessThanOrEqualOp);
379const_eval!(SLessThanOrEqualOp, {
380 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| -> BoolAttr {
381 (lhs <= rhs).into()
382 },
383 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
385 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
386 false => None
387 },
388 custom: |lhs, _| {
390 let ty = lhs?.get_type(ctx);
391 match lhs?.as_const_val(ctx) {
392 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
393 _ => None
394 }
395 },
396 custom: |_, rhs| {
398 let ty = rhs?.get_type(ctx);
399 match rhs?.as_const_val(ctx) {
400 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
401 _ => None
402 }
403 }
404});
405
406cmp_binop!("cmp.u_less_than_or_equal", ULessThanOrEqualOp);
407const_eval!(ULessThanOrEqualOp, {
408 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
409 (lhs <= rhs).into()
410 },
411 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
413 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
414 false => None
415 },
416 custom: |lhs, _| match lhs?.as_int(ctx)?.is_zero() {
418 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
419 false => None
420 },
421 custom: |_, rhs| {
423 let ty = rhs?.get_type(ctx);
424 match rhs?.as_const_val(ctx) {
425 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
426 _ => None
427 }
428 }
429});
430
431cmp_binop!("cmp.f_less_than_or_equal", FLessThanOrEqualOp);
432const_eval!(FLessThanOrEqualOp, {
433 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
434 (lhs <= rhs).into()
435 },
436});
437
438cmp_binop!("cmp.s_greater_than_or_equal", SGreaterThanOrEqualOp);
439const_eval!(SGreaterThanOrEqualOp, {
440 [IndexAttr, IntegerAttr(i8, i16, i32, i64)]: |lhs, rhs| -> BoolAttr {
441 (lhs >= rhs).into()
442 },
443 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
445 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
446 false => None
447 },
448 custom: |_, rhs| {
450 let ty = rhs?.get_type(ctx);
451 match rhs?.as_const_val(ctx) {
452 ConstantValue::Int(val) if is_min_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
453 _ => None
454 }
455 },
456 custom: |lhs, _| {
458 let ty = lhs?.get_type(ctx);
459 match lhs?.as_const_val(ctx) {
460 ConstantValue::Int(val) if is_max_int(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
461 _ => None
462 }
463 }
464});
465
466cmp_binop!("cmp.u_greater_than_or_equal", UGreaterThanOrEqualOp);
467const_eval!(UGreaterThanOrEqualOp, {
468 [IndexAttr, IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
469 (lhs >= rhs).into()
470 },
471 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
473 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
474 false => None
475 },
476 custom: |_, rhs| match rhs?.as_int(ctx)?.is_zero() {
478 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
479 false => None
480 },
481 custom: |lhs, _| {
483 let ty = lhs?.get_type(ctx);
484 match lhs?.as_const_val(ctx) {
485 ConstantValue::UInt(val) if is_max_uint(ctx, ty, val) => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
486 _ => None
487 }
488 }
489});
490
491cmp_binop!("cmp.f_greater_than_or_equal", FGreaterThanOrEqualOp);
492const_eval!(FGreaterThanOrEqualOp, {
493 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
494 (lhs >= rhs).into()
495 },
496});
497
498cmp_binop!("cmp.i_equal", IEqualOp);
499const_eval!(IEqualOp, {
500 [IndexAttr, IntegerAttr(i8, i16, i32, i64), IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
501 (lhs == rhs).into()
502 },
503 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
505 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
506 false => None
507 }
508});
509
510cmp_binop!("cmp.f_equal", FEqualOp);
511const_eval!(FEqualOp, {
512 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
513 (lhs == rhs).into()
514 },
515});
516
517cmp_binop!("cmp.bool_equal", BoolEqualOp);
518const_eval!(BoolEqualOp, {
519 [BoolAttr]: |lhs, rhs| lhs == rhs,
520 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
522 true => BoolAttr::per_lane(ctx, self.get_result(ctx), true),
523 false => None
524 }
525});
526
527cmp_binop!("cmp.i_not_equal", INotEqualOp);
528const_eval!(INotEqualOp, {
529 [IndexAttr, IntegerAttr(i8, i16, i32, i64), IntegerAttr(u8, u16, u32, u64)]: |lhs, rhs| -> BoolAttr {
530 (lhs != rhs).into()
531 },
532 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
534 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
535 false => None
536 },
537});
538
539cmp_binop!("cmp.f_not_equal", FNotEqualOp);
540const_eval!(FNotEqualOp, {
541 [FloatAttr(f16, bf16, f32, f64)]: |lhs, rhs| -> BoolAttr {
542 (lhs != rhs).into()
543 },
544});
545
546cmp_binop!("cmp.bool_not_equal", BoolNotEqualOp);
547const_eval!(BoolNotEqualOp, {
548 [BoolAttr]: |lhs, rhs| lhs != rhs,
549 custom: |_, _| match self.lhs(ctx) == self.rhs(ctx) {
551 true => BoolAttr::per_lane(ctx, self.get_result(ctx), false),
552 false => None
553 }
554});
555
556fn cmp_result_ty(ctx: &Context, lhs: &Value, _: &Value) -> TypeHandle {
557 let vectorization = lhs.vector_size(ctx);
558 let bool = BoolType::get(ctx).into();
559 if vectorization == 1 {
560 bool
561 } else {
562 VectorType::get(ctx, bool, vectorization).into()
563 }
564}
565
566pub(super) fn width(ctx: &Context, ty: TypeHandle) -> usize {
567 ty.size(ctx) * 8
568}
569
570fn is_min_int(ctx: &Context, ty: TypeHandle, val: i64) -> bool {
571 let ty = TypedHandle::<IntegerType>::from_handle(ty, ctx).unwrap();
572 val == min_int(ty.deref(ctx).width() as usize)
573}
574
575pub(super) fn is_max_int(ctx: &Context, ty: TypeHandle, val: i64) -> bool {
576 let ty = TypedHandle::<IntegerType>::from_handle(ty, ctx).unwrap();
577 val == max_int(ty.deref(ctx).width() as usize)
578}
579
580pub(super) fn is_max_uint(ctx: &Context, ty: TypeHandle, val: u64) -> bool {
581 val == max_uint(ty.size_bits(ctx))
582}
583
584fn min_int(width: usize) -> i64 {
585 if width >= 64 {
586 i64::MIN
587 } else {
588 -(1i64 << (width - 1))
589 }
590}
591
592fn max_int(width: usize) -> i64 {
593 if width >= 64 {
594 i64::MAX
595 } else {
596 (1i64 << (width - 1)) - 1
597 }
598}
599
600fn max_uint(width: usize) -> u64 {
601 if width >= 64 {
602 u64::MAX
603 } else {
604 (1u64 << width) - 1
605 }
606}