Register coloring in a small compiler

instruction_select.py

"""
    Select instructions to emit
"""
class Tile(object):
    def __init__(self, cost, pat, code):
        self.cost = cost
        self.pat = pat
        self.code = code

def choose(code, tiles):
    for node in code.postorder([]):
        best = float('inf')
        for pat, base, code in tiles:
            if pat.match(node):
                cost = pat.estimate(node, base)
                if cost < best:
                    best = cost
                    node.tile = Tile(cost, pat, code)

def gen(node, block):
    return node.tile.code(node, block)

def select(code, tiles, block):
    choose(code, tiles)
    return gen(code, block)

output

;   xor vreg1, vreg1          def[vreg1]                             active=[]
    xor ax, ax
;   vreg2 = vreg1                                                    active=[vreg1]
;   sub vreg2, 6              def[vreg2]     use[vreg2]              active=[vreg2]
    sub ax, 6
;   vreg3 = vreg2                                                    active=[vreg2]
;   add vreg3, 2              def[vreg3]     use[vreg3]              active=[vreg3]
    add ax, 2
;   xor vreg4, vreg4          def[vreg4]                             active=[vreg3]
    xor cx, cx
;   vreg5 = vreg4                                                    active=[vreg3, vreg4]
;   sub vreg5, 6              def[vreg5]     use[vreg5]              active=[vreg3, vreg5]
    sub cx, 6
;   vreg6 = vreg5                                                    active=[vreg3, vreg5]
;   add vreg6, 2              def[vreg6]     use[vreg6]              active=[vreg3, vreg6]
    add cx, 2
;   vreg7 = vreg3                                                    active=[vreg6, vreg3]
;   sub vreg7, vreg6          def[vreg7]     use[vreg6, vreg7]       active=[vreg6, vreg7]
    sub ax, cx

patterns.py

"""
    Forms the graphs to map instructions
    into tiles.
"""
from structures import *

class PFunc(object):
    def __init__(self, match):
        self.match = match

    def estimate(self, node, base):
        return base

    def tile_estimate(self, node):
        return node.tile.cost if node.tile is not None else float('inf')

class PConst(object):
    def __init__(self, value):
        self.value = value

    def match(self, node):
        if isinstance(node, Const):
            return node.value == self.value

    def estimate(self, node, base):
        return base

    def tile_estimate(self, node):
        return node.tile.cost if node.tile is not None else float('inf')

class POp(object):
    def __init__(self, name, operands):
        self.name = name
        self.operands = operands
        self.leaf = False

    def __getitem__(self, index):
        return self.operands[index]

    def __len__(self):
        return len(self.operands)

    def match(self, node):
        if not isinstance(node, Op):
            return False
        if node.name != self.name:
            return False
        if len(self) != len(node):
            return False
        return all(pat.match(sn) for pat, sn in zip(self, node))

    def estimate(self, node, base):
        return base + sum(pat.tile_estimate(sn) for pat, sn in zip(self, node))

    def tile_estimate(self, node):
        return self.estimate(node, 0)

Any = PFunc(lambda node: True)
Int = PFunc(lambda node: isinstance(node, Const) and isinstance(node.value, int))

register_alloc.py

"""
    Select location for every virtual register.
"""
from structures import *

def alloc(block):
    # At this point, virtual registers are replaced with real registers.
    registers = {'ax', 'bx', 'cx', 'dx'}
    colors = {}

    # Our graph contains both interferences and coalescing knowledge.
    graph = {}
    def get(vreg):
        if vreg in graph:
            return graph[vreg]
        graph[vreg] = result = (set(), set())
        return result

    active = set()
    for cell in reversed(block):
        cell.active = active = (active ^ cell.defs) | cell.uses
        for vreg in active:
            get(vreg)[0].update(active ^ {vreg})

        if isinstance(cell, Motion):
            get(cell.src)[1].add(cell.dst)
            get(cell.dst)[1].add(cell.src)
    # Going through steps in chaitin-briggs algorithm
    stack = []
    # First simplification...
    while len(graph) > 0:
        for vreg, (interfere, coalesce) in graph.items():
            if len(interfere) < len(registers):
                for other in interfere:
                    graph[other][0].discard(vreg)
                stack.append((vreg, graph.pop(vreg)))
                break
        else:
            # The code compiled doesn't cause this situation yet.
            assert False, "XXX: the next part of coloring"
    # Then an attempt to color, no failure recovery yet.
    while len(stack) > 0:
        vreg, (interfere, coalesce) = stack.pop()
        filled = set(colors[v] for v in interfere if v in colors)
        avail = registers ^ filled
        colors[vreg] = avail.pop()
    return colors

scratch.py

"""
    This is the entry point.
"""
import instruction_select
import register_alloc
from structures import *
from patterns import *
from x86_realmode_tiles import tiles

code = Op('sub', [
    Op('add', [
        Op('sub', [Const(0), Const(6)]),
        Const(2)]),
    Op('add', [
        Op('sub', [Const(0), Const(6)]),
        Const(2)])])

block = []
dst = instruction_select.select(code, tiles, block)
colors = register_alloc.alloc(block)

# Next map the colors while emitting the code.
def colormap(node):
    if isinstance(node, VReg):
        return colors[node]
    return node

print "use16"
print "org 0x0"
for cell in block:
    print ';   {:<64} {}'.format(cell, format_vregs("active=[{}]", cell.active))
    if isinstance(cell, Code):
        form = map(colormap, cell.form)
        print "    " + form[0].format(*form[1:])
    elif isinstance(cell, Motion):
        # there's no coalescing, but this occassionally happens by fortune.
        if colors[cell.dst] != colors[cell.src]:
            print "    mov {}, {}".format(colors[cell.dst], colors[cell.src])
    else:
        assert False

structures.py

"""
    Intermediate representation for
    programs.
"""
class Const(object):
    def __init__(self, value, klass=None):
        self.value = value
        self.klass = klass
        self.tile = None

    def postorder(self, prefix):
        prefix.append(self)
        return prefix

class Op(object):
    def __init__(self, name, operands):
        self.name = name
        self.operands = operands
        self.tile = None

    def __getitem__(self, index):
        return self.operands[index]

    def __len__(self):
        return len(self.operands)

    def postorder(self, prefix):
        for operand in self.operands:
            operand.postorder(prefix)
        prefix.append(self)
        return prefix

class VReg(object):
    next_uid = 1
    def __init__(self, assign=None, klass=None):
        self.assign = assign
        self.klass = klass
        self.uid = VReg.next_uid
        VReg.next_uid += 1

    def __repr__(self):
        return 'vreg{}'.format(self.uid)

class Code(object):
    def __init__(self, form, uses=None, defs=None):
        self.form = form
        self.uses = set() if uses is None else uses
        self.defs = set() if defs is None else defs

    def __str__(self):
        line = self.form[0].format(*self.form[1:])
        if len(self.defs) > 0:
            line = line.ljust(25) + format_vregs(" def[{}]", self.defs)
        if len(self.uses) > 0:
            line = line.ljust(40) + format_vregs(" use[{}]", self.uses)
        return line

class Motion(object):
    def __init__(self, dst, src):
        self.dst = dst
        self.src = src
        self.uses = {src}
        self.defs = {dst}

    def __str__(self):
        return "{} = {}".format(self.dst, self.src)

def format_vregs(form, vregs):
    return form.format(', '.join(map(repr, vregs)))

x86_realmode_tiles.py

from instruction_select import gen
from structures import *
from patterns import *
tiles = []

def tile(pat, cost):
    def _deco_(code):
        tiles.append((pat, cost, code))
        return code
    return _deco_

@tile(PConst(0), 10)
def _(const, block):
    dst = VReg()
    block.append(Code(
        ('xor {}, {}', dst, dst),
        defs={dst}))
    return dst

@tile(Int, 10)
def _(const, block):
    dst = VReg()
    block.append(Code(
        ('mov {}, {}', dst, const.value),
        defs={dst}))
    return dst

@tile(POp('add', [Any, Int]), 15)
def _(op, block):
    src1 = gen(op[0], block)
    src2 = op[1].value
    dst = VReg()
    block.append(Motion(dst, src1))
    block.append(Code(
        ('add {}, {}', dst, src2),
        uses={dst},
        defs={dst}))
    return dst

@tile(POp('add', [Any, Any]), 20)
def _(op, block):
    src1 = gen(op[0], block)
    src2 = gen(op[1], block)
    dst = VReg()
    block.append(Motion(dst, src1))
    block.append(Code(
        ('add {}, {}', dst, src2),
        uses={dst, src2},
        defs={dst}))
    return dst

@tile(POp('sub', [Any, Int]), 15)
def _(op, block):
    src1 = gen(op[0], block)
    src2 = op[1].value
    dst = VReg()
    block.append(Motion(dst, src1))
    block.append(Code(
        ('sub {}, {}', dst, src2),
        uses={dst},
        defs={dst}))
    return dst

@tile(POp('sub', [Any, Any]), 20)
def _(op, block):
    src1 = gen(op[0], block)
    src2 = gen(op[1], block)
    dst = VReg()
    block.append(Motion(dst, src1))
    block.append(Code(
        ('sub {}, {}', dst, src2),
        uses={dst, src2},
        defs={dst}))
    return dst