Check if interval interferes with the call

src/coreclr/jit/lsra.cpp

@@ -703,6 +703,8 @@ LinearScan::LinearScan(Compiler* theCompiler)
     blockSequenceWorkList = nullptr;
     curBBSeqNum           = 0;
     bbSeqCount            = 0;
+    recentRefPosition     = nullptr;
+    recentKillLocation    = 0;
 
     // Information about each block, including predecessor blocks used for variable locations at block entry.
     blockInfo = nullptr;

src/coreclr/jit/lsra.h

@@ -1582,6 +1582,8 @@ class LinearScan : public LinearScanInterface
     // Ordered list of RefPositions
     RefPositionList refPositions;
 
+    RefPosition* recentRefPosition;
+    LsraLocation recentKillLocation;
     // Per-block variable location mappings: an array indexed by block number that yields a
     // pointer to an array of regNumber, one per variable.
     VarToRegMap* inVarToRegMaps;

src/coreclr/jit/lsrabuild.cpp

@@ -186,6 +186,7 @@ RefPosition* LinearScan::newRefPositionRaw(LsraLocation nodeLocation, GenTree* t
     currBuildNode = nullptr;
     newRP->rpNum  = static_cast<unsigned>(refPositions.size() - 1);
 #endif // DEBUG
+    recentRefPosition = newRP;
     return newRP;
 }
 
@@ -356,14 +357,13 @@ void LinearScan::resolveConflictingDefAndUse(Interval* interval, RefPosition* de
 }
 
 //------------------------------------------------------------------------
-// applyCalleeSaveHeuristics: Set register preferences for an interval based on the given RefPosition
+// applyCalleeSaveHeuristics: Check if an interval should be set to prefer callee-save.
+//   This is done by checking if this interval interferes with any call (`recentCallRefPositionLocation`)
+//   and if yes, sets `preferCalleeSave = true`.
 //
 // Arguments:
 //    rp - The RefPosition of interest
 //
-// Notes:
-//    This is slightly more general than its name applies, and updates preferences not just
-//    for callee-save registers.
 //
 void LinearScan::applyCalleeSaveHeuristics(RefPosition* rp)
 {
@@ -375,15 +375,71 @@ void LinearScan::applyCalleeSaveHeuristics(RefPosition* rp)
     }
 #endif // TARGET_AMD64
 
-    Interval* theInterval = rp->getInterval();
+    Interval* theInterval          = rp->getInterval();
+    regMaskTP calleeSavePreference = rp->registerAssignment;
+
+    if (!theInterval->preferCalleeSave && !theInterval->isLocalVar && RefTypeIsUse(rp->refType))
+    {
+        RefPosition* firstPosition = theInterval->firstRefPosition;
+        if ((firstPosition != nullptr) && (firstPosition->nodeLocation > 0))
+        {
+            if (firstPosition->nodeLocation < recentKillLocation)
+            {
+                JITDUMP("Interval %2u: Prefer callee-save because of presence of kills at location %d. First "
+                        "RefPosition is at #%d @%d and current RefPosition is #%d @%d\n",
+                        theInterval->intervalIndex, recentKillLocation, firstPosition->rpNum,
+                        firstPosition->nodeLocation, rp->rpNum, rp->nodeLocation);
+
+                theInterval->preferCalleeSave = true;
+                if (!theInterval->isWriteThru)
+                {
+                    regMaskTP killMask = RBM_CALLEE_TRASH;
+                    // if there is no FP used, we can ignore the FP kills
+                    if (!compiler->compFloatingPointUsed)
+                    {
+                        killMask &= ~RBM_FLT_CALLEE_TRASH;
+                    }
+
+                    //  We are more conservative about allocating callee-saves registers to write-thru vars,
+                    //  since a call only requires reloading after (not spilling before). So we record (above)
+                    //  the fact that we'd prefer a callee-save register, but we don't update the preferences at
+                    //  this point. See the "heuristics for writeThru intervals" in 'buildIntervals()'.
+                    regMaskTP newPreferences = allRegs(theInterval->registerType) & (~killMask);
+#ifdef DEBUG
+                    if (VERBOSE)
+                    {
+                        printf("Interval %2u: Current preference= ", theInterval->intervalIndex);
+                        dumpRegMask(theInterval->registerPreferences);
+                        printf(", New preference= ");
+                        dumpRegMask(newPreferences);
+                        printf(", ");
+                    }
+#endif
+                    if (newPreferences != RBM_NONE)
+                    {
+                        calleeSavePreference = newPreferences;
+                    }
+                }
+            }
+        }
+    }
 
 #ifdef DEBUG
     if (!doReverseCallerCallee())
 #endif // DEBUG
     {
         // Set preferences so that this register set will be preferred for earlier refs
-        theInterval->mergeRegisterPreferences(rp->registerAssignment);
+        theInterval->mergeRegisterPreferences(calleeSavePreference);
+    }
+
+#ifdef DEBUG
+    if (VERBOSE)
+    {
+        printf("Merged preference= ");
+        dumpRegMask(theInterval->registerPreferences);
+        printf("\n");
     }
+#endif
 }
 
 //------------------------------------------------------------------------
@@ -1167,6 +1223,13 @@ bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLo
         // if (!blockSequence[curBBSeqNum]->isRunRarely())
         if (enregisterLocalVars)
         {
+            const bool isCallKill = ((killMask == RBM_INT_CALLEE_TRASH) || (killMask == RBM_CALLEE_TRASH));
+            if (isCallKill)
+            {
+                JITDUMP("Recording location of kill at %d\n", recentRefPosition->nodeLocation);
+                recentKillLocation = recentRefPosition->nodeLocation;
+            }
+
             VarSetOps::Iter iter(compiler, currentLiveVars);
             unsigned        varIndex = 0;
             while (iter.NextElem(&varIndex))
@@ -1187,9 +1250,7 @@ bool LinearScan::buildKillPositionsForNode(GenTree* tree, LsraLocation currentLo
                 {
                     continue;
                 }
-                Interval*  interval   = getIntervalForLocalVar(varIndex);
-                const bool isCallKill = ((killMask == RBM_INT_CALLEE_TRASH) || (killMask == RBM_CALLEE_TRASH));
-
+                Interval* interval = getIntervalForLocalVar(varIndex);
                 if (isCallKill)
                 {
                     interval->preferCalleeSave = true;