From c9d1c068bc578b4f2095b3dc039ed0a2809caad6 Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Tue, 31 Jan 2023 10:04:17 -0800
Subject: [PATCH] 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         | 11 ++++++
 cranelift/codegen/src/opts/algebraic.isle     |  8 ++++-
 cranelift/codegen/src/prelude.isle            |  5 ++-
 .../filetests/filetests/egraph/algebraic.clif |  4 +--
 .../filetests/filetests/egraph/mul-pow-2.clif | 34 +++++++++++++++++++
 cranelift/filetests/src/subtest.rs            |  6 ++++
 6 files changed, 64 insertions(+), 4 deletions(-)
 create mode 100644 cranelift/filetests/filetests/egraph/mul-pow-2.clif

diff --git a/cranelift/codegen/src/isle_prelude.rs b/cranelift/codegen/src/isle_prelude.rs
index 0d69b158db..17bf26eb7d 100644
--- a/cranelift/codegen/src/isle_prelude.rs
+++ b/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 {
diff --git a/cranelift/codegen/src/opts/algebraic.isle b/cranelift/codegen/src/opts/algebraic.isle
index 9262b1b657..865c773bd8 100644
--- a/cranelift/codegen/src/opts/algebraic.isle
+++ 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
diff --git a/cranelift/codegen/src/prelude.isle b/cranelift/codegen/src/prelude.isle
index c8c8496ca8..50edc2b060 100644
--- a/cranelift/codegen/src/prelude.isle
+++ b/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)
-
diff --git a/cranelift/filetests/filetests/egraph/algebraic.clif b/cranelift/filetests/filetests/egraph/algebraic.clif
index 3a852b0ebb..409788ce82 100644
--- a/cranelift/filetests/filetests/egraph/algebraic.clif
+++ b/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
 }
 
diff --git a/cranelift/filetests/filetests/egraph/mul-pow-2.clif b/cranelift/filetests/filetests/egraph/mul-pow-2.clif
new file mode 100644
index 0000000000..e81ae49364
--- /dev/null
+++ b/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
+}
diff --git a/cranelift/filetests/src/subtest.rs b/cranelift/filetests/src/subtest.rs
index f5fa4e0837..b342929171 100644
--- a/cranelift/filetests/src/subtest.rs
+++ b/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)