From 147a88b575b7bb0894bb5a7ac17717911dd6f902 Mon Sep 17 00:00:00 2001
From: Kunal Pathak <Kunal.Pathak@microsoft.com>
Date: Wed, 29 May 2024 15:13:38 -0700
Subject: [PATCH] AddRegsetForType()

---
 src/coreclr/jit/lsra.cpp         |  8 ++++----
 src/coreclr/jit/regMaskTPOps.cpp | 17 +++++++++++++++++
 src/coreclr/jit/target.h         |  1 +
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/coreclr/jit/lsra.cpp b/src/coreclr/jit/lsra.cpp
index 818488eca31ff..6c450334a4161 100644
--- a/src/coreclr/jit/lsra.cpp
+++ b/src/coreclr/jit/lsra.cpp
@@ -5221,13 +5221,13 @@ void LinearScan::allocateRegistersMinimal()
                 }
 #endif // TARGET_ARM
             }
-            regsInUseThisLocation |= currentRefPosition.registerAssignment;
+            regsInUseThisLocation.AddRegsetForType(currentRefPosition.registerAssignment, regRecord->registerType);
             INDEBUG(dumpLsraAllocationEvent(LSRA_EVENT_FIXED_REG, nullptr, currentRefPosition.assignedReg()));
 
 #ifdef SWIFT_SUPPORT
             if (currentRefPosition.delayRegFree)
             {
-                regsInUseNextLocation |= currentRefPosition.registerAssignment;
+                regsInUseNextLocation.AddRegsetForType(currentRefPosition.registerAssignment, regRecord->registerType);
             }
 #endif // SWIFT_SUPPORT
             continue;
@@ -5932,13 +5932,13 @@ void LinearScan::allocateRegisters()
                 }
 #endif // TARGET_ARM
             }
-            regsInUseThisLocation |= currentRefPosition.registerAssignment;
+            regsInUseThisLocation.AddRegsetForType(currentRefPosition.registerAssignment, regRecord->registerType);
             INDEBUG(dumpLsraAllocationEvent(LSRA_EVENT_FIXED_REG, nullptr, currentRefPosition.assignedReg()));
 
 #ifdef SWIFT_SUPPORT
             if (currentRefPosition.delayRegFree)
             {
-                regsInUseNextLocation |= currentRefPosition.registerAssignment;
+                regsInUseNextLocation.AddRegsetForType(currentRefPosition.registerAssignment, regRecord->registerType);
             }
 #endif // SWIFT_SUPPORT
             continue;
diff --git a/src/coreclr/jit/regMaskTPOps.cpp b/src/coreclr/jit/regMaskTPOps.cpp
index 1ea88adbd55e0..43deb7826e3f3 100644
--- a/src/coreclr/jit/regMaskTPOps.cpp
+++ b/src/coreclr/jit/regMaskTPOps.cpp
@@ -108,6 +108,23 @@ void regMaskTP::AddRegNum(regNumber reg, var_types type)
 #endif
 }
 
+//------------------------------------------------------------------------
+// AddRegsetForType: Add regs of `type` in mask.
+//
+// Parameters:
+//  maskToAdd  - Register to check
+//  type       - type of register
+//
+void regMaskTP::AddRegsetForType(SingleTypeRegSet regs, var_types type)
+{
+#ifdef HAS_MORE_THAN_64_REGISTERS
+    int index = getRegisterTypeIndex(type);
+    _registers[index] |= encodeForRegisterIndex(index, regs);
+#else
+    low |= regs;
+#endif
+}
+
 //------------------------------------------------------------------------
 // IsRegNumInMask: Checks if `reg` is in the mask
 //
diff --git a/src/coreclr/jit/target.h b/src/coreclr/jit/target.h
index d145c8fb381c7..ef5f8306ed31b 100644
--- a/src/coreclr/jit/target.h
+++ b/src/coreclr/jit/target.h
@@ -283,6 +283,7 @@ struct regMaskTP
     void AddGprRegs(SingleTypeRegSet gprRegs);
     void AddRegNum(regNumber reg, var_types type);
     void AddRegNumInMask(regNumber reg);
+    void AddRegsetForType(SingleTypeRegSet maskToAdd, var_types type);
     void RemoveRegNum(regNumber reg, var_types type);
     void RemoveRegNumFromMask(regNumber reg);    
     bool IsRegNumInMask(regNumber reg) const;