Instantly share code, notes, and snippets.

# cheery/instruction_select.py Last active Mar 20, 2016

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