def visit_interior_kernel(self, node):
cur_node, ret_node = self.gen_loops(node)
body = cpp_ast.Block()
self.output_index_var = cpp_ast.CName(self.gen_fresh_var())
body.append(cpp_ast.Value("int", self.output_index_var))
body.append(
cpp_ast.Assign(
self.output_index_var,
self.gen_array_macro(self.output_grid_name,
[cpp_ast.CName(x)
for x in self.dim_vars])))
replaced_body = None
for gridname in self.argdict.keys():
replaced_body = [
ast_tools.ASTNodeReplacer(ast.Name(gridname, None),
ast.Name("_my_" + gridname,
None)).visit(x)
for x in node.body
]
body.extend([self.visit(x) for x in replaced_body])
cur_node.body = body
return ret_node
def handle_index(args, converter, kargs):
if len(args) == 0:
return cpp_ast.CName(converter.loopvar[-1])
elif len(args) == 1:
return cpp_ast.CName(converter.loopvar[-1 - args[0]])
else:
raise TypeError("Too many arguments to function 'index': %d" %
len(args))
def gen_array_unpack(self):
ret = [
cpp_ast.Assign(
cpp_ast.Pointer(cpp_ast.Value("npy_double", "_my_" + x)),
cpp_ast.TypeCast(
cpp_ast.Pointer(cpp_ast.Value("npy_double", "")),
cpp_ast.FunctionCall(cpp_ast.CName("PyArray_DATA"),
params=[cpp_ast.CName(x)])))
for x in self.argdict.keys()
]
return ret
def visit_arguments(self, node):
args = map(lambda arg: self.visit(arg), node.args)
defaults = map(lambda tup: self.visit(tup), node.defaults)
annotations = dict()
for tup in defaults:
annotations[tup[0].text] = tuple([elt.text for elt in tup])
ret = []
if len(self._data_model) != len(args):
raise TypeError('Expected %d arguments, received %d' %
(len(args), len(self._data_model)))
foo = zip(args, self._data_model)
for arg, model in foo:
model.name = arg.name
model.should_subsample = not (model.name in annotations
and 'nosubsample'
in annotations[model.name])
if not model.should_subsample and type(model) is not DataModel:
model = DataModel.clone(model)
model.dimensions = tuple([len(model)] + model.dimensions[1:])
ret.append(cpp_ast.Pointer(cpp_ast.Value(model.ctype(), arg.name)))
self.data_model[arg.name] = model
self.arg_model.append(model)
if self.weighted:
ret.append(
cpp_ast.Pointer(
cpp_ast.Value("const unsigned int",
cpp_ast.CName('_blb_weights'))))
return ret
def visit_InputElement(self, node):
index = self.gen_array_macro(
node.grid.name,
map(
lambda x, y: cpp_ast.BinOp(cpp_ast.CName(x), "+",
cpp_ast.CNumber(y)), self.dim_vars,
node.offset_list))
return cpp_ast.Subscript("_my_" + node.grid.name, index)
def handle_matrix(args, converter, kargs):
global anon_vecs
assert len(args) == 3, "matrix takes 3 arguments, %d supplied" % len(args)
assert isinstance(args[0], cpp_ast.CNumber)
assert isinstance(args[1], cpp_ast.CNumber)
assert isinstance(args[2], cpp_ast.CName)
name = "_blb_anon_vec%d" % anon_vecs
anon_vecs += 1
model = DataModel(args[2].name, [args[0].num, args[1].num], None, name)
model.declare = True
converter.data_model[name] = model
return cpp_ast.CName(name)
def inject_registers(self, body):
for model in self.data_model.itervalues():
if model.should_declare():
#TODO: This is an awkward way to do this.
body.contents.insert(
0,
cpp_ast.Call('memset', [
model.ref_name(), 0,
model.size() * model.scalar_t.csize()
]))
body.contents.insert(
0,
cpp_ast.Value(
model.scalar_t.ctype(),
cpp_ast.CName("%s [%d]" %
(model.ref_name(), model.size()))))
def handle_vector(args, converter, kargs):
global anon_vecs
length = 0
if isinstance(args[0], cpp_ast.CNumber):
length = args[0].num
else:
raise TypeError("vector: invalid length argument: %s" % ars[0])
tp = None
if type(args[1]) == cpp_ast.CName:
try:
tp = RobustType(args[1].name)
except KeyError:
raise TypeError("vector: invalid data type argument: %s" % args[1])
else:
raise TypeError("vector: invalid data type argument: %s" % args[1])
name = "_blb_anon_vec%d" % anon_vecs
anon_vecs += 1
model = DataModel(tp, [length], None, name)
model._declare = True
converter.data_model[name] = model
return cpp_ast.CName(name)
def get_register(self, dtype, dim):
""" Returns a pair of name, index representing a register. """
idx = -1
if not (dtype, dim) in self.registers:
self.registers[(dtype, dim)] = []
reglist = self.registers[(dtype, dim)]
for i in range(len(reglist)):
if reglist[i]:
reglist[i] = False
idx = i
break
if idx < 0:
reglist.append(False)
idx = len(reglist) - 1
regname = register_from_spec(dtype, dim, idx)
try:
return self.get_or_create_model(cpp_ast.CName(regname)), idx
except ValueError:
register = DataModel(dtype, [dim], name=regname)
self.data_model[regname] = register
register._declare = True
return register, idx
def declare_child(self, child, idx):
if child.is_scalar():
return cpp_ast.Subscript(self.ref_name(), idx)
else:
return cpp_ast.BinOp(cpp_ast.CName(self.ref_name()), '+',
cpp_ast.BinOp(idx, '*', self.element_size()))
def declare_child(self, child, idx):
return cpp_ast.Subscript(cpp_ast.CName(self.ref_name()), idx)
def visit_DataModel(self, node):
return cpp_ast.CName(node.ref_name())
def visit_For(self, node):
_iters = self.visit(node.iter)
targets = self.visit(node.target)
#are we iterating over a data object?
if type(_iters) is cpp_ast.CName and _iters.name in self.data_model:
_iters = [_iters]
targets = [targets]
#co-iterating over a set of datasets!
if type(_iters) is list:
#set up the loop variable and weight
self.loopvar.append(self.loopvar[-1] + 'i')
loopvar = cpp_ast.CName(self.loopvar[-1])
body = []
weight_loop = self.weighted and reduce(lambda a, b: a or b, [
(_iter.name in self.data_model
and self.data_model[_iter.name].should_subsample)
for _iter in _iters
])
if weight_loop:
self.aggregate_on(loopvar)
# add the temporary children to the data model
for target, _iter in zip(targets, _iters):
if not (type(_iter) is cpp_ast.CName
and _iter.name in self.data_model):
raise ValueError(
"Cannot iterate over unknown data object: '%s'" %
_iter.name)
parent_model = self.data_model[_iter.name]
if type(target) is not cpp_ast.CName:
raise ValueError("Not a valid iteration target: '%s'" %
str(target))
if target.name in self.data_model:
raise ValueError("Cannot reuse iteration target: '%s'" %
target.name)
target_model = self.data_model[
target.name] = parent_model.branch(name=target.name)
ctype = target_model.scalar_t.ctype()
decl = cpp_ast.Value(
ctype, target.name) if target_model.is_scalar(
) else cpp_ast.Pointer(cpp_ast.Value(ctype, target.name))
init = parent_model.declare_child(target_model, loopvar)
#init = cpp_ast.Subscript( _iter, loopvar) if target_model.is_scalar() \
# else cpp_ast.BinOp( _iter, '+', cpp_ast.BinOp( loopvar , '*', parent_model.element_size() ) )
body.append(cpp_ast.Assign(decl, init))
# visit the body of the for
body += [self.visit(x) for x in node.body]
# generate the C for intializer
ret = cpp_ast.RawFor( cpp_ast.Assign( cpp_ast.Value( 'int', loopvar ), cpp_ast.CNumber(0) ), \
cpp_ast.Compare( loopvar, '<', cpp_ast.CNumber( self.data_model[ _iter.name ].dimensions[0] ) ), \
cpp_ast.UnaryOp( '++', loopvar ), cpp_ast.Block( body ) )
# remove temporary children from data model and otherwise clean up
self.loopvar = self.loopvar[:-1]
for target in targets:
del self.data_model[target.name]
if weight_loop:
self.aggregate_off(loopvar)
# return
return ret
#or perhaps, a known range.
elif isinstance(_iters, cpp_ast.CNumber):
return cpp_ast.RawFor( cpp_ast.Assign( cpp_ast.Value( 'int', targets ), cpp_ast.CNumber(0) ), \
cpp_ast.Compare( targets, '<', _iter ), cpp_ast.UnaryOp( '++', target ), \
cpp_ast.Block([ self.visit( x ) for x in node.body ] ) )
else:
raise ValueError(
'Loop iterand "%s" is not a numeric expression or known data object'
% str(_iters))
def gen_array_macro(self, arg, point):
name = "_%s_array_macro" % arg
return cpp_ast.Call(cpp_ast.CName(name), point)
def visit_MathFunction(self, node):
return cpp_ast.FunctionCall(cpp_ast.CName(node.name),
params=map(self.visit, node.args))
def visit_Assign(self, node):
lhs = self.visit(node.targets[0])
rhs = self.visit(node.value)
#find or create target model
#pass pointer into opencall, or do what is necessary
target = None
try:
target = self.get_or_create_model(lhs)
except ValueError:
pass
value = None
try:
value = self.get_or_create_model(rhs)
except TypeError:
pass
if value is not None: #RHS is a data object
if target is not None:
if type(target) != DataModelView:
raise ValueError(
"'%s' refers to a real buffer and cannot be reassigned to '%s'"
% (str(lhs), str(rhs)))
#TODO we allways assume here that the name must be assigned, which is not always the case.
name = target.name
self.data_model[name] = DataModelView(value, name)
return cpp_ast.Assign(cpp_ast.CName(target.ref_name()),
cpp_ast.CName(value.ref_name()))
elif type(lhs) == cpp_ast.CName:
self.data_model[lhs.name] = DataModelView(value, lhs.name)
data_type = value.scalar_t.ctype() if value.is_scalar(
) else value.scalar_t.ctype() + "*"
return cpp_ast.Assign(cpp_ast.Value(data_type, lhs.name),
cpp_ast.CName(value.ref_name()))
else:
raise ValueError(
"could not assign to '%s': not a data object or name" %
str(lhs))
elif type(rhs) == OpenCall:
if target is None and type(lhs) == cpp_ast.CName:
params = rhs.get_output_params()
self.data_model[lhs.name] = target = DataModel(
params['type'], params['shape'], None, lhs.name)
pre = target.declare(self)
post = rhs.write_to(target, self)
return cpp_ast.UnbracedBlock([pre, post])
elif target:
return rhs.write_to(target, self)
else:
raise ValueError(
"could not assign to '%s': not a data object or name" %
str(lhs))
elif type(rhs) == cpp_ast.CNumber:
if target is None and type(lhs) == cpp_ast.CName:
self.data_model[lhs.name] = target = DataModel(
type(rhs.num), [1], None, lhs.name)
return cpp_ast.Initializer(
cpp_ast.Value(target.scalar_t.ctype(), lhs.name), rhs.num)
elif target:
assert target.is_scalar(
), "'%s' is not a scalar data object" % target.name
assert target.scalar_t.matches( type( rhs.num ) ), \
"Type mismatch: '%s' is type '%s', but '%s' is type '%s'" % ( target.name, target.scalar_t, rhs.num, type(rhs.num ) )
return cpp_ast.Initializer(
cpp_ast.Value(taget.scalar_t.ctype(), target.name),
rhs.num)
else:
raise ValueError(
"could not assign to '%s': not a data object or name" %
str(lhs))
else:
raise ValueError("could not assign from '%s'" % str(rhs))
def handle_dtype(args, converter, kargs):
if len(args) < 1:
raise TypeError("dtype takes one argument, %d given" % len(args))
return cpp_ast.CName(
converter.get_or_create_model(args[0]).scalar_t.ctype())
def __init__(self, weight_index, target):
self.left = cpp_ast.Subscript(cpp_ast.CName('_blb_weights'),
weight_index)
self.right = target
self.op = '*'
super(cpp_ast.BinOp, self).__init__()
def visit_Subscript(self, node):
index = self.visit(node.slice)
value = self.visit(node.value)
if type(node.ctx) == ast.Load:
if type(value) == cpp_ast.CName:
target_model = self.get_or_create_model(value)
if target_model.is_scalar():
raise ValueError(
'Invalid target for indexing: %s is scalar' %
value.name)
child_model = target_model.branch(idx=index.name)
if not child_model.weight_index:
child_model.weight_with(index.name)
ret = cpp_ast.Subscript(value, index)
if child_model.is_scalar() and self.weighted:
ret = WeightOp(child_model.weight_index, ret)
self.data_model[ret] = child_model
return ret
else:
self.data_model[ret] = child_model
return ret
elif type(value) == cpp_ast.Subscript:
target_model = self.get_or_create_model(value)
if target_model.is_scalar():
raise ValueError('Invalid target for indexing: %s' %
value.name)
child_model = target_model.branch()
new_index = cpp_ast.CName(
'(((%d)*(%s))+(%s))' %
(target_model.dimensions[0], value.index.generate(),
index.generate()))
ret = cpp_ast.Subscript(value, new_index)
if child_model.is_scalar() and self.weighted:
# This isn't fully general, and in fact only works for two level indexing.
ret = WeightOp(
child_model.weight_index
if child_model.weight_index else index, value)
self.data_model[ret] = child_model
return ret
else:
ret = cpp_ast.Subscript(value.value, new_index)
self.data_model[ret] = child_model
return ret
else:
raise ValueError('Invalid target for indexing: %s' %
str(value))
else:
#Nothing fancy here, just make sure everything gets written out right
if type(value) == cpp_ast.CName:
return cpp_ast.Subscipt(value, index)
elif type(value) == cpp_ast.Subscript:
if not target_model or target_model.is_scalar():
raise ValueError('Invalid target for indexing: %s' %
value.name)
child_model = target_model.branch()
new_index = cpp_ast.CName(
'(((%d)*(%s))+(%s))' %
(target_model.dimensions[0], value.index.generate(),
index.generate()))
return cpp_ast.Subscript(value.value, new_index)
else:
raise ValueError('Invalid target for indexing: %s' %
str(value))
