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)
-