def test_jacobi_variable_field_size():
size = (3, 3, 3)
f = Field.create_generic("f", 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)
body = Block([jacobi])
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)
src_field_c = np.random.rand(*size)
src_field_py = np.copy(src_field_c)
dst_field_c = np.zeros(size)
dst_field_py = np.zeros(size)
for x in range(1, size[0] - 1):
for y in range(1, size[1] - 1):
for z in range(1, size[2] - 1):
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 = 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)
def test_jacobi_fixed_field_size():
size = (30, 20)
src_field_c = np.random.rand(*size)
src_field_py = np.copy(src_field_c)
dst_field_c = np.zeros(size)
dst_field_py = np.zeros(size)
f = Field.create_from_numpy_array("f", src_field_c)
d = Field.create_from_numpy_array("d", dst_field_c)
jacobi = SympyAssignment(d[0, 0],
(f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
body = Block([jacobi])
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)
for x in range(1, size[0] - 1):
for y in range(1, size[1] - 1):
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 = 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)
code_display = show_code(ast_node)
assert 'for' in str(code_display)
assert 'for' in code_display._repr_html_()
def create_kernel(assignments: AssignmentOrAstNodeList,
function_name: str = "kernel",
type_info='double',
split_groups=(),
iteration_slice=None,
ghost_layers=None,
skip_independence_check=False) -> KernelFunction:
"""Creates an abstract syntax tree for a kernel function, by taking a list of update rules.
Loops are created according to the field accesses in the equations.
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
skip_independence_check: don't check that loop iterations are independent. This is needed e.g. for
periodicity kernel, that access the field outside the iteration bounds. Use with care!
Returns:
AST node representing a function, that can be printed as C or CUDA code
"""
def type_symbol(term):
if isinstance(term, Field.Access) or isinstance(term, TypedSymbol):
return term
elif isinstance(term, sp.Symbol):
if not hasattr(type_info, '__getitem__'):
return TypedSymbol(term.name, create_type(type_info))
else:
return TypedSymbol(term.name, type_info[term.name])
else:
raise ValueError("Term has to be field access or symbol")
fields_read, fields_written, assignments = add_types(
assignments, type_info, not skip_independence_check)
all_fields = fields_read.union(fields_written)
read_only_fields = set([f.name for f in fields_read - fields_written])
buffers = set([f for f in all_fields if FieldType.is_buffer(f)])
fields_without_buffers = all_fields - buffers
body = ast.Block(assignments)
loop_order = get_optimal_loop_ordering(fields_without_buffers)
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]
split_inner_loop(ast_node, typed_split_groups)
base_pointer_spec = [['spatialInner0'], ['spatialInner1']
] if len(loop_order) >= 2 else [['spatialInner0']]
base_pointer_info = {
field.name: parse_base_pointer_info(base_pointer_spec, loop_order,
field.spatial_dimensions,
field.index_dimensions)
for field in fields_without_buffers
}
buffer_base_pointer_info = {
field.name: parse_base_pointer_info([['spatialInner0']], [0],
field.spatial_dimensions,
field.index_dimensions)
for field in buffers
}
base_pointer_info.update(buffer_base_pointer_info)
if any(FieldType.is_buffer(f) for f in all_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)
move_constants_before_loop(ast_node)
return ast_node