def oir_iteration_space_computation(stencil: oir.Stencil) -> Dict[int, CartesianIterationSpace]:
iteration_spaces = dict()
offsets: Dict[str, List[CartesianOffset]] = dict()
outputs = set()
access_spaces: Dict[str, CartesianIterationSpace] = dict()
# reversed pre_order traversal is post-order but the last nodes per node come first.
# this is not possible with visitors unless a (reverseed) visit is implemented for every node
for node in reversed(list(iter_tree(stencil, traversal_order=TraversalOrder.PRE_ORDER))):
if isinstance(node, oir.FieldAccess):
if node.name not in offsets:
offsets[node.name] = list()
offsets[node.name].append(node.offset)
elif isinstance(node, oir.AssignStmt):
if node.left.kind == ExprKind.FIELD:
outputs.add(node.left.name)
elif isinstance(node, oir.HorizontalExecution):
iteration_spaces[id(node)] = CartesianIterationSpace.domain()
for name in outputs:
access_space = access_spaces.get(name, CartesianIterationSpace.domain())
iteration_spaces[id(node)] = iteration_spaces[id(node)] | access_space
for name in offsets:
access_space = CartesianIterationSpace.domain()
for offset in offsets[name]:
access_space = access_space | CartesianIterationSpace.from_offset(offset)
accumulated_extent = iteration_spaces[id(node)].compose(access_space)
access_spaces[name] = (
access_spaces.get(name, CartesianIterationSpace.domain()) | accumulated_extent
)
offsets = dict()
outputs = set()
return iteration_spaces
def oir_field_boundary_computation(stencil: oir.Stencil) -> Dict[str, CartesianIterationSpace]:
offsets: Dict[str, List[CartesianOffset]] = dict()
access_spaces: Dict[str, CartesianIterationSpace] = dict()
iteration_spaces = oir_iteration_space_computation(stencil)
for node in iter_tree(stencil, traversal_order=TraversalOrder.POST_ORDER):
if isinstance(node, oir.FieldAccess):
if node.name not in offsets:
offsets[node.name] = list()
offsets[node.name].append(node.offset)
elif isinstance(node, oir.HorizontalExecution):
if iteration_spaces.get(id(node), None) is not None:
for name in offsets:
access_space = CartesianIterationSpace.domain()
for offset in offsets[name]:
access_space = access_space | CartesianIterationSpace.from_offset(offset)
access_spaces[name] = access_spaces.get(
name, CartesianIterationSpace.domain()
) | (iteration_spaces[id(node)].compose(access_space))
offsets = dict()
return access_spaces