Skip to content

Register stack

register_stack

OutOfRegisters

Bases: DiagnosticException

Source code in xdsl/backend/register_stack.py 
15
16
17
class OutOfRegisters(DiagnosticException):
    def __str__(self):
        return "Out of registers."

__str__()

Source code in xdsl/backend/register_stack.py 
16
17
def __str__(self):
    return "Out of registers."

RegisterStack dataclass

LIFO stack of registers available for allocation.

Source code in xdsl/backend/register_stack.py 
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
@dataclass
class RegisterStack:
    """
    LIFO stack of registers available for allocation.
    """

    allocatable_registers: defaultdict[str, set[int]] = field(
        default_factory=lambda: defaultdict(lambda: set[int]())
    )
    """
    Registers that can be used by the register allocator for a given register set.
    """

    next_infinite_indices: defaultdict[str, int] = field(
        default_factory=lambda: defaultdict(lambda: 0)
    )
    """Next index for a given register set."""

    reserved_registers: defaultdict[str, defaultdict[int, int]] = field(
        default_factory=lambda: defaultdict(lambda: defaultdict[int, int](lambda: 0))
    )
    """
    Registers unavailable to be used by the register allocator for a given register set.
    """

    available_registers: defaultdict[str, list[int]] = field(
        default_factory=lambda: defaultdict(list[int])
    )
    """
    Registers from a given register set that values can be allocated to in the current
    context.
    """

    allow_infinite: bool = False
    """
    When there are no more registers, use infinite register set.
    """

    @classmethod
    def default_allocatable_registers(cls) -> Iterable[RegisterType]:
        """
        The default registers to be made available when instantiating the stack.
        """
        return ()

    @classmethod
    def get(
        cls,
        allocatable_registers: Iterable[RegisterType] | None = None,
        *,
        allow_infinite: bool = False,
    ):
        if allocatable_registers is None:
            allocatable_registers = cls.default_allocatable_registers()
        res = cls(allow_infinite=allow_infinite)
        for reg in allocatable_registers:
            res.include_register(reg)
        return res

    def push(self, reg: RegisterType) -> None:
        """
        Return a register to be made available for allocation.
        """
        if not isinstance(reg.index, IntAttr):
            raise ValueError("Cannot push an unallocated register")

        index = reg.index.data
        register_set = reg.name
        if (
            index in self.reserved_registers[register_set]
            or index not in self.allocatable_registers[register_set]
        ) and 0 <= index:
            return

        self.available_registers[register_set].append(index)

    def pop(self, reg_type: type[_T]) -> _T:
        """
        Get the next available register for allocation.
        """
        register_set = reg_type.name
        available_registers = self.available_registers[register_set]

        if available_registers:
            reg = reg_type.from_index(available_registers.pop())
        else:
            if self.allow_infinite:
                reg = reg_type.infinite_register(
                    self.next_infinite_indices[register_set]
                )
                self.next_infinite_indices[register_set] += 1
            else:
                raise OutOfRegisters

        reserved_registers = self.reserved_registers[reg_type.name]

        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

    def reserve_register(self, reg: RegisterType) -> None:
        """
        Increase the reservation count for a register.
        If the reservation count is greater than 0, a register cannot be pushed back
        onto the stack.
        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)
        self.reserved_registers[reg.name][reg.index.data] += 1

    @contextmanager
    def reserve_registers(self, regs: Sequence[RegisterType]):
        for reg in regs:
            self.reserve_register(reg)

        yield

        for reg in regs:
            self.unreserve_register(reg)

    def unreserve_register(self, reg: RegisterType) -> 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)
        reserved_registers = self.reserved_registers[reg.name]
        if reg.index.data not in reserved_registers:
            raise ValueError(f"Cannot unreserve register {reg.register_name}")
        reserved_registers[reg.index.data] -= 1
        if not reserved_registers[reg.index.data]:
            del reserved_registers[reg.index.data]

    def include_register(self, reg: RegisterType) -> None:
        """
        Makes register available for allocation.
        """
        assert not isinstance(reg.index, NoneAttr)
        self.allocatable_registers[reg.name].add(reg.index.data)
        self.push(reg)

    def exclude_register(self, reg: RegisterType) -> None:
        """
        Removes register from available set, if present.
        """
        assert isinstance(reg.index, IntAttr)
        index = reg.index.data
        available_registers = self.available_registers[reg.name]
        allocatable_registers = self.allocatable_registers[reg.name]
        if index in available_registers:
            available_registers.remove(index)
        if index in allocatable_registers:
            allocatable_registers.remove(reg.index.data)

allocatable_registers: defaultdict[str, set[int]] = field(default_factory=(lambda: defaultdict(lambda: set[int]()))) class-attribute instance-attribute

Registers that can be used by the register allocator for a given register set.

next_infinite_indices: defaultdict[str, int] = field(default_factory=(lambda: defaultdict(lambda: 0))) class-attribute instance-attribute

Next index for a given register set.

reserved_registers: defaultdict[str, defaultdict[int, int]] = field(default_factory=(lambda: defaultdict(lambda: defaultdict[int, int](lambda: 0)))) class-attribute instance-attribute

Registers unavailable to be used by the register allocator for a given register set.

available_registers: defaultdict[str, list[int]] = field(default_factory=(lambda: defaultdict(list[int]))) class-attribute instance-attribute

Registers from a given register set that values can be allocated to in the current context.

allow_infinite: bool = False class-attribute instance-attribute

When there are no more registers, use infinite register set.

__init__(allocatable_registers: defaultdict[str, set[int]] = (lambda: defaultdict(lambda: set[int]()))(), next_infinite_indices: defaultdict[str, int] = (lambda: defaultdict(lambda: 0))(), reserved_registers: defaultdict[str, defaultdict[int, int]] = (lambda: defaultdict(lambda: defaultdict[int, int](lambda: 0)))(), available_registers: defaultdict[str, list[int]] = (lambda: defaultdict(list[int]))(), allow_infinite: bool = False) -> None

default_allocatable_registers() -> Iterable[RegisterType] classmethod

The default registers to be made available when instantiating the stack.

Source code in xdsl/backend/register_stack.py 
58
59
60
61
62
63
@classmethod
def default_allocatable_registers(cls) -> Iterable[RegisterType]:
    """
    The default registers to be made available when instantiating the stack.
    """
    return ()

get(allocatable_registers: Iterable[RegisterType] | None = None, *, allow_infinite: bool = False) classmethod

Source code in xdsl/backend/register_stack.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
@classmethod
def get(
    cls,
    allocatable_registers: Iterable[RegisterType] | None = None,
    *,
    allow_infinite: bool = False,
):
    if allocatable_registers is None:
        allocatable_registers = cls.default_allocatable_registers()
    res = cls(allow_infinite=allow_infinite)
    for reg in allocatable_registers:
        res.include_register(reg)
    return res

push(reg: RegisterType) -> None

Return a register to be made available for allocation.

Source code in xdsl/backend/register_stack.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def push(self, reg: RegisterType) -> None:
    """
    Return a register to be made available for allocation.
    """
    if not isinstance(reg.index, IntAttr):
        raise ValueError("Cannot push an unallocated register")

    index = reg.index.data
    register_set = reg.name
    if (
        index in self.reserved_registers[register_set]
        or index not in self.allocatable_registers[register_set]
    ) and 0 <= index:
        return

    self.available_registers[register_set].append(index)

pop(reg_type: type[_T]) -> _T

Get the next available register for allocation.

Source code in xdsl/backend/register_stack.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def pop(self, reg_type: type[_T]) -> _T:
    """
    Get the next available register for allocation.
    """
    register_set = reg_type.name
    available_registers = self.available_registers[register_set]

    if available_registers:
        reg = reg_type.from_index(available_registers.pop())
    else:
        if self.allow_infinite:
            reg = reg_type.infinite_register(
                self.next_infinite_indices[register_set]
            )
            self.next_infinite_indices[register_set] += 1
        else:
            raise OutOfRegisters

    reserved_registers = self.reserved_registers[reg_type.name]

    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

reserve_register(reg: RegisterType) -> None

Increase the reservation count for a register. If the reservation count is greater than 0, a register cannot be pushed back onto the stack. It is invalid to reserve a register that is available, and popping it before unreserving a register will result in an AssertionError.

Source code in xdsl/backend/register_stack.py 
122
123
124
125
126
127
128
129
130
131
def reserve_register(self, reg: RegisterType) -> None:
    """
    Increase the reservation count for a register.
    If the reservation count is greater than 0, a register cannot be pushed back
    onto the stack.
    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)
    self.reserved_registers[reg.name][reg.index.data] += 1

reserve_registers(regs: Sequence[RegisterType])

Source code in xdsl/backend/register_stack.py 
133
134
135
136
137
138
139
140
141
@contextmanager
def reserve_registers(self, regs: Sequence[RegisterType]):
    for reg in regs:
        self.reserve_register(reg)

    yield

    for reg in regs:
        self.unreserve_register(reg)

unreserve_register(reg: RegisterType) -> None

Decrease the reservation count for a register. If the reservation count is 0, make the register available for allocation.

Source code in xdsl/backend/register_stack.py 
143
144
145
146
147
148
149
150
151
152
153
154
def unreserve_register(self, reg: RegisterType) -> 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)
    reserved_registers = self.reserved_registers[reg.name]
    if reg.index.data not in reserved_registers:
        raise ValueError(f"Cannot unreserve register {reg.register_name}")
    reserved_registers[reg.index.data] -= 1
    if not reserved_registers[reg.index.data]:
        del reserved_registers[reg.index.data]

include_register(reg: RegisterType) -> None

Makes register available for allocation.

Source code in xdsl/backend/register_stack.py 
156
157
158
159
160
161
162
def include_register(self, reg: RegisterType) -> None:
    """
    Makes register available for allocation.
    """
    assert not isinstance(reg.index, NoneAttr)
    self.allocatable_registers[reg.name].add(reg.index.data)
    self.push(reg)

exclude_register(reg: RegisterType) -> None

Removes register from available set, if present.

Source code in xdsl/backend/register_stack.py 
164
165
166
167
168
169
170
171
172
173
174
175
def exclude_register(self, reg: RegisterType) -> None:
    """
    Removes register from available set, if present.
    """
    assert isinstance(reg.index, IntAttr)
    index = reg.index.data
    available_registers = self.available_registers[reg.name]
    allocatable_registers = self.allocatable_registers[reg.name]
    if index in available_registers:
        available_registers.remove(index)
    if index in allocatable_registers:
        allocatable_registers.remove(reg.index.data)