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_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 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 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 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 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_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 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 __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 __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 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_gcd(self):
tree = PyCtxScrubber().visit(get_ast(gcd))
self._check(tree)
def test_func_docstring(self):
tree = PyBasicConversions().visit(get_ast(self.foo))
self._check(tree)
def test_identity(self):
tree = PyCtxScrubber().visit(get_ast(identity))
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_class_docstring(self):
tree = PyBasicConversions().visit(get_ast(self.Bar))
self._check(tree)
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):
tree = PyCtxScrubber().visit(get_ast(fib))
self._check(tree)
def test_fib(self):
self.assertIsInstance(get_ast(fib), ast.AST)
def test_func_docstring(self):
tree = PyBasicConversions().visit(get_ast(self.foo))
self._check(tree)
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))
