backend: use register index for allocation (#4047)

Now that the index uniquely identifies the register, we can use it in
the register queue, helping avoid collisions between different names for
the same register.

Author: superlopuh

tests/backend/riscv/test_register_queue.py

@@ -4,6 +4,7 @@
 
 from xdsl.backend.riscv.riscv_register_queue import RiscvRegisterQueue
 from xdsl.dialects import riscv
+from xdsl.dialects.builtin import IntAttr
 
 
 def test_default_reserved_registers():
@@ -51,39 +52,21 @@ def test_push_register():
 def test_reserve_register():
     register_queue = RiscvRegisterQueue()
 
-    register_queue.reserve_register(riscv.IntRegisterType.infinite_register(0))
-    assert (
-        register_queue.reserved_int_registers[
-            riscv.IntRegisterType.infinite_register(0)
-        ]
-        == 1
-    )
+    j0 = riscv.IntRegisterType.infinite_register(0)
+    assert isinstance(j0.index, IntAttr)
 
-    register_queue.reserve_register(riscv.IntRegisterType.infinite_register(0))
-    assert (
-        register_queue.reserved_int_registers[
-            riscv.IntRegisterType.infinite_register(0)
-        ]
-        == 2
-    )
+    register_queue.reserve_register(j0)
+    assert register_queue.reserved_int_registers[j0.index.data] == 1
 
-    register_queue.unreserve_register(riscv.IntRegisterType.infinite_register(0))
-    assert (
-        register_queue.reserved_int_registers[
-            riscv.IntRegisterType.infinite_register(0)
-        ]
-        == 1
-    )
+    register_queue.reserve_register(j0)
+    assert register_queue.reserved_int_registers[j0.index.data] == 2
 
-    register_queue.unreserve_register(riscv.IntRegisterType.infinite_register(0))
-    assert (
-        riscv.IntRegisterType.infinite_register(0)
-        not in register_queue.reserved_int_registers
-    )
-    assert (
-        riscv.IntRegisterType.infinite_register(0)
-        not in register_queue.available_int_registers
-    )
+    register_queue.unreserve_register(j0)
+    assert register_queue.reserved_int_registers[j0.index.data] == 1
+
+    register_queue.unreserve_register(j0)
+    assert j0 not in register_queue.reserved_int_registers
+    assert j0 not in register_queue.available_int_registers
 
     # Check assertion error when reserving an available register
     reg = register_queue.pop(riscv.IntRegisterType)

xdsl/backend/riscv/riscv_register_queue.py

@@ -1,8 +1,9 @@
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import overload
+from typing import cast, overload
 
 from xdsl.backend.register_queue import RegisterQueue
+from xdsl.dialects.builtin import IntAttr
 from xdsl.dialects.riscv import FloatRegisterType, IntRegisterType, Registers
 
 
@@ -30,14 +31,17 @@ class RiscvRegisterQueue(RegisterQueue[IntRegisterType | FloatRegisterType]):
     _fj_idx: int = 0
     """Next `fj` register index."""
 
-    reserved_int_registers: defaultdict[IntRegisterType, int] = field(
-        default_factory=lambda: defaultdict[IntRegisterType, int](lambda: 0)
-        | {r: 1 for r in RiscvRegisterQueue.DEFAULT_RESERVED_REGISTERS}
+    reserved_int_registers: defaultdict[int, int] = field(
+        default_factory=lambda: defaultdict[int, int](lambda: 0)
+        | {
+            cast(IntAttr, r.index).data: 1
+            for r in RiscvRegisterQueue.DEFAULT_RESERVED_REGISTERS
+        }
     )
     "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)
+    reserved_float_registers: defaultdict[int, int] = field(
+        default_factory=lambda: defaultdict[int, int](lambda: 0)
     )
     "Floating-point registers unavailable to be used by the register allocator."
 
@@ -55,13 +59,16 @@ def push(self, reg: IntRegisterType | FloatRegisterType) -> None:
         """
         Return a register to be made available for allocation.
         """
-        if reg in self.reserved_int_registers or reg in self.reserved_float_registers:
-            return
-        if not reg.is_allocated:
+        if not isinstance(reg.index, IntAttr):
             raise ValueError("Cannot push an unallocated register")
+
         if isinstance(reg, IntRegisterType):
+            if reg.index.data in self.reserved_int_registers:
+                return
             self.available_int_registers.append(reg)
         else:
+            if reg.index.data in self.reserved_float_registers:
+                return
             self.available_float_registers.append(reg)
 
     @overload
@@ -97,7 +104,8 @@ def pop(
             else self.reserved_float_registers
         )
 
-        assert reg not in reserved_registers, (
+        assert isinstance(reg.index, IntAttr)
+        assert reg.index.data not in reserved_registers, (
             f"Cannot pop a reserved register ({reg.register_name.data}), it must have been reserved while available."
         )
         return reg
@@ -110,28 +118,30 @@ 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.
         """
+        assert isinstance(reg.index, IntAttr)
         if isinstance(reg, IntRegisterType):
-            self.reserved_int_registers[reg] += 1
+            self.reserved_int_registers[reg.index.data] += 1
         if isinstance(reg, FloatRegisterType):
-            self.reserved_float_registers[reg] += 1
+            self.reserved_float_registers[reg.index.data] += 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.
         """
+        assert isinstance(reg.index, IntAttr)
         if isinstance(reg, IntRegisterType):
-            if reg not in self.reserved_int_registers:
+            if reg.index.data not in self.reserved_int_registers:
                 raise ValueError(f"Cannot unreserve register {reg.register_name}")
-            self.reserved_int_registers[reg] -= 1
-            if not self.reserved_int_registers[reg]:
-                del self.reserved_int_registers[reg]
+            self.reserved_int_registers[reg.index.data] -= 1
+            if not self.reserved_int_registers[reg.index.data]:
+                del self.reserved_int_registers[reg.index.data]
         if isinstance(reg, FloatRegisterType):
-            if reg not in self.reserved_float_registers:
+            if reg.index.data not in self.reserved_float_registers:
                 raise ValueError(f"Cannot unreserve register {reg.register_name}")
-            self.reserved_float_registers[reg] -= 1
-            if not self.reserved_float_registers[reg]:
-                del self.reserved_float_registers[reg]
+            self.reserved_float_registers[reg.index.data] -= 1
+            if not self.reserved_float_registers[reg.index.data]:
+                del self.reserved_float_registers[reg.index.data]
 
     def limit_registers(self, limit: int) -> None:
         """