forked from M-Labs/artiq
wrpll: add simple thls compiler
This commit is contained in:
parent
e4fff390a8
commit
5f461d08cd
|
|
@ -0,0 +1,267 @@
|
|||
import inspect
|
||||
import ast
|
||||
from copy import copy
|
||||
|
||||
|
||||
class Isn:
|
||||
def __init__(self, immediate=None, inputs=None, outputs=None):
|
||||
if inputs is None:
|
||||
inputs = []
|
||||
if outputs is None:
|
||||
outputs = []
|
||||
self.immediate = immediate
|
||||
self.inputs = inputs
|
||||
self.outputs = outputs
|
||||
|
||||
def __repr__(self):
|
||||
r = "<"
|
||||
r += self.__class__.__name__
|
||||
if self.immediate is not None:
|
||||
r += " (" + str(self.immediate) + ")"
|
||||
for inp in self.inputs:
|
||||
r += " r" + str(inp)
|
||||
if self.outputs:
|
||||
r += " ->"
|
||||
for outp in self.outputs:
|
||||
r += " r" + str(outp)
|
||||
r += ">"
|
||||
return r
|
||||
|
||||
|
||||
class NopIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class AddIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class SubIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class MulIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class ShiftIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class CopyIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class InputIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class OutputIsn(Isn):
|
||||
pass
|
||||
|
||||
|
||||
class ASTCompiler:
|
||||
def __init__(self):
|
||||
self.program = []
|
||||
self.data = []
|
||||
self.next_ssa_reg = -1
|
||||
self.constants = dict()
|
||||
self.names = dict()
|
||||
self.globals = dict()
|
||||
|
||||
def get_ssa_reg(self):
|
||||
r = self.next_ssa_reg
|
||||
self.next_ssa_reg -= 1
|
||||
return r
|
||||
|
||||
def add_global(self, name):
|
||||
r = len(self.data)
|
||||
self.data.append(0)
|
||||
self.names[name] = r
|
||||
self.globals[name] = r
|
||||
return r
|
||||
|
||||
def input(self, name):
|
||||
target = self.get_ssa_reg()
|
||||
self.program.append(InputIsn(outputs=[target]))
|
||||
self.names[name] = target
|
||||
|
||||
def emit(self, node):
|
||||
if isinstance(node, ast.BinOp):
|
||||
left = self.emit(node.left)
|
||||
right = self.emit(node.right)
|
||||
if isinstance(node.op, ast.Add):
|
||||
cls = AddIsn
|
||||
elif isinstance(node.op, ast.Sub):
|
||||
cls = SubIsn
|
||||
elif isinstance(node.op, ast.Mult):
|
||||
cls = MulIsn
|
||||
else:
|
||||
raise NotImplementedError
|
||||
output = self.get_ssa_reg()
|
||||
self.program.append(cls(inputs=[left, right], outputs=[output]))
|
||||
return output
|
||||
elif isinstance(node, ast.Num):
|
||||
if node.n in self.constants:
|
||||
return self.constants[node.n]
|
||||
else:
|
||||
r = len(self.data)
|
||||
self.data.append(node.n)
|
||||
self.constants[node.n] = r
|
||||
return r
|
||||
elif isinstance(node, ast.Name):
|
||||
return self.names[node.id]
|
||||
elif isinstance(node, ast.Assign):
|
||||
output = self.emit(node.value)
|
||||
for target in node.targets:
|
||||
assert isinstance(target, ast.Name)
|
||||
self.names[target.id] = output
|
||||
elif isinstance(node, ast.Return):
|
||||
value = self.emit(node.value)
|
||||
self.program.append(OutputIsn(inputs=[value]))
|
||||
elif isinstance(node, ast.Global):
|
||||
pass
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Processor:
|
||||
def __init__(self, multiplier_stages=2):
|
||||
self.multiplier_stages = multiplier_stages
|
||||
|
||||
def get_instruction_latency(self, isn):
|
||||
return {
|
||||
AddIsn: 2,
|
||||
SubIsn: 2,
|
||||
MulIsn: 1 + self.multiplier_stages,
|
||||
ShiftIsn: 2,
|
||||
CopyIsn: 1,
|
||||
InputIsn: 1
|
||||
}[isn.__class__]
|
||||
|
||||
|
||||
class Scheduler:
|
||||
def __init__(self, processor, reserved_data, program):
|
||||
self.processor = processor
|
||||
self.reserved_data = reserved_data
|
||||
self.used_registers = set(range(self.reserved_data))
|
||||
self.exits = dict()
|
||||
self.program = program
|
||||
self.remaining = copy(program)
|
||||
self.output = []
|
||||
|
||||
def allocate_register(self):
|
||||
r = min(set(range(max(self.used_registers) + 2)) - self.used_registers)
|
||||
self.used_registers.add(r)
|
||||
return r
|
||||
|
||||
def free_register(self, r):
|
||||
assert r >= self.reserved_data
|
||||
self.used_registers.discard(r)
|
||||
|
||||
def find_inputs(self, cycle, isn):
|
||||
mapped_inputs = []
|
||||
for inp in isn.inputs:
|
||||
if inp >= 0:
|
||||
mapped_inputs.append(inp)
|
||||
else:
|
||||
found = False
|
||||
for i in range(cycle):
|
||||
if i in self.exits:
|
||||
r, rm = self.exits[i]
|
||||
if r == inp:
|
||||
mapped_inputs.append(rm)
|
||||
found = True
|
||||
break
|
||||
if not found:
|
||||
return None
|
||||
return mapped_inputs
|
||||
|
||||
def schedule_one(self, isn):
|
||||
cycle = len(self.output)
|
||||
mapped_inputs = self.find_inputs(cycle, isn)
|
||||
if mapped_inputs is None:
|
||||
return False
|
||||
|
||||
if isn.outputs:
|
||||
latency = self.processor.get_instruction_latency(isn)
|
||||
exit = cycle + latency
|
||||
if exit in self.exits:
|
||||
return False
|
||||
|
||||
# Instruction can be scheduled
|
||||
|
||||
self.remaining.remove(isn)
|
||||
|
||||
for inp, minp in zip(isn.inputs, mapped_inputs):
|
||||
can_free = inp < 0 and all(inp != rinp for risn in self.remaining for rinp in risn.inputs)
|
||||
if can_free:
|
||||
self.free_register(minp)
|
||||
|
||||
if isn.outputs:
|
||||
assert len(isn.outputs) == 1
|
||||
output = self.allocate_register()
|
||||
self.exits[exit] = (isn.outputs[0], output)
|
||||
self.output.append(isn.__class__(immediate=isn.immediate, inputs=mapped_inputs))
|
||||
|
||||
return True
|
||||
|
||||
def schedule(self):
|
||||
while self.remaining:
|
||||
success = False
|
||||
for isn in self.remaining:
|
||||
if self.schedule_one(isn):
|
||||
success = True
|
||||
break
|
||||
if not success:
|
||||
self.output.append(NopIsn())
|
||||
self.output += [NopIsn()]*(max(self.exits.keys()) - len(self.output) + 1)
|
||||
return self.output
|
||||
|
||||
def pretty_print(self):
|
||||
for cycle, isn in enumerate(self.output):
|
||||
l = "{:4d} {:15}".format(cycle, str(isn))
|
||||
if cycle in self.exits:
|
||||
l += " -> r{}".format(self.exits[cycle][1])
|
||||
print(l)
|
||||
|
||||
|
||||
def compile(function):
|
||||
node = ast.parse(inspect.getsource(function))
|
||||
assert isinstance(node, ast.Module)
|
||||
assert len(node.body) == 1
|
||||
node = node.body[0]
|
||||
assert isinstance(node, ast.FunctionDef)
|
||||
assert len(node.args.args) == 1
|
||||
arg = node.args.args[0].arg
|
||||
body = node.body
|
||||
|
||||
astcompiler = ASTCompiler()
|
||||
for node in body:
|
||||
if isinstance(node, ast.Global):
|
||||
for name in node.names:
|
||||
astcompiler.add_global(name)
|
||||
arg_r = astcompiler.input(arg)
|
||||
for node in body:
|
||||
astcompiler.emit(node)
|
||||
print(astcompiler.data)
|
||||
print(astcompiler.program)
|
||||
|
||||
scheduler = Scheduler(Processor(), len(astcompiler.data), astcompiler.program)
|
||||
scheduler.schedule()
|
||||
scheduler.pretty_print()
|
||||
|
||||
|
||||
a = 0
|
||||
b = 0
|
||||
c = 0
|
||||
|
||||
def foo(x):
|
||||
global a, b, c
|
||||
c = b
|
||||
b = a
|
||||
a = x
|
||||
return 4748*a + 259*b - 155*c
|
||||
|
||||
|
||||
compile(foo)
|
||||
Loading…
Reference in New Issue