def visit_Assign(self, node):
if len(node.targets) > 1:
raise NotImplementedError("cannot process multiple assignment")
if not isinstance(node.targets[0], gast.Name):
raise NotImplementedError("cannot process indexed assignment")
# $lhs = $lhs.update_($rhs, matchbox.EXECUTION_MASK) if (lhs in vars()
# or lhs in globals()) and isinstance($lhs, (matchbox.MaskedBatch,
# matchbox.TENSOR_TYPE)) else $rhs
node.value = gast.IfExp(
gast.BoolOp(
gast.And(),
[
gast.BoolOp(gast.Or(), [
gast.Compare(gast.Str(
node.targets[0].id), [gast.In()], [
gast.Call(gast.Name('vars', gast.Load, None),
[], [])
]),
gast.Compare(gast.Str(
node.targets[0].id), [gast.In()], [
gast.Call(
gast.Name('globals', gast.Load, None), [],
[])
])
]),
# gast.Compare(
# gast.Attribute(
# gast.Name('matchbox', gast.Load(), None),
# gast.Name('EXECUTION_MASK', gast.Load(), None),
# gast.Load()),
# [gast.IsNot()],
# [gast.NameConstant(None)]),
gast.Call(gast.Name('isinstance', gast.Load(), None), [
node.targets[0],
gast.Tuple([
gast.Attribute(
gast.Name('matchbox', gast.Load(), None),
gast.Name('MaskedBatch', gast.Load(), None),
gast.Load()),
gast.Attribute(
gast.Name('matchbox', gast.Load(), None),
gast.Name('TENSOR_TYPE', gast.Load(), None),
gast.Load())
], gast.Load())
], [])
]),
gast.Call(
gast.Attribute(
gast.Name(node.targets[0].id, gast.Load(), None),
gast.Name('_update', gast.Load(), None), gast.Load()), [
node.value,
gast.Attribute(
gast.Name('matchbox', gast.Load(), None),
gast.Name('EXECUTION_MASK', gast.Load(), None),
gast.Load())
], []),
node.value)
return node
def visit_BinOp(self, node):
node = self.generic_visit(node)
if not isinstance(node.op, ast.Mod):
return node
right_range = self.range_values[node.right]
left_range = self.range_values[node.left]
if right_range.low < 0 or isinf(right_range.high):
return node
if left_range.low < -right_range.low:
return node
if left_range.high > right_range.high * 2:
return node
cleft0, cleft1 = deepcopy(node.left), deepcopy(node.left)
cright = deepcopy(node.right)
self.update = True
return ast.IfExp(ast.Compare(node.left, [ast.Lt()], [node.right]),
cleft0, ast.BinOp(cleft1, ast.Sub(), cright))
def visit_BinOp(self, node):
if not isinstance(node.op, ast.Mod):
return self.generic_visit(node)
# check that right is a name defined once outside of loop
# TODO: handle expression instead of names
if not isinstance(node.right, ast.Name):
return self.generic_visit(node)
right_def = self.single_def(node.right)
if not right_def:
return self.generic_visit(node)
if self.range_values[node.right.id].low < 0:
return self.generic_visit(node)
# same for lhs
if not isinstance(node.left, ast.Name):
return self.generic_visit(node)
head = self.single_def(node.left)
if not head:
return self.generic_visit(node)
# check lhs is the actual index of a loop
head = head['name']
loop = self.ancestors[head][-1]
if not isinstance(loop, ast.For):
return self.generic_visit(node)
if not isinstance(loop.iter, ast.Call):
return self.generic_visit(node)
# make sure rhs is defined out of the loop
if loop in self.ancestors[right_def['name']]:
return self.generic_visit(node)
# gather range informations
range_ = None
for alias in self.aliases[loop.iter.func]:
if alias is MODULES['__builtin__']['range']:
range_ = alias
elif alias is MODULES['__builtin__']['xrange']:
range_ = alias
else:
break
if range_ is None:
return self.generic_visit(node)
# everything is setup for the transformation!
new_id = node.left.id + '_m'
i = 0
while new_id in self.identifiers:
new_id = '{}_m{}'.format(node.left.id, i)
i += 1
rargs = range_.args.args
lower = rargs[0] if len(rargs) > 1 else ast.Num(0)
header = ast.Assign([ast.Name(new_id, ast.Store(), None)],
ast.BinOp(
ast.BinOp(deepcopy(lower), ast.Sub(),
ast.Num(1)), ast.Mod(),
deepcopy(node.right)))
incr = ast.BinOp(ast.Name(new_id, ast.Load(), None), ast.Add(),
ast.Num(1))
step = ast.Assign([ast.Name(new_id, ast.Store(), None)],
ast.IfExp(
ast.Compare(incr, [ast.Eq()],
[deepcopy(node.right)]), ast.Num(0),
deepcopy(incr)))
self.loops_mod.setdefault(loop, []).append((header, step))
self.update = True
return ast.Name(new_id, ast.Load(), None)
