def run(func, env=None, return_block=None):
"""
Rewrite 'ret' operations into jumps to a return block and assignments
to a return variable.
"""
b = Builder(func)
return_block = return_block or func.new_block("pykit.return")
# Allocate return variable
if not func.type.restype.is_void:
with b.at_front(func.startblock):
return_var = b.alloca(types.Pointer(func.type.restype), [])
b.store(Undef(func.type.restype), return_var)
else:
return_var = None
# Repace 'ret' instructions with jumps and assignments
for op in func.ops:
if op.opcode == "ret":
b.position_after(op)
if return_var:
b.store(op.args[0], return_var)
b.jump(return_block)
op.delete()
with b.at_end(return_block):
if return_var:
result = b.load(return_var.type.base, [return_var])
else:
result = None
b.ret(result)
def run(func, env=None, return_block=None):
"""
Rewrite 'ret' operations into jumps to a return block and assignments
to a return variable.
"""
b = Builder(func)
return_block = return_block or func.new_block("pykit.return")
# Allocate return variable
if not func.type.restype.is_void:
with b.at_front(func.startblock):
return_var = b.alloca(types.Pointer(func.type.restype))
b.store(Undef(func.type.restype), return_var)
else:
return_var = None
# Repace 'ret' instructions with jumps and assignments
for op in func.ops:
if op.opcode == "ret":
b.position_after(op)
if return_var:
b.store(op.args[0], return_var)
b.jump(return_block)
op.delete()
with b.at_end(return_block):
if return_var:
result = b.load(return_var)
else:
result = None
b.ret(result)
class TestBuilder(unittest.TestCase):
def setUp(self):
self.f = Function("testfunc", ['a'],
types.Function(types.Float32, [types.Int32]))
self.b = Builder(self.f)
self.b.position_at_end(self.f.add_block('entry'))
self.a = self.f.get_arg('a')
def test_basic_builder(self):
v = self.b.alloca(types.Pointer(types.Float32), [])
result = self.b.mul(types.Int32, [self.a, self.a], result='r')
c = self.b.convert(types.Float32, [result])
self.b.store(c, v)
val = self.b.load(types.Float32, [v])
self.b.ret(val)
# print(string(self.f))
self.assertEqual(str(self.f).strip(), basic_expected)
def test_splitblock(self):
old, new = self.b.splitblock('newblock')
with self.b.at_front(old):
self.b.add(types.Int32, [self.a, self.a])
with self.b.at_end(new):
self.b.div(types.Int32, [self.a, self.a])
# print(string(self.f))
self.assertEqual(split_expected, string(self.f))
def test_loop_builder(self):
square = self.b.mul(types.Int32, [self.a, self.a])
c = self.b.convert(types.Float32, [square])
self.b.position_after(square)
_, block = self.b.splitblock('start', terminate=True)
self.b.position_at_end(block)
const = partial(Const, type=types.Int32)
cond, body, exit = self.b.gen_loop(const(5), const(10), const(2))
with self.b.at_front(body):
self.b.print_(c)
with self.b.at_end(exit):
self.b.ret(c)
# print(string(self.f))
# verify.verify(self.f)
# self.assertEqual(loop_expected, string(self.f))
# TestBuilder('test_basic_builder').debug()
# TestBuilder('test_splitblock').debug()
# TestBuilder('test_loop_builder').debug()
# unittest.main()
class TestBuilder(unittest.TestCase):
def setUp(self):
self.f = Function("testfunc", ['a'],
types.Function(types.Float32, [types.Int32]))
self.b = Builder(self.f)
self.b.position_at_end(self.f.add_block('entry'))
self.a = self.f.get_arg('a')
def test_basic_builder(self):
v = self.b.alloca(types.Pointer(types.Float32), [])
result = self.b.mul(types.Int32, [self.a, self.a], result='r')
c = self.b.convert(types.Float32, [result])
self.b.store(c, v)
val = self.b.load(types.Float32, [v])
self.b.ret(val)
# print(string(self.f))
self.assertEqual(str(self.f).strip(), basic_expected)
def test_splitblock(self):
old, new = self.b.splitblock('newblock')
with self.b.at_front(old):
self.b.add(types.Int32, [self.a, self.a])
with self.b.at_end(new):
self.b.div(types.Int32, [self.a, self.a])
# print(string(self.f))
self.assertEqual(split_expected, string(self.f))
def test_loop_builder(self):
square = self.b.mul(types.Int32, [self.a, self.a])
c = self.b.convert(types.Float32, [square])
self.b.position_after(square)
_, block = self.b.splitblock('start', terminate=True)
self.b.position_at_end(block)
const = partial(Const, type=types.Int32)
cond, body, exit = self.b.gen_loop(const(5), const(10), const(2))
with self.b.at_front(body):
self.b.print_(c)
with self.b.at_end(exit):
self.b.ret(c)
# print(string(self.f))
# verify.verify(self.f)
# self.assertEqual(loop_expected, string(self.f))
# TestBuilder('test_basic_builder').debug()
# TestBuilder('test_splitblock').debug()
# TestBuilder('test_loop_builder').debug()
# unittest.main()
def from_expr(graph, expr_context, env):
"""
Map a Blaze expression graph to blaze AIR
Parameters
----------
graph: blaze.expr.Op
Expression graph
expr_context: ExprContext
Context of the expression
ctx: ExecutionContext
"""
inputs = expr_context.params
# -------------------------------------------------
# Types
argtypes = [operand.dshape for operand in inputs]
signature = types.Function(graph.dshape, argtypes, varargs=False)
# -------------------------------------------------
# Setup function
name = "expr"
argnames = ["e%d" % i for i in range(len(inputs))]
f = Function(name, argnames, signature)
builder = Builder(f)
builder.position_at_beginning(f.new_block('entry'))
# -------------------------------------------------
# Generate function
valuemap = dict((expr, f.get_arg("e%d" % i))
for i, expr in enumerate(inputs))
_from_expr(graph, f, builder, valuemap)
retval = valuemap[graph]
builder.ret(retval)
# Update environment with runtime arguments
runtime_args = [expr_context.terms[input] for input in inputs]
env['runtime.args'] = dict(zip(f.args, runtime_args))
return f
def from_expr(graph, expr_context, env):
"""
Map a Blaze expression graph to blaze AIR
Parameters
----------
graph: blaze.expr.Op
Expression graph
expr_context: ExprContext
Context of the expression
ctx: ExecutionContext
"""
inputs = expr_context.params
# -------------------------------------------------
# Types
argtypes = [operand.dshape for operand in inputs]
signature = types.Function(graph.dshape, argtypes, varargs=False)
# -------------------------------------------------
# Setup function
name = "expr"
argnames = ["e%d" % i for i in range(len(inputs))]
f = Function(name, argnames, signature)
builder = Builder(f)
builder.position_at_beginning(f.new_block('entry'))
# -------------------------------------------------
# Generate function
valuemap = dict(
(expr, f.get_arg("e%d" % i)) for i, expr in enumerate(inputs))
_from_expr(graph, f, builder, valuemap)
retval = valuemap[graph]
builder.ret(retval)
# Update environment with runtime arguments
runtime_args = [expr_context.terms[input] for input in inputs]
env['runtime.args'] = dict(zip(f.args, runtime_args))
return f
def from_expr(graph, expr_context, ctx):
"""
Map a Blaze expression graph to blaze AIR
Parameters
----------
graph: blaze.expr.Op
Expression graph
expr_context: ExprContext
Context of the expression
ctx: ExecutionContext
"""
inputs = expr_context.params
# -------------------------------------------------
# Types
argtypes = [operand.dshape for operand in inputs]
signature = types.Function(graph.dshape, argtypes)
# -------------------------------------------------
# Setup function
name = "expr%d" % ctx.incr()
argnames = ["e%d" % i for i in range(len(inputs))]
f = Function(name, argnames, signature)
builder = Builder(f)
builder.position_at_beginning(f.new_block('entry'))
# -------------------------------------------------
# Generate function
values = dict((expr, f.get_arg("e%d" % i))
for i, expr in enumerate(inputs))
_from_expr(graph, f, builder, values)
retval = values[graph]
builder.ret(retval)
return f, values
class TestBuilder(unittest.TestCase):
def setUp(self):
self.f = Function("testfunc", ['a'],
types.Function(types.Float32, [types.Int32]))
self.b = Builder(self.f)
self.b.position_at_end(self.f.new_block('entry'))
self.a = self.f.get_arg('a')
def test_basic_builder(self):
v = self.b.alloca(types.Pointer(types.Float32), [])
result = self.b.mul(types.Int32, [self.a, self.a], result='r')
c = self.b.convert(types.Float32, [result])
self.b.store(c, v)
val = self.b.load(types.Float32, [v])
self.b.ret(val)
# print(string(self.f))
assert interp.run(self.f, args=[10]) == 100
def test_splitblock(self):
old, new = self.b.splitblock('newblock')
with self.b.at_front(old):
self.b.add(types.Int32, [self.a, self.a])
with self.b.at_end(new):
self.b.div(types.Int32, [self.a, self.a])
self.assertEqual(opcodes(self.f), ['add', 'div'])
def test_loop_builder(self):
square = self.b.mul(types.Int32, [self.a, self.a])
c = self.b.convert(types.Float32, [square])
self.b.position_after(square)
_, block = self.b.splitblock('start', terminate=True)
self.b.position_at_end(block)
const = partial(Const, type=types.Int32)
cond, body, exit = self.b.gen_loop(const(5), const(10), const(2))
with self.b.at_front(body):
self.b.print(c)
with self.b.at_end(exit):
self.b.ret(c)
self.assertEqual(interp.run(self.f, args=[10]), 100.0)
class TestBuilder(unittest.TestCase):
def setUp(self):
self.f = Function("testfunc", ['a'],
types.Function(types.Float32, [types.Int32], False))
self.b = Builder(self.f)
self.b.position_at_end(self.f.new_block('entry'))
self.a = self.f.get_arg('a')
def test_basic_builder(self):
v = self.b.alloca(types.Pointer(types.Float32))
result = self.b.mul(self.a, self.a, result='r')
c = self.b.convert(types.Float32, result)
self.b.store(c, v)
val = self.b.load(v)
self.b.ret(val)
# print(string(self.f))
assert interp.run(self.f, args=[10]) == 100
def test_splitblock(self):
old, new = self.b.splitblock('newblock')
with self.b.at_front(old):
self.b.add(self.a, self.a)
with self.b.at_end(new):
self.b.div(self.a, self.a)
self.assertEqual(opcodes(self.f), ['add', 'div'])
def test_loop_builder(self):
square = self.b.mul(self.a, self.a)
c = self.b.convert(types.Float32, square)
self.b.position_after(square)
_, block = self.b.splitblock('start', terminate=True)
self.b.position_at_end(block)
const = partial(Const, type=types.Int32)
cond, body, exit = self.b.gen_loop(const(5), const(10), const(2))
with self.b.at_front(body):
self.b.print(c)
with self.b.at_end(exit):
self.b.ret(c)
self.assertEqual(interp.run(self.f, args=[10]), 100.0)
def test_splitblock_preserve_phis(self):
"""
block1:
%0 = mul a a
jump(newblock)
newblock:
%1 = phi([block1], [%0])
ret %1
"""
square = self.b.mul(self.a, self.a)
old, new = self.b.splitblock('newblock', terminate=True)
with self.b.at_front(new):
phi = self.b.phi(types.Int32, [self.f.startblock], [square])
self.b.ret(phi)
# Now split block1
self.b.position_after(square)
block1, split = self.b.splitblock(terminate=True)
phi, ret = new.ops
blocks, values = phi.args
self.assertEqual(blocks, [split])
class PykitIRVisitor(c_ast.NodeVisitor):
"""
Map pykit IR in the form of polymorphic C to in-memory pykit IR.
int function(float x) {
int i = 0; /* I am a comment */
while (i < 10) { /*: { "unroll": true } :*/
x = call_external("sqrt", x * x);
}
return (int) x;
}
Attributes:
"""
in_function = False
def __init__(self, type_env=None):
self.mod = Module()
self.type_env = type_env or {}
self.func = None
self.builder = None
self.local_vars = None
self.allocas = None
self.global_vars = {}
self.functions = {}
# ______________________________________________________________________
@property
def vars(self):
if self.in_function:
return self.local_vars
else:
return self.global_vars
def enter_func(self):
self.in_function = True
self.local_vars = {}
self.allocas = {}
def leave_func(self):
self.in_function = False
self.mod.add_function(self.func)
self.local_vars = None
self.allocas = None
self.func = None
def visit(self, node, type=None):
"""
Visit a node.
:type: Whether we have a type for this opcode, which is an LHS type
or a cast. E.g.:
(Int) call(...) // cast
result = call(...) // assmnt, assuming 'result' is declared
result = call(..., call(...)) // second 'call' isn't typed
"""
self.type = type
method = 'visit_' + node.__class__.__name__
visitor = getattr(self, method, self.generic_visit)
# if visitor is None:
# raise SyntaxError(
# "Node %s not supported in %s:%s" % (node, node.coord.file,
# node.coord.line))
return visitor(node)
def visitif(self, node):
if node:
return self.visit(node)
def visits(self, node):
return list(map(self.visit, node))
# ______________________________________________________________________
def alloca(self, varname):
if varname not in self.allocas:
# Allocate variable with alloca
with self.builder.at_front(self.func.blocks[0]):
type = self.local_vars[varname]
self.allocas[varname] = self.builder.alloca(type, [], varname)
return self.allocas[varname]
def assign(self, varname, rhs):
if not self.in_function:
error(rhs, "Assignment only allowed in functions")
if varname not in self.allocas:
# Allocate variable with alloca
with self.builder.at_front(self.func.blocks[0]):
type = self.local_vars[varname]
self.allocas[varname] = self.builder.alloca(type, [], varname)
self.builder.store(self.visit(rhs), self.alloca(varname))
# ______________________________________________________________________
def visit_Decl(self, decl):
if decl.name in self.vars:
error(decl, "Var '%s' already declared!" % (decl.name, ))
type = self.visit(decl.type)
self.vars[decl.name] = type
if decl.init:
self.assign(decl.name, decl.init)
return type
def visit_TypeDecl(self, decl):
return self.visit(decl.type)
visit_Typename = visit_TypeDecl
def visit_PtrDecl(self, decl):
return types.Pointer(self.visit(decl.type.type))
def visit_FuncDecl(self, decl):
return types.Function(self.visit(decl.type),
self.visits(decl.args.params))
def visit_IdentifierType(self, node):
name, = node.names
return self.type_env[name]
def visit_Typedef(self, node):
if node.name in ("Type", "_list"):
type = self.type_env[node.name]
else:
type = self.visit(node.type)
if type == types.Type:
type = getattr(types, node.name)
self.type_env[node.name] = type
return type
def visit_Template(self, node):
left = self.visit(node.left)
subtypes = self.visits(node.right)
if left is list:
return list(subtypes)
else:
assert issubclass(left, types.Type)
subtypes = self.visits(node.right)
return left(*subtypes)
# ______________________________________________________________________
def visit_FuncDef(self, node):
assert not node.param_decls
self.enter_func()
name = node.decl.name
type = self.visit(node.decl.type)
argnames = [p.name for p in node.decl.type.args.params]
self.func = Function(name, argnames, type)
self.func.add_block('entry')
self.builder = Builder(self.func)
self.builder.position_at_end(self.func.blocks[0])
self.generic_visit(node.body)
self.leave_func()
# ______________________________________________________________________
def visit_FuncCall(self, node):
name = node.name.name
if not self.in_typed_context:
error(
node, "Expected a type for sub-expression "
"(add a cast or assignment)")
if not hasattr(self.builder, name):
error(node, "No opcode %s" % (name, ))
self.in_typed_context = False
buildop = getattr(self.builder, name)
args = self.visits(node.args.exprs)
return buildop, args
def visit_ID(self, node):
if self.in_function:
if node.name not in self.local_vars:
error(node, "Not a local: %r" % node.name)
result = self.alloca(node.name)
return self.builder.load(result.type, result)
def visit_Cast(self, node):
type = self.visit(node.to_type)
if isinstance(node.expr, c_ast.FuncCall):
self.in_typed_context = True
buildop, args = self.visit(node.expr)
return buildop(type, args, "temp")
else:
result = self.visit(node.expr)
if result.type == type:
return result
return self.builder.convert(type, [result], "temp")
def visit_Assignment(self, node):
if node.op != '=':
error(node, "Only assignment with '=' is supported")
if not isinstance(node.lvalue, c_ast.ID):
error(node, "Canot only assign to a name")
self.assign(node.lvalue.name, node.rvalue)
def visit_Constant(self, node):
type = self.type_env[node.type]
const = types.convert(node.value, type)
return Const(const)
def visit_UnaryOp(self, node):
op = defs.unary_defs[node.op]
buildop = getattr(self.builder, op)
arg = self.visit(node.expr)
type = self.type or arg.type
return buildop(type, [arg])
def visit_BinaryOp(self, node):
op = binary_defs[node.op]
buildop = getattr(self.builder, op)
left, right = self.visits([node.left, node.right])
if not self.type:
assert left.type == right.type, (left, right)
return buildop(self.type or left.type, [left, right], "temp")
def _loop(self, init, cond, next, body):
_, exit_block = self.builder.splitblock("exit")
_, body_block = self.builder.splitblock("body")
_, cond_block = self.builder.splitblock("cond")
self.visitif(init)
self.builder.jump(cond_block)
with self.builder.at_front(cond_block):
cond = self.visit(cond, type=types.Bool)
self.builder.cbranch(cond, cond_block, exit_block)
with self.builder.at_front(body_block):
self.visit(body)
self.visitif(next)
self.builder.jump(cond_block)
self.builder.position_at_end(exit_block)
def visit_While(self, node):
self._loop(None, node.cond, None, node.stmt)
def visit_For(self, node):
self._loop(node.init, node.cond, node.next, node.stmt)
def visit_Return(self, node):
self.builder.ret(self.visit(node.expr))
class PykitIRVisitor(c_ast.NodeVisitor):
"""
Map pykit IR in the form of polymorphic C to in-memory pykit IR.
int function(float x) {
int i = 0; /* I am a comment */
while (i < 10) { /*: { "unroll": true } :*/
x = call_external("sqrt", x * x);
}
return (int) x;
}
Attributes:
"""
in_function = False
def __init__(self, type_env=None):
self.mod = Module()
self.type_env = type_env or {}
self.func = None
self.builder = None
self.local_vars = None
self.allocas = None
self.global_vars = {}
self.functions = {}
# ______________________________________________________________________
@property
def vars(self):
if self.in_function:
return self.local_vars
else:
return self.global_vars
def enter_func(self):
self.in_function = True
self.local_vars = {}
self.allocas = {}
def leave_func(self):
self.in_function = False
self.mod.add_function(self.func)
self.local_vars = None
self.allocas = None
self.func = None
def visit(self, node, type=None):
"""
Visit a node.
:type: Whether we have a type for this opcode, which is an LHS type
or a cast. E.g.:
(Int) call(...) // cast
result = call(...) // assmnt, assuming 'result' is declared
result = call(..., call(...)) // second 'call' isn't typed
"""
self.type = type
method = 'visit_' + node.__class__.__name__
visitor = getattr(self, method, self.generic_visit)
# if visitor is None:
# raise SyntaxError(
# "Node %s not supported in %s:%s" % (node, node.coord.file,
# node.coord.line))
return visitor(node)
def visitif(self, node):
if node:
return self.visit(node)
def visits(self, node):
return list(map(self.visit, node))
# ______________________________________________________________________
def alloca(self, varname):
if varname not in self.allocas:
# Allocate variable with alloca
with self.builder.at_front(self.func.blocks[0]):
type = self.local_vars[varname]
self.allocas[varname] = self.builder.alloca(type, [], varname)
return self.allocas[varname]
def assign(self, varname, rhs):
if not self.in_function:
error(rhs, "Assignment only allowed in functions")
if varname not in self.allocas:
# Allocate variable with alloca
with self.builder.at_front(self.func.blocks[0]):
type = self.local_vars[varname]
self.allocas[varname] = self.builder.alloca(type, [], varname)
self.builder.store(self.visit(rhs), self.alloca(varname))
# ______________________________________________________________________
def visit_Decl(self, decl):
if decl.name in self.vars:
error(decl, "Var '%s' already declared!" % (decl.name,))
type = self.visit(decl.type)
self.vars[decl.name] = type
if decl.init:
self.assign(decl.name, decl.init)
return type
def visit_TypeDecl(self, decl):
return self.visit(decl.type)
visit_Typename = visit_TypeDecl
def visit_PtrDecl(self, decl):
return types.Pointer(self.visit(decl.type.type))
def visit_FuncDecl(self, decl):
return types.Function(self.visit(decl.type),
self.visits(decl.args.params))
def visit_IdentifierType(self, node):
name, = node.names
return self.type_env[name]
def visit_Typedef(self, node):
if node.name in ("Type", "_list"):
type = self.type_env[node.name]
else:
type = self.visit(node.type)
if type == types.Type:
type = getattr(types, node.name)
self.type_env[node.name] = type
return type
def visit_Template(self, node):
left = self.visit(node.left)
subtypes = self.visits(node.right)
if left is list:
return list(subtypes)
else:
assert issubclass(left, types.Type)
subtypes = self.visits(node.right)
return left(*subtypes)
# ______________________________________________________________________
def visit_FuncDef(self, node):
assert not node.param_decls
self.enter_func()
name = node.decl.name
type = self.visit(node.decl.type)
argnames = [p.name for p in node.decl.type.args.params]
self.func = Function(name, argnames, type)
self.func.add_block('entry')
self.builder = Builder(self.func)
self.builder.position_at_end(self.func.blocks[0])
self.generic_visit(node.body)
self.leave_func()
# ______________________________________________________________________
def visit_FuncCall(self, node):
name = node.name.name
if not self.in_typed_context:
error(node, "Expected a type for sub-expression "
"(add a cast or assignment)")
if not hasattr(self.builder, name):
error(node, "No opcode %s" % (name,))
self.in_typed_context = False
buildop = getattr(self.builder, name)
args = self.visits(node.args.exprs)
return buildop, args
def visit_ID(self, node):
if self.in_function:
if node.name not in self.local_vars:
error(node, "Not a local: %r" % node.name)
result = self.alloca(node.name)
return self.builder.load(result.type, result)
def visit_Cast(self, node):
type = self.visit(node.to_type)
if isinstance(node.expr, c_ast.FuncCall):
self.in_typed_context = True
buildop, args = self.visit(node.expr)
return buildop(type, args, "temp")
else:
result = self.visit(node.expr)
if result.type == type:
return result
return self.builder.convert(type, [result], "temp")
def visit_Assignment(self, node):
if node.op != '=':
error(node, "Only assignment with '=' is supported")
if not isinstance(node.lvalue, c_ast.ID):
error(node, "Canot only assign to a name")
self.assign(node.lvalue.name, node.rvalue)
def visit_Constant(self, node):
type = self.type_env[node.type]
const = types.convert(node.value, type)
return Const(const)
def visit_UnaryOp(self, node):
op = defs.unary_defs[node.op]
buildop = getattr(self.builder, op)
arg = self.visit(node.expr)
type = self.type or arg.type
return buildop(type, [arg])
def visit_BinaryOp(self, node):
op = binary_defs[node.op]
buildop = getattr(self.builder, op)
left, right = self.visits([node.left, node.right])
if not self.type:
assert left.type == right.type, (left, right)
return buildop(self.type or left.type, [left, right], "temp")
def _loop(self, init, cond, next, body):
_, exit_block = self.builder.splitblock("exit")
_, body_block = self.builder.splitblock("body")
_, cond_block = self.builder.splitblock("cond")
self.visitif(init)
self.builder.jump(cond_block)
with self.builder.at_front(cond_block):
cond = self.visit(cond, type=types.Bool)
self.builder.cbranch(cond, cond_block, exit_block)
with self.builder.at_front(body_block):
self.visit(body)
self.visitif(next)
self.builder.jump(cond_block)
self.builder.position_at_end(exit_block)
def visit_While(self, node):
self._loop(None, node.cond, None, node.stmt)
def visit_For(self, node):
self._loop(node.init, node.cond, node.next, node.stmt)
def visit_Return(self, node):
self.builder.ret(self.visit(node.expr))
class TestBuilder(unittest.TestCase):
def setUp(self):
self.f = Function("testfunc", ['a'],
types.Function(types.Float32, [types.Int32], False))
self.b = Builder(self.f)
self.b.position_at_end(self.f.new_block('entry'))
self.a = self.f.get_arg('a')
def test_basic_builder(self):
v = self.b.alloca(types.Pointer(types.Float32))
result = self.b.mul(self.a, self.a, result='r')
c = self.b.convert(types.Float32, result)
self.b.store(c, v)
val = self.b.load(v)
self.b.ret(val)
# print(string(self.f))
assert interp.run(self.f, args=[10]) == 100
def test_splitblock(self):
old, new = self.b.splitblock('newblock')
with self.b.at_front(old):
self.b.add(self.a, self.a)
with self.b.at_end(new):
self.b.div(self.a, self.a)
self.assertEqual(opcodes(self.f), ['add', 'div'])
def test_loop_builder(self):
square = self.b.mul(self.a, self.a)
c = self.b.convert(types.Float32, square)
self.b.position_after(square)
_, block = self.b.splitblock('start', terminate=True)
self.b.position_at_end(block)
const = partial(Const, type=types.Int32)
cond, body, exit = self.b.gen_loop(const(5), const(10), const(2))
with self.b.at_front(body):
self.b.print(c)
with self.b.at_end(exit):
self.b.ret(c)
self.assertEqual(interp.run(self.f, args=[10]), 100.0)
def test_splitblock_preserve_phis(self):
"""
block1:
%0 = mul a a
jump(newblock)
newblock:
%1 = phi([block1], [%0])
ret %1
"""
square = self.b.mul(self.a, self.a)
old, new = self.b.splitblock('newblock', terminate=True)
with self.b.at_front(new):
phi = self.b.phi(types.Int32, [self.f.startblock], [square])
self.b.ret(phi)
# Now split block1
self.b.position_after(square)
block1, split = self.b.splitblock(terminate=True)
phi, ret = new.ops
blocks, values = phi.args
self.assertEqual(blocks, [split])
