Cranelift: constant propagate shifts (#5671)

Thanks to Souper for pointing out we weren't doing this!

cranelift/codegen/src/isle_prelude.rs

@@ -80,6 +80,30 @@ macro_rules! isle_common_prelude_methods {
             x ^ y
         }
 
+        #[inline]
+        fn u64_shl(&mut self, x: u64, y: u64) -> u64 {
+            x << y
+        }
+
+        #[inline]
+        fn imm64_ushr(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let shift = u32::checked_sub(64, ty.bits()).unwrap_or(0);
+            let mask = u64::MAX >> shift;
+            let x = (x.bits() as u64) & mask;
+            let y = (y.bits() as u64) & mask;
+            Imm64::new((x >> y) as i64)
+        }
+
+        #[inline]
+        fn imm64_sshr(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            let shift = u32::checked_sub(64, ty.bits()).unwrap_or(0);
+            let mask = u64::MAX >> shift;
+            let x = (x.bits() as u64) & mask;
+            let y = (y.bits() as u64) & mask;
+            let sext = |v| ((v << shift) as i64) >> shift;
+            Imm64::new(sext(x) >> sext(y))
+        }
+
         #[inline]
         fn u64_not(&mut self, x: u64) -> u64 {
             !x

cranelift/codegen/src/opts/cprop.isle

@@ -55,6 +55,21 @@
             (iconst ty (u64_from_imm64 k))))
       (subsume (iconst ty (imm64_masked ty (u64_not k)))))
 
+(rule (simplify (ishl (fits_in_64 ty)
+                      (iconst ty (u64_from_imm64 k1))
+                      (iconst ty (u64_from_imm64 k2))))
+      (subsume (iconst ty (imm64_masked ty (u64_shl k1 k2)))))
+
+(rule (simplify (ushr (fits_in_64 ty)
+                      (iconst ty k1)
+                      (iconst ty k2)))
+      (subsume (iconst ty (imm64_ushr ty k1 k2))))
+
+(rule (simplify (sshr (fits_in_64 ty)
+                      (iconst ty k1)
+                      (iconst ty k2)))
+      (subsume (iconst ty (imm64_sshr ty k1 k2))))
+
 ;; Canonicalize via commutativity: push immediates to the right.
 ;;
 ;;   (op k x) --> (op x k)

cranelift/codegen/src/prelude.isle

@@ -129,6 +129,15 @@
 (decl pure u64_xor (u64 u64) u64)
 (extern constructor u64_xor u64_xor)
 
+(decl pure u64_shl (u64 u64) u64)
+(extern constructor u64_shl u64_shl)
+
+(decl pure imm64_ushr (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_ushr imm64_ushr)
+
+(decl pure imm64_sshr (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_sshr imm64_sshr)
+
 (decl pure u64_not (u64) u64)
 (extern constructor u64_not u64_not)
 

cranelift/filetests/filetests/egraph/cprop.clif

@@ -41,3 +41,35 @@ block0(v1: i8):
 ; check: v3 = iconst.i8 0
 ; check: return v3
 
+function %ishl() -> i8 {
+block0:
+    v0 = iconst.i8 1
+    v1 = iconst.i8 2
+    v2 = ishl v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 4
+; check: return v3
+
+function %ushr() -> i8 {
+block0:
+    v0 = iconst.i8 -1
+    v1 = iconst.i8 2
+    v2 = ushr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 63
+; check: return v3
+
+function %sshr() -> i8 {
+block0:
+    v0 = iconst.i8 0xf0
+    v1 = iconst.i8 2
+    v2 = sshr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 -4
+; check: return v3