Instantly share code, notes, and snippets.

Embed
What would you like to do?
Register coloring in a small compiler
"""
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)
; 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
"""
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))
"""
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
"""
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
"""
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)))
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment