Fix test and perform removals after fixpoint

Author: jacob-shops

zjit/src/hir.rs

@@ -6743,6 +6743,13 @@ impl Function {
         // The cache is pruned when loads and stores can alias between objects.
         let mut cached_insns: Vec<HashMap<Key, InsnId>> = vec![HashMap::new(); cfi.num_blocks];
 
+        // An available expression is an expression that has already been computed with values that have not changed since the last computation.
+        // In this case, available expressions == redundant loads and stores
+        // To compute the available expressions, we need to know which expressions have been evaluated, and when effectful operations interrupt the optimization.
+        // This information is stored in cached_insns and is used to populate available_expressions.
+        // Once the fixpoint loop is finished, we elide all available expressions, concluding the load store forwarding pass.
+        let mut available_expressions: InsnSet = InsnSet::with_capacity(self.insns.len());
+
         loop {
             for (rpo_index, &block_id) in rpo.iter().enumerate() {
                 // Set block_cache equal to the intersection of cached insns of all predecessors
@@ -6767,22 +6774,17 @@ impl Function {
 
                     }
                 };
-                let old_insns = std::mem::take(&mut self.blocks[block_id].insns);
-                let mut new_insns = Vec::with_capacity(old_insns.len());
-                for insn_id in old_insns {
-                    let replacement_insn: InsnId = match self.resolve(insn_id).insn(self) {
+                let insns = std::mem::take(&mut self.blocks[block_id].insns);
+                for insn_id in insns.iter() {
+                    match self.resolve(*insn_id).insn(self) {
                         &Insn::StoreField { recv, offset, val, .. } => {
                             let key = Key { id: self.chase_insn(recv), offset };
                             let heap_entry = block_cache.get(&key).copied();
-                            // TODO(Jacob): Switch from actual to partial equality
                             if Some(val) == heap_entry {
-                                // If the value is already stored, short circuit and don't add an instruction to the block
-                                continue
+                                available_expressions.insert(*insn_id);
                             }
-                            // TODO(Jacob): Add type based alias analysis to avoid removing so many entries
                             block_cache.retain(|key, _| key.offset != offset);
                             block_cache.insert(key, val);
-                            insn_id
                         },
                         &Insn::LoadField { recv, offset, return_type, .. } => {
                             let key = Key { id: self.chase_insn(recv), offset };
@@ -6803,38 +6805,32 @@ impl Function {
                                     if can_forward_cached_insn {
                                         // If the value is stored already, we should short circuit.
                                         // However, we need to replace insn_id with its representative in the SSA union.
-                                        self.make_equal_to(insn_id, cached_insn);
-                                        continue
+                                        self.make_equal_to(*insn_id, cached_insn);
+                                        available_expressions.insert(*insn_id);
                                     }
                                 }
                                 Entry::Vacant(entry) => {
                                     // If the value has not been accessed, cache a copy to optimize future loads or stores.
-                                    entry.insert(insn_id);
+                                    entry.insert(*insn_id);
                                 }
                             }
-                            insn_id
                         }
                         &Insn::WriteBarrier { .. } => {
                             // Currently, WriteBarrier write effects are Allocator and Memory when we'd really like them to be flags.
                             // We don't use LoadField for mark bits so we can ignore them for now.
                             // But flags does not exist in our effects abstract heap modeling and we don't want to add special casing to effects.
                             // This special casing in this pass here should be removed once we refine our effects system to provide greater granularity for WriteBarrier.
-                            // TODO: use type based alias analysis
                             let offset = RUBY_OFFSET_RBASIC_FLAGS;
                             block_cache.retain(|key, _| key.offset != offset);
-                            insn_id
                         },
                         insn => {
                             if insn.effects_of().includes(Effect::write(abstract_heaps::Memory)) {
                                 block_cache.clear();
                             }
-                            insn_id
                         }
                     };
-                    new_insns.push(replacement_insn);
                 }
-
-                self.blocks[block_id].insns = new_insns;
+                self.blocks[block_id].insns = insns;
 
                 if cached_insns[block_id] == block_cache {
                     changed = false;
@@ -6848,6 +6844,10 @@ impl Function {
                 break;
             }
         }
+
+        for &block_id in rpo {
+            self.blocks[block_id].insns.retain(|&insn_id| !available_expressions.get(insn_id));
+        }
     }
 
     /// Fold a binary operator on fixnums.

zjit/src/hir/opt_tests.rs

@@ -26253,7 +26253,8 @@ mod hir_opt_tests {
              def read(val)
                even = val.even?
                a = @value
-               even ? a : @value + 1
+               b = even ? 0 : 1
+               a + b
              end
            end
 
@@ -26274,34 +26275,37 @@ mod hir_opt_tests {
           Jump bb3(v1, v3)
         bb2():
           EntryPoint JIT(0)
-          v8:BasicObject = LoadArg :self@0
-          v9:BasicObject = LoadArg :val@1
-          Jump bb3(v8, v9)
-        bb3(v13:BasicObject, v14:BasicObject):
-          v74:NilClass = Const Value(nil)
-          v73:NilClass = Const Value(nil)
+          v9:BasicObject = LoadArg :self@0
+          v10:BasicObject = LoadArg :val@1
+          Jump bb3(v9, v10)
+        bb3(v15:BasicObject, v16:BasicObject):
+          v85:NilClass = Const Value(nil)
+          v84:NilClass = Const Value(nil)
+          v83:NilClass = Const Value(nil)
           PatchPoint MethodRedefined(Integer@0x1008, even?@0x1010, cme:0x1018)
-          v67:Fixnum = GuardType v14, Fixnum recompile
-          v68:BoolExact = InvokeBuiltin leaf <inline_expr>, v67
+          v77:Fixnum = GuardType v16, Fixnum recompile
+          v78:BoolExact = InvokeBuiltin leaf <inline_expr>, v77
           PatchPoint NoEPEscape(read)
-          v27:HeapBasicObject = GuardType v13, HeapBasicObject
-          v28:CShape = LoadField v27, :shape_id@0x1040
-          v29:CShape[0x1041] = GuardBitEquals v28, CShape(0x1041) recompile
-          v30:BasicObject = LoadField v27, :@value@0x1042
-          v37:CBool = Test v68
-          v38:FalseClass = RefineType v68, Falsy
-          CondBranch v37, bb5(), bb4()
-        bb5():
-          v40:TrueClass = RefineType v68, Truthy
-          CheckInterrupts
-          Return v30
+          v30:HeapBasicObject = GuardType v15, HeapBasicObject
+          v31:CShape = LoadField v30, :shape_id@0x1040
+          v32:CShape[0x1041] = GuardBitEquals v31, CShape(0x1041) recompile
+          v33:BasicObject = LoadField v30, :@value@0x1042
+          v40:CBool = Test v78
+          v41:FalseClass = RefineType v78, Falsy
+          CondBranch v40, bb6(), bb4()
+        bb6():
+          v43:TrueClass = RefineType v78, Truthy
+          v45:Fixnum[0] = Const Value(0)
+          Jump bb5(v43, v45)
         bb4():
-          v57:Fixnum[1] = Const Value(1)
+          v55:Fixnum[1] = Const Value(1)
+          Jump bb5(v41, v55)
+        bb5(v59:BoolExact, v62:Fixnum):
           PatchPoint MethodRedefined(Integer@0x1008, +@0x1043, cme:0x1048)
-          v71:Fixnum = GuardType v30, Fixnum recompile
-          v72:Fixnum = FixnumAdd v71, v57
+          v81:Fixnum = GuardType v33, Fixnum recompile
+          v82:Fixnum = FixnumAdd v81, v62
           CheckInterrupts
-          Return v72
+          Return v82
         "
         );
     }