Commit b9e53581 authored by Martin Bauer's avatar Martin Bauer
Browse files

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
parent 690a0ed8
...@@ -162,18 +162,30 @@ class KernelFunction(Node): ...@@ -162,18 +162,30 @@ class KernelFunction(Node):
def field_name(self): def field_name(self):
return self.fields[0].name 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__() super(KernelFunction, self).__init__()
self._body = body self._body = body
body.parent = self body.parent = self
self.function_name = function_name self.function_name = function_name
self._body.parent = self self._body.parent = self
self.compile = None
self.ghost_layers = ghost_layers 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 # these variables are assumed to be global, so no automatic parameter is generated for them
self.global_variables = set() 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 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 @property
def symbols_defined(self): def symbols_defined(self):
...@@ -194,7 +206,7 @@ class KernelFunction(Node): ...@@ -194,7 +206,7 @@ class KernelFunction(Node):
@property @property
def args(self): def args(self):
return [self._body] return self._body,
@property @property
def fields_accessed(self) -> Set['ResolvedFieldAccess']: def fields_accessed(self) -> Set['ResolvedFieldAccess']:
...@@ -231,6 +243,11 @@ class KernelFunction(Node): ...@@ -231,6 +243,11 @@ class KernelFunction(Node):
params = [p.symbol for p in self.get_parameters()] params = [p.symbol for p in self.get_parameters()]
return '{0} {1}({2})'.format(type(self).__name__, self.function_name, params) 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): class SkipIteration(Node):
@property @property
......
import sympy as sp import sympy as sp
from functools import partial
from pystencils.astnodes import SympyAssignment, Block, LoopOverCoordinate, KernelFunction from pystencils.astnodes import SympyAssignment, Block, LoopOverCoordinate, KernelFunction
from pystencils.transformations import resolve_buffer_accesses, resolve_field_accesses, make_loop_over_domain, \ 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, \ 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 ...@@ -61,9 +60,10 @@ def create_kernel(assignments: AssignmentOrAstNodeList, function_name: str = "ke
body = ast.Block(assignments) body = ast.Block(assignments)
loop_order = get_optimal_loop_ordering(fields_without_buffers) loop_order = get_optimal_loop_ordering(fields_without_buffers)
ast_node = make_loop_over_domain(body, function_name, iteration_slice=iteration_slice, loop_node, ghost_layer_info = make_loop_over_domain(body, iteration_slice=iteration_slice,
ghost_layers=ghost_layers, loop_order=loop_order) ghost_layers=ghost_layers, loop_order=loop_order)
ast_node.target = 'cpu' ast_node = KernelFunction(loop_node, 'cpu', 'c', compile_function=make_python_function,
ghost_layers=ghost_layer_info, function_name=function_name)
if split_groups: if split_groups:
typed_split_groups = [[type_symbol(s) for s in split_group] for split_group in 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 ...@@ -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_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) resolve_field_accesses(ast_node, read_only_fields, field_to_base_pointer_info=base_pointer_info)
move_constants_before_loop(ast_node) move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node return ast_node
...@@ -141,14 +140,14 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu ...@@ -141,14 +140,14 @@ def create_indexed_kernel(assignments: AssignmentOrAstNodeList, index_fields, fu
loop_body.append(assignment) loop_body.append(assignment)
function_body = Block([loop_node]) 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} 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]) 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) resolve_field_accesses(ast_node, read_only_fields, field_to_fixed_coordinates=fixed_coordinate_mapping)
move_constants_before_loop(ast_node) move_constants_before_loop(ast_node)
ast_node.compile = partial(make_python_function, ast_node)
return ast_node return ast_node
......
from functools import partial
from pystencils.gpucuda.indexing import BlockIndexing from pystencils.gpucuda.indexing import BlockIndexing
from pystencils.transformations import resolve_field_accesses, add_types, parse_base_pointer_info, \ 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 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 ...@@ -55,7 +53,7 @@ def create_cuda_kernel(assignments, function_name="kernel", type_info=None, inde
block = indexing.guard(block, common_shape) block = indexing.guard(block, common_shape)
unify_shape_symbols(block, common_shape=common_shape, fields=fields_without_buffers) 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) ast.global_variables.update(indexing.index_variables)
base_pointer_spec = [['spatialInner0']] base_pointer_spec = [['spatialInner0']]
...@@ -84,7 +82,6 @@ def create_cuda_kernel(assignments, function_name="kernel", type_info=None, inde ...@@ -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.body.insert_front(SympyAssignment(loop_counter, indexing.coordinates[i]))
ast.indexing = indexing ast.indexing = indexing
ast.compile = partial(make_python_function, ast)
return ast return ast
...@@ -124,7 +121,7 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel ...@@ -124,7 +121,7 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel
function_body = Block(coordinate_symbol_assignments + assignments) function_body = Block(coordinate_symbol_assignments + assignments)
function_body = indexing.guard(function_body, get_common_shape(index_fields)) 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) ast.global_variables.update(indexing.index_variables)
coord_mapping = indexing.coordinates coord_mapping = indexing.coordinates
...@@ -141,5 +138,4 @@ def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel ...@@ -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 # add the function which determines #blocks and #threads as additional member to KernelFunction node
# this is used by the jit # this is used by the jit
ast.indexing = indexing ast.indexing = indexing
ast.compile = partial(make_python_function, ast)
return ast return ast
from pystencils.transformations import insert_casts from pystencils.transformations import insert_casts
from functools import partial
from pystencils.llvm.llvmjit import make_python_function from pystencils.llvm.llvmjit import make_python_function
...@@ -10,16 +9,17 @@ def create_kernel(assignments, function_name="kernel", type_info=None, split_gro ...@@ -10,16 +9,17 @@ def create_kernel(assignments, function_name="kernel", type_info=None, split_gro
Loops are created according to the field accesses in the equations. Loops are created according to the field accesses in the equations.
:param assignments: list of sympy equations, containing accesses to :class:`pystencils.field.Field`. Args:
assignments: list of sympy equations, containing accesses to :class:`pystencils.field.Field`.
Defining the update rules of the kernel Defining the update rules of the kernel
:param function_name: name of the generated function - only important if generated code is written out 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 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 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' . 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 split_groups: Specification on how to split up inner loop into multiple loops. For details see
transformation :func:`pystencils.transformation.split_inner_loop` transformation :func:`pystencils.transformation.split_inner_loop`
:param iteration_slice: if not None, iteration is done only over this slice of the field 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 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 if None, the number of ghost layers is determined automatically and assumed to be equal for a
all dimensions all dimensions
...@@ -27,6 +27,7 @@ def create_kernel(assignments, function_name="kernel", type_info=None, split_gro ...@@ -27,6 +27,7 @@ def create_kernel(assignments, function_name="kernel", type_info=None, split_gro
""" """
from pystencils.cpu import create_kernel from pystencils.cpu import create_kernel
code = create_kernel(assignments, function_name, type_info, split_groups, iteration_slice, ghost_layers) code = create_kernel(assignments, function_name, type_info, split_groups, iteration_slice, ghost_layers)
code = insert_casts(code) code.body = insert_casts(code.body)
code.compile = partial(make_python_function, code) code._compile_function = make_python_function
code._backend = 'llvm'
return code return code
...@@ -143,12 +143,11 @@ def get_common_shape(field_set): ...@@ -143,12 +143,11 @@ def get_common_shape(field_set):
return shape 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. """Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.
Args: Args:
body: Block object with inner loop contents 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 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 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 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 ...@@ -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) loop_order: loop ordering from outer to inner loop (optimal ordering is same as layout)
Returns: 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 # find correct ordering by inspecting participating FieldAccesses
field_accesses = body.atoms(AbstractField.AbstractAccess) field_accesses = body.atoms(AbstractField.AbstractAccess)
...@@ -199,8 +198,7 @@ def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layer ...@@ -199,8 +198,7 @@ def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layer
sp.sympify(slice_component)) sp.sympify(slice_component))
current_body.insert_front(assignment) current_body.insert_front(assignment)
ast_node = ast.KernelFunction(current_body, ghost_layers=ghost_layers, function_name=function_name, backend='cpu') return current_body, ghost_layers
return ast_node
def create_intermediate_base_pointer(field_access, coordinates, previous_ptr): def create_intermediate_base_pointer(field_access, coordinates, previous_ptr):
......
...@@ -2,7 +2,7 @@ import numpy as np ...@@ -2,7 +2,7 @@ import numpy as np
from pystencils import show_code from pystencils import show_code
from pystencils.transformations import move_constants_before_loop, make_loop_over_domain, resolve_field_accesses from pystencils.transformations import move_constants_before_loop, make_loop_over_domain, resolve_field_accesses
from pystencils.field import Field 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 from pystencils.cpu import make_python_function
...@@ -19,7 +19,8 @@ def test_jacobi_fixed_field_size(): ...@@ -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) jacobi = SympyAssignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
body = Block([jacobi]) 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) resolve_field_accesses(ast_node)
move_constants_before_loop(ast_node) move_constants_before_loop(ast_node)
...@@ -28,7 +29,7 @@ def test_jacobi_fixed_field_size(): ...@@ -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] + 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]) 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) kernel(f=src_field_c, d=dst_field_c)
error = np.sum(np.abs(dst_field_py - dst_field_c)) error = np.sum(np.abs(dst_field_py - dst_field_c))
np.testing.assert_allclose(error, 0.0, atol=1e-13) np.testing.assert_allclose(error, 0.0, atol=1e-13)
...@@ -44,7 +45,8 @@ def test_jacobi_variable_field_size(): ...@@ -44,7 +45,8 @@ def test_jacobi_variable_field_size():
d = Field.create_generic("d", 3) 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) 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]) 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) resolve_field_accesses(ast_node)
move_constants_before_loop(ast_node) move_constants_before_loop(ast_node)
...@@ -59,7 +61,7 @@ def test_jacobi_variable_field_size(): ...@@ -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] + 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]) 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) kernel(f=src_field_c, d=dst_field_c)
error = np.sum(np.abs(dst_field_py-dst_field_c)) error = np.sum(np.abs(dst_field_py-dst_field_c))
np.testing.assert_allclose(error, 0.0, atol=1e-13) np.testing.assert_allclose(error, 0.0, atol=1e-13)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment