egraphs: Remove extends and reduces to the shift input in shift instr…

…uctions

We support variations of this instruction with different inputs, so lets take advantage of this where possible.

cranelift/codegen/src/opts/shifts.isle

@@ -102,3 +102,31 @@
 (rule (simplify (ineg ty (ushr ty x sconst @ (iconst ty (u64_from_imm64 shift_amt)))))
       (if-let $true (u64_eq shift_amt (u64_sub (ty_bits ty) 1)))
       (sshr ty x sconst))
+
+;; Shifts and rotates allow a different type for the shift amount, so we
+;; can remove any extend/reduce operations on the shift amount.
+;;
+;; (op x (ireduce y)) == (op x y)
+;; (op x (uextend y)) == (op x y)
+;; (op x (sextend y)) == (op x y)
+;;
+;; where `op` is one of ishl, ushr, sshr, rotl, rotr
+;;
+;; TODO: This rule is restricted to <=64 bits for ireduce since the x86
+;; backend doesn't support SIMD shifts with 128-bit shift amounts.
+
+(rule (simplify (ishl ty x (ireduce _ y @ (value_type (fits_in_64 _))))) (ishl ty x y))
+(rule (simplify (ishl ty x (uextend _ y))) (ishl ty x y))
+(rule (simplify (ishl ty x (sextend _ y))) (ishl ty x y))
+(rule (simplify (ushr ty x (ireduce _ y @ (value_type (fits_in_64 _))))) (ushr ty x y))
+(rule (simplify (ushr ty x (uextend _ y))) (ushr ty x y))
+(rule (simplify (ushr ty x (sextend _ y))) (ushr ty x y))
+(rule (simplify (sshr ty x (ireduce _ y @ (value_type (fits_in_64 _))))) (sshr ty x y))
+(rule (simplify (sshr ty x (uextend _ y))) (sshr ty x y))
+(rule (simplify (sshr ty x (sextend _ y))) (sshr ty x y))
+(rule (simplify (rotr ty x (ireduce _ y @ (value_type (fits_in_64 _))))) (rotr ty x y))
+(rule (simplify (rotr ty x (uextend _ y))) (rotr ty x y))
+(rule (simplify (rotr ty x (sextend _ y))) (rotr ty x y))
+(rule (simplify (rotl ty x (ireduce _ y @ (value_type (fits_in_64 _))))) (rotl ty x y))
+(rule (simplify (rotl ty x (uextend _ y))) (rotl ty x y))
+(rule (simplify (rotl ty x (sextend _ y))) (rotl ty x y))

cranelift/filetests/filetests/egraph/shifts.clif

@@ -315,3 +315,157 @@ block0(v0: i64):
     ; check: v8 = uextend.i64 v7
     ; check: return v8
 }
+
+function %ishl_amt_type_ireduce(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = ireduce.i8 v1
+    v3 = ishl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ishl v0, v1
+; check: return v4
+
+function %ishl_amt_type_sextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = sextend.i32 v1
+    v3 = ishl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ishl v0, v1
+; check: return v4
+
+function %ishl_amt_type_uextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = uextend.i32 v1
+    v3 = ishl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ishl v0, v1
+; check: return v4
+
+
+function %ushr_amt_type_ireduce(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = ireduce.i8 v1
+    v3 = ushr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ushr v0, v1
+; check: return v4
+
+function %ushr_amt_type_sextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = sextend.i32 v1
+    v3 = ushr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ushr v0, v1
+; check: return v4
+
+function %ushr_amt_type_uextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = uextend.i32 v1
+    v3 = ushr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = ushr v0, v1
+; check: return v4
+
+
+function %sshr_amt_type_ireduce(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = ireduce.i8 v1
+    v3 = sshr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = sshr v0, v1
+; check: return v4
+
+function %sshr_amt_type_sextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = sextend.i32 v1
+    v3 = sshr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = sshr v0, v1
+; check: return v4
+
+function %sshr_amt_type_uextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = uextend.i32 v1
+    v3 = sshr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = sshr v0, v1
+; check: return v4
+
+
+function %rotr_amt_type_ireduce(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = ireduce.i8 v1
+    v3 = rotr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotr v0, v1
+; check: return v4
+
+function %rotr_amt_type_sextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = sextend.i32 v1
+    v3 = rotr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotr v0, v1
+; check: return v4
+
+function %rotr_amt_type_uextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = uextend.i32 v1
+    v3 = rotr.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotr v0, v1
+; check: return v4
+
+
+function %rotl_amt_type_ireduce(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = ireduce.i8 v1
+    v3 = rotl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotl v0, v1
+; check: return v4
+
+function %rotl_amt_type_sextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = sextend.i32 v1
+    v3 = rotl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotl v0, v1
+; check: return v4
+
+function %rotl_amt_type_uextend(i8, i16) -> i8 {
+block0(v0: i8, v1: i16):
+    v2 = uextend.i32 v1
+    v3 = rotl.i8 v0, v2
+    return v3
+}
+
+; check: v4 = rotl v0, v1
+; check: return v4