def test_subtree_docstrings(self): tree = PyBasicConversions().visit( ast.Module([ get_ast(self.foo), get_ast(self.Bar), ])) self._check(tree)
def test_boolop(self): def fn(): 1 or 2 or 3 py_ast = get_ast(fn).body[0].body[0] c_ast = Or(Or(Constant(1), Constant(2)), Constant(3)) self._check(py_ast, c_ast)
def test_compare(self): def fn(): 0 < a < 3 py_ast = get_ast(fn).body[0].body[0] c_ast = And(Gt(0, SymbolRef("a")), Lt(SymbolRef("a"), Constant(3))) self._check(py_ast, c_ast)
def specialize(fn=None, output=None): """ Specializes function fn using SpecalizedFn. """ if fn is None: def wrapped(fn): return specialize(fn, output) return wrapped frame = inspect.stack()[1][0] symbol_table = frame.f_locals for i in range(3): frame = frame.f_back if frame is None: # End of stack break symbol_table.update(frame.f_locals) # FIXME: symbol_table prints out a huge dict, why?? # TODO: We grab the last two frames, what to do if there's more? spec_fn = SpecializedFn(get_ast(fn), symbol_table, output) @wraps(fn) def fn(*args, **kwargs): return spec_fn(*args, **kwargs) fn._specializer = spec_fn return fn
def test_fuse_3(self): a = numpy.random.rand(100, 100).astype(numpy.float32) * 100 b = numpy.random.rand(100, 100).astype(numpy.float32) * 100 def f(a, b): c = array_mul(a, b) d = array_sub(a, c) e = array_add(c, d) return e tree = get_ast(f) blocks = [] BlockBuilder(blocks).visit(tree) fuser = Fuser(blocks, dict(locals(), **globals())) fused = fuser._fuse([blocks[0], blocks[1], blocks[2]]).value actual_c, actual_d, actual_e = fuser._symbol_table[fused.func.id]( *(fuser._symbol_table[arg.id] if isinstance(arg, ast.Name) else arg.n for arg in fused.args) ) expected_c = a * b expected_d = a - expected_c expected_e = expected_c + expected_d try: testing.assert_array_almost_equal(actual_c, expected_c) testing.assert_array_almost_equal(actual_d, expected_d) testing.assert_array_almost_equal(actual_e, expected_e) except Exception as e: self.fail("Arrays not almost equal: {0}".format(e))
def emit(cls, sources, sinks, keywords, symbol_table): tree = get_ast(cls.fn) tree = PyBasicConversions().visit(tree) body = tree.body[0].defn mapping = {arg.name: source for arg, source in zip(tree.body[0].params, sources)} visitor = MapTransformer(mapping, sinks[0]) body = [visitor.visit(s) for s in body] return "\n".join([str(s) + ";" for s in body])
def test_simple(self): def f(a): return array_mul(a) tree = get_ast(f) blocks = [] BlockBuilder(blocks).visit(tree) self.assertEqual(len(blocks), 1) self.assertIsInstance(blocks[0], ast.Return)
def meta(fn): frame = inspect.stack()[1][0].f_back tree = get_ast(fn) symbol_table = frame.f_locals spec = MetaSpecialized(tree, symbol_table) def fn(*args, **kwargs): return spec(*args, **kwargs) return fn
def test_is_fusable_false(self): def f(a, b): c = array_mul(a, b) return c tree = get_ast(f) blocks = [] BlockBuilder(blocks).visit(tree) fuser = Fuser(blocks, dict(locals(), **globals())) self.assertFalse(fuser._is_fusable(blocks[0], blocks[1]))
def __init__(self, py_ast=None, sub_dir=None, backend_name="default"): if py_ast is not None and \ self.apply is not LazySpecializedFunction.apply: raise TypeError('Cannot define apply and pass py_ast') self.original_tree = py_ast or \ (get_ast(self.apply) if self.apply is not LazySpecializedFunction.apply else None) self.concrete_functions = {} # config -> callable map self._tuner = self.get_tuning_driver() self.sub_dir = sub_dir or \ self.NameExtractor().visit(self.original_tree) or \ hex(hash(self))[2:] self.backend_name = backend_name
def test_multiline(self): def f(a, b): c = array_mul(a, b) d = array_sub(a, c) e = d * c return e / 4 tree = get_ast(f) blocks = [] BlockBuilder(blocks).visit(tree) self.assertEqual(len(blocks), 4) self.assertIsInstance(blocks[0], ast.Assign) self.assertIsInstance(blocks[1], ast.Assign) self.assertIsInstance(blocks[2], ast.Assign) self.assertIsInstance(blocks[3], ast.Return)
def blasc(fn): frame = inspect.stack()[1][0] symbol_table = frame.f_locals symbol_table.update(frame.f_globals) symbol_table.update(frame.f_back.f_locals) symbol_table.update(frame.f_back.f_globals) tree = get_ast(fn) def wrapped(*args, **kwargs): print(args) print([arg.id for arg in tree.body[0].args.args]) for index, arg in enumerate(tree.body[0].args.args): symbol_table[arg.id] = args[index] # Process tree return get_callable(tree, tree.body[0].name, symbol_table)(*args, **kwargs) return wrapped
def get_semantic_node(self, arg_names, *args): class StencilCall(ast.AST): _fields = ['params', 'body'] def __init__(self, function_decl, input_grids, output_grid, kernel): self.params = function_decl.params self.body = function_decl.defn self.function_decl = function_decl self.input_grids = input_grids self.output_grid = output_grid self.kernel = kernel def label(self): return "" def to_dot(self): return "digraph mytree {\n%s}" % self._to_dot() def _to_dot(self): from ctree.dotgen import DotGenVisitor return DotGenVisitor().visit(self) def add_undef(self): self.function_decl.defn[0].add_undef() def remove_types_from_decl(self): self.function_decl.defn[1].remove_types_from_decl() def backend_transform(self, block_padding, local_input): return StencilOclTransformer( self.input_grids, self.output_grid, self.kernel, block_padding).visit(self.function_decl) def backend_semantic_transform(self, fusion_padding): self.function_decl = StencilOclSemanticTransformer( self.input_grids, self.output_grid, self.kernel, fusion_padding).visit(self.function_decl) self.body = self.function_decl.defn self.params = self.function_decl.params func_decl = PythonToStencilModel(arg_names).visit(get_ast( self.model)).files[0].body[0] return StencilCall(func_decl, args[:-1], args[-1], self)
def from_function(cls, func, folder_name=''): class Replacer(ast.NodeTransformer): def visit_Module(self, node): return MultiNode(body=[self.visit(i) for i in node.body]) def visit_FunctionDef(self, node): if node.name == func.__name__: node.name = 'apply' node.body = [self.visit(item) for item in node.body] return node def visit_Name(self, node): if node.id == func.__name__: node.id = 'apply' return node func_ast = Replacer().visit(get_ast(func)) return cls(py_ast=func_ast, sub_dir=folder_name or func.__name__)
def __init__(self, stencil_kernel, backend_name, boundary_handling=""): """ Initializes an instance of a SpecializedStencil. This function inherits from ctree's LazySpecializedFunction. When the specialized function is called, it will either load a cached version, or generate a new version using the kernel method's AST and the passed parameters . The tuning configurations are defined in get_tuning_driver. The arguments to the specialized function call are passed to args_to_subconfig where they can be processed to a form usable by the specializer. For more information consult the ctree docs. :param stencil_kernel: The Kernel object containing the kernel function. :param backend_name: the type of specialized kernel to generate \ """ self.kernel = stencil_kernel self.backend = self.backend_dict[backend_name] self.output = None self.args = None backend_key = "{}_{}".format(backend_name, boundary_handling) super(SpecializedStencil, self).__init__(get_ast(stencil_kernel.kernel), backend_name=backend_key)
def visit_FunctionCall(self, node): if node.func.name in {'min', 'max'}: node.func.name = "f" + node.func.name # TODO: Add support for all math funcs self.includes.add("math.h") return super(Backend, self).generic_visit(node) # FIXME: This is specific for handling a map function # do we have to generalize? node.args = [self.visit(arg) for arg in node.args] func_tree = get_ast(self.symbol_table[node.func.name]) func_tree = PyBasicConversions().visit(func_tree).body[0] func_tree = self.visit(func_tree) func_tree.name = C.SymbolRef(node.func.name) func_tree.set_static() func_tree.set_inline() self.defns.append(func_tree) # FIXME: Infer type for p in func_tree.params: p.type = ct.c_float() func_tree.return_type = ct.c_float() return node
def test_fuse_with_return(self): a = numpy.random.rand(100, 100).astype(numpy.float32) * 100 b = numpy.random.rand(100, 100).astype(numpy.float32) * 100 def f(a, b): c = array_mul(a, b) d = scalar_array_mul(4, c) return array_sub(c, d) orig_f = f tree = get_ast(f) blocks = get_blocks(tree) fuser = Fuser(blocks, dict(locals(), **globals())) fused_blocks = fuser.do_fusion() tree.body[0].body = fused_blocks tree = ast.fix_missing_locations(tree) exec(compile(tree, '', 'exec')) in fuser._symbol_table try: testing.assert_array_almost_equal(fuser._symbol_table['f'](a, b), orig_f(a, b)) except Exception as e: self.fail("Arrays not almost equal: {0}".format(e))
def __init__(self, kernel, backend, testing=False): """ Initializes an instance of a SpecializedStencil. This function inherits from ctree's LazySpecializedFunction. When the specialized function is called, it will either load a cached version, or generate a new version using the kernel method's AST and the passed parameters . The tuning configurations are defined in get_tuning_driver. The arguments to the specialized function call are passed to args_to_subconfig where they can be processed to a form usable by the specializer. For more information consult the ctree docs. :param func: Stencil Kernel function to be specialized. :param input_grids: List of input grids passed as arguments to stencil. :param output_grid: Output grid passed to stencil. :param kernel: The Kernel object containing the kernel function. :param testing: Optional - whether or not we are in testing mode """ self.testing = testing self.kernel = kernel self.backend = self.backend_dict[backend] self.output = None super(SpecializedStencil, self).__init__(get_ast(kernel.kernel))
def test_fib_py(self): self._check_raises(get_ast(fib))
def test_gcd(self): tree = PyCtxScrubber().visit(get_ast(gcd)) self._check(tree)
def __init__(self, py_ast=None): py_ast = py_ast or get_ast(self.apply) super(LazySlimmy, self).__init__(py_ast)
def test_class_docstring(self): tree = PyBasicConversions().visit(get_ast(self.Bar)) self._check(tree)
def test_identity(self): tree = PyCtxScrubber().visit(get_ast(identity)) self._check(tree)
def test_py_identity(self): self.assertNotEqual(to_dot(get_ast(identity)), "")
def test_py_fib(self): self.assertNotEqual(to_dot(get_ast(fib)), "")
def test_func_docstring(self): tree = PyBasicConversions().visit(get_ast(self.foo)) self._check(tree)
def test_fib(self): tree = PyCtxScrubber().visit(get_ast(fib)) self._check(tree)
def specialize_model(sim): """ Take out the model, feed it to the specializer and replace the dfun with specialized version. Simulation is expected to be already configured. """ py_ast = get_ast(sim.model.dfun) sim.model.dfun = nes.CModelDfun(py_ast, sim)
def __init__(self, func, entry_point, input_grids, output_grid, constants): self.input_grids = input_grids self.output_grid = output_grid self.constants = constants super(StencilConvert, self).__init__(get_ast(func))
def test_gcd_py(self): self._check_raises(get_ast(gcd))
def test_py_gcd(self): self.assertNotEqual(get_ast(gcd).to_dot(), "")
def test_py_fib(self): self.assertNotEqual(get_ast(fib).to_dot(), "")
def test_subtree_docstrings(self): tree = PyBasicConversions().visit(ast.Module([ get_ast(self.foo), get_ast(self.Bar), ])) self._check(tree)
def test_py_gcd(self): self.assertNotEqual(to_dot(get_ast(gcd)), "")
def __init__(self, fn): """Instantiate translator.""" self.jit = OpTranslator(get_ast(fn))
def __init__(self): """Instantiate translator.""" from ctree.frontend import get_ast self.c_apply_all = OpTranslator(get_ast(self.apply), "apply_all")
def __init__(self): """Instantiate translator.""" self.c_apply_all = OpTranslator(get_ast(self.apply), "apply_all")
def test_gcd(self): self.assertIsInstance(get_ast(gcd), ast.AST)
def test_identity(self): self.assertIsInstance(get_ast(identity), ast.AST)
def test_fib(self): self.assertIsInstance(get_ast(fib), ast.AST)
def test_py_identity(self): self.assertNotEqual(get_ast(identity).to_dot(), "")
def __init__(self): """Instantiate translator.""" self.c_apply_all = OpTranslator(get_ast(self.apply))
def test_identity_py(self): self._check_raises(get_ast(identity))