Cranelift: Add egraph rule to rewrite x * C ==> x << log2(C) when C is a power of two (#5647)

cranelift/codegen/src/isle_prelude.rs

@@ -343,6 +343,17 @@ macro_rules! isle_common_prelude_methods {
             imm.bits() as u64
         }
 
+        #[inline]
+        fn imm64_power_of_two(&mut self, x: Imm64) -> Option<u64> {
+            let x = i64::from(x);
+            let x = u64::try_from(x).ok()?;
+            if x.is_power_of_two() {
+                Some(x.trailing_zeros().into())
+            } else {
+                None
+            }
+        }
+
         #[inline]
         fn u64_from_bool(&mut self, b: bool) -> u64 {
             if b {

cranelift/codegen/src/opts/algebraic.isle

@@ -144,6 +144,12 @@
 (rule (simplify (imul ty (iconst _ (simm32 2)) x))
       (iadd ty x x))
 
+;; x*c == x<<log2(c) when c is a power of two.
+(rule (simplify (imul ty x (iconst _ (imm64_power_of_two c))))
+      (ishl ty x (iconst ty (imm64 c))))
+(rule (simplify (imul ty (iconst _ (imm64_power_of_two c)) x))
+      (ishl ty x (iconst ty (imm64 c))))
+
 ;; x<<32>>32: uextend/sextend 32->64.
 (rule (simplify (ushr $I64 (ishl $I64 (uextend $I64 x @ (value_type $I32)) (iconst _ (simm32 32))) (iconst _ (simm32 32))))
       (uextend $I64 x))
@@ -151,7 +157,7 @@
 (rule (simplify (sshr $I64 (ishl $I64 (uextend $I64 x @ (value_type $I32)) (iconst _ (simm32 32))) (iconst _ (simm32 32))))
       (sextend $I64 x))
 
-;; TODO: strength reduction: mul/div to shifts
+;; TODO: strength reduction: div to shifts
 ;; TODO: div/rem by constants -> magic multiplications
 
 ;; Rematerialize ALU-op-with-imm and iconsts in each block where they're

cranelift/codegen/src/prelude.isle

@@ -350,6 +350,10 @@
 (decl nonzero_u64_from_imm64 (u64) Imm64)
 (extern extractor nonzero_u64_from_imm64 nonzero_u64_from_imm64)
 
+;; If the given `Imm64` is a power-of-two, extract its log2 value.
+(decl imm64_power_of_two (u64) Imm64)
+(extern extractor imm64_power_of_two imm64_power_of_two)
+
 ;; Create a new Imm64.
 (decl pure imm64 (u64) Imm64)
 (extern constructor imm64 imm64)
@@ -452,4 +456,3 @@
 ;;;; Automatic conversions ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (convert Offset32 u32 offset32_to_u32)
-

cranelift/filetests/filetests/egraph/algebraic.clif

@@ -7,8 +7,8 @@ function %f0(i32) -> i32 {
 block0(v0: i32):
     v1 = iconst.i32 2
     v2 = imul v0, v1
-    ; check: v3 = iadd v0, v0
-    ; check: return v3
+    ; check: v5 = ishl v0, v4  ; v4 = 1
+    ; check: return v5
     return v2
 }
 

cranelift/filetests/filetests/egraph/mul-pow-2.clif

@@ -0,0 +1,34 @@
+test optimize
+set opt_level=speed
+set use_egraphs=true
+target x86_64
+
+function %f0(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 4
+    v2 = imul v0, v1
+    ; check:  v3 = iconst.i32 2
+    ; nextln: v4 = ishl v0, v3
+    ; check:  return v4
+    return v2
+}
+
+function %f1(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 8
+    v2 = imul v0, v1
+    ; check:  v3 = iconst.i32 3
+    ; nextln: v4 = ishl v0, v3
+    ; check:  return v4
+    return v2
+}
+
+function %f2(i32) -> i32 {
+block0(v0: i32):
+    v1 = iconst.i32 16
+    v2 = imul v0, v1
+    ; check:  v3 = iconst.i32 4
+    ; nextln: v4 = ishl v0, v3
+    ; check:  return v4
+    return v2
+}

cranelift/filetests/src/subtest.rs

@@ -107,6 +107,12 @@ pub trait SubTest {
 
 /// Run filecheck on `text`, using directives extracted from `context`.
 pub fn run_filecheck(text: &str, context: &Context) -> anyhow::Result<()> {
+    log::debug!(
+        "Filecheck Input:\n\
+         =======================\n\
+         {text}\n\
+         ======================="
+    );
     let checker = build_filechecker(context)?;
     if checker
         .check(text, NO_VARIABLES)