[wasm] Reuse constrained input for same vreg

src/compiler/register-allocator.cc

Index: src/compiler/register-allocator.cc
diff --git a/src/compiler/register-allocator.cc b/src/compiler/register-allocator.cc
index 403c344aee09209934a6581f0d5695aee71aeffd..75fc1c5f3266fbb1b38806791491ab4fa219a19e 100644
--- a/src/compiler/register-allocator.cc
+++ b/src/compiler/register-allocator.cc
@@ -1705,6 +1705,7 @@ void ConstraintBuilder::MeetConstraintsAfter(int instr_index) {
   // Handle constant/fixed output operands.
   for (size_t i = 0; i < first->OutputCount(); i++) {
     InstructionOperand* output = first->OutputAt(i);
+    if (output->IsAllocated()) continue;
     if (output->IsConstant()) {
       int output_vreg = ConstantOperand::cast(output)->virtual_register();
       TopLevelLiveRange* range = data()->GetOrCreateLiveRangeFor(output_vreg);
@@ -1751,8 +1752,10 @@ void ConstraintBuilder::MeetConstraintsAfter(int instr_index) {
 
 void ConstraintBuilder::MeetConstraintsBefore(int instr_index) {
   Instruction* second = code()->InstructionAt(instr_index);
+  ZoneMap<int, InstructionOperand> constraints(allocation_zone());
+
   // Handle fixed input operands of second instruction.
-  for (size_t i = 0; i < second->InputCount(); i++) {
+  for (size_t i = 0; i < second->InputCount(); ++i) {
     InstructionOperand* input = second->InputAt(i);
     if (input->IsImmediate() || input->IsExplicit()) {
       continue;  // Ignore immediates and explicitly reserved registers.
@@ -1763,11 +1766,19 @@ void ConstraintBuilder::MeetConstraintsBefore(int instr_index) {
       UnallocatedOperand input_copy(UnallocatedOperand::ANY, input_vreg);
       bool is_tagged = code()->IsReference(input_vreg);
       AllocateFixed(cur_input, instr_index, is_tagged);
+      // TODO(mtrofin): refactor AllocateFixed, its behavior is
+      // not the most maintainable: cur_input is mutated.
+      if (cur_input->IsAnyRegister()) {
+        // It may be that the vreg is constrained to be in different registers.
+        // We just want one of them for the other potential non-fixed register
+        // uses.
+        constraints.insert(std::make_pair(input_vreg, *cur_input));
+      }
       data()->AddGapMove(instr_index, Instruction::END, input_copy, *cur_input);
     }
   }
   // Handle "output same as input" for second instruction.
-  for (size_t i = 0; i < second->OutputCount(); i++) {
+  for (size_t i = 0; i < second->OutputCount(); ++i) {
     InstructionOperand* output = second->OutputAt(i);
     if (!output->IsUnallocated()) continue;
     UnallocatedOperand* second_output = UnallocatedOperand::cast(output);
@@ -1777,24 +1788,53 @@ void ConstraintBuilder::MeetConstraintsBefore(int instr_index) {
         UnallocatedOperand::cast(second->InputAt(0));
     int output_vreg = second_output->virtual_register();
     int input_vreg = cur_input->virtual_register();
-    UnallocatedOperand input_copy(UnallocatedOperand::ANY, input_vreg);
-    cur_input->set_virtual_register(second_output->virtual_register());
-    MoveOperands* gap_move = data()->AddGapMove(instr_index, Instruction::END,
-                                                input_copy, *cur_input);
-    if (code()->IsReference(input_vreg) && !code()->IsReference(output_vreg)) {
-      if (second->HasReferenceMap()) {
-        RegisterAllocationData::DelayedReference delayed_reference = {
-            second->reference_map(), &gap_move->source()};
-        data()->delayed_references().push_back(delayed_reference);
-      }
-    } else if (!code()->IsReference(input_vreg) &&
-               code()->IsReference(output_vreg)) {
-      // The input is assumed to immediately have a tagged representation,
-      // before the pointer map can be used. I.e. the pointer map at the
-      // instruction will include the output operand (whose value at the
-      // beginning of the instruction is equal to the input operand). If
-      // this is not desired, then the pointer map at this instruction needs
-      // to be adjusted manually.
+    if (constraints.count(input_vreg) > 0) {
+      InstructionOperand reg_operand = constraints[input_vreg];
+      InstructionOperand::ReplaceWith(second_output, &reg_operand);
+      if (data()->code()->instructions().size() >
+          static_cast<size_t>(instr_index) + 1) {
+        UnallocatedOperand new_output_op(UnallocatedOperand::ANY, output_vreg);
+        data()->AddGapMove(instr_index + 1, Instruction::START, reg_operand,
+                           new_output_op);
+        TopLevelLiveRange* range = data()->GetOrCreateLiveRangeFor(output_vreg);
+        range->RecordSpillLocation(allocation_zone(), instr_index + 1,
+                                   second_output);
+        range->SetSpillStartIndex(instr_index + 1);
+      }
+    } else {
+      UnallocatedOperand input_copy(UnallocatedOperand::ANY, input_vreg);
+      cur_input->set_virtual_register(second_output->virtual_register());
+      MoveOperands* gap_move = data()->AddGapMove(instr_index, Instruction::END,
+                                                  input_copy, *cur_input);
+      if (code()->IsReference(input_vreg) &&
+          !code()->IsReference(output_vreg)) {
+        if (second->HasReferenceMap()) {
+          RegisterAllocationData::DelayedReference delayed_reference = {
+              second->reference_map(), &gap_move->source()};
+          data()->delayed_references().push_back(delayed_reference);
+        }
+      } else if (!code()->IsReference(input_vreg) &&
+                 code()->IsReference(output_vreg)) {
+        // The input is assumed to immediately have a tagged representation,
+        // before the pointer map can be used. I.e. the pointer map at the
+        // instruction will include the output operand (whose value at the
+        // beginning of the instruction is equal to the input operand). If
+        // this is not desired, then the pointer map at this instruction needs
+        // to be adjusted manually.
+      }
+    }
+  }
+
+  for (size_t i = 0; i < second->InputCount(); ++i) {
+    InstructionOperand* input = second->InputAt(i);
+    if (!input->IsUnallocated()) continue;
+    UnallocatedOperand* cur_input = UnallocatedOperand::cast(input);
+    if (cur_input->HasFixedPolicy() || cur_input->HasSlotPolicy()) {
+      continue;
+    }
+    int vreg = cur_input->virtual_register();
+    if (constraints.count(vreg) > 0) {
+      InstructionOperand::ReplaceWith(input, &constraints.at(vreg));
     }
   }
 }