Get load_store optimization working

Author: jacob-shops

zjit/src/hir.rs

@@ -1184,7 +1184,7 @@ impl Insn {
             Insn::LoadSelf { .. } => Effect::read_write(abstract_heaps::Frame, abstract_heaps::Empty),
             Insn::LoadField { .. } => Effect::read_write(abstract_heaps::Memory, abstract_heaps::Empty),
             Insn::StoreField { .. } => effects::Any,
-            Insn::WriteBarrier { .. } => effects::Any,
+            Insn::WriteBarrier { .. } => effects::Allocator,
             Insn::SetLocal { .. } => effects::Any,
             Insn::GetSpecialSymbol { .. } => effects::Any,
             Insn::GetSpecialNumber { .. } => effects::Any,
@@ -1285,6 +1285,9 @@ impl Insn {
     /// Note: These are restrictions on the `write` `EffectSet` only. Even instructions with
     /// `read: effects::Any` could potentially be omitted.
     fn is_elidable(&self) -> bool {
+        if let Insn::WriteBarrier { .. } = self {
+            return false
+        }
         abstract_heaps::Allocator.includes(self.effects_of().write_bits())
     }
 }
@@ -4616,41 +4619,35 @@ impl Function {
         // The key for the hashmap should be type and offset, with a value of value
         // This lets us to index in with both load and store fields since insn_ids are probably always going to be different and we can't easily match on that
         // So... how do we match against and store the enum label without all the data?? not sure yet :/
-        eprintln!("{}", FunctionPrinter::with_snapshot(self));
         let mut compile_time_heap: HashMap<(InsnId, i32), InsnId>  = HashMap::new();
         for block in self.rpo() {
             let old_insns = std::mem::take(&mut self.blocks[block.0].insns);
             let mut new_insns = vec![];
             for insn_id in old_insns {
                 let replacement_insn: InsnId = match self.find(insn_id) {
                     Insn::StoreField { recv, offset, val, .. } => {
-                        let key = (recv, offset);
-                        eprintln!("hi jane");
+                        let key = (self.chase_insn(recv), offset);
                         let heap_entry = compile_time_heap.get(&key).copied();
-                        eprintln!("{heap_entry:?}");
                         // TODO(Jacob): Switch from actual to partial equality
                         if Some(val) == heap_entry {
-                            eprintln!("Matched value {val}, {heap_entry:?}");
                             // TODO(Jacob): Add TBAA to avoid removing so many entries
-                            eprintln!("Erasing aliasing offsets {offset}");
                             compile_time_heap.retain(|(_, off), _| *off != offset);
                             // If the value is already stored, short circuit and don't add an instruction to the block
                             continue
                         }
                         // TODO(Jacob): Add TBAA to avoid removing so many entries
-                        eprintln!("Erasing aliasing offsets {offset}");
                         compile_time_heap.retain(|(_, off), _| *off != offset);
                         compile_time_heap.insert(key, val);
-                        eprintln!("Inserted into heap {key:?}, {val}");
                         insn_id
                     },
                     Insn::LoadField { recv, offset, .. } => {
-                        let key = (recv, offset);
+                        let key = (self.chase_insn(recv), offset);
                         match compile_time_heap.entry(key) {
                             std::collections::hash_map::Entry::Occupied(entry) => {
                                 // If the value is already saved, we can't short circuit like we can with a store.
                                 // However, we can avoid the load with a reference to the representative to the union from SSA.
                                 self.make_equal_to(insn_id, *entry.get());
+                                continue
                             }
                             std::collections::hash_map::Entry::Vacant(_) => {
                                 // TODO(Jacob): Make sure this is correct, could be wrong?
@@ -4662,7 +4659,6 @@ impl Function {
                     insn => {
                         // If an instruction affects memory and we haven't modeled it, the compile_time_heap is invalidated
                         if insn.effects_of().includes(Effect::write(abstract_heaps::Memory)) {
-                            eprintln!("Clearing... {insn:?}");
                             compile_time_heap.clear();
                         }
                         insn_id

zjit/src/hir/opt_tests.rs

@@ -5205,8 +5205,6 @@ mod hir_opt_tests {
           v14:HeapBasicObject = RefineType v6, HeapBasicObject
           v17:Fixnum[2] = Const Value(2)
           PatchPoint SingleRactorMode
-          v36:CShape = LoadField v14, :_shape_id@0x1000
-          v37:CShape[0x1003] = GuardBitEquals v36, CShape(0x1003)
           StoreField v14, :@bar@0x1004, v17
           WriteBarrier v14, v17
           v40:CShape[0x1005] = Const CShape(0x1005)
@@ -8761,7 +8759,7 @@ mod hir_opt_tests {
             end
             test([])
         ");
-        assert_snapshot!(hir_string("test"), @"
+        assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:3:
         bb1():
           EntryPoint interpreter
@@ -8782,8 +8780,7 @@ mod hir_opt_tests {
           v33:ArrayExact = GuardType v10, ArrayExact
           v34:CUInt64 = LoadField v33, :_rbasic_flags@0x1040
           v35:CUInt64 = GuardNoBitsSet v34, RUBY_FL_FREEZE=CUInt64(2048)
-          v36:CUInt64 = LoadField v33, :_rbasic_flags@0x1040
-          v37:CUInt64 = GuardNoBitsSet v36, RUBY_ELTS_SHARED=CUInt64(4096)
+          v37:CUInt64 = GuardNoBitsSet v34, RUBY_ELTS_SHARED=CUInt64(4096)
           v38:CInt64[1] = UnboxFixnum v17
           v39:CInt64 = ArrayLength v33
           v40:CInt64[1] = GuardLess v38, v39
@@ -8805,7 +8802,7 @@ mod hir_opt_tests {
             end
             test([], 0, 1)
         ");
-        assert_snapshot!(hir_string("test"), @"
+        assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:3:
         bb1():
           EntryPoint interpreter
@@ -8829,8 +8826,7 @@ mod hir_opt_tests {
           v38:Fixnum = GuardType v15, Fixnum
           v39:CUInt64 = LoadField v37, :_rbasic_flags@0x1040
           v40:CUInt64 = GuardNoBitsSet v39, RUBY_FL_FREEZE=CUInt64(2048)
-          v41:CUInt64 = LoadField v37, :_rbasic_flags@0x1040
-          v42:CUInt64 = GuardNoBitsSet v41, RUBY_ELTS_SHARED=CUInt64(4096)
+          v42:CUInt64 = GuardNoBitsSet v39, RUBY_ELTS_SHARED=CUInt64(4096)
           v43:CInt64 = UnboxFixnum v38
           v44:CInt64 = ArrayLength v37
           v45:CInt64 = GuardLess v43, v44
@@ -13515,17 +13511,13 @@ mod hir_opt_tests {
           v14:HeapBasicObject = RefineType v6, HeapBasicObject
           v17:Fixnum[2] = Const Value(2)
           PatchPoint SingleRactorMode
-          v50:CShape = LoadField v14, :_shape_id@0x1000
-          v51:CShape[0x1003] = GuardBitEquals v50, CShape(0x1003)
           StoreField v14, :@b@0x1004, v17
           WriteBarrier v14, v17
           v54:CShape[0x1005] = Const CShape(0x1005)
           StoreField v14, :_shape_id@0x1000, v54
           v21:HeapBasicObject = RefineType v14, HeapBasicObject
           v24:Fixnum[3] = Const Value(3)
           PatchPoint SingleRactorMode
-          v57:CShape = LoadField v21, :_shape_id@0x1000
-          v58:CShape[0x1005] = GuardBitEquals v57, CShape(0x1005)
           StoreField v21, :@c@0x1006, v24
           WriteBarrier v21, v24
           v61:CShape[0x1007] = Const CShape(0x1007)
@@ -13735,9 +13727,6 @@ mod hir_opt_tests {
           v20:HeapBasicObject = RefineType v8, HeapBasicObject
           PatchPoint NoEPEscape(initialize)
           PatchPoint SingleRactorMode
-          v43:CShape = LoadField v20, :_shape_id@0x1000
-          v44:CShape[0x1003] = GuardBitEquals v43, CShape(0x1003)
-          StoreField v20, :@a@0x1002, v13
           WriteBarrier v20, v13
           CheckInterrupts
           Return v13