Commit 23a8f33d authored by Michael Kuron's avatar Michael Kuron
Browse files

create_staggered_kernel: re-add gpu_exclusive_conditions

parent 68da7457
...@@ -191,7 +191,7 @@ def create_indexed_kernel(assignments, ...@@ -191,7 +191,7 @@ def create_indexed_kernel(assignments,
raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,)) raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,))
def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwargs): def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=False, **kwargs):
"""Kernel that updates a staggered field. """Kernel that updates a staggered field.
.. image:: /img/staggered_grid.svg .. image:: /img/staggered_grid.svg
...@@ -205,7 +205,9 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg ...@@ -205,7 +205,9 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg
regular fields are passed through to `create_kernel`. Multiple different staggered fields can be regular fields are passed through to `create_kernel`. Multiple different staggered fields can be
used, but they all need to use the same stencil (i.e. the same number of staggered points) and used, but they all need to use the same stencil (i.e. the same number of staggered points) and
shape. shape.
gpu_exclusive_conditions: whether to use nested conditionals instead of multiple conditionals target: 'cpu', 'llvm' or 'gpu'
gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then
handled in an else branch.
kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed
Returns: Returns:
...@@ -277,7 +279,28 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg ...@@ -277,7 +279,28 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg
return sp.And(*conditions) return sp.And(*conditions)
if gpu_exclusive_conditions: if gpu_exclusive_conditions:
raise NotImplementedError('gpu_exclusive_conditions is not implemented yet') outer_assignment = None
for assignment in assignments:
direction = assignment.lhs.field.staggered_stencil[assignment.lhs.index[0]]
assignment = SympyAssignment(assignment.lhs, assignment.rhs)
outer_assignment = Conditional(condition(direction), Block([assignment]), outer_assignment)
inner_assignment = []
for assignment in assignments:
direction = assignment.lhs.field.staggered_stencil[assignment.lhs.index[0]]
inner_assignment.append(SympyAssignment(assignment.lhs, assignment.rhs))
last_conditional = Conditional(sp.And(*[condition(d) for d in stencil]),
Block(inner_assignment), outer_assignment)
final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
[SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
[last_conditional]
if target == 'cpu':
from pystencils.cpu import create_kernel as create_kernel_cpu
ast = create_kernel_cpu(final_assignments, ghost_layers=ghost_layers, **kwargs)
else:
ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, **kwargs)
return ast
for assignment in assignments: for assignment in assignments:
direction = assignment.lhs.field.staggered_stencil[assignment.lhs.index[0]] direction = assignment.lhs.field.staggered_stencil[assignment.lhs.index[0]]
...@@ -290,6 +313,6 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg ...@@ -290,6 +313,6 @@ def create_staggered_kernel(assignments, gpu_exclusive_conditions=False, **kwarg
remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers]) remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers])
prepend_optimizations = [lambda ast: remove_conditionals_in_staggered_kernel(ast, remove_start_conditional), prepend_optimizations = [lambda ast: remove_conditionals_in_staggered_kernel(ast, remove_start_conditional),
move_constants_before_loop] move_constants_before_loop]
ast = create_kernel(final_assignments, ghost_layers=ghost_layers, cpu_prepend_optimizations=prepend_optimizations, ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target,
**kwargs) cpu_prepend_optimizations=prepend_optimizations, **kwargs)
return ast return ast
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