Last active
March 20, 2016 05:42
-
-
Save cheery/af3136905c698967cd14 to your computer and use it in GitHub Desktop.
Register coloring in a small compiler
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
; 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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 file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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))) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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