Skip to content

Register allocation

register_allocation

RegisterAllocatorLivenessBlockNaive

Bases: BlockNaiveAllocator

Source code in xdsl/backend/riscv/register_allocation.py 
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
class RegisterAllocatorLivenessBlockNaive(BlockNaiveAllocator):
    def __init__(self, available_registers: RegisterStack) -> None:
        super().__init__(available_registers, RISCVRegisterType)

    def new_type_for_value(self, reg: SSAValue) -> RegisterType | None:
        if (
            isinstance(reg.type, self.register_base_class)
            and not reg.type.is_allocated
            and (val := get_constant_value(reg)) is not None
            and val.value.data == 0
        ):
            return Registers.ZERO
        return super().new_type_for_value(reg)

    def allocate_func(
        self, func: riscv_func.FuncOp, *, add_regalloc_stats: bool = False
    ) -> None:
        """
        Allocates values in function passed in to registers.
        The whole function must have been lowered to the relevant riscv dialects
        and it must contain no unrealized casts.
        If `add_regalloc_stats` is set to `True`, then a comment op will be inserted
        before the function op passed in with a json containing the relevant data.
        """
        if not func.body.blocks:
            # External function declaration
            return

        if len(func.body.blocks) != 1:
            raise NotImplementedError(
                f"Cannot register allocate func with {len(func.body.blocks)} blocks."
            )

        preallocated = RegisterAllocatableOperation.all_used_registers(func.body)
        excluded = RegisterAllocatableOperation.all_excluded_registers(func.body)

        for pa_reg in preallocated | excluded:
            self.available_registers.exclude_register(pa_reg)

        block = func.body.block

        self.live_ins_per_block = live_ins_per_block(block)
        assert not self.live_ins_per_block[block]

        self.allocate_block(block)

        if add_regalloc_stats:
            preallocated_stats = reg_types_by_name(preallocated)
            excluded_stats = reg_types_by_name(excluded)
            allocated_stats = reg_types_by_name(
                val.type
                for op in block.walk()
                for vals in (op.results, op.operands)
                for val in vals
                if isinstance(val.type, RISCVRegisterType)
            )
            stats = {
                "preallocated_float": sorted(preallocated_stats["riscv.freg"]),
                "preallocated_int": sorted(preallocated_stats["riscv.reg"]),
                "excluded_float": sorted(excluded_stats["riscv.freg"]),
                "excluded_int": sorted(excluded_stats["riscv.reg"]),
                "allocated_float": sorted(allocated_stats["riscv.freg"]),
                "allocated_int": sorted(allocated_stats["riscv.reg"]),
            }

            stats_str = json.dumps(stats)

            Rewriter.insert_op(
                riscv.CommentOp(f"Regalloc stats: {stats_str}"),
                InsertPoint.before(func),
            )

__init__(available_registers: RegisterStack) -> None

Source code in xdsl/backend/riscv/register_allocation.py 
28
29
def __init__(self, available_registers: RegisterStack) -> None:
    super().__init__(available_registers, RISCVRegisterType)

new_type_for_value(reg: SSAValue) -> RegisterType | None

Source code in xdsl/backend/riscv/register_allocation.py 
31
32
33
34
35
36
37
38
39
def new_type_for_value(self, reg: SSAValue) -> RegisterType | None:
    if (
        isinstance(reg.type, self.register_base_class)
        and not reg.type.is_allocated
        and (val := get_constant_value(reg)) is not None
        and val.value.data == 0
    ):
        return Registers.ZERO
    return super().new_type_for_value(reg)

allocate_func(func: riscv_func.FuncOp, *, add_regalloc_stats: bool = False) -> None

Allocates values in function passed in to registers. The whole function must have been lowered to the relevant riscv dialects and it must contain no unrealized casts. If add_regalloc_stats is set to True, then a comment op will be inserted before the function op passed in with a json containing the relevant data.

Source code in xdsl/backend/riscv/register_allocation.py 
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
def allocate_func(
    self, func: riscv_func.FuncOp, *, add_regalloc_stats: bool = False
) -> None:
    """
    Allocates values in function passed in to registers.
    The whole function must have been lowered to the relevant riscv dialects
    and it must contain no unrealized casts.
    If `add_regalloc_stats` is set to `True`, then a comment op will be inserted
    before the function op passed in with a json containing the relevant data.
    """
    if not func.body.blocks:
        # External function declaration
        return

    if len(func.body.blocks) != 1:
        raise NotImplementedError(
            f"Cannot register allocate func with {len(func.body.blocks)} blocks."
        )

    preallocated = RegisterAllocatableOperation.all_used_registers(func.body)
    excluded = RegisterAllocatableOperation.all_excluded_registers(func.body)

    for pa_reg in preallocated | excluded:
        self.available_registers.exclude_register(pa_reg)

    block = func.body.block

    self.live_ins_per_block = live_ins_per_block(block)
    assert not self.live_ins_per_block[block]

    self.allocate_block(block)

    if add_regalloc_stats:
        preallocated_stats = reg_types_by_name(preallocated)
        excluded_stats = reg_types_by_name(excluded)
        allocated_stats = reg_types_by_name(
            val.type
            for op in block.walk()
            for vals in (op.results, op.operands)
            for val in vals
            if isinstance(val.type, RISCVRegisterType)
        )
        stats = {
            "preallocated_float": sorted(preallocated_stats["riscv.freg"]),
            "preallocated_int": sorted(preallocated_stats["riscv.reg"]),
            "excluded_float": sorted(excluded_stats["riscv.freg"]),
            "excluded_int": sorted(excluded_stats["riscv.reg"]),
            "allocated_float": sorted(allocated_stats["riscv.freg"]),
            "allocated_int": sorted(allocated_stats["riscv.reg"]),
        }

        stats_str = json.dumps(stats)

        Rewriter.insert_op(
            riscv.CommentOp(f"Regalloc stats: {stats_str}"),
            InsertPoint.before(func),
        )

reg_types_by_name(regs: Iterable[RegisterType]) -> dict[str, set[str]]

Groups register types by name.

Source code in xdsl/backend/riscv/register_allocation.py 
17
18
19
20
21
22
23
24
def reg_types_by_name(regs: Iterable[RegisterType]) -> dict[str, set[str]]:
    """
    Groups register types by name.
    """
    res = defaultdict[str, set[str]](set)
    for reg in regs:
        res[reg.name].add(reg.register_name)
    return res