Combine free and busy register allocation

src/coreclr/src/jit/codegenarmarch.cpp

@@ -1405,7 +1405,7 @@ void CodeGen::genMultiRegStoreToSIMDLocal(GenTreeLclVar* lclNode)
     for (int i = regCount - 1; i >= 0; --i)
     {
         var_types type = op1->gtSkipReloadOrCopy()->GetRegTypeByIndex(i);
-        regNumber reg  = op1->GetRegByIndex(i);
+        regNumber reg  = actualOp1->GetRegByIndex(i);
         if (op1->IsCopyOrReload())
         {
             // GT_COPY/GT_RELOAD will have valid reg for those positions

src/coreclr/src/jit/lsra.h

@@ -471,12 +471,11 @@ class RegRecord : public Referenceable
 public:
     RegRecord()
     {
-        assignedInterval    = nullptr;
-        previousInterval    = nullptr;
-        regNum              = REG_NA;
-        isCalleeSave        = false;
-        registerType        = IntRegisterType;
-        isBusyUntilNextKill = false;
+        assignedInterval = nullptr;
+        previousInterval = nullptr;
+        regNum           = REG_NA;
+        isCalleeSave     = false;
+        registerType     = IntRegisterType;
     }
 
     void init(regNumber reg)
@@ -511,8 +510,6 @@ class RegRecord : public Referenceable
     void tinyDump();
 #endif // DEBUG
 
-    bool isFree();
-
     // RefPosition   * getNextRefPosition();
     // LsraLocation    getNextRefLocation();
 
@@ -528,15 +525,10 @@ class RegRecord : public Referenceable
     // assignedInterval becomes inactive.
     Interval* previousInterval;
 
-    regNumber    regNum;
-    bool         isCalleeSave;
-    RegisterType registerType;
-    // This register must be considered busy until the next time it is explicitly killed.
-    // This is used so that putarg_reg can avoid killing its lclVar source, while avoiding
-    // the problem with the reg becoming free if the last-use is encountered before the call.
-    bool isBusyUntilNextKill;
-
-    bool conflictingFixedRegReference(RefPosition* refPosition);
+    regNumber     regNum;
+    bool          isCalleeSave;
+    RegisterType  registerType;
+    unsigned char regOrder;
 };
 
 inline bool leafInRange(GenTree* leaf, int lower, int upper)
@@ -976,18 +968,16 @@ class LinearScan : public LinearScanInterface
     bool isSecondHalfReg(RegRecord* regRec, Interval* interval);
     RegRecord* getSecondHalfRegRec(RegRecord* regRec);
     RegRecord* findAnotherHalfRegRec(RegRecord* regRec);
-    bool canSpillDoubleReg(RegRecord*            physRegRecord,
-                           LsraLocation          refLocation,
-                           BasicBlock::weight_t* recentAssignedRefWeight);
+    bool canSpillDoubleReg(RegRecord* physRegRecord, LsraLocation refLocation);
     void unassignDoublePhysReg(RegRecord* doubleRegRecord);
 #endif
     void updateAssignedInterval(RegRecord* reg, Interval* interval, RegisterType regType);
     void updatePreviousInterval(RegRecord* reg, Interval* interval, RegisterType regType);
     bool canRestorePreviousInterval(RegRecord* regRec, Interval* assignedInterval);
     bool isAssignedToInterval(Interval* interval, RegRecord* regRec);
     bool isRefPositionActive(RefPosition* refPosition, LsraLocation refLocation);
-    bool canSpillReg(RegRecord* physRegRecord, LsraLocation refLocation, BasicBlock::weight_t* recentAssignedRefWeight);
-    bool isRegInUse(RegRecord* regRec, RefPosition* refPosition);
+    bool canSpillReg(RegRecord* physRegRecord, LsraLocation refLocation);
+    float getSpillWeight(RegRecord* physRegRecord);
 
     // insert refpositions representing prolog zero-inits which will be added later
     void insertZeroInitRefPositions();
@@ -1056,11 +1046,7 @@ class LinearScan : public LinearScanInterface
     regMaskTP allSIMDRegs();
     regMaskTP internalFloatRegCandidates();
 
-    bool registerIsFree(regNumber regNum, RegisterType regType);
-    bool registerIsAvailable(RegRecord*    physRegRecord,
-                             LsraLocation  currentLoc,
-                             LsraLocation* nextRefLocationPtr,
-                             RegisterType  regType);
+    void makeRegisterInactive(RegRecord* physRegRecord);
     void freeRegister(RegRecord* physRegRecord);
     void freeRegisters(regMaskTP regsToFree);
 
@@ -1143,15 +1129,12 @@ class LinearScan : public LinearScanInterface
      * Register management
      ****************************************************************************/
     RegisterType getRegisterType(Interval* currentInterval, RefPosition* refPosition);
-    regNumber tryAllocateFreeReg(Interval* current, RefPosition* refPosition);
-    regNumber allocateBusyReg(Interval* current, RefPosition* refPosition, bool allocateIfProfitable);
+
+    regNumber allocateReg(Interval* current, RefPosition* refPosition);
     regNumber assignCopyReg(RefPosition* refPosition);
 
     bool isMatchingConstant(RegRecord* physRegRecord, RefPosition* refPosition);
-    bool isSpillCandidate(Interval*     current,
-                          RefPosition*  refPosition,
-                          RegRecord*    physRegRecord,
-                          LsraLocation& nextLocation);
+    bool isSpillCandidate(Interval* current, RefPosition* refPosition, RegRecord* physRegRecord);
     void checkAndAssignInterval(RegRecord* regRec, Interval* interval);
     void assignPhysReg(RegRecord* regRec, Interval* interval);
     void assignPhysReg(regNumber reg, Interval* interval)
@@ -1160,7 +1143,6 @@ class LinearScan : public LinearScanInterface
     }
 
     bool isAssigned(RegRecord* regRec ARM_ARG(RegisterType newRegType));
-    bool isAssigned(RegRecord* regRec, LsraLocation lastLocation ARM_ARG(RegisterType newRegType));
     void checkAndClearInterval(RegRecord* regRec, RefPosition* spillRefPosition);
     void unassignPhysReg(RegRecord* regRec ARM_ARG(RegisterType newRegType));
     void unassignPhysReg(RegRecord* regRec, RefPosition* spillRefPosition);
@@ -1176,6 +1158,59 @@ class LinearScan : public LinearScanInterface
 
     void spillGCRefs(RefPosition* killRefPosition);
 
+    /*****************************************************************************
+    * Register selection
+    ****************************************************************************/
+    regMaskTP getFreeCandidates(regMaskTP candidates, var_types regType)
+    {
+        regMaskTP result = candidates & m_AvailableRegs;
+#ifdef TARGET_ARM
+        // For TYP_DOUBLE on ARM, we can only use register for which the odd half is
+        // also available.
+        if (regType == TYP_DOUBLE)
+        {
+            result &= (m_AvailableRegs >> 1);
+        }
+#endif // TARGET_ARM
+        return result;
+    }
+
+    struct registerSelector
+    {
+        regMaskTP candidates;
+        int       score;
+#ifdef TARGET_ARM
+        var_types regType;
+#endif // TARGET_ARM
+
+        // Apply a simple mask-based selection heuristic, and return 'true' if we now have a single candidate.
+        bool applySelection(int selectionScore, regMaskTP selectionCandidates)
+        {
+            regMaskTP newCandidates = candidates & selectionCandidates;
+            if (newCandidates != RBM_NONE)
+            {
+                score += selectionScore;
+                candidates = newCandidates;
+                return isSingleRegister(candidates);
+            }
+            return false;
+        }
+
+        // Select a single register, if it is in the candidate set.
+        // Return true if so.
+        bool applySingleRegSelection(int selectionScore, regMaskTP selectionCandidate)
+        {
+            assert(isSingleRegister(selectionCandidate));
+            regMaskTP newCandidates = candidates & selectionCandidate;
+            if (newCandidates != RBM_NONE)
+            {
+                candidates = newCandidates;
+                return true;
+            }
+            return false;
+        }
+    };
+
     /*****************************************************************************
      * For Resolution phase
      ****************************************************************************/
@@ -1316,8 +1351,7 @@ class LinearScan : public LinearScanInterface
         // Allocation decisions
         LSRA_EVENT_FIXED_REG, LSRA_EVENT_EXP_USE, LSRA_EVENT_ZERO_REF, LSRA_EVENT_NO_ENTRY_REG_ALLOCATED,
         LSRA_EVENT_KEPT_ALLOCATION, LSRA_EVENT_COPY_REG, LSRA_EVENT_MOVE_REG, LSRA_EVENT_ALLOC_REG,
-        LSRA_EVENT_ALLOC_SPILLED_REG, LSRA_EVENT_NO_REG_ALLOCATED, LSRA_EVENT_RELOAD, LSRA_EVENT_SPECIAL_PUTARG,
-        LSRA_EVENT_REUSE_REG,
+        LSRA_EVENT_NO_REG_ALLOCATED, LSRA_EVENT_RELOAD, LSRA_EVENT_SPECIAL_PUTARG, LSRA_EVENT_REUSE_REG,
     };
     void dumpLsraAllocationEvent(LsraDumpEvent event,
                                  Interval*     interval     = nullptr,
@@ -1488,6 +1522,154 @@ class LinearScan : public LinearScanInterface
     VARSET_TP largeVectorCalleeSaveCandidateVars;
 #endif // FEATURE_PARTIAL_SIMD_CALLEE_SAVE
 
+    //-----------------------------------------------------------------------
+    // Register status
+    //-----------------------------------------------------------------------
+
+    regMaskTP m_AvailableRegs;
+    regNumber getRegForType(regNumber reg, var_types regType)
+    {
+#ifdef TARGET_ARM
+        if ((regType == TYP_DOUBLE) && !genIsValidDoubleReg(reg))
+        {
+            reg = REG_PREV(reg);
+        }
+#endif // TARGET_ARM
+        return reg;
+    }
+
+    regMaskTP getRegMask(regNumber reg, var_types regType)
+    {
+        reg               = getRegForType(reg, regType);
+        regMaskTP regMask = genRegMask(reg);
+#ifdef TARGET_ARM
+        if (regType == TYP_DOUBLE)
+        {
+            assert(genIsValidDoubleReg(reg));
+            regMask |= (regMask << 1);
+        }
+#endif // TARGET_ARM
+        return regMask;
+    }
+
+    void resetAvailableRegs()
+    {
+        m_AvailableRegs          = (availableIntRegs | availableFloatRegs);
+        m_RegistersWithConstants = RBM_NONE;
+    }
+
+    bool isRegAvailable(regNumber reg, var_types regType)
+    {
+        regMaskTP regMask = getRegMask(reg, regType);
+        return (m_AvailableRegs & regMask) == regMask;
+    }
+    void setRegsInUse(regMaskTP regMask)
+    {
+        m_AvailableRegs &= ~regMask;
+    }
+    void setRegInUse(regNumber reg, var_types regType)
+    {
+        regMaskTP regMask = getRegMask(reg, regType);
+        setRegsInUse(regMask);
+    }
+    void makeRegsAvailable(regMaskTP regMask)
+    {
+        m_AvailableRegs |= regMask;
+    }
+    void makeRegAvailable(regNumber reg, var_types regType)
+    {
+        regMaskTP regMask = getRegMask(reg, regType);
+        makeRegsAvailable(regMask);
+    }
+
+    void clearNextIntervalRef(regNumber reg, var_types regType);
+    void updateNextIntervalRef(regNumber reg, Interval* interval);
+
+    void clearSpillCost(regNumber reg, var_types regType);
+    void updateSpillCost(regNumber reg, Interval* interval);
+
+    regMaskTP m_RegistersWithConstants;
+    void clearConstantReg(regNumber reg, var_types regType)
+    {
+        m_RegistersWithConstants &= ~getRegMask(reg, regType);
+    }
+    void setConstantReg(regNumber reg, var_types regType)
+    {
+        m_RegistersWithConstants |= getRegMask(reg, regType);
+    }
+    bool isRegConstant(regNumber reg, var_types regType)
+    {
+        reg               = getRegForType(reg, regType);
+        regMaskTP regMask = getRegMask(reg, regType);
+        return (m_RegistersWithConstants & regMask) == regMask;
+    }
+    regMaskTP getMatchingConstants(regMaskTP mask, Interval* currentInterval, RefPosition* refPosition);
+
+    regMaskTP    fixedRegs;
+    LsraLocation nextFixedRef[REG_COUNT];
+    void updateNextFixedRef(RegRecord* regRecord, RefPosition* nextRefPosition);
+    LsraLocation getNextFixedRef(regNumber regNum, var_types regType)
+    {
+        LsraLocation loc = nextFixedRef[regNum];
+#ifdef TARGET_ARM
+        if (regType == TYP_DOUBLE)
+        {
+            loc = Min(loc, nextFixedRef[regNum + 1]);
+        }
+#endif
+        return loc;
+    }
+
+    LsraLocation nextIntervalRef[REG_COUNT];
+    LsraLocation getNextIntervalRef(regNumber regNum, var_types regType)
+    {
+        LsraLocation loc = nextIntervalRef[regNum];
+#ifdef TARGET_ARM
+        if (regType == TYP_DOUBLE)
+        {
+            loc = Min(loc, nextIntervalRef[regNum + 1]);
+        }
+#endif
+        return loc;
+    }
+    float spillCost[REG_COUNT];
+
+    regMaskTP regsBusyUntilKill;
+    regMaskTP regsInUseThisLocation;
+    regMaskTP regsInUseNextLocation;
+    bool isRegBusy(regNumber reg, var_types regType)
+    {
+        regMaskTP regMask = getRegMask(reg, regType);
+        return (regsBusyUntilKill & regMask) != RBM_NONE;
+    }
+    void setRegBusyUntilKill(regNumber reg, var_types regType)
+    {
+        regsBusyUntilKill |= getRegMask(reg, regType);
+    }
+    void clearRegBusyUntilKill(regNumber reg)
+    {
+        regsBusyUntilKill &= ~genRegMask(reg);
+    }
+
+    bool isRegInUse(regNumber reg, var_types regType)
+    {
+        regMaskTP regMask = getRegMask(reg, regType);
+        return (regsInUseThisLocation & regMask) != RBM_NONE;
+    }
+
+    void resetRegState()
+    {
+        resetAvailableRegs();
+        regsBusyUntilKill = RBM_NONE;
+    }
+
+    bool conflictingFixedRegReference(regNumber regNum, RefPosition* refPosition);
+
+    // This method should not be used and is here to retain old behavior.
+    // It should be replaced by isRegAvailable().
+    // See comment in allocateReg();
+    bool isFree(RegRecord* regRecord);
+
     //-----------------------------------------------------------------------
     // Build methods
     //-----------------------------------------------------------------------
@@ -1551,7 +1733,7 @@ class LinearScan : public LinearScanInterface
 
     int BuildSimple(GenTree* tree);
     int BuildOperandUses(GenTree* node, regMaskTP candidates = RBM_NONE);
-    int BuildDelayFreeUses(GenTree* node, regMaskTP candidates = RBM_NONE);
+    int BuildDelayFreeUses(GenTree* node, GenTree* rmwNode = nullptr, regMaskTP candidates = RBM_NONE);
     int BuildIndirUses(GenTreeIndir* indirTree, regMaskTP candidates = RBM_NONE);
     int BuildAddrUses(GenTree* addr, regMaskTP candidates = RBM_NONE);
     void HandleFloatVarArgs(GenTreeCall* call, GenTree* argNode, bool* callHasFloatRegArgs);
@@ -1728,7 +1910,7 @@ class Interval : public Referenceable
     // True if this interval is defined by a putArg, whose source is a non-last-use lclVar.
     // During allocation, this flag will be cleared if the source is not already in the required register.
     // Othewise, we will leave the register allocated to the lclVar, but mark the RegRecord as
-    // isBusyUntilNextKill, so that it won't be reused if the lclVar goes dead before the call.
+    // isBusyUntilKill, so that it won't be reused if the lclVar goes dead before the call.
     bool isSpecialPutArg : 1;
 
     // True if this interval interferes with a call.

src/coreclr/src/jit/lsraarm64.cpp

@@ -1098,8 +1098,8 @@ int LinearScan::BuildHWIntrinsic(GenTreeHWIntrinsic* intrinsicTree)
         {
             if (isRMW)
             {
-                srcCount += BuildDelayFreeUses(intrin.op2);
-                srcCount += BuildDelayFreeUses(intrin.op3, RBM_ASIMD_INDEXED_H_ELEMENT_ALLOWED_REGS);
+                srcCount += BuildDelayFreeUses(intrin.op2, nullptr);
+                srcCount += BuildDelayFreeUses(intrin.op3, nullptr, RBM_ASIMD_INDEXED_H_ELEMENT_ALLOWED_REGS);
             }
             else
             {