Cranelift: Correctly wrap shifts in constant propagation (#5695)

Fixes #5690
Fixes #5696

Co-authored-by: Jamey Sharp <jsharp@fastly.com>

cranelift/codegen/src/isle_prelude.rs

@@ -81,27 +81,42 @@ macro_rules! isle_common_prelude_methods {
         }
 
         #[inline]
-        fn u64_shl(&mut self, x: u64, y: u64) -> u64 {
-            x << y
+        fn imm64_shl(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as u64;
+            let y = (y.bits() as u64) & shift_mask;
+
+            // Mask the result to `ty` bits.
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x.bits() << y) & ty_mask)
         }
 
         #[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;
+            let ty_mask = self.ty_mask(ty);
+            let x = (x.bits() as u64) & ty_mask;
+
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as u64;
+            let y = (y.bits() as u64) & shift_mask;
+
+            // NB: No need to mask off high bits because they are already zero.
             Imm64::new((x >> y) as i64)
         }
 
         #[inline]
         fn imm64_sshr(&mut self, ty: Type, x: Imm64, y: Imm64) -> Imm64 {
+            // Sign extend `x` from `ty.bits()`-width to the full 64 bits.
             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))
+            let x = (x.bits() << shift) >> shift;
+
+            // Mask off any excess shift bits.
+            let shift_mask = (ty.bits() - 1) as i64;
+            let y = y.bits() & shift_mask;
+
+            // Mask off sign bits that aren't part of `ty`.
+            let ty_mask = self.ty_mask(ty) as i64;
+            Imm64::new((x >> y) & ty_mask)
         }
 
         #[inline]
@@ -179,14 +194,12 @@ macro_rules! isle_common_prelude_methods {
 
         #[inline]
         fn ty_mask(&mut self, ty: Type) -> u64 {
-            match ty.bits() {
-                1 => 1,
-                8 => 0xff,
-                16 => 0xffff,
-                32 => 0xffff_ffff,
-                64 => 0xffff_ffff_ffff_ffff,
-                _ => unimplemented!(),
-            }
+            let ty_bits = ty.bits();
+            debug_assert_ne!(ty_bits, 0);
+            let shift = 64_u64
+                .checked_sub(ty_bits.into())
+                .expect("unimplemented for > 64 bits");
+            u64::MAX >> shift
         }
 
         fn fits_in_16(&mut self, ty: Type) -> Option<Type> {

cranelift/codegen/src/opts/algebraic.isle

@@ -191,15 +191,15 @@
 ;; `(x >> k) << k` is the same as masking off the bottom `k` bits (regardless if
 ;; this is a signed or unsigned shift right).
 (rule (simplify (ishl (fits_in_64 ty)
-                      (ushr ty x (iconst _ (u64_from_imm64 k)))
-                      (iconst _ (u64_from_imm64 k))))
-      (let ((mask u64 (u64_shl 0xFFFFFFFFFFFFFFFF k)))
-        (band ty x (iconst ty (imm64_masked ty mask)))))
+                      (ushr ty x (iconst _ k))
+                      (iconst _ k)))
+      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
+        (band ty x (iconst ty mask))))
 (rule (simplify (ishl (fits_in_64 ty)
-                      (sshr ty x (iconst _ (u64_from_imm64 k)))
-                      (iconst _ (u64_from_imm64 k))))
-      (let ((mask u64 (u64_shl 0xFFFFFFFFFFFFFFFF k)))
-        (band ty x (iconst ty (imm64_masked ty mask)))))
+                      (sshr ty x (iconst _ k))
+                      (iconst _ k)))
+      (let ((mask Imm64 (imm64_shl ty (imm64 0xFFFF_FFFF_FFFF_FFFF) k)))
+        (band ty x (iconst ty mask))))
 
 ;; Rematerialize ALU-op-with-imm and iconsts in each block where they're
 ;; used. This is neutral (add-with-imm) or positive (iconst) for

cranelift/codegen/src/opts/cprop.isle

@@ -56,18 +56,18 @@
       (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)))))
+                      (iconst ty k1)
+                      (iconst _ k2)))
+      (subsume (iconst ty (imm64_shl ty k1 k2))))
 
 (rule (simplify (ushr (fits_in_64 ty)
                       (iconst ty k1)
-                      (iconst ty k2)))
+                      (iconst _ k2)))
       (subsume (iconst ty (imm64_ushr ty k1 k2))))
 
 (rule (simplify (sshr (fits_in_64 ty)
                       (iconst ty k1)
-                      (iconst ty k2)))
+                      (iconst _ k2)))
       (subsume (iconst ty (imm64_sshr ty k1 k2))))
 
 (rule (simplify

cranelift/codegen/src/prelude.isle

@@ -129,8 +129,8 @@
 (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_shl (Type Imm64 Imm64) Imm64)
+(extern constructor imm64_shl imm64_shl)
 
 (decl pure imm64_ushr (Type Imm64 Imm64) Imm64)
 (extern constructor imm64_ushr imm64_ushr)

cranelift/filetests/filetests/egraph/algebraic.clif

@@ -87,6 +87,17 @@ block0(v0: i64):
     ; return v5
 }
 
+function %signed_shift_right_shift_left_i8_mask_rhs(i8) -> i8 {
+block0(v0: i8):
+    v1 = iconst.i8 0xf5
+    v2 = sshr v0, v1
+    v3 = ishl v2, v1
+    return v3
+    ; check: v4 = iconst.i8 224
+    ; check: v5 = band v0, v4
+    ; return v5
+}
+
 function %or_and_y_with_not_y_i8(i8, i8) -> i8 {
 block0(v0: i8, v1: i8):
     v2 = band v0, v1

cranelift/filetests/filetests/egraph/cprop.clif

@@ -34,6 +34,28 @@ block0:
 ; check: v3 = iconst.i8 4
 ; check: return v3
 
+function %ishl_i8_i16() -> i8 {
+block0:
+    v0 = iconst.i8 1
+    v1 = iconst.i16 0xf2
+    v2 = ishl v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 4
+; check: return v3
+
+function %ishl_i16_i8() -> i16 {
+block0:
+    v0 = iconst.i16 1
+    v1 = iconst.i8 0xf2
+    v2 = ishl v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i16 4
+; check: return v3
+
 function %ushr() -> i8 {
 block0:
     v0 = iconst.i8 -1
@@ -45,6 +67,28 @@ block0:
 ; check: v3 = iconst.i8 63
 ; check: return v3
 
+function %ushr_i8_i16() -> i8 {
+block0:
+    v0 = iconst.i8 -1
+    v1 = iconst.i16 0xf2
+    v2 = ushr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 63
+; check: return v3
+
+function %ushr_i16_i8() -> i16 {
+block0:
+    v0 = iconst.i16 -1
+    v1 = iconst.i8 0xf2
+    v2 = ushr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i16 0x3fff
+; check: return v3
+
 function %sshr() -> i8 {
 block0:
     v0 = iconst.i8 0xf0
@@ -53,7 +97,29 @@ block0:
     return v2
 }
 
-; check: v3 = iconst.i8 -4
+; check: v3 = iconst.i8 252
+; check: return v3
+
+function %sshr_i8_i16() -> i8 {
+block0:
+    v0 = iconst.i8 0xf0
+    v1 = iconst.i16 0xf2
+    v2 = sshr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i8 252
+; check: return v3
+
+function %sshr_i16_i8() -> i16 {
+block0:
+    v0 = iconst.i16 0xfff0
+    v1 = iconst.i8 0xf2
+    v2 = sshr v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i16 0xfffc
 ; check: return v3
 
 function %icmp_eq_i32() -> i8 {

cranelift/filetests/filetests/runtests/issue-5690.clif

@@ -0,0 +1,29 @@
+test interpret
+test run
+set opt_level=speed
+set enable_simd=true
+set enable_safepoints=true
+set unwind_info=false
+set preserve_frame_pointers=true
+set machine_code_cfg_info=true
+set enable_table_access_spectre_mitigation=false
+target aarch64
+target x86_64
+
+function %u1() -> i64 sext, f64, i8, i8 sext, i8 sext system_v {
+block0:
+    v0 = f64const 0x1.8373638ff3738p-124
+    v1 = iconst.i8 53
+    v2 = iconst.i64 0x4445_00ff_ffff_ffff
+    v3 = iconst.i8 0
+    v4 = iconst.i16 0
+    v5 = iconst.i32 0
+    v6 = iconst.i64 0
+    v7 = uextend.i128 v6
+    v8 = ishl v2, v2
+    v9 = rotr v1, v1
+    nop
+    return v8, v0, v9, v9, v9
+}
+
+; run: %u1() == [-9223372036854775808, 0x1.8373638ff3738p-124, -87, -87, -87]

cranelift/filetests/filetests/wasm/issue-5696.wat

@@ -0,0 +1,20 @@
+;;! target = "x86_64"
+;;! optimize = true
+;;! settings = ["opt_level=speed"]
+
+(module
+  (func (;0;) (param i64) (result i64)
+    i64.const 32
+    i64.const -19
+    i64.shr_u
+    ;; call 0
+  )
+)
+;; function u0:0(i64, i64 vmctx) -> i64 fast {
+;;                                 block0(v0: i64, v1: i64):
+;; @001e                               jump block1
+;;
+;;                                 block1:
+;;                                     v6 = iconst.i64 0
+;; @001e                               return v6  ; v6 = 0
+;; }

cranelift/filetests/src/test_wasm.rs

@@ -24,6 +24,7 @@ pub fn run(path: &Path, wat: &str) -> Result<()> {
 
     let config: TestConfig =
         toml::from_str(&config_text).context("failed to parse the test configuration")?;
+    log::debug!("Wasm test config = {config:#?}");
 
     config
         .validate()

cranelift/wasm/src/func_translator.rs

@@ -117,6 +117,7 @@ impl FuncTranslator {
         parse_function_body(validator, reader, &mut builder, &mut self.state, environ)?;
 
         builder.finalize();
+        log::trace!("translated Wasm to CLIF:\n{}", func.display());
         Ok(())
     }
 }