Skip to content
Snippets Groups Projects

Reimplement create_staggered_kernel

Merged Michael Kuron requested to merge staggered_kernel into master
1 file
+ 28
5
Compare changes
  • Side-by-side
  • Inline
@@ -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
Loading