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