[maglev] Reuse unused stack slots

This CL splits the stack frame into tagged and untagged slots, and allows reuse
of slots for values that are also tagged/untagged.

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

src/maglev/maglev-code-gen-state.h

@@ -7,6 +7,7 @@
 
 #include "src/codegen/assembler.h"
 #include "src/codegen/label.h"
+#include "src/codegen/machine-type.h"
 #include "src/codegen/macro-assembler.h"
 #include "src/codegen/safepoint-table.h"
 #include "src/common/globals.h"
@@ -37,7 +38,8 @@ class MaglevCodeGenState {
         safepoint_table_builder_(safepoint_table_builder),
         masm_(isolate(), CodeObjectRequired::kNo) {}
 
-  void SetVregSlots(int slots) { vreg_slots_ = slots; }
+  void set_tagged_slots(int slots) { tagged_slots_ = slots; }
+  void set_untagged_slots(int slots) { untagged_slots_ = slots; }
 
   void PushDeferredCode(DeferredCodeInfo* deferred_code) {
     deferred_code_.push_back(deferred_code);
@@ -73,7 +75,7 @@ class MaglevCodeGenState {
     return compilation_unit_->graph_labeller();
   }
   MacroAssembler* masm() { return &masm_; }
-  int vreg_slots() const { return vreg_slots_; }
+  int stack_slots() const { return untagged_slots_ + tagged_slots_; }
   SafepointTableBuilder* safepoint_table_builder() const {
     return safepoint_table_builder_;
   }
@@ -87,7 +89,38 @@ class MaglevCodeGenState {
     return found_unsupported_code_paths_;
   }
 
+  inline int GetFramePointerOffsetForStackSlot(
+      const compiler::AllocatedOperand& operand) {
+    int index = operand.index();
+    if (operand.representation() != MachineRepresentation::kTagged) {
+      index += tagged_slots_;
+    }
+    return GetFramePointerOffsetForStackSlot(index);
+  }
+
+  inline MemOperand GetStackSlot(const compiler::AllocatedOperand& operand) {
+    return MemOperand(rbp, GetFramePointerOffsetForStackSlot(operand));
+  }
+
+  inline MemOperand ToMemOperand(const compiler::InstructionOperand& operand) {
+    return GetStackSlot(compiler::AllocatedOperand::cast(operand));
+  }
+
+  inline MemOperand ToMemOperand(const ValueLocation& location) {
+    return ToMemOperand(location.operand());
+  }
+
+  inline MemOperand TopOfStack() {
+    return MemOperand(rbp,
+                      GetFramePointerOffsetForStackSlot(stack_slots() - 1));
+  }
+
  private:
+  inline constexpr int GetFramePointerOffsetForStackSlot(int index) {
+    return StandardFrameConstants::kExpressionsOffset -
+           index * kSystemPointerSize;
+  }
+
   MaglevCompilationUnit* const compilation_unit_;
   SafepointTableBuilder* const safepoint_table_builder_;
 
@@ -95,7 +128,8 @@ class MaglevCodeGenState {
   std::vector<DeferredCodeInfo*> deferred_code_;
   std::vector<EagerDeoptInfo*> eager_deopts_;
   std::vector<LazyDeoptInfo*> lazy_deopts_;
-  int vreg_slots_ = 0;
+  int untagged_slots_ = 0;
+  int tagged_slots_ = 0;
 
   // Allow marking some codegen paths as unsupported, so that we can test maglev
   // incrementally.
@@ -106,30 +140,12 @@ class MaglevCodeGenState {
 // Some helpers for codegen.
 // TODO(leszeks): consider moving this to a separate header.
 
-inline constexpr int GetFramePointerOffsetForStackSlot(int index) {
-  return StandardFrameConstants::kExpressionsOffset -
-         index * kSystemPointerSize;
-}
-
-inline int GetFramePointerOffsetForStackSlot(
-    const compiler::AllocatedOperand& operand) {
-  return GetFramePointerOffsetForStackSlot(operand.index());
-}
-
 inline int GetSafepointIndexForStackSlot(int i) {
   // Safepoint tables also contain slots for all fixed frame slots (both
   // above and below the fp).
   return StandardFrameConstants::kFixedSlotCount + i;
 }
 
-inline MemOperand GetStackSlot(int index) {
-  return MemOperand(rbp, GetFramePointerOffsetForStackSlot(index));
-}
-
-inline MemOperand GetStackSlot(const compiler::AllocatedOperand& operand) {
-  return GetStackSlot(operand.index());
-}
-
 inline Register ToRegister(const compiler::InstructionOperand& operand) {
   return compiler::AllocatedOperand::cast(operand).GetRegister();
 }
@@ -138,24 +154,10 @@ inline Register ToRegister(const ValueLocation& location) {
   return ToRegister(location.operand());
 }
 
-inline MemOperand ToMemOperand(const compiler::InstructionOperand& operand) {
-  return GetStackSlot(compiler::AllocatedOperand::cast(operand));
-}
-
-inline MemOperand ToMemOperand(const ValueLocation& location) {
-  return ToMemOperand(location.operand());
-}
-
 inline void MaglevCodeGenState::DefineSafepointStackSlots(
     SafepointTableBuilder::Safepoint& safepoint) const {
-  DCHECK_EQ(compilation_unit()->stack_value_repr().size(), vreg_slots());
-  int stack_slot = 0;
-  for (ValueRepresentation repr : compilation_unit()->stack_value_repr()) {
-    if (repr == ValueRepresentation::kTagged) {
-      safepoint.DefineTaggedStackSlot(
-          GetSafepointIndexForStackSlot(stack_slot));
-    }
-    stack_slot++;
+  for (int stack_slot = 0; stack_slot < tagged_slots_; stack_slot++) {
+    safepoint.DefineTaggedStackSlot(GetSafepointIndexForStackSlot(stack_slot));
   }
 }
 

src/maglev/maglev-code-generator.cc

@@ -64,16 +64,18 @@ class MaglevCodeGeneratingNodeProcessor {
     __ Push(kJavaScriptCallArgCountRegister);  // Actual argument count.
 
     // Extend rsp by the size of the frame.
-    code_gen_state_->SetVregSlots(graph->stack_slots());
-    __ subq(rsp, Immediate(code_gen_state_->vreg_slots() * kSystemPointerSize));
+    code_gen_state_->set_untagged_slots(graph->untagged_stack_slots());
+    code_gen_state_->set_tagged_slots(graph->tagged_stack_slots());
+    __ subq(rsp,
+            Immediate(code_gen_state_->stack_slots() * kSystemPointerSize));
 
     // Initialize stack slots.
     // TODO(jgruber): Update logic once the register allocator is further along.
     {
       ASM_CODE_COMMENT_STRING(masm(), "Initializing stack slots");
       __ Move(rax, Immediate(0));
-      __ Move(rcx, Immediate(code_gen_state_->vreg_slots()));
-      __ leaq(rdi, GetStackSlot(code_gen_state_->vreg_slots() - 1));
+      __ Move(rcx, Immediate(code_gen_state_->stack_slots()));
+      __ leaq(rdi, code_gen_state_->TopOfStack());
       __ repstosq();
     }
 
@@ -123,7 +125,8 @@ class MaglevCodeGeneratingNodeProcessor {
         if (!source.IsStackSlot()) {
           if (FLAG_code_comments) __ RecordComment("--   Spill:");
           DCHECK(!source.IsStackSlot());
-          __ movq(GetStackSlot(value_node->spill_slot()), ToRegister(source));
+          __ movq(code_gen_state_->GetStackSlot(value_node->spill_slot()),
+                  ToRegister(source));
         } else {
           // Otherwise, the result source stack slot should be equal to the
           // spill slot.
@@ -179,7 +182,8 @@ class MaglevCodeGeneratingNodeProcessor {
   void EmitStackToRegisterGapMove(compiler::InstructionOperand source,
                                   Register target) {
     if (!source.IsAllocated()) return;
-    __ movq(target, GetStackSlot(compiler::AllocatedOperand::cast(source)));
+    __ movq(target, code_gen_state_->GetStackSlot(
+                        compiler::AllocatedOperand::cast(source)));
   }
 
   void RecordGapMove(compiler::AllocatedOperand source, Register target_reg,
@@ -215,10 +219,10 @@ class MaglevCodeGeneratingNodeProcessor {
     // clobbered by reg->reg or stack->reg, so emit them immediately.
     if (source.IsRegister()) {
       Register source_reg = ToRegister(source);
-      __ movq(GetStackSlot(target), source_reg);
+      __ movq(code_gen_state_->GetStackSlot(target), source_reg);
     } else {
-      __ movq(kScratchRegister, GetStackSlot(source));
-      __ movq(GetStackSlot(target), kScratchRegister);
+      __ movq(kScratchRegister, code_gen_state_->GetStackSlot(source));
+      __ movq(code_gen_state_->GetStackSlot(target), kScratchRegister);
     }
   }
 
@@ -304,15 +308,6 @@ class MaglevCodeGeneratingNodeProcessor {
   MaglevCodeGenState* code_gen_state_;
 };
 
-constexpr int DeoptStackSlotIndexFromFPOffset(int offset) {
-  return 1 - offset / kSystemPointerSize;
-}
-
-int DeoptStackSlotFromStackSlot(const compiler::AllocatedOperand& operand) {
-  return DeoptStackSlotIndexFromFPOffset(
-      GetFramePointerOffsetForStackSlot(operand));
-}
-
 }  // namespace
 
 class MaglevCodeGeneratorImpl final {
@@ -467,6 +462,15 @@ class MaglevCodeGeneratorImpl final {
     }
   }
 
+  constexpr int DeoptStackSlotIndexFromFPOffset(int offset) {
+    return 1 - offset / kSystemPointerSize;
+  }
+
+  int DeoptStackSlotFromStackSlot(const compiler::AllocatedOperand& operand) {
+    return DeoptStackSlotIndexFromFPOffset(
+        code_gen_state_.GetFramePointerOffsetForStackSlot(operand));
+  }
+
   void EmitDeoptFrameValues(
       const MaglevCompilationUnit& compilation_unit,
       const CompactInterpreterFrameState* checkpoint_state,
@@ -635,7 +639,7 @@ class MaglevCodeGeneratorImpl final {
     return data;
   }
 
-  int stack_slot_count() const { return code_gen_state_.vreg_slots(); }
+  int stack_slot_count() const { return code_gen_state_.stack_slots(); }
   int stack_slot_count_with_fixed_frame() const {
     return stack_slot_count() + StandardFrameConstants::kFixedSlotCount;
   }

src/maglev/maglev-compilation-unit.cc

@@ -22,8 +22,7 @@ MaglevCompilationUnit::MaglevCompilationUnit(MaglevCompilationInfo* info,
       bytecode_analysis_(bytecode_.object(), zone(), BytecodeOffset::None(),
                          true),
       register_count_(bytecode_.register_count()),
-      parameter_count_(bytecode_.parameter_count()),
-      stack_value_repr_(info->zone()) {}
+      parameter_count_(bytecode_.parameter_count()) {}
 
 compiler::JSHeapBroker* MaglevCompilationUnit::broker() const {
   return info_->broker();

src/maglev/maglev-compilation-unit.h

@@ -49,14 +49,6 @@ class MaglevCompilationUnit : public ZoneObject {
 
   void RegisterNodeInGraphLabeller(const Node* node);
 
-  const ZoneVector<ValueRepresentation>& stack_value_repr() const {
-    return stack_value_repr_;
-  }
-
-  void push_stack_value_repr(ValueRepresentation r) {
-    stack_value_repr_.push_back(r);
-  }
-
  private:
   MaglevCompilationInfo* const info_;
   const compiler::SharedFunctionInfoRef shared_function_info_;
@@ -65,10 +57,6 @@ class MaglevCompilationUnit : public ZoneObject {
   const compiler::BytecodeAnalysis bytecode_analysis_;
   const int register_count_;
   const int parameter_count_;
-
-  // TODO(victorgomes): Compress these values, if only tagged/untagged, we could
-  // use a binary vector? We might also want to deal with safepoints properly.
-  ZoneVector<ValueRepresentation> stack_value_repr_;
 };
 
 }  // namespace maglev

src/maglev/maglev-graph.h

@@ -41,15 +41,22 @@ class Graph final : public ZoneObject {
 
   void Add(BasicBlock* block) { blocks_.push_back(block); }
 
-  uint32_t stack_slots() const { return stack_slots_; }
-  void set_stack_slots(uint32_t stack_slots) {
-    DCHECK_EQ(kMaxUInt32, stack_slots_);
+  uint32_t tagged_stack_slots() const { return tagged_stack_slots_; }
+  uint32_t untagged_stack_slots() const { return untagged_stack_slots_; }
+  void set_tagged_stack_slots(uint32_t stack_slots) {
+    DCHECK_EQ(kMaxUInt32, tagged_stack_slots_);
     DCHECK_NE(kMaxUInt32, stack_slots);
-    stack_slots_ = stack_slots;
+    tagged_stack_slots_ = stack_slots;
+  }
+  void set_untagged_stack_slots(uint32_t stack_slots) {
+    DCHECK_EQ(kMaxUInt32, untagged_stack_slots_);
+    DCHECK_NE(kMaxUInt32, stack_slots);
+    untagged_stack_slots_ = stack_slots;
   }
 
  private:
-  uint32_t stack_slots_ = kMaxUInt32;
+  uint32_t tagged_stack_slots_ = kMaxUInt32;
+  uint32_t untagged_stack_slots_ = kMaxUInt32;
   ZoneVector<BasicBlock*> blocks_;
 };
 

src/maglev/maglev-ir.cc

@@ -100,7 +100,7 @@ void PushInput(MaglevCodeGenState* code_gen_state, const Input& input) {
     __ Push(operand.GetRegister());
   } else {
     DCHECK(operand.IsStackSlot());
-    __ Push(GetStackSlot(operand));
+    __ Push(code_gen_state->GetStackSlot(operand));
   }
 }
 
@@ -628,15 +628,15 @@ void GapMove::GenerateCode(MaglevCodeGenState* code_gen_state,
     if (target().IsAnyRegister()) {
       __ movq(ToRegister(target()), source_reg);
     } else {
-      __ movq(ToMemOperand(target()), source_reg);
+      __ movq(code_gen_state->ToMemOperand(target()), source_reg);
     }
   } else {
-    MemOperand source_op = ToMemOperand(source());
+    MemOperand source_op = code_gen_state->ToMemOperand(source());
     if (target().IsAnyRegister()) {
       __ movq(ToRegister(target()), source_op);
     } else {
       __ movq(kScratchRegister, source_op);
-      __ movq(ToMemOperand(target()), kScratchRegister);
+      __ movq(code_gen_state->ToMemOperand(target()), kScratchRegister);
     }
   }
 }

src/maglev/maglev-regalloc.cc

@@ -6,6 +6,7 @@
 
 #include "src/base/bits.h"
 #include "src/base/logging.h"
+#include "src/codegen/machine-type.h"
 #include "src/codegen/register.h"
 #include "src/codegen/reglist.h"
 #include "src/compiler/backend/instruction.h"
@@ -82,7 +83,8 @@ StraightForwardRegisterAllocator::StraightForwardRegisterAllocator(
     : compilation_unit_(compilation_unit) {
   ComputePostDominatingHoles(graph);
   AllocateRegisters(graph);
-  graph->set_stack_slots(top_of_stack_);
+  graph->set_tagged_stack_slots(tagged_.top);
+  graph->set_untagged_stack_slots(untagged_.top);
 }
 
 StraightForwardRegisterAllocator::~StraightForwardRegisterAllocator() = default;
@@ -318,6 +320,20 @@ void StraightForwardRegisterAllocator::UpdateUse(
 
   // If a value is dead, make sure it's cleared.
   FreeRegisters(node);
+
+  // If the stack slot is a local slot, free it so it can be reused.
+  if (node->is_spilled()) {
+    compiler::AllocatedOperand slot = node->spill_slot();
+    if (slot.index() > 0) {
+      SpillSlots& slots =
+          slot.representation() == MachineRepresentation::kTagged ? tagged_
+                                                                  : untagged_;
+      DCHECK_IMPLIES(
+          slots.free_slots.size() > 0,
+          slots.free_slots.back().freed_at_position <= node->live_range().end);
+      slots.free_slots.emplace_back(slot.index(), node->live_range().end);
+    }
+  }
 }
 
 void StraightForwardRegisterAllocator::UpdateUse(
@@ -651,11 +667,30 @@ void StraightForwardRegisterAllocator::SpillAndClearRegisters() {
 
 void StraightForwardRegisterAllocator::AllocateSpillSlot(ValueNode* node) {
   DCHECK(!node->is_spilled());
-  uint32_t free_slot = top_of_stack_++;
-  compilation_unit_->push_stack_value_repr(node->value_representation());
+  uint32_t free_slot;
+  bool is_tagged = node->value_representation() == ValueRepresentation::kTagged;
+  SpillSlots& slots = is_tagged ? tagged_ : untagged_;
+  MachineRepresentation representation = is_tagged
+                                             ? MachineRepresentation::kTagged
+                                             : MachineRepresentation::kWord64;
+  if (slots.free_slots.empty()) {
+    free_slot = slots.top++;
+  } else {
+    NodeIdT start = node->live_range().start;
+    auto it =
+        std::upper_bound(slots.free_slots.begin(), slots.free_slots.end(),
+                         start, [](NodeIdT s, const SpillSlotInfo& slot_info) {
+                           return slot_info.freed_at_position < s;
+                         });
+    if (it != slots.free_slots.end()) {
+      free_slot = it->slot_index;
+      slots.free_slots.erase(it);
+    } else {
+      free_slot = slots.top++;
+    }
+  }
   node->Spill(compiler::AllocatedOperand(compiler::AllocatedOperand::STACK_SLOT,
-                                         MachineRepresentation::kTagged,
-                                         free_slot));
+                                         representation, free_slot));
 }
 
 void StraightForwardRegisterAllocator::FreeSomeRegister() {

src/maglev/maglev-regalloc.h

@@ -25,15 +25,27 @@ class StraightForwardRegisterAllocator {
                                    Graph* graph);
   ~StraightForwardRegisterAllocator();
 
-  int stack_slots() const { return top_of_stack_; }
-
  private:
   std::vector<int> future_register_uses_[Register::kNumRegisters];
   ValueNode* register_values_[Register::kNumRegisters];
 
-  int top_of_stack_ = 0;
   RegList free_registers_ = kAllocatableGeneralRegisters;
 
+  struct SpillSlotInfo {
+    SpillSlotInfo(uint32_t slot_index, NodeIdT freed_at_position)
+        : slot_index(slot_index), freed_at_position(freed_at_position) {}
+    uint32_t slot_index;
+    NodeIdT freed_at_position;
+  };
+  struct SpillSlots {
+    int top = 0;
+    // Sorted from earliest freed_at_position to latest freed_at_position.
+    std::vector<SpillSlotInfo> free_slots;
+  };
+
+  SpillSlots untagged_;
+  SpillSlots tagged_;
+
   RegList used_registers() const {
     // Only allocatable registers should be free.
     DCHECK_EQ(free_registers_, free_registers_ & kAllocatableGeneralRegisters);