[Cranelift] add type-aware imm64 constant folding operations (#12826)

* [Cranelift] add imm64 operations

* [Cranelift] add imm64 partial ops

* [Cranelift] implement udiv

* [Cranelift] implement urem and use imm64 ops

* [Cranelift] use imm64 ops

* [Cranelift] `(-X) * C = X * (-C)`

* [Cranelift] update register numbering in filetests

Author: bongjunj

cranelift/codegen/src/isle_prelude.rs

@@ -74,6 +74,17 @@ macro_rules! isle_common_prelude_methods {
             Some(Imm64::new(result).mask_to_width(type_width))
         }
 
+        #[inline]
+        fn imm64_udiv(&mut self, ty: Type, x: Imm64, y: Imm64) -> Option<Imm64> {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+            let mask = self.ty_mask(ty);
+            let x = (x.bits() as u64) & mask;
+            let y = (y.bits() as u64) & mask;
+            let result = x.checked_div(y)?;
+            Some(Imm64::new(result as i64).mask_to_width(type_width))
+        }
+
         #[inline]
         fn imm64_srem(&mut self, ty: Type, x: Imm64, y: Imm64) -> Option<Imm64> {
             // Sign extend `x` and `y`.
@@ -89,6 +100,125 @@ macro_rules! isle_common_prelude_methods {
             Some(Imm64::new(result).mask_to_width(type_width))
         }
 
+        #[inline]
+        fn imm64_urem(&mut self, ty: Type, x: Imm64, y: Imm64) -> Option<Imm64> {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+            let mask = self.ty_mask(ty);
+            let x = (x.bits() as u64) & mask;
+            let y = (y.bits() as u64) & mask;
+            let result = x.checked_rem(y)?;
+            Some(Imm64::new(result as i64).mask_to_width(type_width))
+        }
+
+        #[inline]
+        fn imm64_add(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new(x.bits().wrapping_add(y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_sub(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new(x.bits().wrapping_sub(y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_mul(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new(x.bits().wrapping_mul(y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_and(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x.bits() & y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_or(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x.bits() | y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_xor(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x.bits() ^ y.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_not(&mut self, ty: Type, x: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((!x.bits()) & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_neg(&mut self, ty: Type, x: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new(x.bits().wrapping_neg() & ty_mask)
+        }
+
+        #[inline]
+        fn imm64_abs(&mut self, ty: Type, x: Imm64) -> Option<Imm64> {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+
+            let x = x.sign_extend_from_width(type_width).bits();
+            let shift = 64 - type_width;
+            let min = ((self.ty_smin(ty) as i64) << shift) >> shift;
+            if x == min {
+                return None;
+            }
+
+            Some(Imm64::new(x.abs()).mask_to_width(type_width))
+        }
+
+        #[inline]
+        fn imm64_ilog2(&mut self, ty: Type, x: Imm64) -> Option<Imm64> {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+            let masked = (x.bits() as u64) & self.ty_mask(ty);
+            let result = masked.checked_ilog2()?;
+            Some(Imm64::new(result.into()))
+        }
+
+        #[inline]
+        fn imm64_umin(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty);
+            let x_u = (x.bits() as u64) & ty_mask;
+            let y_u = (y.bits() as u64) & ty_mask;
+            Imm64::new((if x_u <= y_u { x_u } else { y_u }) as i64)
+        }
+
+        #[inline]
+        fn imm64_umax(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let ty_mask = self.ty_mask(ty);
+            let x_u = (x.bits() as u64) & ty_mask;
+            let y_u = (y.bits() as u64) & ty_mask;
+            Imm64::new((if x_u >= y_u { x_u } else { y_u }) as i64)
+        }
+
+        #[inline]
+        fn imm64_smin(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+            let x_s = x.sign_extend_from_width(type_width).bits();
+            let y_s = y.sign_extend_from_width(type_width).bits();
+            let selected = if x_s <= y_s { x_s } else { y_s };
+            Imm64::new(selected).mask_to_width(type_width)
+        }
+
+        #[inline]
+        fn imm64_smax(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let type_width = ty.bits();
+            assert!(type_width <= 64);
+            let x_s = x.sign_extend_from_width(type_width).bits();
+            let y_s = y.sign_extend_from_width(type_width).bits();
+            let selected = if x_s >= y_s { x_s } else { y_s };
+            Imm64::new(selected).mask_to_width(type_width)
+        }
+
         #[inline]
         fn imm64_shl(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
             // Mask off any excess shift bits.

cranelift/codegen/src/opts/arithmetic.isle

@@ -226,10 +226,10 @@
 ;; or(x, C) + (-C)  -->  and(x, ~C)
 (rule
   (simplify (iadd ty
-              (bor ty x (iconst_s ty n))
-              (iconst_s ty m)))
-  (if-let m (i64_checked_neg n))
-  (band ty x (iconst ty (imm64_masked ty (i64_cast_unsigned (i64_not n))))))
+              (bor ty x (iconst ty n))
+              (iconst ty m)))
+  (if-let m (imm64_neg ty n))
+  (band ty x (iconst ty (imm64_not ty n))))
 
 ;; (x + y) - (x | y) --> x & y
 (rule (simplify (isub ty (iadd ty x y) (bor ty x y))) (band ty x y))
@@ -500,3 +500,6 @@
 (rule (simplify (ugt ty x (umax ty y x))) (iconst_u ty 0))
 (rule (simplify (ult ty (umax ty x y) x)) (iconst_u ty 0))
 (rule (simplify (ult ty (umax ty y x) x)) (iconst_u ty 0))
+
+;; (-X) * C = X * (-C)
+(rule (simplify (imul (fits_in_64 ty) (ineg ty x) (iconst ty y))) (imul ty x (iconst ty (imm64_neg ty y))))

cranelift/codegen/src/opts/cprop.isle

@@ -13,21 +13,21 @@
 
 (rule (simplify
        (iadd (fits_in_64 ty)
-             (iconst ty (u64_from_imm64 k1))
-             (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_wrapping_add k1 k2)))))
+             (iconst ty k1)
+             (iconst ty k2)))
+      (subsume (iconst ty (imm64_add ty k1 k2))))
 
 (rule (simplify
        (isub (fits_in_64 ty)
-             (iconst ty (u64_from_imm64 k1))
-             (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_wrapping_sub k1 k2)))))
+             (iconst ty k1)
+             (iconst ty k2)))
+      (subsume (iconst ty (imm64_sub ty k1 k2))))
 
 (rule (simplify
        (imul (fits_in_64 ty)
-             (iconst ty (u64_from_imm64 k1))
-             (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_wrapping_mul k1 k2)))))
+             (iconst ty k1)
+             (iconst ty k2)))
+      (subsume (iconst ty (imm64_mul ty k1 k2))))
 
 (rule (simplify_skeleton
        (sdiv (iconst ty k1)
@@ -42,39 +42,39 @@
       (iconst ty d))
 
 (rule (simplify_skeleton
-       (udiv (iconst_u ty k1)
-             (iconst_u ty k2)))
-      (if-let d (u64_checked_div k1 k2))
-      (iconst ty (imm64_masked ty d)))
+       (udiv (iconst ty k1)
+             (iconst ty k2)))
+      (if-let d (imm64_udiv ty k1 k2))
+      (iconst ty d))
 
 (rule (simplify_skeleton
-       (urem (iconst_u ty k1)
-             (iconst_u ty k2)))
-      (if-let d (u64_checked_rem k1 k2))
-      (iconst ty (imm64_masked ty d)))
+       (urem (iconst ty k1)
+             (iconst ty k2)))
+      (if-let d (imm64_urem ty k1 k2))
+      (iconst ty d))
 
 (rule (simplify
        (bor (fits_in_64 ty)
-            (iconst ty (u64_from_imm64 k1))
-            (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_or k1 k2)))))
+            (iconst ty k1)
+            (iconst ty k2)))
+      (subsume (iconst ty (imm64_or ty k1 k2))))
 
 (rule (simplify
        (band (fits_in_64 ty)
-            (iconst ty (u64_from_imm64 k1))
-            (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_and k1 k2)))))
+            (iconst ty k1)
+            (iconst ty k2)))
+      (subsume (iconst ty (imm64_and ty k1 k2))))
 
 (rule (simplify
        (bxor (fits_in_64 ty)
-            (iconst ty (u64_from_imm64 k1))
-            (iconst ty (u64_from_imm64 k2))))
-      (subsume (iconst ty (imm64_masked ty (u64_xor k1 k2)))))
+            (iconst ty k1)
+            (iconst ty k2)))
+      (subsume (iconst ty (imm64_xor ty k1 k2))))
 
 (rule (simplify
        (bnot (fits_in_64 ty)
-            (iconst ty (u64_from_imm64 k))))
-      (subsume (iconst ty (imm64_masked ty (u64_not k)))))
+            (iconst ty k)))
+      (subsume (iconst ty (imm64_not ty k))))
 
 (rule (simplify (ishl (fits_in_64 ty)
                       (iconst ty k1)
@@ -166,25 +166,25 @@
 ;; - (add (sub x k1) k2) -> (add x (sub k2 k1))
 ;; - (add (sub k1 x) k2) -> (sub (add k1 k2) x)
 (rule (simplify (isub ty
-                      (isub ty x (iconst ty (u64_from_imm64 k1)))
-                      (iconst ty (u64_from_imm64 k2))))
-      (isub ty x (iconst ty (imm64_masked ty (u64_wrapping_add k1 k2)))))
+                      (isub ty x (iconst ty k1))
+                      (iconst ty k2)))
+      (isub ty x (iconst ty (imm64_add ty k1 k2))))
 (rule (simplify (isub ty
-                      (isub ty (iconst ty (u64_from_imm64 k1)) x)
-                      (iconst ty (u64_from_imm64 k2))))
-      (isub ty (iconst ty (imm64_masked ty (u64_wrapping_sub k1 k2))) x))
+                      (isub ty (iconst ty k1) x)
+                      (iconst ty k2)))
+      (isub ty (iconst ty (imm64_sub ty k1 k2)) x))
 (rule (simplify (isub ty
-                      (iadd ty x (iconst ty (u64_from_imm64 k1)))
-                      (iconst ty (u64_from_imm64 k2))))
-      (isub ty x (iconst ty (imm64_masked ty (u64_wrapping_sub k2 k1)))))
+                      (iadd ty x (iconst ty k1))
+                      (iconst ty k2)))
+      (isub ty x (iconst ty (imm64_sub ty k2 k1))))
 (rule (simplify (iadd ty
-                      (isub ty x (iconst ty (u64_from_imm64 k1)))
-                      (iconst ty (u64_from_imm64 k2))))
-      (iadd ty x (iconst ty (imm64_masked ty (u64_wrapping_sub k2 k1)))))
+                      (isub ty x (iconst ty k1))
+                      (iconst ty k2)))
+      (iadd ty x (iconst ty (imm64_sub ty k2 k1))))
 (rule (simplify (iadd ty
-                      (isub ty (iconst ty (u64_from_imm64 k1)) x)
-                      (iconst ty (u64_from_imm64 k2))))
-      (isub ty (iconst ty (imm64_masked ty (u64_wrapping_add k1 k2))) x))
+                      (isub ty (iconst ty k1) x)
+                      (iconst ty k2)))
+      (isub ty (iconst ty (imm64_add ty k1 k2)) x))
 
 (rule (simplify
        (imul ty (imul ty x k1 @ (iconst ty _)) k2 @ (iconst ty _)))
@@ -231,7 +231,7 @@
 (rule (simplify (isub ty x (iconst_s ty k)))
       (if-let true (u64_lt (i64_cast_unsigned (i64_wrapping_neg k))
                            (i64_cast_unsigned k)))
-      (iadd ty x (iconst ty (imm64_masked ty (i64_cast_unsigned (i64_wrapping_neg k))))))
+      (iadd ty x (iconst ty (imm64_neg ty (imm64 (i64_cast_unsigned k))))))
 
 ;; A splat of a constant can become a direct `vconst` with the appropriate bit
 ;; pattern.

cranelift/codegen/src/opts/selects.isle

@@ -94,10 +94,10 @@
 
 ;; fold add-select to select-add
 (rule (simplify
-        (iadd ty (select ty c (iconst_u ty x) (iconst_u ty y)) (iconst_u ty z)))
+        (iadd ty (select ty c (iconst ty x) (iconst ty y)) (iconst ty z)))
         (select ty c
-            (iconst ty (imm64_masked ty (u64_wrapping_add x z)))
-            (iconst ty (imm64_masked ty (u64_wrapping_add y z)))))
+            (iconst ty (imm64_add ty x z))
+            (iconst ty (imm64_add ty y z))))
 
 (rule (simplify (select ty d a (select ty d _ y))) (select ty d a y))
 (rule (simplify (select ty d (select ty d x _) a)) (select ty d x a))

cranelift/codegen/src/prelude.isle

@@ -73,9 +73,57 @@
 (decl pure partial imm64_sdiv (Type Imm64 Imm64) Imm64)
 (extern constructor imm64_sdiv imm64_sdiv)
 
+(decl pure partial imm64_udiv (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_udiv imm64_udiv)
+
 (decl pure partial imm64_srem (Type Imm64 Imm64) Imm64)
 (extern constructor imm64_srem imm64_srem)
 
+(decl pure partial imm64_urem (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_urem imm64_urem)
+
+(decl pure imm64_add (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_add imm64_add)
+
+(decl pure imm64_sub (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_sub imm64_sub)
+
+(decl pure imm64_mul (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_mul imm64_mul)
+
+(decl pure imm64_and (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_and imm64_and)
+
+(decl pure imm64_or (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_or imm64_or)
+
+(decl pure imm64_xor (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_xor imm64_xor)
+
+(decl pure imm64_not (Type Imm64) Imm64)
+(extern constructor imm64_not imm64_not)
+
+(decl pure imm64_neg (Type Imm64) Imm64)
+(extern constructor imm64_neg imm64_neg)
+
+(decl pure partial imm64_abs (Type Imm64) Imm64)
+(extern constructor imm64_abs imm64_abs)
+
+(decl pure partial imm64_ilog2 (Type Imm64) Imm64)
+(extern constructor imm64_ilog2 imm64_ilog2)
+
+(decl pure imm64_umin (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_umin imm64_umin)
+
+(decl pure imm64_umax (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_umax imm64_umax)
+
+(decl pure imm64_smin (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_smin imm64_smin)
+
+(decl pure imm64_smax (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_smax imm64_smax)
+
 (decl pure imm64_shl (Type Imm64 Imm64) Imm64)
 (extern constructor imm64_shl imm64_shl)
 

cranelift/filetests/filetests/egraph/issue-10409.clif

@@ -23,6 +23,6 @@ function u0:1(i64) -> i64 system_v {
         return v4
 }
 
-; check: v5 = ireduce.i32 v1
-; next: v6 = uextend.i64 v5
-; next: return v6
+; check: v7 = ireduce.i32 v1
+; next: v8 = uextend.i64 v7
+; next: return v8