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)