Check if interval interfers with the call

Author: kunalspathak

src/coreclr/jit/clrjit.natvis

@@ -153,7 +153,7 @@ Documentation for VS debugger format specifiers: https://docs.microsoft.com/en-u
   </Type>
 
   <Type Name="RefPosition">
-    <DisplayString>[#{rpNum,d} - {refType,en}]</DisplayString>
+    <DisplayString>[{nodeLocation,d}. #{rpNum,d} - {refType,en}]</DisplayString>
     <Expand>
         <Item Name="Referent" Condition="this->isPhysRegRef">(RegRecord*)this-&gt;referent</Item>
         <Item Name="Referent" Condition="!this->isPhysRegRef">(Interval*)this-&gt;referent</Item>

src/coreclr/jit/lsra.cpp

@@ -624,6 +624,8 @@ LinearScan::LinearScan(Compiler* theCompiler)
     , intervals(theCompiler->getAllocator(CMK_LSRA_Interval))
     , allocationPassComplete(false)
     , refPositions(theCompiler->getAllocator(CMK_LSRA_RefPosition))
+    , callRefPositionLocations(theCompiler->getAllocator(CMK_LSRA))
+    , callKillMasks(theCompiler->getAllocator(CMK_LSRA))
     , listNodePool(theCompiler)
 {
     regSelector  = new (theCompiler, CMK_LSRA) RegisterSelection(this);
@@ -703,6 +705,8 @@ LinearScan::LinearScan(Compiler* theCompiler)
     blockSequenceWorkList = nullptr;
     curBBSeqNum           = 0;
     bbSeqCount            = 0;
+    callRefPositionCount  = 0;
+    recentRefPosition     = nullptr;
 
     // Information about each block, including predecessor blocks used for variable locations at block entry.
     blockInfo = nullptr;

src/coreclr/jit/lsra.h

@@ -1582,6 +1582,15 @@ class LinearScan : public LinearScanInterface
     // Ordered list of RefPositions
     RefPositionList refPositions;
 
+    // Most recent refpostion assigned
+    RefPosition* recentRefPosition;
+
+    // Ordered list of locations where CALL happens
+    jitstd::list<LsraLocation> callRefPositionLocations;
+    // Equivalent list of kill masks for those calls
+    jitstd::list<regMaskTP> callKillMasks;
+    int                     callRefPositionCount;
+
     // 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;
 }
 
@@ -1141,8 +1142,9 @@ regMaskTP LinearScan::getKillSetForNode(GenTree* tree)
 //
 // Notes:
 //    The return value is needed because if we have any kills, we need to make sure that
-//    all defs are located AFTER the kills.  On the other hand, if there aren't kills,
-//    the multiple defs for a regPair are in different locations.
+//    all defs are located AFTER the kills. TODO: This is always the case.
+//    On the other hand, if there aren't kills, the multiple defs for a regPair are in
+//    different locations.
 //    If we generate any kills, we will mark all currentLiveVars as being preferenced
 //    to avoid the killed registers.  This is somewhat conservative.
 //
@@ -1167,6 +1169,18 @@ 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)
+            {
+                assert(recentRefPosition);
+                if (recentRefPosition->refType != RefTypeBB)
+                {
+                    callRefPositionLocations.push_back(recentRefPosition->nodeLocation);
+                    callKillMasks.push_back(killMask);
+                    callRefPositionCount++;
+                }
+            }
+
             VarSetOps::Iter iter(compiler, currentLiveVars);
             unsigned        varIndex = 0;
             while (iter.NextElem(&varIndex))
@@ -1187,9 +1201,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;
@@ -2613,6 +2625,74 @@ void LinearScan::buildIntervals()
             }
         }
 
+        // Adjust preferCalleeSave
+        if (callRefPositionCount > 0)
+        {
+            LsraLocation firstCallLocation = callRefPositionLocations.front();
+            LsraLocation lastCallLocation  = callRefPositionLocations.back();
+
+            JITDUMP("\n\nCHECKING if Intervals can use callee-save registers");
+            JITDUMP("\n==============================\n");
+            for (Interval& interval : intervals)
+            {
+                if (interval.isLocalVar)
+                {
+                    // We already handle them in BuildKillPositionsForNode.
+                    continue;
+                }
+
+                LsraLocation firstRef = interval.firstRefPosition->nodeLocation;
+                LsraLocation lastRef  = interval.lastRefPosition->nodeLocation;
+
+                if ((firstRef > lastCallLocation) || (lastRef < firstCallLocation))
+                {
+                    // Current interval doesn't interfer with any calls. Skip it.
+                    continue;
+                }
+
+                jitstd::list<regMaskTP> callKillMasksLocal = callKillMasks;
+                for (LsraLocation& callLocation : callRefPositionLocations)
+                {
+                    regMaskTP killMask = callKillMasksLocal.front();
+                    callKillMasksLocal.pop_front();
+
+                    if ((firstRef <= callLocation) && (callLocation < lastRef))
+                    {
+                        interval.preferCalleeSave = true;
+                        if (!interval.isWriteThru)
+                        {
+                            // 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(interval.registerType) & ~killMask;
+
+                            if (newPreferences != RBM_NONE)
+                            {
+                                if (VERBOSE)
+                                {
+                                    printf("Interval %2u: Update preferences (callee-save) from ",
+                                           interval.intervalIndex);
+                                    dumpRegMask(interval.registerPreferences);
+                                    printf(" to ");
+                                    dumpRegMask(newPreferences);
+                                    printf("\n");
+                                }
+                                interval.updateRegisterPreferences(newPreferences);
+                            }
+                        }
+                        break;
+                    }
+
+                    if (lastRef < callLocation)
+                    {
+                        // We passed lastRef, no point in checking further callRefPositions.
+                        break;
+                    }
+                }
+            }
+        }
+
 #ifdef DEBUG
         if (getLsraExtendLifeTimes())
         {