def visit_StencilDefinition(self, node: StencilDefinition) -> gtir.Stencil:
field_params = {f.name: self.visit(f) for f in node.api_fields}
scalar_params = {p.name: self.visit(p) for p in node.parameters}
vertical_loops = [self.visit(c) for c in node.computations if c.body.stmts]
return gtir.Stencil(
name=node.name,
params=[
self.visit(f, all_params={**field_params, **scalar_params})
for f in node.api_signature
],
vertical_loops=vertical_loops,
)
def visit_StencilDefinition(self, node: StencilDefinition) -> gtir.Stencil:
field_params = {f.name: self.visit(f) for f in node.api_fields}
scalar_params = {p.name: self.visit(p) for p in node.parameters}
self._scalar_params = scalar_params
vertical_loops = [self.visit(c) for c in node.computations if c.body.stmts]
return gtir.Stencil(
name=node.name.split(".")[
-1
], # TODO probably definition IR should not contain '.' in the name
params=[
self.visit(f, all_params={**field_params, **scalar_params})
for f in node.api_signature
],
vertical_loops=vertical_loops,
)
def prune_unused_parameters(node: gtir.Stencil) -> gtir.Stencil:
"""
Remove unused parameters from the gtir signature.
(Maybe this pass should go into a later stage. If you need to touch this pass,
e.g. when the definition_ir gets removed, consider moving it to a more appropriate
level. Maybe to the backend IR?)
"""
assert isinstance(node, gtir.Stencil)
used_variables = (node.iter_tree().if_isinstance(
gtir.FieldAccess, gtir.ScalarAccess).getattr("name").to_list())
used_params = list(
filter(lambda param: param.name in used_variables, node.params))
return gtir.Stencil(name=node.name,
params=used_params,
vertical_loops=node.vertical_loops)
def visit_StencilDefinition(self, node: StencilDefinition) -> gtir.Stencil:
field_params = {f.name: self.visit(f) for f in node.api_fields}
scalar_params = {p.name: self.visit(p) for p in node.parameters}
vertical_loops = [
self.visit(c) for c in node.computations if c.body.stmts
]
if node.externals is not None:
externals = {
name: _make_literal(value)
for name, value in node.externals.items()
if isinstance(value, numbers.Number)
}
else:
externals = {}
return gtir.Stencil(
name=node.name,
api_signature=[
gtir.Argument(
name=f.name,
is_keyword=f.is_keyword,
default=str(f.default)
if not isinstance(f.default, type(Empty)) else "",
) for f in node.api_signature
],
params=[
self.visit(f, all_params={
**field_params,
**scalar_params
}) for f in node.api_signature
],
vertical_loops=vertical_loops,
externals=externals,
sources=node.sources or "",
docstring=node.docstring,
loc=location_to_source_location(node.loc),
)
