Merge pull request #1632 from ltratt/opt_ptr_adds

Further optimise pointer adds
ykjit · Feb 27, 2025 · a797bdf · a797bdf
2 parents 6eac949 + 4c1d86c
commit a797bdf
Showing 1 changed file with 64 additions and 10 deletions.
diff --git a/ykrt/src/compile/jitc_yk/opt/mod.rs b/ykrt/src/compile/jitc_yk/opt/mod.rs
@@ -580,8 +580,9 @@ impl Opt {
                     Operand::Const(cidx) => self.m.replace(iidx, Inst::Const(cidx)),
                 }
             } else {
-                self.m
-                    .replace(iidx, Inst::PtrAdd(PtrAddInst::new(inst.ptr(&self.m), off)));
+                let pa_inst = PtrAddInst::new(inst.ptr(&self.m), off);
+                self.m.replace(iidx, Inst::PtrAdd(pa_inst));
+                self.opt_ptradd(iidx, pa_inst)?;
             }
         }
 
@@ -751,12 +752,30 @@ impl Opt {
         Ok(())
     }
 
-    fn opt_ptradd(&mut self, iidx: InstIdx, inst: PtrAddInst) -> Result<(), CompilationError> {
-        match self.an.op_map(&self.m, inst.ptr(&self.m)) {
-            Operand::Const(_) => todo!(),
-            Operand::Var(op_iidx) => {
-                if inst.off() == 0 {
-                    self.m.replace(iidx, Inst::Copy(op_iidx));
+    fn opt_ptradd(
+        &mut self,
+        iidx: InstIdx,
+        mut pa_inst: PtrAddInst,
+    ) -> Result<(), CompilationError> {
+        let mut off = 0;
+        loop {
+            off += pa_inst.off();
+            match self.an.op_map(&self.m, pa_inst.ptr(&self.m)) {
+                Operand::Const(_) => todo!(),
+                Operand::Var(op_iidx) => {
+                    if let Inst::PtrAdd(x) = self.m.inst(op_iidx) {
+                        pa_inst = x;
+                    } else {
+                        if off == 0 {
+                            self.m.replace(iidx, Inst::Copy(op_iidx));
+                        } else {
+                            self.m.replace(
+                                iidx,
+                                Inst::PtrAdd(PtrAddInst::new(Operand::Var(op_iidx), off)),
+                            );
+                        }
+                        break;
+                    }
                 }
             }
         }
@@ -1721,7 +1740,7 @@ mod test {
     }
 
     #[test]
-    fn opt_ptradd_zero() {
+    fn opt_ptradd() {
         Module::assert_ir_transform_eq(
             "
           entry:
@@ -1737,10 +1756,28 @@ mod test {
             black_box %0
         ",
         );
+
+        Module::assert_ir_transform_eq(
+            "
+          entry:
+            %0: ptr = param reg
+            %1: ptr = ptr_add %0, 4
+            %2: ptr = ptr_add %1, 4
+            %3: ptr = ptr_add %2, -8
+            black_box %3
+        ",
+            |m| opt(m).unwrap(),
+            "
+          ...
+          entry:
+            %0: ptr = param ...
+            black_box %0
+        ",
+        );
     }
 
     #[test]
-    fn opt_dynptradd_const() {
+    fn opt_dynptradd() {
         Module::assert_ir_transform_eq(
             "
           entry:
@@ -1771,6 +1808,23 @@ mod test {
             black_box 0x1234
         ",
         );
+
+        Module::assert_ir_transform_eq(
+            "
+          entry:
+            %0: ptr = param reg
+            %1: ptr = ptr_add %0, -4
+            %2: ptr = dyn_ptr_add %1, 1i64, 4
+            black_box %2
+        ",
+            |m| opt(m).unwrap(),
+            "
+          ...
+          entry:
+            %0: ptr = param ...
+            black_box %0
+        ",
+        );
     }
 
     #[test]