egraphs: Normalize shift and rotate amounts

cranelift/codegen/src/opts/shifts.isle

@@ -286,3 +286,25 @@
                       (ishl ty x k @ (iconst _ (u64_from_imm64 k1)))))
       (if-let $true (u64_eq k2 (u64_sub (ty_bits_u64 (lane_type ty)) k1)))
       (rotl ty x k))
+
+;; Normalize the shift amount. Some rules can't fire unless the shift amount
+;; is normalized. This also helps us materialize fewer and smaller constants.
+;;
+;; (op x k) == (op x (and k (ty_shift_mask ty)))
+;;
+;; where `op` is one of ishl, ushr, sshr, rotl, rotr
+(rule (simplify (ishl ty x (iconst kty (u64_from_imm64 k))))
+      (if-let $false (u64_eq k (u64_and k (ty_shift_mask ty))))
+      (ishl ty x (iconst kty (imm64 (u64_and k (ty_shift_mask ty))))))
+(rule (simplify (ushr ty x (iconst kty (u64_from_imm64 k))))
+      (if-let $false (u64_eq k (u64_and k (ty_shift_mask ty))))
+      (ushr ty x (iconst kty (imm64 (u64_and k (ty_shift_mask ty))))))
+(rule (simplify (sshr ty x (iconst kty (u64_from_imm64 k))))
+      (if-let $false (u64_eq k (u64_and k (ty_shift_mask ty))))
+      (sshr ty x (iconst kty (imm64 (u64_and k (ty_shift_mask ty))))))
+(rule (simplify (rotr ty x (iconst kty (u64_from_imm64 k))))
+      (if-let $false (u64_eq k (u64_and k (ty_shift_mask ty))))
+      (rotr ty x (iconst kty (imm64 (u64_and k (ty_shift_mask ty))))))
+(rule (simplify (rotl ty x (iconst kty (u64_from_imm64 k))))
+      (if-let $false (u64_eq k (u64_and k (ty_shift_mask ty))))
+      (rotl ty x (iconst kty (imm64 (u64_and k (ty_shift_mask ty))))))

cranelift/filetests/filetests/egraph/shifts.clif

@@ -74,9 +74,9 @@ block0(v0: i8):
     v2 = sshr v0, v1
     v3 = ishl v2, v1
     return v3
-    ; check: v4 = iconst.i8 224
-    ; check: v5 = band v0, v4
-    ; check: return v5
+    ; check: v8 = iconst.i8 224
+    ; check: v9 = band v0, v8
+    ; check: return v9
 }
 
 function %sextend_shift_32_64_unsigned(i32) -> i64 {
@@ -596,9 +596,9 @@ block0(v0: i128):
     return v3
 }
 
-; check: v4 = iconst.i64 0
-; nextln: v5 = uextend.i128 v4  ; v4 = 0
-; nextln: return v5
+; check: v7 = iconst.i64 0
+; nextln: v8 = uextend.i128 v7
+; nextln: return v8
 
 
 function %i128_ishl_overflow_becomes_0(i128) -> i128 {
@@ -609,9 +609,9 @@ block0(v0: i128):
     return v3
 }
 
-; check: v4 = iconst.i64 0
-; nextln: v5 = uextend.i128 v4  ; v4 = 0
-; nextln: return v5
+; check: v7 = iconst.i64 0
+; nextln: v8 = uextend.i128 v7
+; nextln: return v8
 
 function %simd_shift_does_not_panic(i8x16) -> i8x16 {
 block0(v0: i8x16):
@@ -636,11 +636,10 @@ block0(v0: i8x16):
     return v5
 }
 
-; check: v1 = iconst.i16 30
-; nextln: v3 = sshr v0, v1  ; v1 = 30
-; nextln: v6 = iconst.i16 6
-; nextln: v7 = sshr v3, v6  ; v6 = 6
-; check: return v7
+; check: v6 = iconst.i16 6
+; check: v7 = sshr v0, v6
+; check: v9 = sshr v7, v6
+; check: return v9
 
 
 function %rotr_rotr_prop(i8) -> i8 {
@@ -801,3 +800,14 @@ block0(v0: i64):
 ; check: v2 = iconst.i32 61
 ; nextln: v6 = rotl v0, v2  ; v2 = 61
 ; check: return v6
+
+function %ishl_normalize_amount(i64) -> i64 {
+block0(v0: i64):
+    v1 = iconst.i32 99
+    v2 = ishl v0, v1
+    return v2
+}
+
+; check: v3 = iconst.i32 35
+; check: v4 = ishl v0, v3
+; check: return v4