Skip to content

Register allocator

register_allocator

ValueAllocator

Base class for register allocators.

Source code in xdsl/backend/register_allocator.py 
 14
 15
 16
 17
 18
 19
 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
class ValueAllocator:
    """
    Base class for register allocators.
    """

    available_registers: RegisterStack
    register_base_class: type[RegisterType]
    new_value_by_old_value: dict[SSAValue, SSAValue]

    def __init__(
        self,
        available_registers: RegisterStack,
        register_base_class: type[RegisterType],
    ) -> None:
        self.available_registers = available_registers
        self.register_base_class = register_base_class
        self.new_value_by_old_value = {}

    def new_type_for_value(self, reg: SSAValue) -> RegisterType | None:
        if isinstance(reg.type, self.register_base_class) and not reg.type.is_allocated:
            return self.available_registers.pop(type(reg.type))

    def _replace_value_with_new_type(
        self, val: SSAValue, new_type: Attribute
    ) -> SSAValue:
        new_val = Rewriter.replace_value_with_new_type(val, new_type)
        self.new_value_by_old_value[val] = new_val
        return new_val

    def allocate_value(self, val: SSAValue) -> SSAValue | None:
        """
        Allocate a register if not already allocated.
        """
        if val in self.new_value_by_old_value:
            return
        new_type = self.new_type_for_value(val)
        if new_type is not None:
            return self._replace_value_with_new_type(val, new_type)

    def allocate_values_same_reg(self, vals: Sequence[SSAValue]) -> bool:
        """
        Allocates the values passed in to the same register.
        If some of the values are already allocated, they must be allocated to the same
        register, and unallocated values are then allocated to this register.
        If the values passed in are already allocated to differing registers, a
        `DiagnosticException` is raised.
        """
        reg_types = set(val.type for val in vals)
        assert all(
            isinstance(reg_type, self.register_base_class) for reg_type in reg_types
        )
        reg_types = cast(set[RegisterType], reg_types)

        match len(reg_types):
            case 0:
                # No inputs, nothing to do
                return False
            case 1:
                # Single input, may already be allocated
                (reg_type,) = reg_types
                if reg_type.is_allocated:
                    return False
                else:
                    reg_type = self.available_registers.pop(type(reg_type))
            case 2:
                # Two inputs, either one is allocated or two
                reg_type_0, reg_type_1 = reg_types
                if reg_type_0.is_allocated:
                    if reg_type_1.is_allocated:
                        reg_names = [f"{reg_type}" for reg_type in reg_types]
                        reg_names.sort()
                        raise DiagnosticException(
                            f"Cannot allocate registers to the same register {reg_names}"
                        )
                    else:
                        reg_type = reg_type_0
                else:
                    reg_type = reg_type_1
            case _:
                # More than one input is allocated, meaning we can't allocate them to be
                # the same, error.
                reg_names = [f"{reg_type}" for reg_type in reg_types]
                reg_names.sort()
                raise DiagnosticException(
                    f"Cannot allocate registers to the same register {reg_names}"
                )

        did_allocate = False

        for val in vals:
            if val.type != reg_type:
                self._replace_value_with_new_type(val, reg_type)
                did_allocate = True

        return did_allocate

    def free_value(self, val: SSAValue) -> None:
        if isinstance(val.type, self.register_base_class) and val.type.is_allocated:
            self.available_registers.push(val.type)

available_registers: RegisterStack = available_registers instance-attribute

register_base_class: type[RegisterType] = register_base_class instance-attribute

new_value_by_old_value: dict[SSAValue, SSAValue] = {} instance-attribute

__init__(available_registers: RegisterStack, register_base_class: type[RegisterType]) -> None

Source code in xdsl/backend/register_allocator.py 
23
24
25
26
27
28
29
30
def __init__(
    self,
    available_registers: RegisterStack,
    register_base_class: type[RegisterType],
) -> None:
    self.available_registers = available_registers
    self.register_base_class = register_base_class
    self.new_value_by_old_value = {}

new_type_for_value(reg: SSAValue) -> RegisterType | None

Source code in xdsl/backend/register_allocator.py 
32
33
34
def new_type_for_value(self, reg: SSAValue) -> RegisterType | None:
    if isinstance(reg.type, self.register_base_class) and not reg.type.is_allocated:
        return self.available_registers.pop(type(reg.type))

allocate_value(val: SSAValue) -> SSAValue | None

Allocate a register if not already allocated.

Source code in xdsl/backend/register_allocator.py 
43
44
45
46
47
48
49
50
51
def allocate_value(self, val: SSAValue) -> SSAValue | None:
    """
    Allocate a register if not already allocated.
    """
    if val in self.new_value_by_old_value:
        return
    new_type = self.new_type_for_value(val)
    if new_type is not None:
        return self._replace_value_with_new_type(val, new_type)

allocate_values_same_reg(vals: Sequence[SSAValue]) -> bool

Allocates the values passed in to the same register. If some of the values are already allocated, they must be allocated to the same register, and unallocated values are then allocated to this register. If the values passed in are already allocated to differing registers, a DiagnosticException is raised.

Source code in xdsl/backend/register_allocator.py 
 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
def allocate_values_same_reg(self, vals: Sequence[SSAValue]) -> bool:
    """
    Allocates the values passed in to the same register.
    If some of the values are already allocated, they must be allocated to the same
    register, and unallocated values are then allocated to this register.
    If the values passed in are already allocated to differing registers, a
    `DiagnosticException` is raised.
    """
    reg_types = set(val.type for val in vals)
    assert all(
        isinstance(reg_type, self.register_base_class) for reg_type in reg_types
    )
    reg_types = cast(set[RegisterType], reg_types)

    match len(reg_types):
        case 0:
            # No inputs, nothing to do
            return False
        case 1:
            # Single input, may already be allocated
            (reg_type,) = reg_types
            if reg_type.is_allocated:
                return False
            else:
                reg_type = self.available_registers.pop(type(reg_type))
        case 2:
            # Two inputs, either one is allocated or two
            reg_type_0, reg_type_1 = reg_types
            if reg_type_0.is_allocated:
                if reg_type_1.is_allocated:
                    reg_names = [f"{reg_type}" for reg_type in reg_types]
                    reg_names.sort()
                    raise DiagnosticException(
                        f"Cannot allocate registers to the same register {reg_names}"
                    )
                else:
                    reg_type = reg_type_0
            else:
                reg_type = reg_type_1
        case _:
            # More than one input is allocated, meaning we can't allocate them to be
            # the same, error.
            reg_names = [f"{reg_type}" for reg_type in reg_types]
            reg_names.sort()
            raise DiagnosticException(
                f"Cannot allocate registers to the same register {reg_names}"
            )

    did_allocate = False

    for val in vals:
        if val.type != reg_type:
            self._replace_value_with_new_type(val, reg_type)
            did_allocate = True

    return did_allocate

free_value(val: SSAValue) -> None

Source code in xdsl/backend/register_allocator.py 
110
111
112
def free_value(self, val: SSAValue) -> None:
    if isinstance(val.type, self.register_base_class) and val.type.is_allocated:
        self.available_registers.push(val.type)

BlockAllocator

Bases: ValueAllocator, ABC

Abstract base class for allocators that can process blocks at a time.

Source code in xdsl/backend/register_allocator.py 
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
class BlockAllocator(ValueAllocator, abc.ABC):
    """
    Abstract base class for allocators that can process blocks at a time.
    """

    live_ins_per_block: dict[Block, OrderedSet[SSAValue]]

    def __init__(
        self,
        available_registers: RegisterStack,
        register_base_class: type[RegisterType],
    ) -> None:
        self.live_ins_per_block = {}
        super().__init__(available_registers, register_base_class)

    @abc.abstractmethod
    def allocate_block(self, block: Block):
        """
        For each operation in the block, allocate the registers.
        """

live_ins_per_block: dict[Block, OrderedSet[SSAValue]] = {} instance-attribute

__init__(available_registers: RegisterStack, register_base_class: type[RegisterType]) -> None

Source code in xdsl/backend/register_allocator.py 
158
159
160
161
162
163
164
def __init__(
    self,
    available_registers: RegisterStack,
    register_base_class: type[RegisterType],
) -> None:
    self.live_ins_per_block = {}
    super().__init__(available_registers, register_base_class)

allocate_block(block: Block) abstractmethod

For each operation in the block, allocate the registers.

Source code in xdsl/backend/register_allocator.py 
166
167
168
169
170
@abc.abstractmethod
def allocate_block(self, block: Block):
    """
    For each operation in the block, allocate the registers.
    """

live_ins_per_block(block: Block) -> dict[Block, OrderedSet[SSAValue]]

Returns a mapping from a block to the set of values used in it but defined outside of it.

Source code in xdsl/backend/register_allocator.py 
141
142
143
144
145
146
147
148
def live_ins_per_block(block: Block) -> dict[Block, OrderedSet[SSAValue]]:
    """
    Returns a mapping from a block to the set of values used in it but defined outside of
    it.
    """
    res: dict[Block, OrderedSet[SSAValue]] = {}
    _ = _live_ins_per_block(block, res)
    return res