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")
