backend: (riscv) add tied operands/results to register allocation

tests/backend/riscv/test_register_allocation.py

@@ -0,0 +1,166 @@
+from typing import Annotated, TypeAlias
+
+import pytest
+
+from xdsl.backend.riscv.register_allocation import RegisterAllocatorLivenessBlockNaive
+from xdsl.builder import Builder
+from xdsl.dialects import builtin, riscv, riscv_func
+from xdsl.dialects.riscv import (
+    AssemblyInstructionArg,
+    IntRegisterType,
+    RISCVInstruction,
+)
+from xdsl.ir import OpResult
+from xdsl.irdl import (
+    ConstraintVar,
+    Operand,
+    irdl_op_definition,
+    operand_def,
+    result_def,
+)
+from xdsl.utils.exceptions import VerifyException
+
+
+# We use the infamous post-increment load instruction
+# provided by several architectures (ARM, PULP, ...) that presents
+# the classical form of a 2-address instruction: the first operand
+# is incremented and returned in a brand new SSA value as the second result.
+# Both tied SSA values must be allocated on the same register, otherwise the
+# generated code is broken.
+@irdl_op_definition
+class PostIncrementLoad(RISCVInstruction):
+    name = "postload"
+
+    SameIntRegisterType: TypeAlias = Annotated[IntRegisterType, ConstraintVar("T")]
+
+    rd: OpResult = result_def(IntRegisterType)  # loaded value
+    off: OpResult = result_def(
+        SameIntRegisterType
+    )  # new value of the incremented offset
+    rs1: Operand = operand_def(SameIntRegisterType)  # offset
+    rs2: Operand = operand_def(IntRegisterType)  # base adddress
+
+    def assembly_line_args(self) -> tuple[AssemblyInstructionArg, ...]:
+        return self.rd, self.rs1, self.rs2
+
+
+def test_allocate_2address():
+    # riscv_func.func @function() {
+    #     %0 = riscv.li 1 : !riscv.reg<t1>
+    #     %1 = riscv.li 2 : !riscv.reg<t0>
+    #        tied -------- tied                  tied ---------------------------------------------- tied
+    #         |             |                     |                                                   |
+    #     %2, %3 = postload %0, %1 : (!riscv.reg<t1>, !riscv.reg<t0>) -> (!riscv.reg<t0>, !riscv.reg<t1>)
+    #
+    #                this add forces all results of the postload to be allocated first
+    #                 |
+    #     %4 = riscv.add %2, %3 : (!riscv.reg<t0>, !riscv.reg<t1>) -> !riscv.reg<t0>
+    #     %5 = builtin.unrealized_conversion_cast %4 : !riscv.reg<t0> to i32
+    #     riscv_func.return
+    # }
+    @Builder.implicit_region
+    def region():
+        base = riscv.LiOp(1)
+        offset = riscv.LiOp(2)
+        load = PostIncrementLoad(
+            operands=[base.rd, offset.rd],
+            result_types=[
+                IntRegisterType.unallocated(),
+                IntRegisterType.unallocated(),
+            ],
+        )
+        add = riscv.AddOp(load.rd, load.off, rd=IntRegisterType.unallocated())
+        builtin.UnrealizedConversionCastOp.get((add.rd,), (builtin.i32,))
+        riscv_func.ReturnOp()
+
+    func = riscv_func.FuncOp("function", region, ((), ()))
+    func.verify()
+
+    allocator = RegisterAllocatorLivenessBlockNaive()
+    allocator.allocate_func(func)
+    func.verify()
+
+    body = list(func.body.block.ops)
+    # load.rs1 == load.off
+    assert body[2].operands[0].type == body[2].results[1].type
+    # load.rs1 != load.rd
+    assert body[2].operands[0].type != body[2].results[0].type
+    # load.rs1 != load.rs2
+    assert body[2].operands[0].type != body[2].operands[1].type
+
+
+def test_allocate_2address_no_riscv_successor():
+    # riscv_func.func @function() {
+    #     %0 = riscv.li 1 : !riscv.reg<t1>
+    #     %1 = riscv.li 2 : !riscv.reg<t2>
+    #        tied -------- tied                  tied ---------------------------------------------- tied
+    #         |             |                     |                                                   |
+    #     %2, %3 = postload %0, %1 : (!riscv.reg<t1>, !riscv.reg<t2>) -> (!riscv.reg<t0>, !riscv.reg<t1>)
+    #     %4 = builtin.unrealized_conversion_cast %2 : !riscv.reg<t0> to i32
+    #     %5 = builtin.unrealized_conversion_cast %3 : !riscv.reg<t1> to i32
+    #     riscv_func.return
+    # }
+    @Builder.implicit_region
+    def region():
+        base = riscv.LiOp(1)
+        offset = riscv.LiOp(2)
+        load = PostIncrementLoad(
+            operands=[base.rd, offset.rd],
+            result_types=[
+                IntRegisterType.unallocated(),
+                IntRegisterType.unallocated(),
+            ],
+        )
+        builtin.UnrealizedConversionCastOp.get((load.rd,), (builtin.i32,))
+        builtin.UnrealizedConversionCastOp.get((load.off,), (builtin.i32,))
+        riscv_func.ReturnOp()
+
+    func = riscv_func.FuncOp("function", region, ((), ()))
+    func.verify()
+
+    allocator = RegisterAllocatorLivenessBlockNaive()
+    allocator.allocate_func(func)
+    func.verify()
+
+    body = list(func.body.block.ops)
+    # load.rs1 == load.off
+    assert body[2].operands[0].type == body[2].results[1].type
+    # load.rs1 != load.rd
+    assert body[2].operands[0].type != body[2].results[0].type
+    # load.rs1 != load.rs2
+    assert body[2].operands[0].type != body[2].operands[1].type
+
+
+def test_allocate_2address_preallocated():
+    @Builder.implicit_region
+    def region():
+        base = riscv.LiOp(1)
+        offset = riscv.LiOp(2)
+        load = PostIncrementLoad(
+            operands=[base.rd, offset.rd],
+            result_types=[
+                IntRegisterType.unallocated(),
+                riscv.Registers.A7,
+            ],
+        )
+        builtin.UnrealizedConversionCastOp.get((load.rd,), (builtin.i32,))
+        builtin.UnrealizedConversionCastOp.get((load.off,), (builtin.i32,))
+        riscv_func.ReturnOp()
+
+    func = riscv_func.FuncOp("function", region, ((), ()))
+    with pytest.raises(VerifyException):
+        func.verify()
+
+    allocator = RegisterAllocatorLivenessBlockNaive()
+    allocator.exclude_preallocated = True
+    allocator.allocate_func(func)
+    func.verify()
+
+    body = list(func.body.block.ops)
+    # load.rs1 == load.off
+    assert body[2].operands[0].type == riscv.Registers.A7
+    assert body[2].results[1].type == riscv.Registers.A7
+    # load.rs1 != load.rd
+    assert body[2].operands[0].type != body[2].results[0].type
+    # load.rs1 != load.rs2
+    assert body[2].operands[0].type != body[2].operands[1].type

tests/backend/riscv/test_register_allocation_constraints.py

@@ -0,0 +1,177 @@
+from abc import ABC
+from typing import Annotated, Generic, TypeAlias, TypeVar
+
+import pytest
+
+from xdsl.backend.riscv.register_allocation_constraints import OpTieConstraints
+from xdsl.dialects.riscv import (
+    AssemblyInstructionArg,
+    IntRegisterType,
+    Registers,
+    RISCVInstruction,
+    RISCVRegisterType,
+)
+from xdsl.ir import OpResult
+from xdsl.irdl import (
+    ConstraintVar,
+    Operand,
+    irdl_op_definition,
+    operand_def,
+    result_def,
+)
+from xdsl.utils.exceptions import VerifyException
+from xdsl.utils.test_value import TestSSAValue
+
+RD1InvT = TypeVar("RD1InvT", bound=RISCVRegisterType)
+RD2InvT = TypeVar("RD2InvT", bound=RISCVRegisterType)
+RS1InvT = TypeVar("RS1InvT", bound=RISCVRegisterType)
+RS2InvT = TypeVar("RS2InvT", bound=RISCVRegisterType)
+
+
+class TestOp4(Generic[RD1InvT, RD2InvT, RS1InvT, RS2InvT], RISCVInstruction, ABC):
+
+    rd1: OpResult = result_def(RD1InvT)
+    rd2: OpResult = result_def(RD2InvT)
+    rs1: Operand = operand_def(RS1InvT)
+    rs2: Operand = operand_def(RS2InvT)
+
+    def assembly_line_args(self) -> tuple[AssemblyInstructionArg, ...]:
+        return self.rd1, self.rd2, self.rs1, self.rs2
+
+
+SameIntRegisterType: TypeAlias = Annotated[IntRegisterType, ConstraintVar("T")]
+
+
+def test_tie_constrained_single():
+    @irdl_op_definition
+    class ConstrainedOp(
+        TestOp4[SameIntRegisterType, IntRegisterType, IntRegisterType, IntRegisterType]
+    ):
+        name = "constrained_op"
+
+    op = ConstrainedOp(
+        result_types=[
+            IntRegisterType.unallocated(),  # rd1
+            IntRegisterType.unallocated(),  # rd2
+        ],
+        operands=[
+            TestSSAValue(IntRegisterType.unallocated()),  # rs1
+            TestSSAValue(IntRegisterType.unallocated()),  # rs2
+        ],
+    )
+    op.verify()
+
+    constr = OpTieConstraints.from_op(op)
+    assert constr.result_has_constraints(0)
+    assert not constr.result_has_constraints(1)
+    assert not constr.operand_has_constraints(0)
+    assert not constr.operand_has_constraints(1)
+
+    assert constr.result_is_constrained_to(0) is None
+    assert constr.result_is_constrained_to(1) is None
+    assert constr.operand_is_constrained_to(2) is None
+    assert constr.operand_is_constrained_to(3) is None
+
+    # Allocate rd1
+    op.rd1.type = Registers.A1
+    constr.result_satisfy_constraint(0, Registers.A1)
+
+    assert constr.result_is_constrained_to(0) == Registers.A1
+    assert constr.result_is_constrained_to(1) is None
+    assert constr.operand_is_constrained_to(2) is None
+    assert constr.operand_is_constrained_to(3) is None
+
+
+def test_tie_constrained_all_unallocated():
+    @irdl_op_definition
+    class ConstrainedOp(
+        TestOp4[
+            SameIntRegisterType,
+            SameIntRegisterType,
+            SameIntRegisterType,
+            SameIntRegisterType,
+        ]
+    ):
+        name = "constrained_op"
+
+    op = ConstrainedOp(
+        result_types=[
+            IntRegisterType.unallocated(),  # rd1
+            IntRegisterType.unallocated(),  # rd2
+        ],
+        operands=[
+            TestSSAValue(IntRegisterType.unallocated()),  # rs1
+            TestSSAValue(IntRegisterType.unallocated()),  # rs2
+        ],
+    )
+    op.verify()
+
+    constr = OpTieConstraints.from_op(op)
+    assert constr.result_has_constraints(0)
+    assert constr.result_has_constraints(1)
+    assert constr.operand_has_constraints(0)
+    assert constr.operand_has_constraints(1)
+
+    assert constr.result_is_constrained_to(0) is None
+    assert constr.result_is_constrained_to(1) is None
+    assert constr.operand_is_constrained_to(0) is None
+    assert constr.operand_is_constrained_to(1) is None
+
+    # Allocate rs2
+    op.rs2.type = Registers.A1
+    constr.operand_satisfy_constraint(1, Registers.A1)
+
+    assert constr.result_is_constrained_to(0) == Registers.A1
+    assert constr.result_is_constrained_to(1) == Registers.A1
+    assert constr.operand_is_constrained_to(0) == Registers.A1
+    assert constr.operand_is_constrained_to(1) == Registers.A1
+
+
+def test_tie_constrained():
+    @irdl_op_definition
+    class ConstrainedOp(
+        TestOp4[
+            SameIntRegisterType,
+            SameIntRegisterType,
+            SameIntRegisterType,
+            SameIntRegisterType,
+        ]
+    ):
+        name = "constrained_op"
+
+    op = ConstrainedOp(
+        result_types=[
+            IntRegisterType.unallocated(),  # rd1
+            IntRegisterType.unallocated(),  # rd2
+        ],
+        operands=[
+            TestSSAValue(Registers.A1),  # rs1
+            TestSSAValue(IntRegisterType.unallocated()),  # rs2
+        ],
+    )
+    with pytest.raises(VerifyException):
+        op.verify()
+
+    constr = OpTieConstraints.from_op(op)
+    assert constr.result_has_constraints(0)
+    assert constr.result_has_constraints(1)
+    assert constr.operand_has_constraints(0)
+    assert constr.operand_has_constraints(1)
+
+    assert constr.result_is_constrained_to(0) == Registers.A1
+    assert constr.result_is_constrained_to(1) == Registers.A1
+    assert constr.operand_is_constrained_to(0) == Registers.A1
+    assert constr.operand_is_constrained_to(1) == Registers.A1
+
+    op.rd1.type = Registers.A1
+    op.rd2.type = Registers.A1
+    op.rs1.type = Registers.A1
+    op.rs2.type = Registers.A1
+    op.verify()
+
+    op.rs2.type = Registers.A2
+    with pytest.raises(VerifyException):
+        op.verify()
+    constr.operand_satisfy_constraint(0, Registers.A1)
+    with pytest.raises(ValueError):
+        constr.operand_satisfy_constraint(0, Registers.A2)