backend: (riscv) split out reserved int and float registers in regalloc

tests/backend/riscv/test_register_queue.py

@@ -46,16 +46,16 @@ def test_reserve_register():
     register_queue = RiscvRegisterQueue()
 
     register_queue.reserve_register(riscv.IntRegisterType("j0"))
-    assert register_queue.reserved_registers[riscv.IntRegisterType("j0")] == 1
+    assert register_queue.reserved_int_registers[riscv.IntRegisterType("j0")] == 1
 
     register_queue.reserve_register(riscv.IntRegisterType("j0"))
-    assert register_queue.reserved_registers[riscv.IntRegisterType("j0")] == 2
+    assert register_queue.reserved_int_registers[riscv.IntRegisterType("j0")] == 2
 
     register_queue.unreserve_register(riscv.IntRegisterType("j0"))
-    assert register_queue.reserved_registers[riscv.IntRegisterType("j0")] == 1
+    assert register_queue.reserved_int_registers[riscv.IntRegisterType("j0")] == 1
 
     register_queue.unreserve_register(riscv.IntRegisterType("j0"))
-    assert riscv.IntRegisterType("j0") not in register_queue.reserved_registers
+    assert riscv.IntRegisterType("j0") not in register_queue.reserved_int_registers
     assert riscv.IntRegisterType("j0") not in register_queue.available_int_registers
 
     # Check assertion error when reserving an available register

xdsl/backend/riscv/riscv_register_queue.py

@@ -30,13 +30,16 @@ class RiscvRegisterQueue(RegisterQueue[IntRegisterType | FloatRegisterType]):
     _fj_idx: int = 0
     """Next `fj` register index."""
 
-    reserved_registers: defaultdict[IntRegisterType | FloatRegisterType, int] = field(
-        default_factory=lambda: defaultdict[IntRegisterType | FloatRegisterType, int](
-            lambda: 0
-        )
+    reserved_int_registers: defaultdict[IntRegisterType, int] = field(
+        default_factory=lambda: defaultdict[IntRegisterType, int](lambda: 0)
         | {r: 1 for r in RiscvRegisterQueue.DEFAULT_RESERVED_REGISTERS}
     )
-    "Registers unavailable to be used by the register allocator."
+    "Integer registers unavailable to be used by the register allocator."
+
+    reserved_float_registers: defaultdict[FloatRegisterType, int] = field(
+        default_factory=lambda: defaultdict[FloatRegisterType, int](lambda: 0)
+    )
+    "Floating-point registers unavailable to be used by the register allocator."
 
     available_int_registers: list[IntRegisterType] = field(
         default_factory=lambda: list(RiscvRegisterQueue.DEFAULT_INT_REGISTERS)
@@ -52,7 +55,7 @@ def push(self, reg: IntRegisterType | FloatRegisterType) -> None:
         """
         Return a register to be made available for allocation.
         """
-        if reg in self.reserved_registers:
+        if reg in self.reserved_int_registers or reg in self.reserved_float_registers:
             return
         if not reg.is_allocated:
             raise ValueError("Cannot push an unallocated register")
@@ -88,7 +91,13 @@ def pop(
                 reg = reg_type(f"fj{self._fj_idx}")
                 self._fj_idx += 1
 
-        assert reg not in self.reserved_registers, (
+        reserved_registers = (
+            self.reserved_int_registers
+            if issubclass(reg_type, IntRegisterType)
+            else self.reserved_float_registers
+        )
+
+        assert reg not in reserved_registers, (
             f"Cannot pop a reserved register ({reg.register_name}), it must have been reserved while available."
         )
         return reg
@@ -101,18 +110,28 @@ def reserve_register(self, reg: IntRegisterType | FloatRegisterType) -> None:
         It is invalid to reserve a register that is available, and popping it before
         unreserving a register will result in an AssertionError.
         """
-        self.reserved_registers[reg] += 1
+        if isinstance(reg, IntRegisterType):
+            self.reserved_int_registers[reg] += 1
+        if isinstance(reg, FloatRegisterType):
+            self.reserved_float_registers[reg] += 1
 
     def unreserve_register(self, reg: IntRegisterType | FloatRegisterType) -> None:
         """
         Decrease the reservation count for a register. If the reservation count is 0, make
         the register available for allocation.
         """
-        if reg not in self.reserved_registers:
-            raise ValueError("Cannot unreserve an unreserved register")
-        self.reserved_registers[reg] -= 1
-        if not self.reserved_registers[reg]:
-            del self.reserved_registers[reg]
+        if isinstance(reg, IntRegisterType):
+            if reg not in self.reserved_int_registers:
+                raise ValueError(f"Cannot unreserve register {reg.spelling}")
+            self.reserved_int_registers[reg] -= 1
+            if not self.reserved_int_registers[reg]:
+                del self.reserved_int_registers[reg]
+        if isinstance(reg, FloatRegisterType):
+            if reg not in self.reserved_float_registers:
+                raise ValueError(f"Cannot unreserve register {reg.spelling}")
+            self.reserved_float_registers[reg] -= 1
+            if not self.reserved_float_registers[reg]:
+                del self.reserved_float_registers[reg]
 
     def limit_registers(self, limit: int) -> None:
         """