Merge rust-lang#68

68: Fix multiple locations in stackmaps r=ltratt a=ptersilie



Co-authored-by: Lukas Diekmann <[email protected]>

llvm/include/llvm/CodeGen/StackMaps.h

@@ -195,10 +195,11 @@ class StackMaps {
     unsigned Size = 0;
     unsigned Reg = 0;
     int64_t Offset = 0;
+    unsigned ExtraReg = 0;
 
     Location() = default;
-    Location(LocationType Type, unsigned Size, unsigned Reg, int64_t Offset)
-        : Type(Type), Size(Size), Reg(Reg), Offset(Offset) {}
+    Location(LocationType Type, unsigned Size, unsigned Reg, int64_t Offset, unsigned ExtraReg = 0)
+        : Type(Type), Size(Size), Reg(Reg), Offset(Offset), ExtraReg(ExtraReg) {}
   };
 
   struct LiveOutReg {

llvm/lib/CodeGen/StackMaps.cpp

@@ -286,21 +286,46 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
 
 
     signed Offset = 0;
-    // Check if there are any mappings for this register in the spillmap. If so,
-    // encode the additional location in the offset field of the stackmap record
-    // (which is unused for register locations). Note that this assumes that
-    // there can only be one additional location for each value, which may turn
-    // out to be false.
+    int RHS = 0;
+    // Check for any additional mappings in the spillmap and add them to his
+    // location to be later encoded into stackmaps. A typical case where we need
+    // to record multiple sources is the following (annotated with SpillMap
+    // info):
+    //
+    // %rcx = load %rbp, -8             // SpillMap[rcx] = -8
+    // %rbx = %rcx                      // SpillMap[rbx] = rcx
+    // STACKMAP(%rbx) or STACKMAP(%rcx)
+    //
+    // First a value is loaded from the stack into %rcx and then immediately
+    // moved into %rbx. This means there are three sources for the same value,
+    // and during deoptimisation we need to make sure we write the value back to
+    // each one of them. Note, that the stackmap may track either of %rbx or
+    // %rcx, resulting in different ways below to retrieve the mappings.
+    int ExtraReg = 0;
     if (MOI->isReg()) {
       Register R = MOI->getReg();
       if (SpillOffsets.count(R) > 0) {
-        Offset = SpillOffsets[R];
+        RHS = SpillOffsets[R];
         assert(SpillOffsets[R] != 0);
-        if (Offset > 0) {
-          // If the additional location is another register encode its DWARF id.
-          // Also temporarily add 1 since 0 is used to mean there is no
-          // additional location.
-          Offset = getDwarfRegNum(Offset, TRI) + 1;
+        if (RHS > 0) {
+          // The additional location is another register: encode its DWARF id.
+          // Also temporarily add 1 since 0 means there is no additional
+          // location.
+          ExtraReg = getDwarfRegNum(RHS, TRI) + 1;
+          // Check if the other register also has a mapping and add it.
+          if (SpillOffsets.count(RHS) > 0) {
+            Offset = SpillOffsets[RHS];
+          }
+        }  else {
+          // The other location is an offset.
+          Offset = RHS;
+          // Find any additional mappings where the current register appears on
+          // the right hand side.
+          for (auto I = SpillOffsets.begin(); I != SpillOffsets.end(); I++) {
+            if (I->second == R) {
+              ExtraReg = getDwarfRegNum(I->first, TRI) + 1;
+            }
+          }
         }
       }
     }
@@ -312,7 +337,7 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
       Offset = TRI->getSubRegIdxOffset(SubRegIdx);
 
     Locs.emplace_back(Location::Register, TRI->getSpillSize(*RC), DwarfRegNum,
-                      Offset);
+                      Offset, ExtraReg);
     return ++MOI;
   }
 
@@ -813,7 +838,7 @@ void StackMaps::emitCallsiteEntries(MCStreamer &OS) {
         OS.emitIntValue(0, 1); // Reserved
         OS.emitInt16(Loc.Size);
         OS.emitInt16(Loc.Reg);
-        OS.emitInt16(0); // Reserved
+        OS.emitInt16(Loc.ExtraReg); // Reserved
         OS.emitInt32(Loc.Offset);
       }
       LiveIdx++;

llvm/lib/Target/X86/X86AsmPrinter.cpp

@@ -97,14 +97,7 @@ void processInstructions(
       const MachineOperand Rhs = Instr.getOperand(1);
       assert(Lhs.isReg() && "Is register.");
       assert(Rhs.isReg() && "Is register.");
-      if (SpillMap.count(Rhs.getReg()) > 0) {
-        // If the RHS has a mapping, apply the same mapping to the new register.
-        // This means, that stack spills moved into one register and then into
-        // another will continued to be tracked.
-        SpillMap[Lhs.getReg()] = SpillMap[Rhs.getReg()];
-      } else {
-        SpillMap[Lhs.getReg()] = Rhs.getReg();
-      }
+      SpillMap[Lhs.getReg()] = Rhs.getReg();
       // YKFIXME: If the `mov` instruction has a killed-flag, remove the
       // register from the map.
 

llvm/lib/Transforms/Yk/ControlPoint.cpp

@@ -231,13 +231,13 @@ class YkControlPoint : public ModulePass {
     Builder.CreateCall(YKFR, {NewCtrlPointCallInst});
     Builder.CreateUnreachable();
 
-    // To do so we need to first split up the current block and then
-    // insert a conditional branch that either continues or returns.
-
+    // Split up the current block and then insert a conditional branch that
+    // either continues after the control point or invokes frame reconstruction.
     BasicBlock *BB = NewCtrlPointCallInst->getParent();
-    BasicBlock *ContBB = BB->splitBasicBlock(New);
-
+    BasicBlock *ContBB = BB->splitBasicBlock(New->getNextNonDebugInstruction());
     Instruction &OldBr = BB->back();
+    assert(OldBr.getPrevNonDebugInstruction() == New &&
+           "Split block at the wrong spot.");
     OldBr.eraseFromParent();
     Builder.SetInsertPoint(BB);
     // The return value of the control point tells us whether a guard has failed