def visit_Stencil(self, node: oir.Stencil, **kwargs):
from gtc.passes.oir_optimizations.utils import compute_horizontal_block_extents
block_extents = compute_horizontal_block_extents(node)
library_nodes = [
self.visit(vl, stencil=node, block_extents=block_extents, **kwargs)
for vl in node.vertical_loops
]
return StencilOirSDFGBuilder.build(node.name, node, library_nodes)
def visit_Stencil(self, node: oir.Stencil, **kwargs: Any) -> cuir.Program:
block_extents = compute_horizontal_block_extents(node)
accessed_fields: Set[str] = set()
kernels = self.visit(
node.vertical_loops,
new_symbol_name=symbol_name_creator(set(kwargs["symtable"])),
accessed_fields=accessed_fields,
block_extents=block_extents,
**kwargs,
)
temporaries = [
self.visit(d) for d in node.declarations
if d.name in accessed_fields
]
return cuir.Program(
name=node.name,
params=self.visit(node.params),
temporaries=temporaries,
kernels=kernels,
)
def visit_Stencil(self, node: oir.Stencil, **kwargs: Any) -> cuir.Program:
block_extents = compute_horizontal_block_extents(node)
ctx = self.Context(
new_symbol_name=symbol_name_creator(collect_symbol_names(node)))
kernels = self.visit(
node.vertical_loops,
ctx=ctx,
block_extents=block_extents,
**kwargs,
)
temporaries = [
self.visit(d) for d in node.declarations
if d.name in ctx.accessed_fields
]
return cuir.Program(
name=node.name,
params=self.visit(node.params),
positionals=list(ctx.positionals.values()),
temporaries=temporaries,
kernels=kernels,
)
def test_stencil_extents_region(mask, offset, access_extent):
testee = StencilFactory(
vertical_loops__0__sections__0__horizontal_executions=[
HorizontalExecutionFactory(body=[
AssignStmtFactory(left__name="tmp", right__name="input")
]),
HorizontalExecutionFactory(body=[
HorizontalRestrictionFactory(
mask=mask,
body=[
AssignStmtFactory(left__name="tmp",
right__name="input",
right__offset__i=offset)
],
),
]),
HorizontalExecutionFactory(body=[
AssignStmtFactory(
left__name="output", right__name="tmp", right__offset__i=1)
]),
],
declarations=[TemporaryFactory(name="tmp")],
)
block_extents = compute_horizontal_block_extents(testee)
hexecs = testee.vertical_loops[0].sections[0].horizontal_executions
mask_read_accesses = AccessCollector.apply(hexecs[1].body[0])
input_access = next(
iter(acc for acc in mask_read_accesses.ordered_accesses()
if acc.field == "input"))
block_extent = ((0, 1), (0, 0))
assert block_extents[id(hexecs[1])] == block_extent
if access_extent is not None:
assert input_access.to_extent(Extent(block_extent)) == access_extent
else:
assert input_access.to_extent(Extent(block_extent)) is None
def visit_Stencil(self, node: oir.Stencil) -> Dict[str, AccessKind]:
access: Dict[str, AccessKind] = collections.defaultdict(
lambda: AccessKind.NONE)
block_extents = compute_horizontal_block_extents(node)
self.generic_visit(node, access=access, block_extents=block_extents)
return access
def visit_Stencil(self, node: oir.Stencil) -> oir.Stencil:
block_extents = compute_horizontal_block_extents(node)
return self.generic_visit(node, block_extents=block_extents)