[RegAlloc] Fix "ran out of regs" with uses in statepoint

[skatkov]

Summary

Statepoint instruction is known to have a variable and big number of operands.
It is possible that Register Allocator will split live intervals in the way that all
physical registers are occupied by "zero-length" live intervals which are marked
as not-spillable.
While intervals are marked as not-spillable in the moment of creation when they are
really zero-length it is possible that in future as part of re-materialization there will
need for physical register between def and use of such tiny interval (the use is not
related to this interval at all).
As all physical registers are assigned to not-spillable intervals there is not avaialbe
registers and RA reports an error.

The idea of the fix is avoid marking tiny live intervals where there is a use in statepoint
instruction in var args section. Such interval may be perfectly spilled and folded to
operand of statepoint.

[reames]

LGTM w/comments address. Please hold submit for a couple of days for other reviewers to comment if desired.

For those reading along, Serguei and I have talked about alternatives extensively in offline discussion. I really wanted a patch which could be generic (i.e. not specific to statepoints here), but every variant we tried hit either crashes or infinite loops in the register allocator. The "basic" register allocator in particular is extraordinarily sensitive to this piece of code and relies on short live ranges being marked not spillable for forward progress. I view this patch as a workaround rather than a true fix, but I can't continue to block a workaround without an idea of how to fix something the "right way".

The risk with this patch is that the register allocator may fail to make progress on the short live range intervals. When exploring offline (with a variant, of a variant, of a variant of this patch), I saw at least one case where allowing the registers created by spiltting to be assigned normally seemed to cause an infinite split loop. I'd traced that one far enough to be reasonably convinced it was simply a latent bug in last chance recoloring, but am not sure of my diagnosis.

llvm/include/llvm/CodeGen/CalcSpillWeights.h

53

This comment tries to explain too much and is simply confusing. I'd describe what this function does here, and then the purpose at the call site.

llvm/include/llvm/CodeGen/CalcSpillWeights.h

class VirtRegMap;
   /// spill weight and allocation hint.
   class VirtRegAuxInfo {
     MachineFunction &MF;
     LiveIntervals &LIS;
     const VirtRegMap &VRM;
     const MachineLoopInfo &Loops;
     const MachineBlockFrequencyInfo &MBFI;
 
+    /// Returns true if Reg of live interval LI is used in instruction with many
+    /// operands like STATEPOINT.
+    bool isLiveAtStatepointVarArg(LiveInterval &LI);
+
   public:
     VirtRegAuxInfo(MachineFunction &MF, LiveIntervals &LIS,
                    const VirtRegMap &VRM, const MachineLoopInfo &Loops,
                    const MachineBlockFrequencyInfo &MBFI)
         : MF(MF), LIS(LIS), VRM(VRM), Loops(Loops), MBFI(MBFI) {}
 
     virtual ~VirtRegAuxInfo() = default;
 

llvm/lib/CodeGen/CalcSpillWeights.cpp

 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/CodeGen/VirtRegMap.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/CodeGen/StackMaps.h"
 #include <cassert>
 #include <tuple>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "calcspillweights"
 
 void VirtRegAuxInfo::calculateSpillWeightsAndHints() {
for (LiveInterval::const_vni_iterator I = LI.vni_begin(), E = LI.vni_end();
     }
 
     if (!TII.isTriviallyReMaterializable(*MI, LIS.getAliasAnalysis()))
       return false;
   }
   return true;
 }
 
+bool VirtRegAuxInfo::isLiveAtStatepointVarArg(LiveInterval &LI) {
+  return any_of(VRM.getRegInfo().reg_operands(LI.reg()),
+                [](MachineOperand &MO) {
+    MachineInstr *MI = MO.getParent();
+    if (MI->getOpcode() != TargetOpcode::STATEPOINT)
+      return false;
+    return StatepointOpers(MI).getVarIdx() <= MI->getOperandNo(&MO);
+  });
+}
+
 void VirtRegAuxInfo::calculateSpillWeightAndHint(LiveInterval &LI) {
   float Weight = weightCalcHelper(LI);
   // Check if unspillable.
   if (Weight < 0)
     return;
   LI.setWeight(Weight);
 }
 
float VirtRegAuxInfo::weightCalcHelper(LiveInterval &LI, SlotIndex *Start,
   }
 
   // If the live interval was already unspillable, leave it that way.
   if (!IsSpillable)
     return -1.0;
 
   // Mark li as unspillable if all live ranges are tiny and the interval
   // is not live at any reg mask.  If the interval is live at a reg mask
-  // spilling may be required.
+  // spilling may be required. If li is live as use in statepoint instruction
+  // spilling may be required due to if we mark interval with use in statepoint
+  // as not spillable we are risky to end up with no register to allocate.
+  // At the same time STATEPOINT instruction is perfectly fine to have this
+  // operand on stack, so spilling such interval and folding its load from stack
+  // into instruction itself makes perfect sense.
   if (ShouldUpdateLI && LI.isZeroLength(LIS.getSlotIndexes()) &&
-      !LI.isLiveAtIndexes(LIS.getRegMaskSlots())) {
+      !LI.isLiveAtIndexes(LIS.getRegMaskSlots()) &&
+      !isLiveAtStatepointVarArg(LI)) {
     LI.markNotSpillable();
     return -1.0;
   }
 
   // If all of the definitions of the interval are re-materializable,
   // it is a preferred candidate for spilling.
   // FIXME: this gets much more complicated once we support non-trivial
   // re-materialization.
   if (isRematerializable(LI, LIS, VRM, *MF.getSubtarget().getInstrInfo()))
     TotalWeight *= 0.5F;
 
   if (IsLocalSplitArtifact)
     return normalize(TotalWeight, Start->distance(*End), NumInstr);
   return normalize(TotalWeight, LI.getSize(), NumInstr);
 }

llvm/test/CodeGen/X86/statepoint-ra.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -verify-machineinstrs -O3 -use-registers-for-deopt-values -restrict-statepoint-remat=true < %s 2>&1 | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK-NOT: error: ran out of registers during register allocation
+
+define void @barney(i8 addrspace(1)* %arg, double %arg1, double %arg2, double %arg3, double %arg4, double %arg5, double %arg6, double %arg7, double %arg8, double %arg9, double %arg10, double %arg11, double %arg12) gc "statepoint-example" personality i32* ()* @widget {
+bb:
+  %tmp = call coldcc token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* nonnull @blam, i32 0, i32 0, i32 0, i32 0) [ "deopt"(i32 0, i32 1, i32 0, i32 0, i32 0, i32 26, i32 0, i32 0, i8 addrspace(1)* %arg, i32 4, double %arg1, i32 7, i8* null, i32 4, double %arg2, i32 7, i8* null, i32 4, double %arg3, i32 7, i8* null, i32 4, double %arg4, i32 7, i8* null, i32 4, double %arg5, i32 7, i8* null, i32 4, double %arg6, i32 7, i8* null, i32 4, double %arg7, i32 7, i8* null, i32 4, double %arg8, i32 7, i8* null, i32 4, double %arg9, i32 7, i8* null, i32 4, double %arg10, i32 7, i8* null, i32 4, double %arg11, i32 7, i8* null, i32 4, double %arg12, i32 7, i8* null, i32 7, i8* null), "gc-live"(i8 addrspace(1)* %arg) ]
+  br i1 undef, label %bb13, label %bb15
+
+bb13:                                             ; preds = %bb
+  %tmp14 = call token (i64, i32, i32 (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double)*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_i32p1i8f64f64f64f64f64f64f64f64f64f(i64 2, i32 5, i32 (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double)* nonnull @quux, i32 10, i32 0, i8 addrspace(1)* nonnull null, double %arg1, double %arg2, double %arg3, double %arg5, double %arg6, double %arg7, double %arg9, double %arg10, double %arg11, i32 0, i32 0) [ "deopt"(i32 0, i32 2, i32 0, i32 70, i32 0, i32 26, i32 0, i32 0, i8 addrspace(1)* null, i32 4, double %arg1, i32 7, i8* null, i32 4, double %arg2, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 4, double %arg4, i32 7, i8* null, i32 4, double %arg5, i32 7, i8* null, i32 4, double %arg6, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 4, double %arg8, i32 7, i8* null, i32 4, double %arg9, i32 7, i8* null, i32 4, double %arg10, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 4, double %arg12, i32 7, i8* null, i32 7, i8* null), "gc-live"(i8 addrspace(1)* null) ]
+  br label %bb15
+
+bb15:                                             ; preds = %bb13, %bb
+  %tmp16 = phi double [ %arg4, %bb13 ], [ 1.000000e+00, %bb ]
+  %tmp17 = phi double [ %arg8, %bb13 ], [ 1.000000e+00, %bb ]
+  %tmp18 = phi double [ %arg12, %bb13 ], [ 1.000000e+00, %bb ]
+  br i1 undef, label %bb25, label %bb19
+
+bb19:                                             ; preds = %bb15
+  %tmp20 = invoke token (i64, i32, i32 (i32, i8 addrspace(1)*, i32)*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_i32i32p1i8i32f(i64 1, i32 16, i32 (i32, i8 addrspace(1)*, i32)* nonnull @eggs, i32 3, i32 0, i32 undef, i8 addrspace(1)* nonnull undef, i32 0, i32 0, i32 0) [ "deopt"(i32 0, i32 2, i32 0, i32 97, i32 0, i32 26, i32 0, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 0, i32 2, i32 3, i32 0, i32 20, i32 0, i32 0, i8 addrspace(1)* undef, i32 4, double %arg1, i32 7, i8* null, i32 4, double %arg2, i32 7, i8* null, i32 4, double %tmp16, i32 7, i8* null, i32 4, double %arg5, i32 7, i8* null, i32 4, double %arg6, i32 7, i8* null, i32 4, double %tmp17, i32 7, i8* null, i32 4, double %arg9, i32 7, i8* null, i32 4, double %arg10, i32 7, i8* null, i32 4, double %tmp18, i32 7, i8* null, i32 7, i8* null), "gc-live"(i8 addrspace(1)* undef) ]
+          to label %bb21 unwind label %bb23
+
+bb21:                                             ; preds = %bb19
+  %tmp22 = call token (i64, i32, void (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double, i32)*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidp1i8f64f64f64f64f64f64f64f64f64i32f(i64 2, i32 5, void (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double, i32)* nonnull @ham, i32 11, i32 0, i8 addrspace(1)* nonnull undef, double %arg1, double %arg2, double %tmp16, double %arg5, double %arg6, double %tmp17, double %arg9, double %arg10, double %tmp18, i32 51, i32 0, i32 0) [ "deopt"(i32 0, i32 2, i32 0, i32 97, i32 0, i32 26, i32 0, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 7, i8* null, i32 2, i32 2, i32 46, i32 0, i32 20, i32 0, i32 0, i8 addrspace(1)* undef, i32 4, double %arg1, i32 7, i8* null, i32 4, double %arg2, i32 7, i8* null, i32 4, double %tmp16, i32 7, i8* null, i32 4, double %arg5, i32 7, i8* null, i32 4, double %arg6, i32 7, i8* null, i32 4, double %tmp17, i32 7, i8* null, i32 4, double %arg9, i32 7, i8* null, i32 4, double %arg10, i32 7, i8* null, i32 4, double %tmp18, i32 7, i8* null, i32 3, i32 51), "gc-live"(i8 addrspace(1)* undef) ]
+  unreachable
+
+bb23:                                             ; preds = %bb19
+  %tmp24 = landingpad token
+          cleanup
+  ret void
+
+bb25:                                             ; preds = %bb15
+  ret void
+}
+
+declare i32* @widget()
+declare i32 @quux(i8 addrspace(1)*, double, double, double, double, double, double, double, double, double)
+declare void @blam()
+declare i32 @eggs(i32, i8 addrspace(1)*, i32)
+declare void @ham(i8 addrspace(1)*, double, double, double, double, double, double, double, double, double, i32)
+declare token @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 , i32 , void ()*, i32 , i32 , ...)
+declare token @llvm.experimental.gc.statepoint.p0f_i32p1i8f64f64f64f64f64f64f64f64f64f(i64 , i32 , i32 (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double)*, i32 , i32 , ...)
+declare token @llvm.experimental.gc.statepoint.p0f_i32i32p1i8i32f(i64 , i32 , i32 (i32, i8 addrspace(1)*, i32)*, i32 , i32 , ...)
+declare token @llvm.experimental.gc.statepoint.p0f_isVoidp1i8f64f64f64f64f64f64f64f64f64i32f(i64 , i32 , void (i8 addrspace(1)*, double, double, double, double, double, double, double, double, double, i32)*, i32 , i32 , ...)