Restructured KernelFunction node to get rid of monkey-patching

- backend, target and compile are now normal members of the
  KernelFunction node and populated in constructor

pystencils/astnodes.py

@@ -162,18 +162,30 @@ class KernelFunction(Node):
         def field_name(self):
             return self.fields[0].name
 
-    def __init__(self, body, ghost_layers=None, function_name="kernel", backend=""):
+    def __init__(self, body, target, backend, compile_function, ghost_layers, function_name="kernel"):
         super(KernelFunction, self).__init__()
         self._body = body
         body.parent = self
         self.function_name = function_name
         self._body.parent = self
-        self.compile = None
         self.ghost_layers = ghost_layers
+        self._target = target
+        self._backend = backend
         # these variables are assumed to be global, so no automatic parameter is generated for them
         self.global_variables = set()
-        self.backend = backend
         self.instruction_set = None  # used in `vectorize` function to tell the backend which i.s. (SSE,AVX) to use
+        # function that compiles the node to a Python callable, is set by the backends
+        self._compile_function = compile_function
+
+    @property
+    def target(self):
+        """Currently either 'cpu' or 'gpu' """
+        return self._target
+
+    @property
+    def backend(self):
+        """Backend for generating the code e.g. 'llvm', 'c', 'cuda' """
+        return self._backend
 
     @property
     def symbols_defined(self):
@@ -194,7 +206,7 @@ class KernelFunction(Node):
 
     @property
     def args(self):
-        return [self._body]
+        return self._body,
 
     @property
     def fields_accessed(self) -> Set['ResolvedFieldAccess']:
@@ -231,6 +243,11 @@ class KernelFunction(Node):
         params = [p.symbol for p in self.get_parameters()]
         return '{0} {1}({2})'.format(type(self).__name__, self.function_name, params)
 
+    def compile(self, *args, **kwargs):
+        if self._compile_function is None:
+            raise ValueError("No compile-function provided for this KernelFunction node")
+        return self._compile_function(self, *args, **kwargs)
+
 
 class SkipIteration(Node):
     @property

pystencils/cpu/kernelcreation.py

 import sympy as sp
-from functools import partial
 from pystencils.astnodes import SympyAssignment, Block, LoopOverCoordinate, KernelFunction
 from pystencils.transformations import resolve_buffer_accesses, resolve_field_accesses, make_loop_over_domain, \
     add_types, get_optimal_loop_ordering, parse_base_pointer_info, move_constants_before_loop, \
@@ -61,9 +60,10 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
 
     body = ast.Block(assignments)
     loop_order = get_optimal_loop_ordering(fields_without_buffers)
-    ast_node = make_loop_over_domain(body, function_name, iteration_slice=iteration_slice,
-                                     ghost_layers=ghost_layers, loop_order=loop_order)
-    ast_node.target = 'cpu'
+    loop_node, ghost_layer_info = make_loop_over_domain(body, iteration_slice=iteration_slice,
+                                                        ghost_layers=ghost_layers, loop_order=loop_order)
+    ast_node = KernelFunction(loop_node, 'cpu', 'c', compile_function=make_python_function,
+                              ghost_layers=ghost_layer_info, function_name=function_name)
 
     if split_groups:
         typed_split_groups = [[type_symbol(s) for s in split_group] for split_group in split_groups]
@@ -83,7 +83,6 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
         resolve_buffer_accesses(ast_node, get_base_buffer_index(ast_node), read_only_fields)
     resolve_field_accesses(ast_node, read_only_fields, field_to_base_pointer_info=base_pointer_info)
     move_constants_before_loop(ast_node)
-    ast_node.compile = partial(make_python_function, ast_node)
     return ast_node
 
 
@@ -141,14 +140,14 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu
         loop_body.append(assignment)
 
     function_body = Block([loop_node])
-    ast_node = KernelFunction(function_body, backend="cpu", function_name=function_name)
+    ast_node = KernelFunction(function_body, "cpu", "c", make_python_function,
+                              ghost_layers=None, function_name=function_name)
 
     fixed_coordinate_mapping = {f.name: coordinate_typed_symbols for f in non_index_fields}
 
     read_only_fields = set([f.name for f in fields_read - fields_written])
     resolve_field_accesses(ast_node, read_only_fields, field_to_fixed_coordinates=fixed_coordinate_mapping)
     move_constants_before_loop(ast_node)
-    ast_node.compile = partial(make_python_function, ast_node)
     return ast_node
 
 

pystencils/gpucuda/kernelcreation.py

-from functools import partial
-
 from pystencils.gpucuda.indexing import BlockIndexing
 from pystencils.transformations import resolve_field_accesses, add_types, parse_base_pointer_info, \
     get_common_shape, resolve_buffer_accesses, unify_shape_symbols, get_base_buffer_index
@@ -55,7 +53,7 @@ def create_cuda_kernel(assignments, function_name="kernel", type_info=None, inde
     block = indexing.guard(block, common_shape)
     unify_shape_symbols(block, common_shape=common_shape, fields=fields_without_buffers)
 
-    ast = KernelFunction(block, function_name=function_name, ghost_layers=ghost_layers, backend='gpucuda')
+    ast = KernelFunction(block, 'gpu', 'gpucuda', make_python_function, ghost_layers, function_name)
     ast.global_variables.update(indexing.index_variables)
 
     base_pointer_spec = [['spatialInner0']]
@@ -84,7 +82,6 @@ def create_cuda_kernel(assignments, function_name="kernel", type_info=None, inde
         ast.body.insert_front(SympyAssignment(loop_counter, indexing.coordinates[i]))
 
     ast.indexing = indexing
-    ast.compile = partial(make_python_function, ast)
     return ast
 
 
@@ -124,7 +121,7 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel
 
     function_body = Block(coordinate_symbol_assignments + assignments)
     function_body = indexing.guard(function_body, get_common_shape(index_fields))
-    ast = KernelFunction(function_body, function_name=function_name, backend='gpucuda')
+    ast = KernelFunction(function_body, 'gpu', 'gpucuda', make_python_function, function_name)
     ast.global_variables.update(indexing.index_variables)
 
     coord_mapping = indexing.coordinates
@@ -141,5 +138,4 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel
     # add the function which determines #blocks and #threads as additional member to KernelFunction node
     # this is used by the jit
     ast.indexing = indexing
-    ast.compile = partial(make_python_function, ast)
     return ast

pystencils/llvm/kernelcreation.py

 from pystencils.transformations import insert_casts
-from functools import partial
 from pystencils.llvm.llvmjit import make_python_function
 
 
@@ -10,23 +9,25 @@ def create_kernel(assignments, function_name="kernel", type_info=None, split_gro
 
     Loops are created according to the field accesses in the equations.
 
-    :param assignments: list of sympy equations, containing accesses to :class:`pystencils.field.Field`.
-           Defining the update rules of the kernel
-    :param function_name: name of the generated function - only important if generated code is written out
-    :param type_info: a map from symbol name to a C type specifier. If not specified all symbols are assumed to
-           be of type 'double' except symbols which occur on the left hand side of equations where the
-           right hand side is a sympy Boolean which are assumed to be 'bool' .
-    :param split_groups: Specification on how to split up inner loop into multiple loops. For details see
-           transformation :func:`pystencils.transformation.split_inner_loop`
-    :param iteration_slice: if not None, iteration is done only over this slice of the field
-    :param ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
-                        if None, the number of ghost layers is determined automatically and assumed to be equal for a
-                        all dimensions
+    Args:
+    assignments: list of sympy equations, containing accesses to :class:`pystencils.field.Field`.
+                 Defining the update rules of the kernel
+    function_name: name of the generated function - only important if generated code is written out
+                   type_info: a map from symbol name to a C type specifier. If not specified all symbols are assumed to
+                   be of type 'double' except symbols which occur on the left hand side of equations where the
+                   right hand side is a sympy Boolean which are assumed to be 'bool' .
+    split_groups: Specification on how to split up inner loop into multiple loops. For details see
+                  transformation :func:`pystencils.transformation.split_inner_loop`
+    iteration_slice: if not None, iteration is done only over this slice of the field
+    ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                 if None, the number of ghost layers is determined automatically and assumed to be equal for a
+                 all dimensions
 
     :return: :class:`pystencils.ast.KernelFunction` node
     """
     from pystencils.cpu import create_kernel
     code = create_kernel(assignments, function_name, type_info, split_groups, iteration_slice, ghost_layers)
-    code = insert_casts(code)
-    code.compile = partial(make_python_function, code)
+    code.body = insert_casts(code.body)
+    code._compile_function = make_python_function
+    code._backend = 'llvm'
     return code

pystencils/transformations.py

@@ -143,12 +143,11 @@ def get_common_shape(field_set):
     return shape
 
 
-def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layers=None, loop_order=None):
+def make_loop_over_domain(body, iteration_slice=None, ghost_layers=None, loop_order=None):
     """Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.
 
     Args:
         body: Block object with inner loop contents
-        function_name: name of generated C function
         iteration_slice: if not None, iteration is done only over this slice of the field
         ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
              if None, the number of ghost layers is determined automatically and assumed to be equal for a
@@ -156,7 +155,7 @@ def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layer
         loop_order: loop ordering from outer to inner loop (optimal ordering is same as layout)
 
     Returns:
-        :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
+        tuple of loop-node, ghost_layer_info
     """
     # find correct ordering by inspecting participating FieldAccesses
     field_accesses = body.atoms(AbstractField.AbstractAccess)
@@ -199,8 +198,7 @@ def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layer
                                                  sp.sympify(slice_component))
                 current_body.insert_front(assignment)
 
-    ast_node = ast.KernelFunction(current_body, ghost_layers=ghost_layers, function_name=function_name, backend='cpu')
-    return ast_node
+    return current_body, ghost_layers
 
 
 def create_intermediate_base_pointer(field_access, coordinates, previous_ptr):

pystencils_tests/test_jacobi_cbackend.py

@@ -2,7 +2,7 @@ import numpy as np
 from pystencils import show_code
 from pystencils.transformations import move_constants_before_loop, make_loop_over_domain, resolve_field_accesses
 from pystencils.field import Field
-from pystencils.astnodes import SympyAssignment, Block
+from pystencils.astnodes import SympyAssignment, Block, KernelFunction
 from pystencils.cpu import make_python_function
 
 
@@ -19,7 +19,8 @@ def test_jacobi_fixed_field_size():
 
     jacobi = SympyAssignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
     body = Block([jacobi])
-    ast_node = make_loop_over_domain(body, "kernel")
+    loop_node, gl_info = make_loop_over_domain(body)
+    ast_node = KernelFunction(loop_node, 'cpu', 'c', make_python_function, ghost_layers=gl_info)
     resolve_field_accesses(ast_node)
     move_constants_before_loop(ast_node)
 
@@ -28,7 +29,7 @@ def test_jacobi_fixed_field_size():
             dst_field_py[x, y] = 0.25 * (src_field_py[x - 1, y] + src_field_py[x + 1, y] +
                                          src_field_py[x, y - 1] + src_field_py[x, y + 1])
 
-    kernel = make_python_function(ast_node)
+    kernel = ast_node.compile()
     kernel(f=src_field_c, d=dst_field_c)
     error = np.sum(np.abs(dst_field_py - dst_field_c))
     np.testing.assert_allclose(error, 0.0, atol=1e-13)
@@ -44,7 +45,8 @@ def test_jacobi_variable_field_size():
     d = Field.create_generic("d", 3)
     jacobi = SympyAssignment(d[0, 0, 0], (f[1, 0, 0] + f[-1, 0, 0] + f[0, 1, 0] + f[0, -1, 0]) / 4)
     body = Block([jacobi])
-    ast_node = make_loop_over_domain(body, "kernel")
+    loop_node, gl_info = make_loop_over_domain(body)
+    ast_node = KernelFunction(loop_node, 'cpu', 'c', make_python_function, ghost_layers=gl_info)
     resolve_field_accesses(ast_node)
     move_constants_before_loop(ast_node)
 
@@ -59,7 +61,7 @@ def test_jacobi_variable_field_size():
                 dst_field_py[x, y, z] = 0.25 * (src_field_py[x - 1, y, z] + src_field_py[x + 1, y, z] +
                                                 src_field_py[x, y - 1, z] + src_field_py[x, y + 1, z])
 
-    kernel = make_python_function(ast_node)
+    kernel = ast_node.compile()
     kernel(f=src_field_c, d=dst_field_c)
     error = np.sum(np.abs(dst_field_py-dst_field_c))
     np.testing.assert_allclose(error, 0.0, atol=1e-13)