def test_jacobi_variable_field_size():
size = (3, 3, 3)
f = Field.create_generic("f", 3)
d = Field.create_generic("d", 3)
jacobi = Assignment(
d[0, 0, 0], (f[1, 0, 0] + f[-1, 0, 0] + f[0, 1, 0] + f[0, -1, 0]) / 4)
ast = create_kernel([jacobi])
src_field_llvm = np.random.rand(*size)
src_field_py = np.copy(src_field_llvm)
dst_field_llvm = 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 = make_python_function(ast, {
'f': src_field_llvm,
'd': dst_field_llvm
})
kernel()
error = np.sum(np.abs(dst_field_py - dst_field_llvm))
np.testing.assert_almost_equal(error, 0.0)
def test_variable_sized_fields():
src_field = Field.create_generic('src', spatial_dimensions=2)
dst_field = Field.create_generic('dst', spatial_dimensions=2)
update_rule = Assignment(dst_field[0, 0],
(src_field[0, 1] + src_field[0, -1] +
src_field[1, 0] + src_field[-1, 0]) / 4)
ast = create_cuda_kernel(sympy_cse_on_assignment_list([update_rule]))
kernel = make_python_function(ast)
size = (3, 3)
src_arr = np.random.rand(*size)
src_arr = add_ghost_layers(src_arr)
dst_arr = np.zeros_like(src_arr)
gpu_src_arr = gpuarray.to_gpu(src_arr)
gpu_dst_arr = gpuarray.to_gpu(dst_arr)
kernel(src=gpu_src_arr, dst=gpu_dst_arr)
gpu_dst_arr.get(dst_arr)
stencil = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]) / 4.0
reference = convolve(remove_ghost_layers(src_arr),
stencil,
mode='constant',
cval=0.0)
reference = add_ghost_layers(reference)
np.testing.assert_almost_equal(reference, dst_arr)
def __init__(self, boundary, method, pdf_field_sparse):
full_pdf_field = Field.create_generic('pdfFull', spatial_dimensions=method.dim, index_dimensions=1)
additional_data_field = Field.create_generic('additionalData', spatial_dimensions=1,
dtype=boundary.additional_data)
boundary_eqs = boundary(full_pdf_field, self.DIR_SYMBOL, method, additional_data_field)
neighbor_offsets = {fa.offsets for eq in boundary_eqs for fa in eq.atoms(Field.Access)}
neighbor_offsets = list(neighbor_offsets)
neighbor_offsets_dtype = [(self.NEIGHBOR_IDX_NAME.format(i), np.uint32)
for i in range(len(neighbor_offsets))]
index_field_dtype = np.dtype([('dir', np.uint32),
*neighbor_offsets_dtype,
*boundary.additional_data])
index_field = Field.create_generic('indexField', spatial_dimensions=1, dtype=index_field_dtype)
boundary_eqs = boundary(full_pdf_field, self.DIR_SYMBOL, method, index_field)
offset_subs = {off: sp.Symbol(self.NEIGHBOR_IDX_NAME.format(i)) for i, off in enumerate(neighbor_offsets)}
new_boundary_eqs = []
for eq in boundary_eqs:
substitutions = {
fa: pdf_field_sparse.absolute_access([index_field(offset_subs[fa.offsets].name)], fa.index)
for fa in eq.atoms(Field.Access)
if fa.field == full_pdf_field
}
new_boundary_eqs.append(eq.subs(substitutions))
self.boundary_eqs = new_boundary_eqs
self.boundary_eqs_orig = boundary_eqs
self.method = method
self.index_field_dtype = index_field_dtype
self.neighbor_offsets = neighbor_offsets
self.index_field = index_field
def test_full_scalar_field():
"""Tests fully (un)packing a scalar field (from)to a buffer."""
fields = _generate_fields()
for (src_arr, dst_arr, buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field", src_arr)
dst_field = Field.create_from_numpy_array("dst_field", dst_arr)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = [Assignment(buffer.center(), src_field.center())]
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=buffer_arr, src_field=src_arr)
unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(dst_field=dst_arr, buffer=buffer_arr)
np.testing.assert_equal(src_arr, dst_arr)
def create_copy_kernel(domain_size,
from_slice,
to_slice,
index_dimensions=0,
index_dim_shape=1,
dtype=np.float64):
"""Copies a rectangular part of an array to another non-overlapping part"""
f = Field.create_generic("pdfs",
len(domain_size),
index_dimensions=index_dimensions,
dtype=dtype)
normalized_from_slice = normalize_slice(from_slice, f.spatial_shape)
normalized_to_slice = normalize_slice(to_slice, f.spatial_shape)
offset = [
s1.start - s2.start
for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
]
assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
"Slices have to have same size"
update_eqs = []
if index_dimensions < 2:
index_dim_shape = [index_dim_shape]
for i in product(*[range(d) for d in index_dim_shape]):
eq = Assignment(f(*i), f[tuple(offset)](*i))
update_eqs.append(eq)
ast = create_cuda_kernel(update_eqs,
iteration_slice=to_slice,
skip_independence_check=True)
return ast
def create_copy_kernel(domain_size,
from_slice,
to_slice,
index_dimensions=0,
index_dim_shape=1,
dtype=np.float64):
"""Copies a rectangular part of an array to another non-overlapping part"""
if index_dimensions not in (0, 1):
raise NotImplementedError(
"Works only for one or zero index coordinates")
f = Field.create_generic("pdfs",
len(domain_size),
index_dimensions=index_dimensions,
dtype=dtype)
normalized_from_slice = normalize_slice(from_slice, f.spatial_shape)
normalized_to_slice = normalize_slice(to_slice, f.spatial_shape)
offset = [
s1.start - s2.start
for s1, s2 in zip(normalized_from_slice, normalized_to_slice)
]
assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
"Slices have to have same size"
update_eqs = []
for i in range(index_dim_shape):
eq = Assignment(f(i), f[tuple(offset)](i))
update_eqs.append(eq)
ast = create_cuda_kernel(update_eqs,
iteration_slice=to_slice,
skip_independence_check=True)
return make_python_function(ast)
def test_full_scalar_field():
"""Tests fully (un)packing a scalar field (from)to a GPU buffer."""
fields = _generate_fields()
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field", src_arr)
dst_field = Field.create_from_numpy_array("dst_field", src_arr)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = [Assignment(buffer.center(), src_field.center())]
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs, type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(dst_field=gpu_dst_arr, buffer=gpu_buffer_arr)
dst_arr = gpu_dst_arr.get()
np.testing.assert_equal(src_arr, dst_arr)
def test_two_variable_shaped_fields():
src = np.zeros((20, 21, 9))
dst = np.zeros((22, 21, 9))
sym_src = Field.create_generic("src",
spatial_dimensions=2,
index_dimensions=1)
sym_dst = Field.create_generic("dst",
spatial_dimensions=2,
index_dimensions=1)
update_rule = Assignment(sym_dst(0), sym_src[-1, 1](1) + sym_src[1, -1](2))
ast = create_kernel([update_rule])
func = ast.compile()
with pytest.raises(TypeError) as e:
func(src=src, dst=dst)
assert 'must have same' in str(e.value)
def test_two_variable_shaped_fields():
src = np.zeros((20, 21, 9))
dst = np.zeros((22, 21, 9))
sym_src = Field.create_generic("src",
spatial_dimensions=2,
index_dimensions=1)
sym_dst = Field.create_generic("dst",
spatial_dimensions=2,
index_dimensions=1)
update_rule = Assignment(sym_dst(0), sym_src[-1, 1](1) + sym_src[1, -1](2))
ast = create_kernel([update_rule])
func = make_python_function(ast)
try:
func(src=src, dst=dst)
assert False, "Expected ValueError because fields with different sized where passed"
except ValueError:
pass
def test_variable_sized_field():
for order in ('f', 'c'):
for align in (True, False):
dt = np.dtype([('e1', np.float32), ('e2', np.double),
('e3', np.double)],
align=align)
f = Field.create_generic("f", 2, dt, layout=order)
d = Field.create_generic("d", 2, dt, layout=order)
update_rules = [Assignment(d[0, 0]['e2'], f[0, 0]['e3'])]
arr = np.zeros((3, 2), dtype=dt, order=order)
result = arr.copy(order=order)
arr['e2'] = 0
arr['e3'] = np.random.rand(3, 2)
kernel = create_kernel(update_rules).compile()
kernel(f=arr, d=result)
np.testing.assert_almost_equal(result['e2'], arr['e3'])
np.testing.assert_equal(arr['e2'], np.zeros((3, 2)))
def test_subset_cell_values():
"""Tests (un)packing a subset of cell values of the a field (from)to a buffer."""
num_cell_values = 7
# Cell indices of the field to be (un)packed (from)to the buffer
cell_indices = [1, 3, 5, 6]
fields = _generate_fields(stencil_directions=num_cell_values)
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field",
gpu_src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
gpu_src_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=gpu_src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(buffer(buffer_idx), src_field(cell_idx))
pack_eqs.append(eq)
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = []
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(dst_field(cell_idx), buffer(buffer_idx))
unpack_eqs.append(eq)
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs, type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
dst_arr = gpu_dst_arr.get()
mask_arr = np.ma.masked_where((src_arr - dst_arr) != 0, src_arr)
np.testing.assert_equal(dst_arr, mask_arr.filled(int(0)))
def test_setting_value():
arr_cpu = np.arange(25, dtype=np.float64).reshape(5, 5)
arr_gpu = gpuarray.to_gpu(arr_cpu)
iteration_slice = make_slice[:, :]
f = Field.create_generic("f", 2)
update_rule = [Assignment(f(0), sp.Symbol("value"))]
ast = create_cuda_kernel(update_rule,
iteration_slice=iteration_slice,
indexing_creator=LineIndexing)
kernel = make_python_function(ast)
kernel(f=arr_gpu, value=np.float64(42.0))
np.testing.assert_equal(arr_gpu.get(), np.ones((5, 5)) * 42.0)
def test_fixed_and_variable_field_check():
"""Create kernel with two variable sized fields - calling them with different sizes"""
src = np.zeros((20, 21, 9))
sym_src = Field.create_from_numpy_array("src", src, index_dimensions=1)
sym_dst = Field.create_generic("dst",
spatial_dimensions=2,
index_dimensions=1)
update_rule = Assignment(sym_dst(0), sym_src[-1, 1](1) + sym_src[1, -1](2))
with pytest.raises(ValueError) as e:
create_kernel(update_rule)
assert 'Mixing fixed-shaped and variable-shape fields' in str(e.value)
def test_subset_cell_values():
"""Tests (un)packing a subset of cell values of the a field (from)to a buffer."""
num_cell_values = 19
# Cell indices of the field to be (un)packed (from)to the buffer
cell_indices = [1, 5, 7, 8, 10, 12, 13]
fields = _generate_fields(num_directions=num_cell_values)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(buffer(buffer_idx), src_field(cell_idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for buffer_idx, cell_idx in enumerate(cell_indices):
eq = Assignment(dst_field(cell_idx), buffer(buffer_idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
mask_arr = np.ma.masked_where((src_arr - dst_arr) != 0, src_arr)
np.testing.assert_equal(dst_arr, mask_arr.filled(int(0)))
def test_all_cell_values():
"""Tests (un)packing all cell values of the a field (from)to a buffer."""
num_cell_values = 7
fields = _generate_fields(stencil_directions=num_cell_values)
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field",
gpu_src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
gpu_src_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=gpu_src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs, type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
dst_arr = gpu_dst_arr.get()
np.testing.assert_equal(src_arr, dst_arr)
def test_fixed_and_variable_field_check():
"""Create kernel with two variable sized fields - calling them with different sizes"""
src = np.zeros((20, 21, 9))
sym_src = Field.create_from_numpy_array("src", src, index_dimensions=1)
sym_dst = Field.create_generic("dst",
spatial_dimensions=2,
index_dimensions=1)
update_rule = Assignment(sym_dst(0), sym_src[-1, 1](1) + sym_src[1, -1](2))
try:
create_kernel([update_rule])
assert False, "Expected ValueError because fields with different sized where passed"
except ValueError:
pass
def test_all_cell_values():
"""Tests (un)packing all cell values of the a field (from)to a buffer."""
num_cell_values = 19
fields = _generate_fields(num_directions=num_cell_values)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
np.testing.assert_equal(src_arr, dst_arr)
def test_field_layouts():
num_cell_values = 7
for layout_str in ['numpy', 'fzyx', 'zyxf', 'reverse_numpy']:
fields = _generate_fields(stencil_directions=num_cell_values,
layout=layout_str)
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
src_field = Field.create_from_numpy_array("src_field",
gpu_src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
gpu_src_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr, src_field=gpu_src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs,
type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr, dst_field=gpu_dst_arr)
def test_field_layouts():
num_cell_values = 27
for layout_str in ['numpy', 'fzyx', 'zyxf', 'reverse_numpy']:
fields = _generate_fields(num_directions=num_cell_values,
layout=layout_str)
for (src_arr, dst_arr, bufferArr) in fields:
src_field = Field.create_from_numpy_array("src_field",
src_arr,
index_dimensions=1)
dst_field = Field.create_from_numpy_array("dst_field",
dst_arr,
index_dimensions=1)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
pack_code = create_kernel(pack_eqs,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
def test_field_slice():
"""Tests (un)packing slices of a scalar field (from)to a buffer."""
fields = _generate_fields()
for d in ['N', 'S', 'NW', 'SW', 'TNW', 'B']:
for (src_arr, gpu_src_arr, gpu_dst_arr, gpu_buffer_arr) in fields:
# Extract slice from N direction of the field
slice_dir = direction_string_to_offset(d, dim=len(src_arr.shape))
pack_slice = get_slice_before_ghost_layer(slice_dir)
unpack_slice = get_ghost_region_slice(slice_dir)
src_field = Field.create_from_numpy_array("src_field",
src_arr[pack_slice])
dst_field = Field.create_from_numpy_array("dst_field",
src_arr[unpack_slice])
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = [Assignment(buffer.center(), src_field.center())]
pack_types = {
'src_field': gpu_src_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
pack_code = create_cuda_kernel(pack_eqs, type_info=pack_types)
pack_kernel = make_python_function(pack_code)
pack_kernel(buffer=gpu_buffer_arr,
src_field=gpu_src_arr[pack_slice])
# Unpack into ghost layer of dst_field in N direction
unpack_eqs = [Assignment(dst_field.center(), buffer.center())]
unpack_types = {
'dst_field': gpu_dst_arr.dtype,
'buffer': gpu_buffer_arr.dtype
}
unpack_code = create_cuda_kernel(unpack_eqs,
type_info=unpack_types)
unpack_kernel = make_python_function(unpack_code)
unpack_kernel(buffer=gpu_buffer_arr,
dst_field=gpu_dst_arr[unpack_slice])
dst_arr = gpu_dst_arr.get()
np.testing.assert_equal(src_arr[pack_slice], dst_arr[unpack_slice])
def _add_inplace_boundary(self, boundary_obj, flag=None):
if boundary_obj not in self._boundary_object_to_boundary_info:
sym_index_field = Field.create_generic(
'indexField',
spatial_dimensions=1,
dtype=numpy_data_type_for_boundary_object(
boundary_obj, self.dim))
ast_even = self._create_boundary_kernel(
self._data_handling.fields[self._field_name], sym_index_field,
boundary_obj, Timestep.EVEN)
ast_odd = self._create_boundary_kernel(
self._data_handling.fields[self._field_name], sym_index_field,
boundary_obj, Timestep.ODD)
kernels = [ast_even.compile(), ast_odd.compile()]
if flag is None:
flag = self.flag_interface.reserve_next_flag()
boundary_info = self.InplaceStreamingBoundaryInfo(
self, boundary_obj, flag, kernels)
self._boundary_object_to_boundary_info[
boundary_obj] = boundary_info
return self._boundary_object_to_boundary_info[boundary_obj].flag
def visualize_pdf_field_accessor(pdf_field_accessor,
title=True,
read_plot_params=None,
write_plot_params=None,
figure=None):
if write_plot_params is None:
write_plot_params = {}
if read_plot_params is None:
read_plot_params = {}
if figure is None:
import matplotlib.pyplot as plt
figure = plt.gcf()
stencil = LBStencil(Stencil.D2Q9)
figure.patch.set_facecolor('white')
field = Field.create_generic('f', spatial_dimensions=2, index_dimensions=1)
pre_collision_accesses = pdf_field_accessor.read(field, stencil)
post_collision_accesses = pdf_field_accessor.write(field, stencil)
ax_left = figure.add_subplot(1, 2, 1)
ax_right = figure.add_subplot(1, 2, 2)
if 'color' not in read_plot_params:
read_plot_params['color'] = 'k'
if 'color' not in write_plot_params:
write_plot_params['color'] = 'r'
visualize_field_mapping(ax_left, stencil, pre_collision_accesses,
**read_plot_params)
visualize_field_mapping(ax_right, stencil, post_collision_accesses,
**write_plot_params)
if title:
ax_left.set_title("Read")
ax_right.set_title("Write")
def generate_boundary(generation_context, class_name, boundary_object,
lb_method, **create_kernel_params):
struct_name = "IndexInfo"
boundary_object.name = class_name
create_kernel_params = default_create_kernel_parameters(
generation_context, create_kernel_params)
target = create_kernel_params['target']
index_struct_dtype = numpy_data_type_for_boundary_object(
boundary_object, lb_method.dim)
pdf_field = Field.create_generic(
'pdfs',
lb_method.dim,
np.float64 if generation_context.double_accuracy else np.float32,
index_dimensions=1,
layout='fzyx',
index_shape=[len(lb_method.stencil)])
index_field = Field('indexVector',
FieldType.INDEXED,
index_struct_dtype,
layout=[0],
shape=(TypedSymbol("indexVectorSize",
create_type(np.int64)), 1),
strides=(1, 1))
kernel = create_lattice_boltzmann_boundary_kernel(
pdf_field,
index_field,
lb_method,
boundary_object,
target=target,
openmp=generation_context.openmp)
kernel.function_name = "boundary_" + boundary_object.name
# waLBerla is a 3D framework. Therefore, a zero for the z index has to be added if we work in 2D
if lb_method.dim == 2:
stencil = ()
for d in lb_method.stencil:
d = d + (0, )
stencil = stencil + (d, )
else:
stencil = lb_method.stencil
stencil_info = [(i, ", ".join([str(e) for e in d]))
for i, d in enumerate(stencil)]
context = {
'class_name':
boundary_object.name,
'StructName':
struct_name,
'StructDeclaration':
struct_from_numpy_dtype(struct_name, index_struct_dtype),
'kernel':
KernelInfo(kernel),
'stencil_info':
stencil_info,
'dim':
lb_method.dim,
'target':
target,
'namespace':
'lbm',
}
env = Environment(loader=PackageLoader('lbmpy_walberla'),
undefined=StrictUndefined)
add_pystencils_filters_to_jinja_env(env)
header = env.get_template('Boundary.tmpl.h').render(**context)
source = env.get_template('Boundary.tmpl.cpp').render(**context)
source_extension = "cpp" if create_kernel_params.get(
"target", "cpu") == "cpu" else "cu"
generation_context.write_file("{}.h".format(class_name), header)
generation_context.write_file("{}.{}".format(class_name, source_extension),
source)
def generate_pack_info(generation_context,
class_name: str,
directions_to_pack_terms: Dict[Tuple[Tuple],
Sequence[Field.Access]],
namespace='pystencils',
**create_kernel_params):
"""Generates a waLBerla GPU PackInfo
Args:
generation_context: see documentation of `generate_sweep`
class_name: name of the generated class
directions_to_pack_terms: maps tuples of directions to read field accesses, specifying which values have to be
packed for which direction
namespace: inner namespace of the generated class
**create_kernel_params: remaining keyword arguments are passed to `pystencils.create_kernel`
"""
items = [(e[0], sorted(e[1], key=lambda x: str(x)))
for e in directions_to_pack_terms.items()]
items = sorted(items, key=lambda e: e[0])
directions_to_pack_terms = OrderedDict(items)
create_kernel_params = default_create_kernel_parameters(
generation_context, create_kernel_params)
target = create_kernel_params.get('target', 'cpu')
template_name = "CpuPackInfo.tmpl" if target == 'cpu' else 'GpuPackInfo.tmpl'
fields_accessed = set()
for terms in directions_to_pack_terms.values():
for term in terms:
assert isinstance(
term, Field.Access) # and all(e == 0 for e in term.offsets)
fields_accessed.add(term)
field_names = {fa.field.name for fa in fields_accessed}
data_types = {fa.field.dtype for fa in fields_accessed}
if len(data_types) == 0:
raise ValueError("No fields to pack!")
if len(data_types) != 1:
err_detail = "\n".join(" - {} [{}]".format(f.name, f.dtype)
for f in fields_accessed)
raise NotImplementedError(
"Fields of different data types are used - this is not supported.\n"
+ err_detail)
dtype = data_types.pop()
pack_kernels = OrderedDict()
unpack_kernels = OrderedDict()
all_accesses = set()
elements_per_cell = OrderedDict()
for direction_set, terms in directions_to_pack_terms.items():
for d in direction_set:
if not all(abs(i) <= 1 for i in d):
raise NotImplementedError("Only first neighborhood supported")
buffer = Field.create_generic('buffer',
spatial_dimensions=1,
field_type=FieldType.BUFFER,
dtype=dtype.numpy_dtype,
index_shape=(len(terms), ))
direction_strings = tuple(
offset_to_direction_string(d) for d in direction_set)
all_accesses.update(terms)
pack_assignments = [
Assignment(buffer(i), term) for i, term in enumerate(terms)
]
pack_ast = create_kernel(pack_assignments,
**create_kernel_params,
ghost_layers=0)
pack_ast.function_name = 'pack_{}'.format("_".join(direction_strings))
unpack_assignments = [
Assignment(term, buffer(i)) for i, term in enumerate(terms)
]
unpack_ast = create_kernel(unpack_assignments,
**create_kernel_params,
ghost_layers=0)
unpack_ast.function_name = 'unpack_{}'.format(
"_".join(direction_strings))
pack_kernels[direction_strings] = KernelInfo(pack_ast)
unpack_kernels[direction_strings] = KernelInfo(unpack_ast)
elements_per_cell[direction_strings] = len(terms)
fused_kernel = create_kernel(
[Assignment(buffer.center, t) for t in all_accesses],
**create_kernel_params)
jinja_context = {
'class_name': class_name,
'pack_kernels': pack_kernels,
'unpack_kernels': unpack_kernels,
'fused_kernel': KernelInfo(fused_kernel),
'elements_per_cell': elements_per_cell,
'headers': get_headers(fused_kernel),
'target': target,
'dtype': dtype,
'field_name': field_names.pop(),
'namespace': namespace,
}
env = Environment(loader=PackageLoader('pystencils_walberla'),
undefined=StrictUndefined)
add_pystencils_filters_to_jinja_env(env)
header = env.get_template(template_name + ".h").render(**jinja_context)
source = env.get_template(template_name + ".cpp").render(**jinja_context)
source_extension = "cpp" if target == "cpu" else "cu"
generation_context.write_file("{}.h".format(class_name), header)
generation_context.write_file("{}.{}".format(class_name, source_extension),
source)
def test_iteration_slices():
num_cell_values = 19
dt = np.uint64
fields = _generate_fields(dt=dt, num_directions=num_cell_values)
for (src_arr, dst_arr, bufferArr) in fields:
spatial_dimensions = len(src_arr.shape) - 1
# src_field = Field.create_from_numpy_array("src_field", src_arr, index_dimensions=1)
# dst_field = Field.create_from_numpy_array("dst_field", dst_arr, index_dimensions=1)
src_field = Field.create_generic("src_field",
spatial_dimensions,
index_shape=(num_cell_values, ),
dtype=dt)
dst_field = Field.create_generic("dst_field",
spatial_dimensions,
index_shape=(num_cell_values, ),
dtype=dt)
buffer = Field.create_generic("buffer",
spatial_dimensions=1,
index_dimensions=1,
field_type=FieldType.BUFFER,
dtype=src_arr.dtype)
pack_eqs = []
# Since we are packing all cell values for all cells, then
# the buffer index is equivalent to the field index
for idx in range(num_cell_values):
eq = Assignment(buffer(idx), src_field(idx))
pack_eqs.append(eq)
dim = src_field.spatial_dimensions
# Pack only the leftmost slice, only every second cell
pack_slice = (slice(None, None, 2), ) * (dim - 1) + (0, )
# Fill the entire array with data
src_arr[(slice(None, None, 1), ) * dim] = np.arange(num_cell_values)
dst_arr.fill(0)
pack_code = create_kernel(pack_eqs,
iteration_slice=pack_slice,
data_type={
'src_field': src_arr.dtype,
'buffer': buffer.dtype
})
pack_kernel = pack_code.compile()
pack_kernel(buffer=bufferArr, src_field=src_arr)
unpack_eqs = []
for idx in range(num_cell_values):
eq = Assignment(dst_field(idx), buffer(idx))
unpack_eqs.append(eq)
unpack_code = create_kernel(unpack_eqs,
iteration_slice=pack_slice,
data_type={
'dst_field': dst_arr.dtype,
'buffer': buffer.dtype
})
unpack_kernel = unpack_code.compile()
unpack_kernel(buffer=bufferArr, dst_field=dst_arr)
# Check if only every second entry of the leftmost slice has been copied
np.testing.assert_equal(dst_arr[pack_slice], src_arr[pack_slice])
np.testing.assert_equal(
dst_arr[(slice(1, None, 2), ) * (dim - 1) + (0, )], 0)
np.testing.assert_equal(
dst_arr[(slice(None, None, 1), ) * (dim - 1) + (slice(1, None), )],
0)
