egraphs: Add some select optimizations (#6843)

* egraphs: Remove extends before selects when possible

* egraphs: Rewrite `select+{i,fcmp}` when select arms match cmp result

* egraphs: Simplify double `bmask`

* egraphs: Simplify double `bswap`

* egraphs: Modularize select optimizations

Co-Authored-By: Jamey Sharp <jsharp@fastly.com>

* egraphs: Remove `bmask+{i,f}cmp` optimizations

These were unsound since we only represent the results of `{i,f}cmp` as
`-1` for vectors, and not for scalars, scalars get the output of `1`

* egraphs: Optimize `bitrev` in `select` and `bmask`

* egraphs: Optimize `popcnt` in `select` and `bmask`

* egraphs: Optimize double `bitrev`

* egraphs: Add `thruthy` expression matcher

Cleans up the previous logic a bit.

This also adds `rotl`/`rotr` as thruth preserving expressions. As well as
`select` when it's branches match certain constants.

Co-Authored-By: Jamey Sharp <jsharp@fastly.com>

* egraphs: Make `icmp ne x, 0` a thruthy expression

* egraphs: Add `icmp ne x, 0` as an expression that can delete thruthy inputs

* egraphs: Add tests for comutative versions of the `icmp ne` thruthy rules

* egraphs: Restrict `icmp ne` optimization to arms with the

This also cleans up some merge conflicts that happened when rebasing.

* egraphs: Restore some accidentally deleted tests

* egraphs: Use `t` and `f` as names for select arms

Co-authored-by: Jamey Sharp <jamey@minilop.net>

* egraphs: Allow `icmp ne` opt to match arms of different types

Co-authored-by: Jamey Sharp <jamey@minilop.net>

---------

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

Author: afonso360

cranelift/codegen/src/opts/bitops.isle

@@ -104,3 +104,44 @@
 (rule (simplify (sshr ty (bor ty (ineg ty x) x) (iconst ty (u64_from_imm64 shift_amt))))
       (if-let $true (u64_eq shift_amt (u64_sub (ty_bits ty) 1)))
       (bmask ty x))
+
+;; Matches any expressions that preserve "truthiness".
+;; i.e. If the input is zero it remains zero, and if it is nonzero it can have
+;; a different value as long as it is still nonzero.
+(decl pure multi truthy (Value) Value)
+(rule (truthy (sextend _ x)) x)
+(rule (truthy (uextend _ x)) x)
+(rule (truthy (bmask _ x)) x)
+(rule (truthy (ineg _ x)) x)
+(rule (truthy (bswap _ x)) x)
+(rule (truthy (bitrev _ x)) x)
+(rule (truthy (popcnt _ x)) x)
+(rule (truthy (rotl _ x _)) x)
+(rule (truthy (rotr _ x _)) x)
+(rule (truthy (select _ x (iconst _ (u64_from_imm64 (u64_nonzero _))) (iconst _ (u64_from_imm64 0)))) x)
+;; (ne ty (iconst 0) v) is also canonicalized into this form via another rule
+(rule (truthy (ne _ x (iconst _ (u64_from_imm64 0)))) x)
+
+;; All of these expressions don't care about their input as long as it is truthy.
+;; so we can remove expressions that preserve that property from the input.
+(rule (simplify (bmask ty v)) (if-let x (truthy v)) (bmask ty x))
+(rule (simplify (select ty v t f)) (if-let c (truthy v)) (select ty c t f))
+;; (ne ty (iconst 0) v) is also canonicalized into this form via another rule
+(rule (simplify (ne cty v (iconst _ (u64_from_imm64 0))))
+      (if-let c (truthy v))
+      (if-let (value_type ty) c)
+      (ne cty c (iconst ty (imm64 0))))
+
+
+
+;; (sextend (bmask x)) can be replaced with (bmask x) since bmask
+;; supports any size of output type, regardless of input.
+;; Same with `ireduce`
+(rule (simplify (sextend ty (bmask _ x))) (bmask ty x))
+(rule (simplify (ireduce ty (bmask _ x))) (bmask ty x))
+
+;; (bswap (bswap x)) == x
+(rule (simplify (bswap ty (bswap ty x))) (subsume x))
+
+;; (bitrev (bitrev x)) == x
+(rule (simplify (bitrev ty (bitrev ty x))) (subsume x))

cranelift/filetests/filetests/egraph/bitops.clif

@@ -178,3 +178,197 @@ block0(v0: i64):
 
 ; check: v5 = bmask.i64 v0
 ; check: return v5
+
+
+function %double_bmask(i16) -> i64 {
+block0(v0: i16):
+    v1 = bmask.i32 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_sextend(i16) -> i64 {
+block0(v0: i16):
+    v1 = sextend.i32 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_uextend(i16) -> i64 {
+block0(v0: i16):
+    v1 = uextend.i32 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_ineg(i16) -> i64 {
+block0(v0: i16):
+    v1 = ineg.i16 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_bswap(i16) -> i64 {
+block0(v0: i16):
+    v1 = bswap.i16 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_bitrev(i16) -> i64 {
+block0(v0: i16):
+    v1 = bitrev.i16 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_popcnt(i16) -> i64 {
+block0(v0: i16):
+    v1 = popcnt.i16 v0
+    v2 = bmask.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %bmask_rotl(i16, i16) -> i64 {
+block0(v0: i16, v1: i16):
+    v2 = rotl.i16 v0, v1
+    v3 = bmask.i64 v2
+    return v3
+}
+
+; check: v4 = bmask.i64 v0
+; check: return v4
+
+function %bmask_rotr(i16, i16) -> i64 {
+block0(v0: i16, v1: i16):
+    v2 = rotr.i16 v0, v1
+    v3 = bmask.i64 v2
+    return v3
+}
+
+; check: v4 = bmask.i64 v0
+; check: return v4
+
+function %bmask_select_non_zero(i16) -> i64 {
+block0(v0: i16):
+    v1 = iconst.i16 123
+    v2 = iconst.i16 0
+    v3 = select v0, v1, v2
+    v4 = bmask.i64 v3
+    return v4
+}
+
+; check: v5 = bmask.i64 v0
+; check: return v5
+
+function %bmask_icmp_ne_zero(i16) -> i64 {
+block0(v0: i16):
+    v1 = iconst.i16 0
+    v2 = icmp ne v0, v1
+    v3 = bmask.i64 v2
+    return v3
+}
+
+; check: v4 = bmask.i64 v0
+; check: return v4
+
+function %icmp_ne_zero_deletes_thruthy_input(i16) -> i8 {
+block0(v0: i16):
+    v1 = popcnt v0
+    v2 = iconst.i16 0
+    v3 = icmp ne v1, v2
+    return v3
+}
+
+; check: v2 = iconst.i16 0
+; check: v4 = icmp ne v0, v2
+; check: return v4
+
+
+function %icmp_ne_zero_deletes_thruthy_input_const_lhs(i16) -> i8 {
+block0(v0: i16):
+    v1 = iconst.i16 0
+    v2 = popcnt v0
+    v3 = icmp ne v1, v2
+    return v3
+}
+
+; check: v1 = iconst.i16 0
+; check: v5 = icmp ne v0, v1
+; check: return v5
+
+
+function %icmp_ne_matches_arms_of_different_types(i16) -> i8 {
+block0(v0: i16):
+    v1 = iconst.i64 0
+    v2 = uextend.i64 v0
+    v3 = icmp ne v1, v2
+    return v3
+}
+
+; check: v5 = iconst.i16 0
+; check: v6 = icmp ne v0, v5
+; check: return v6
+
+
+
+function %sextend_bmask(i16) -> i64 {
+block0(v0: i16):
+    v1 = bmask.i32 v0
+    v2 = sextend.i64 v1
+    return v2
+}
+
+; check: v3 = bmask.i64 v0
+; check: return v3
+
+function %ireduce_bmask(i16) -> i8 {
+block0(v0: i16):
+    v1 = bmask.i32 v0
+    v2 = ireduce.i8 v1
+    return v2
+}
+
+; check: v3 = bmask.i8 v0
+; check: return v3
+
+
+function %double_bswap(i64) -> i64 {
+block0(v0: i64):
+    v1 = bswap.i64 v0
+    v2 = bswap.i64 v1
+    return v2
+}
+
+; check: return v0
+
+
+function %double_bitrev(i64) -> i64 {
+block0(v0: i64):
+    v1 = bitrev.i64 v0
+    v2 = bitrev.i64 v1
+    return v2
+}
+
+; check: return v0