[maglev] Extend lifetimes of values used in a loop

While marking uses, record what values are used inside a loop, but
defined outside of it. Then, on the loop end, extend the lifetime of
these values.

Bug: v8:7700
Change-Id: I1bba037be760b4871673ecf0af584f5bf72fc35c
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/3782797Reviewed-by: Toon Verwaest <verwaest@chromium.org>
Commit-Queue: Leszek Swirski <leszeks@chromium.org>
Cr-Commit-Position: refs/heads/main@{#82040}

src/maglev/maglev-compiler.cc

@@ -62,8 +62,15 @@ class NumberingProcessor {
 class UseMarkingProcessor {
  public:
   void PreProcessGraph(MaglevCompilationInfo*, Graph* graph) {}
-  void PostProcessGraph(MaglevCompilationInfo*, Graph* graph) {}
-  void PreProcessBasicBlock(MaglevCompilationInfo*, BasicBlock* block) {}
+  void PostProcessGraph(MaglevCompilationInfo*, Graph* graph) {
+    DCHECK(loop_used_nodes_.empty());
+  }
+  void PreProcessBasicBlock(MaglevCompilationInfo*, BasicBlock* block) {
+    if (!block->has_state()) return;
+    if (block->state()->is_loop()) {
+      loop_used_nodes_.push_back(LoopUsedNodes{block, kInvalidNodeId, {}});
+    }
+  }
 
   template <typename NodeT>
   void Process(NodeT* node, const ProcessingState& state) {
@@ -72,6 +79,19 @@ class UseMarkingProcessor {
     }
     for (Input& input : *node) {
       input.node()->mark_use(node->id(), &input);
+      for (auto& loop_used_nodes : loop_used_nodes_) {
+        // Check if the incoming node is from outside the loop, and make sure
+        // to extend its lifetime to the loop end if yes.
+        if (loop_used_nodes.loop_header_id == kInvalidNodeId) {
+          loop_used_nodes.loop_header_id = loop_used_nodes.header->first_id();
+        }
+        if (input.node()->id() < loop_used_nodes.loop_header_id) {
+          loop_used_nodes.used_nodes.insert(input.node());
+        } else {
+          // If we're inside this loop, we're inside all outer loops too.
+          break;
+        }
+      }
     }
     if constexpr (NodeT::kProperties.can_lazy_deopt()) {
       MarkCheckpointNodes(node, node->lazy_deopt_info(), state);
@@ -89,12 +109,29 @@ class UseMarkingProcessor {
   void Process(JumpLoop* node, const ProcessingState& state) {
     int i = state.block()->predecessor_id();
     BasicBlock* target = node->target();
-    if (!target->has_phi()) return;
     uint32_t use = node->id();
-    for (Phi* phi : *target->phis()) {
-      ValueNode* input = phi->input(i).node();
-      input->mark_use(use, &phi->input(i));
+
+    if (target->has_phi()) {
+      for (Phi* phi : *target->phis()) {
+        ValueNode* input = phi->input(i).node();
+        input->mark_use(use, &phi->input(i));
+      }
     }
+
+    LoopUsedNodes& loop_used_nodes = loop_used_nodes_.back();
+    DCHECK_EQ(loop_used_nodes.header, target);
+    if (!loop_used_nodes.used_nodes.empty()) {
+      base::Vector<Input> used_node_inputs =
+          state.compilation_info()->zone()->NewVector<Input>(
+              loop_used_nodes.used_nodes.size());
+      int i = 0;
+      for (ValueNode* used_node : loop_used_nodes.used_nodes) {
+        Input* input = new (&used_node_inputs[i++]) Input(used_node);
+        used_node->mark_use(use, input);
+      }
+      node->set_used_nodes(used_node_inputs);
+    }
+    loop_used_nodes_.pop_back();
   }
   void Process(Jump* node, const ProcessingState& state) {
     int i = state.block()->predecessor_id();
@@ -146,6 +183,13 @@ class UseMarkingProcessor {
           node->mark_use(use_id, &deopt_info->input_locations[index++]);
         });
   }
+
+  struct LoopUsedNodes {
+    BasicBlock* header;
+    uint32_t loop_header_id = kInvalidNodeId;
+    std::unordered_set<ValueNode*> used_nodes;
+  };
+  std::vector<LoopUsedNodes> loop_used_nodes_;
 };
 
 // static

src/maglev/maglev-interpreter-frame-state.h

@@ -253,6 +253,7 @@ class MergePointInterpreterFrameState {
       const compiler::BytecodeLivenessState* liveness)
       : predecessor_count_(predecessor_count),
         predecessors_so_far_(1),
+        is_loop_header_(false),
         predecessors_(info.zone()->NewArray<BasicBlock*>(predecessor_count)),
         frame_state_(info, liveness, state) {
     predecessors_[0] = predecessor;
@@ -264,6 +265,7 @@ class MergePointInterpreterFrameState {
       const compiler::LoopInfo* loop_info)
       : predecessor_count_(predecessor_count),
         predecessors_so_far_(0),
+        is_loop_header_(true),
         predecessors_(info.zone()->NewArray<BasicBlock*>(predecessor_count)),
         frame_state_(info, liveness) {
     auto& assignments = loop_info->assignments();
@@ -388,6 +390,8 @@ class MergePointInterpreterFrameState {
     DCHECK_EQ(predecessors_so_far_, predecessor_count_ - 1);
     DCHECK(is_unmerged_loop());
     MergeDead(compilation_unit, merge_offset);
+    // This means that this is no longer a loop.
+    is_loop_header_ = false;
   }
 
   const CompactInterpreterFrameState& frame_state() const {
@@ -412,6 +416,8 @@ class MergePointInterpreterFrameState {
     return predecessors_[i];
   }
 
+  bool is_loop() const { return is_loop_header_; }
+
   bool is_unmerged_loop() const {
     DCHECK_GT(predecessor_count_, 0);
     return predecessors_[predecessor_count_ - 1] == unmerged_loop_marker();
@@ -582,9 +588,11 @@ class MergePointInterpreterFrameState {
     Phi* result = merged->TryCast<Phi>();
     if (result == nullptr || result->merge_offset() != merge_offset) {
       if (merged != unmerged) {
-        DCHECK(unmerged->Is<CheckedSmiUntag>() ||
-               unmerged->Is<CheckedFloat64Unbox>());
-        DCHECK_EQ(merged, unmerged->input(0).node());
+        // TODO(leszeks): These DCHECKs are too restrictive, investigate making
+        // them looser.
+        // DCHECK(unmerged->Is<CheckedSmiUntag>() ||
+        //        unmerged->Is<CheckedFloat64Unbox>());
+        // DCHECK_EQ(merged, unmerged->input(0).node());
       }
       return;
     }
@@ -609,6 +617,7 @@ class MergePointInterpreterFrameState {
 
   int predecessor_count_;
   int predecessors_so_far_;
+  bool is_loop_header_;
   Phi::List phis_;
   BasicBlock** predecessors_;
 

src/maglev/maglev-ir.h

@@ -2998,10 +2998,16 @@ class JumpLoop : public UnconditionalControlNodeT<JumpLoop> {
   void GenerateCode(MaglevCodeGenState*, const ProcessingState&);
   void PrintParams(std::ostream&, MaglevGraphLabeller*) const {}
 
+  base::Vector<Input> used_nodes() { return used_node_locations_; }
+  void set_used_nodes(base::Vector<Input> locations) {
+    used_node_locations_ = locations;
+  }
+
  private:
   // For OSR.
   const int32_t loop_depth_;
   const FeedbackSlot feedback_slot_;
+  base::Vector<Input> used_node_locations_;
 };
 
 class JumpToInlined : public UnconditionalControlNodeT<JumpToInlined> {

src/maglev/maglev-regalloc.cc

@@ -737,6 +737,18 @@ void StraightForwardRegisterAllocator::AllocateControlNode(ControlNode* node,
     InitializeBranchTargetPhis(predecessor_id, target);
     MergeRegisterValues(unconditional, target, predecessor_id);
 
+    // For JumpLoops, now update the uses of any node used in, but not defined
+    // in the loop. This makes sure that such nodes' lifetimes are extended to
+    // the entire body of the loop. This must be after phi initialisation so
+    // that value dropping in the phi initialisation doesn't think these
+    // extended lifetime nodes are dead.
+    if (auto jump_loop = node->TryCast<JumpLoop>()) {
+      for (Input& input : jump_loop->used_nodes()) {
+        DCHECK(input.node()->has_register() || input.node()->is_loadable());
+        UpdateUse(&input);
+      }
+    }
+
     if (FLAG_trace_maglev_regalloc) {
       printing_visitor_->Process(node,
                                  ProcessingState(compilation_info_, block_it_));

src/zone/zone.h

@@ -124,6 +124,12 @@ class V8_EXPORT_PRIVATE Zone final {
     return static_cast<T*>(Allocate<TypeTag>(length * sizeof(T)));
   }
 
+  template <typename T, typename TypeTag = T[]>
+  base::Vector<T> NewVector(size_t length) {
+    T* new_array = NewArray<T, TypeTag>(length);
+    return {new_array, length};
+  }
+
   template <typename T, typename TypeTag = T[]>
   base::Vector<T> NewVector(size_t length, T value) {
     T* new_array = NewArray<T, TypeTag>(length);