def test_fib(self): py_ast = get_ast(fib).body[0] c_ast = PyBasicConversions().visit(py_ast) filled_ast = DeclarationFiller().visit(c_ast) print(filled_ast) expected = """ void fib(n) { double a = 0.0; double b = 1.0; char* k = "hello"; while (n > 0) { double ____temp__a = b; double ____temp__b = a + b; a = ____temp__a; b = ____temp__b; n -= 1; } return a; }""" stripped_actual = str(filled_ast).replace(" ", "").replace("\n", "") stripped_expected = expected.replace(" ", "").replace("\n", "") self.assertEqual(stripped_actual, stripped_expected)
def transform(self, py_ast, program_cfg): arg_cfg, tune_cfg = program_cfg tree = PyBasicConversions().visit(py_ast) param_dict = {} tree.body[0].params.append(C.SymbolRef("retval", arg_cfg[0]())) # Annotate arguments for param, type in zip(tree.body[0].params, arg_cfg): param.type = type() param_dict[param.name] = type._dtype_ length = np.prod(arg_cfg[0]._shape_) transformer = MapTransformer("i", param_dict, "retval") body = list(map(transformer.visit, tree.body[0].defn)) tree.body[0].defn = [C.For( C.Assign(C.SymbolRef("i", ct.c_int()), C.Constant(0)), C.Lt(C.SymbolRef("i"), C.Constant(length)), C.PostInc(C.SymbolRef("i")), body=body, pragma="ivdep" )] tree = DeclarationFiller().visit(tree) defns = [] tree = HwachaVectorize(param_dict, defns).visit(tree) file_body = [ StringTemplate("#include <stdlib.h>"), StringTemplate("#include <stdint.h>"), StringTemplate("#include <assert.h>"), StringTemplate("extern \"C\" void __hwacha_body(void);"), ] file_body.extend(defns) file_body.append(tree) return [CFile("generated", file_body)]
def transform(self, tree, program_cfg): arg_cfg, tune_cfg = program_cfg tree = PyBasicConversions().visit(tree) tree = Backend(arg_cfg, self.symbol_table).visit(tree) tree = ConstantFold().visit(tree) tree.name = self.original_tree.body[0].name return tree
def transform(self, tree, program_config): """Convert the Python AST to a C AST.""" param_types = [] for arg in program_config[0]: param_types.append(NdPointer(arg[1], arg[2], arg[3])) kernel_sig = FuncType(Void(), param_types) tune_cfg = program_config[1] # block_factor = 2**tune_cfg['block_factor'] unroll_factor = 2**tune_cfg['unroll_factor'] for transformer in [StencilTransformer(self.input_grids, self.output_grid, self.constants ), PyBasicConversions()]: tree = transformer.visit(tree) first_For = tree.find(For) inner_For = FindInnerMostLoop().find(first_For) # self.block(inner_For, first_For, block_factor) self.unroll(inner_For, unroll_factor) # remove self param # TODO: Better way to do this? params = tree.find(FunctionDecl, name="kernel").params params.pop(0) self.gen_array_macro_definition(tree, params) entry_point = tree.find(FunctionDecl, name="kernel") entry_point.set_typesig(kernel_sig) return tree, entry_point.get_type().as_ctype()
def transform(self, tree, program_config): arg_types = program_config[0]['arg_typesig'] tree = PyBasicConversions().visit(tree.body[0]) tree.return_type = arg_types[0]() for param, ty in zip(tree.params, arg_types): param.type = ty() return [CFile(tree.name, [tree])]
def eval_with_loop(self, elts): new_elts = [] for elt in elts: elt = self.replace_loopvars_as_constants(copy.deepcopy(elt)) elt = PyBasicConversions().visit(elt) elt = ConstantFold().visit(elt) new_elts.append(elt.value) return tuple(new_elts)
def test_multiple_assign_constant(self): node = ast.Assign([ ast.Tuple(elts=(ast.Name(id="x", ctx=None), ast.Name(id="y", ctx=None))) ], ast.Tuple(elts=(Constant(1), Constant(2)))) transformed_node = PyBasicConversions().visit(node) self.assertEqual(str(transformed_node), "\nx = 1;\ny = 2;\n")
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 visit_AugAssign(self, node): # TODO: Handle all types? value = self.visit(node.value) # HACK to get this to work, PyBasicConversions will skip this AugAssign node # TODO Figure out why value = PyBasicConversions().visit(value) if type(node.op) is ast.Add: return AddAssign(self.visit(node.target), value) if type(node.op) is ast.Sub: return SubAssign(self.visit(node.target), value)
def transform(self, tree, program_cfg): arg_cfg, tune_cfg = program_cfg # tree = Desugar().visit(tree) inliner = InlineEnvironment(self.symbol_table) tree = inliner.visit(tree) tree = PyBasicConversions().visit(tree) tree.body = inliner.files + tree.body # tree.find(C.For).pragma = 'omp parallel for' tree.name = self.original_tree.body[0].name tree.body.insert(0, StringTemplate("#include <math.h>")) # print(tree) return [tree]
def transform(self, tree, program_config): args_subconfig, tuning_config = program_config function = tree.body[0] c_func = PyBasicConversions().visit(function) #print(c_func) c_func.defn = c_func.defn[1:] #print(c_func) c_func.params[0].type = ctree.types.get_ctype(args_subconfig) c_func.params[1].type = c_func.params[0].type c_func.params.append(SymbolRef('arr', ctypes.POINTER(ctypes.c_int32)())) #print(c_func) return CFile(body=[c_func])
def test_multiple_assign_dependent(self): node = ast.Assign([ ast.Tuple(elts=(ast.Name(id="x", ctx=None), ast.Name(id="y", ctx=None))) ], ast.Tuple(elts=(ast.Name(id="y", ctx=None), ast.Name(id="x", ctx=None)))) transformed_node = PyBasicConversions().visit(node) self.assertEqual( str(transformed_node), "\n____temp__x = x;\n____temp__y = y;\ny = ____temp__y;\nx = ____temp__x;\n" )
def test_multiple_assign_dependent(self): node = ast.Assign( [ ast.Tuple(elts=(ast.Name(id="x", ctx=None), ast.Name(id="y", ctx=None))) ], ast.Tuple( elts=(FunctionCall(func='square', args=[Constant(5), Constant(5)]), FunctionCall(func='square', args=[Constant(5), Constant(5)])))) transformed_node = PyBasicConversions().visit(node) self.assertEqual( str(transformed_node), "\n____temp__x = square(5, 5);\n____temp__y = square(5, 5);\nx = ____temp__x;\ny = ____temp__y;\n" )
def mini_transform(self, node): """ This method acts as a simulation of a specializer's transform() method. It's the bare minimum required of a transform() method by the specializer writer. :param node: the node to transform :return: the node transformed through PyBasicConversions into a rough C-AST. """ transformed_node = PyBasicConversions().visit(node) transformed_node.name = "apply" transformed_node.return_type = ct.c_float() for param in transformed_node.params: param.type = ct.c_float() return transformed_node
def test_fmin(self): def func(): a = 3.0 b = 4.0 c = fmax(a + b, 0.0) return c py_ast = get_ast(func).body[0] c_ast = PyBasicConversions().visit(py_ast) filled_ast = DeclarationFiller().visit(c_ast) expected = """ void func() { double a = 3.0; double b = 4.0; double c = fmax(a + b, 0.0); return c; }""" stripped_actual = str(filled_ast).replace(" ", "").replace("\n", "") stripped_expected = expected.replace(" ", "").replace("\n", "") self.assertEqual(stripped_actual, stripped_expected)
def transform(self, py_ast, program_config): # Get the initial data input_data = program_config[0] length = np.prod(input_data.size) pointer = np.ctypeslib.ndpointer(input_data.dtype, input_data.ndim, input_data.shape) data_type = get_c_type_from_numpy_dtype(input_data.dtype)() scalar_data_type = get_c_type_from_numpy_dtype( np.dtype(input_data.scalar_type))() apply_one = PyBasicConversions().visit(py_ast.body[0]) apply_one.name = 'apply' apply_one.params[0].type = data_type apply_one.params[1].type = scalar_data_type apply_one.return_type = data_type # TODO: figure out which data type to actually preserve # TODO: MAKE A CLASS THAT HANDLES SUPPORTED TYPES (INT, FLOAT, DOUBLE) array_add_template = StringTemplate( r""" #pragma omp parallel for for (int i = 0; i < $length; i++) { output[i] = apply(arr[i], scalar); } """, {'length': Constant(length)}) array_op = CFile("generated", [ CppInclude("omp.h"), CppInclude("stdio.h"), apply_one, FunctionDecl(None, FUNC_NAME, params=[ SymbolRef("arr", pointer()), SymbolRef("scalar", scalar_data_type), SymbolRef("output", pointer()) ], defn=[array_add_template]) ], 'omp') return [array_op]
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_simple_transform(self): class Kernel(StencilKernel): def kernel(self, in_img, out_img): for x in out_img.interior_points(): for y in in_img.neighbors(x, 1): out_img[x] += in_img[y] kernel = Kernel() kernel.should_unroll = False out_grid = StencilGrid([5]) out_grid.ghost_depth = radius in_grid = StencilGrid([5]) in_grid.ghost_depth = radius for x in range(0, 5): in_grid.data[x] = 1 tree1 = ctree.get_ast(Kernel.kernel) tree2 = PyBasicConversions().visit(tree1) actual = StencilOmpTransformer([in_grid], out_grid, kernel).visit(tree2) self.assertEqual(actual, second)
def transform(self, py_ast, program_config): # Get the initial data input_data = program_config[0] length = np.prod(input_data.size) pointer = np.ctypeslib.ndpointer(input_data.dtype, input_data.ndim, input_data.shape) data_type = get_c_type_from_numpy_dtype(input_data.dtype)() apply_one = PyBasicConversions().visit(py_ast.body[0]) apply_one.name = 'apply' apply_one.params[0].type = data_type apply_one.params[1].type = data_type apply_one.return_type = data_type array_add_template = StringTemplate(r""" #pragma omp parallel for for (int i = 0; i < $length; i++) { output[i] = apply(input1[i], input2[i]); } """, { 'length': Constant(length) }) array_op = CFile("generated", [ CppInclude("omp.h"), CppInclude("stdio.h"), apply_one, FunctionDecl(None, FUNC_NAME, params=[ SymbolRef("input1", pointer()), SymbolRef("input2", pointer()), SymbolRef("output", pointer()) ], defn=[ array_add_template ]) ], 'omp') return [array_op]
def visit_Assign(self, node): target = PyBasicConversions().visit(self.visit(node.targets[0])) value = PyBasicConversions().visit(self.visit(node.value)) return Assign(target, value)
def test_minus(self): op = ast.parse("- foo") op = PyBasicConversions().visit(op).find(UnaryOp) self._check(op, "- foo")
def test_not(self): op = ast.parse("not foo") op = PyBasicConversions().visit(op).find(UnaryOp) self._check(op, "! foo")
def test_CUnaryOp(self): op = Not(SymbolRef("foo")) op = PyBasicConversions().visit(op).find(UnaryOp) self._check(str(op), "! foo")
def test_LessThan(self): comp = ast.parse("5 < foo < 6") comp = PyBasicConversions().visit(comp).find(BinaryOp) self.assertEqual(str(comp), "5 < foo && foo < 6")
def test_Equals(self): comp = ast.parse("5 == foo == 6") comp = PyBasicConversions().visit(comp).find(BinaryOp) self.assertEqual(str(comp), "5 == foo && foo == 6")
def test_List(self): array = ast.parse("[1, 5, 7, 3]") array = PyBasicConversions().visit(array).find(Array) self.assertEqual(str(array), "{1, 5, 7, 3}")
__author__ = 'nzhang-dev' import ast from ctree.transforms import DeclarationFiller from ctree.transformations import PyBasicConversions from ctree.frontend import * from ctree.c.nodes import MultiNode code = ["a = 1", "a,b = 1,1", "a = b = 1", """a,b = 1,1 \na,b = b,a"""] def fib(n): a, b = 0, 1 while n > 0: n -= 1 a, b = b, a + b return a asts = [] for c in code: parsed = ast.parse(c) asts.append(MultiNode(body=parsed.body)) asts.append(get_ast(fib).body[0]) processed = [ DeclarationFiller().visit(PyBasicConversions().visit(a)) for a in asts ]
def test_GreaterThan(self): comp = ast.parse("5 >= foo >= 6") comp = PyBasicConversions().visit(comp).find(BinaryOp) self.assertEqual(str(comp), "5 >= foo && foo >= 6")