src/bigint_presentation_code/test_register_allocator.py

   1 import unittest
   2
   3 from bigint_presentation_code.compiler_ir import (FixedGPRRangeType, Fn, GPRRange,
   4                                                   GPRType, GlobalMem, Op, OpAddSubE,
   5                                                   OpClearCY, OpConcat, OpCopy,
   6                                                   OpFuncArg, OpInputMem, OpLI,
   7                                                   OpLoad, OpStore, XERBit)
   8 from bigint_presentation_code.register_allocator import (
   9     AllocationFailed, allocate_registers, MergedRegSet,
  10     try_allocate_registers_without_spilling)
  11
  12
  13 class TestMergedRegSet(unittest.TestCase):
  14     maxDiff = None
  15
  16     def test_from_equality_constraint(self):
  17         fn = Fn()
  18         op0 = OpLI(fn, 0, length=1)
  19         op1 = OpLI(fn, 0, length=2)
  20         op2 = OpLI(fn, 0, length=3)
  21         self.assertEqual(MergedRegSet.from_equality_constraint([
  22             op0.out,
  23             op1.out,
  24             op2.out,
  25         ]), MergedRegSet({
  26             op0.out: 0,
  27             op1.out: 1,
  28             op2.out: 3,
  29         }.items()))
  30         self.assertEqual(MergedRegSet.from_equality_constraint([
  31             op1.out,
  32             op0.out,
  33             op2.out,
  34         ]), MergedRegSet({
  35             op1.out: 0,
  36             op0.out: 2,
  37             op2.out: 3,
  38         }.items()))
  39
  40
  41 class TestRegisterAllocator(unittest.TestCase):
  42     maxDiff = None
  43
  44     def test_try_alloc_fail(self):
  45         fn = Fn()
  46         op0 = OpLI(fn, 0, length=52)
  47         op1 = OpLI(fn, 0, length=64)
  48         op2 = OpConcat(fn, [op0.out, op1.out])
  49
  50         reg_assignments = try_allocate_registers_without_spilling(fn.ops)
  51         self.assertEqual(
  52             repr(reg_assignments),
  53             "AllocationFailed("
  54             "node=IGNode(#0, merged_reg_set=MergedRegSet(["
  55             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)]), "
  56             "edges={}, reg=None), "
  57             "live_intervals=LiveIntervals("
  58             "live_intervals={"
  59             "MergedRegSet([(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)]): "
  60             "LiveInterval(first_write=0, last_use=2)}, "
  61             "merged_reg_sets=MergedRegSets(data={"
  62             "<#0.out>: MergedRegSet(["
  63             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)]), "
  64             "<#1.out>: MergedRegSet(["
  65             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)]), "
  66             "<#2.dest>: MergedRegSet(["
  67             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)])}), "
  68             "reg_sets_live_after={"
  69             "0: OFSet([MergedRegSet(["
  70             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)])]), "
  71             "1: OFSet([MergedRegSet(["
  72             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)])]), "
  73             "2: OFSet()}), "
  74             "interference_graph=InterferenceGraph(nodes={"
  75             "...: IGNode(#0, "
  76             "merged_reg_set=MergedRegSet(["
  77             "(<#2.dest>, 0), (<#0.out>, 0), (<#1.out>, 52)]), "
  78             "edges={}, reg=None)}))"
  79         )
  80
  81     def test_try_alloc_bigint_inc(self):
  82         fn = Fn()
  83         op0 = OpFuncArg(fn, FixedGPRRangeType(GPRRange(3)))
  84         op1 = OpCopy(fn, op0.out, GPRType())
  85         arg = op1.dest
  86         op2 = OpInputMem(fn)
  87         mem = op2.out
  88         op3 = OpLoad(fn, arg, offset=0, mem=mem, length=32)
  89         a = op3.RT
  90         op4 = OpLI(fn, 1)
  91         b_0 = op4.out
  92         op5 = OpLI(fn, 0, length=31)
  93         b_rest = op5.out
  94         op6 = OpConcat(fn, [b_0, b_rest])
  95         b = op6.dest
  96         op7 = OpClearCY(fn)
  97         cy = op7.out
  98         op8 = OpAddSubE(fn, a, b, cy, is_sub=False)
  99         s = op8.RT
 100         op9 = OpStore(fn, s, arg, offset=0, mem_in=mem)
 101         mem = op9.mem_out
 102
 103         reg_assignments = try_allocate_registers_without_spilling(fn.ops)
 104
 105         expected_reg_assignments = {
 106             op0.out: GPRRange(start=3, length=1),
 107             op1.dest: GPRRange(start=3, length=1),
 108             op2.out: GlobalMem.GlobalMem,
 109             op3.RT: GPRRange(start=78, length=32),
 110             op4.out: GPRRange(start=46, length=1),
 111             op5.out: GPRRange(start=47, length=31),
 112             op6.dest: GPRRange(start=46, length=32),
 113             op7.out: XERBit.CY,
 114             op8.RT: GPRRange(start=14, length=32),
 115             op8.CY_out: XERBit.CY,
 116             op9.mem_out: GlobalMem.GlobalMem,
 117         }
 118
 119         self.assertEqual(reg_assignments, expected_reg_assignments)
 120
 121     def tst_try_alloc_concat(self, expected_regs, expected_dest_reg):
 122         # type: (list[GPRRange], GPRRange) -> None
 123         fn = Fn()
 124         li_ops = [OpLI(fn, i, r.length) for i, r in enumerate(expected_regs)]
 125         concat = OpConcat(fn, [i.out for i in li_ops])
 126
 127         reg_assignments = try_allocate_registers_without_spilling(fn.ops)
 128
 129         expected_reg_assignments = {concat.dest: expected_dest_reg}
 130         for li_op, reg in zip(li_ops, expected_regs):
 131             expected_reg_assignments[li_op.out] = reg
 132
 133         self.assertEqual(reg_assignments, expected_reg_assignments)
 134
 135     def test_try_alloc_concat_1(self):
 136         self.tst_try_alloc_concat([GPRRange(3)], GPRRange(3))
 137
 138     def test_try_alloc_concat_3(self):
 139         self.tst_try_alloc_concat([GPRRange(3, 3)], GPRRange(3, 3))
 140
 141     def test_try_alloc_concat_3_5(self):
 142         self.tst_try_alloc_concat([GPRRange(3, 3), GPRRange(6, 5)],
 143                                   GPRRange(3, 8))
 144
 145     def test_try_alloc_concat_5_3(self):
 146         self.tst_try_alloc_concat([GPRRange(3, 5), GPRRange(8, 3)],
 147                                   GPRRange(3, 8))
 148
 149     def test_try_alloc_concat_1_2_3_4_5_6(self):
 150         self.tst_try_alloc_concat([
 151             GPRRange(14, 1),
 152             GPRRange(15, 2),
 153             GPRRange(17, 3),
 154             GPRRange(20, 4),
 155             GPRRange(24, 5),
 156             GPRRange(29, 6),
 157         ], GPRRange(14, 21))
 158
 159
 160 if __name__ == "__main__":
 161     unittest.main()