def visit_HorizontalExecution(self, node: oir.HorizontalExecution, *,
ctx: Context) -> None:
results = AccessCollector.apply(node)
horizontal_extent = functools.reduce(
lambda ext, name: ext | ctx.fields.get(name, self.zero_extent),
results.write_fields(),
self.zero_extent,
)
ctx.blocks[id(node)] = horizontal_extent
for name, accesses in results.read_offsets().items():
extent = functools.reduce(
lambda ext, off: ext | Extent.from_offset(off[:2]),
accesses,
Extent.from_offset(accesses.pop()[:2]),
)
total_extent = horizontal_extent + extent
ctx.fields.setdefault(name, total_extent)
ctx.fields[name] |= total_extent
for name in results.write_fields():
ctx.fields.setdefault(name, horizontal_extent)
def to_extent(self, horizontal_extent: Extent) -> Optional[Extent]:
"""
Convert the access to an extent provided a horizontal extent for the access.
This returns None if no overlap exists between the horizontal mask and interval.
"""
offset_as_extent = Extent.from_offset(
cast(Tuple[int, int, int], self.offset)[:2])
zeros = Extent.zeros(ndims=2)
if self.horizontal_mask:
if dist_from_edge := mask_overlap_with_extent(
self.horizontal_mask, horizontal_extent):
return ((horizontal_extent - dist_from_edge) +
offset_as_extent) | zeros
else:
return None
def visit_HorizontalExecution(self, node: oir.HorizontalExecution, *,
ctx: Context) -> None:
results = AccessCollector.apply(node).cartesian_accesses()
horizontal_extent = functools.reduce(
lambda ext, name: ext | ctx.field_extents.get(
name, Extent.zeros(ndims=2)),
results.write_fields(),
Extent.zeros(ndims=2),
)
ctx.block_extents[id(node)] = horizontal_extent
for name, accesses in results.read_offsets().items():
extent = functools.reduce(
lambda ext, off: ext | Extent.from_offset(off[:2]), accesses,
Extent.zeros(ndims=2))
ctx.field_extents[name] = ctx.field_extents.get(
name, Extent.zeros(ndims=2)).union(horizontal_extent + extent)
def visit_FieldRef(self, node: gt_ir.FieldRef):
extent = Extent.from_offset([node.offset.get(ax, 0) for ax in self.data.axes_names])
result = [(node.name, extent)]
return result
