From f99b3652dc402ff1b9e374171dcec5a39a12d85a Mon Sep 17 00:00:00 2001
From: Laurence Tratt <laurie@tratt.net>
Date: Mon, 18 Nov 2024 17:24:08 +0000
Subject: [PATCH] Skip intermediate `PtrAdd`s.

A common pattern we see is:

```
%7: ptr = ptr_add %0, 2
%8: ptr = ptr_add %7, 3
```

This can be annoying because `%7` sometimes isn't really used other than as a
vehicle for `%8`. This commit rewrites `%8` to:

```
%8: ptr = ptr_add %0, 5
```

This can reduce the pressure on the register allocator a little bit. Has
a small (0.5%), but measurable effect on big_loop.
---
 ykrt/src/compile/jitc_yk/opt/mod.rs | 39 +++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/ykrt/src/compile/jitc_yk/opt/mod.rs b/ykrt/src/compile/jitc_yk/opt/mod.rs
index 1e7c4a7f0..9d64f8c6a 100644
--- a/ykrt/src/compile/jitc_yk/opt/mod.rs
+++ b/ykrt/src/compile/jitc_yk/opt/mod.rs
@@ -326,6 +326,21 @@ impl Opt {
                     Operand::Var(op_iidx) => {
                         if x.off() == 0 {
                             self.m.replace(iidx, Inst::Copy(op_iidx));
+                        } else {
+                            // Skip intermediate `PtrAdds`. For example given:
+                            //   %7: ptr = ptr_add %0, 2
+                            //   %8: ptr = ptr_add %7, 3
+                            // we can replace the latter with:
+                            //   %8: ptr = ptr_add %0, 5
+                            if let Some(Inst::PtrAdd(back_inst)) = self.m.inst_nocopy(op_iidx) {
+                                self.m.replace(
+                                    iidx,
+                                    Inst::PtrAdd(PtrAddInst::new(
+                                        back_inst.ptr(&self.m),
+                                        back_inst.off().checked_add(x.off()).unwrap(),
+                                    )),
+                                );
+                            }
                         }
                     }
                 },
@@ -1022,4 +1037,28 @@ mod test {
         ",
         );
     }
+
+    #[test]
+    fn opt_ptr_adds() {
+        Module::assert_ir_transform_eq(
+            "
+          entry:
+            %0: ptr = load_ti 0
+            %1: ptr = ptr_add %0, 2
+            %2: ptr = ptr_add %1, 3
+            black_box %1
+            black_box %2
+        ",
+            |m| opt(m).unwrap(),
+            "
+          ...
+          entry:
+            %0: ptr = load_ti ...
+            %1: ptr = ptr_add %0, 2
+            %2: ptr = ptr_add %0, 5
+            black_box %1
+            black_box %2
+        ",
+        );
+    }
 }