def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
data = ModuleData()
node = pipeline.full()
field_extents = compute_legacy_extents(node)
write_fields = (
node.iter_tree()
.if_isinstance(gtir.ParAssignStmt)
.getattr("left")
.if_isinstance(gtir.FieldAccess)
.getattr("name")
.to_set()
)
read_fields: Set[str] = set()
for expr in node.iter_tree().if_isinstance(gtir.ParAssignStmt).getattr("right"):
read_fields |= expr.iter_tree().if_isinstance(gtir.FieldAccess).getattr("name").to_set()
referenced_field_params = [
param.name for param in node.params if isinstance(param, gtir.FieldDecl)
]
for name in sorted(referenced_field_params):
access = AccessKind.NONE
if name in read_fields:
access |= AccessKind.READ
if name in write_fields:
access |= AccessKind.WRITE
data.field_info[name] = FieldInfo(
access=access,
boundary=field_extents[name].to_boundary(),
axes=tuple(dimension_flags_to_names(node.symtable_[name].dimensions).upper()),
data_dims=tuple(node.symtable_[name].data_dims),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = [
param.name for param in node.params if param.name not in referenced_field_params
]
for name in sorted(referenced_scalar_params):
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower())
)
unref_params = get_unused_params_from_gtir(pipeline)
for param in sorted(unref_params, key=lambda decl: decl.name):
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = [*sorted(param.name for param in unref_params)]
return data
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
"""
Compute module data containing information about stencil arguments from gtir.
This is no longer compatible with the legacy backends.
"""
data = ModuleData()
# NOTE: pipeline.gtir has not had prune_unused_parameters applied.
all_params = pipeline.gtir.params
node = pipeline.full()
oir = gtir_to_oir.GTIRToOIR().visit(node)
field_extents = compute_fields_extents(oir)
accesses = compute_access_kinds(oir)
for decl in (param for param in all_params
if isinstance(param, gtir.FieldDecl)):
access = accesses[decl.name]
dtype = numpy.dtype(decl.dtype.name.lower())
if access != AccessKind.NONE:
k_boundary = compute_k_boundary(node)[decl.name]
boundary = Boundary(*field_extents[decl.name].to_boundary()[0:2],
k_boundary)
else:
boundary = Boundary.zeros(ndims=3)
data.field_info[decl.name] = FieldInfo(
access=access,
boundary=boundary,
axes=tuple(dimension_flags_to_names(decl.dimensions).upper()),
data_dims=tuple(decl.data_dims),
dtype=dtype,
)
for decl in (param for param in all_params
if isinstance(param, gtir.ScalarDecl)):
access = accesses[decl.name]
dtype = numpy.dtype(decl.dtype.name.lower())
data.parameter_info[decl.name] = ParameterInfo(access=access,
dtype=dtype)
data.unreferenced = [
*sorted(name for name in accesses if accesses[name] == AccessKind.NONE)
]
return data
def make_args_data_from_gtir(pipeline: GtirPipeline) -> ModuleData:
data = ModuleData()
node = pipeline.full()
field_extents = compute_legacy_extents(node)
write_fields = (node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("left").if_isinstance(
gtir.FieldAccess).getattr("name").to_list())
referenced_field_params = {
param.name
for param in node.params if isinstance(param, gtir.FieldDecl)
}
for name in referenced_field_params:
data.field_info[name] = FieldInfo(
access=AccessKind.READ_WRITE
if name in write_fields else AccessKind.READ_ONLY,
boundary=field_extents[name].to_boundary(),
axes=list(
dimension_flags_to_names(
node.symtable_[name].dimensions).upper()),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = set(
node.param_names).difference(referenced_field_params)
for name in referenced_scalar_params:
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()))
unref_params = get_unused_params_from_gtir(pipeline)
for param in unref_params:
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = {param.name for param in unref_params}
return data
def gtir_is_not_emtpy(pipeline: GtirPipeline) -> bool:
node = pipeline.full()
return bool(node.iter_tree().if_isinstance(gtir.ParAssignStmt).to_list())
def make_args_data_from_gtir(pipeline: GtirPipeline,
legacy=False) -> ModuleData:
"""
Compute module data containing information about stencil arguments from gtir.
Use `legacy` parameter to ensure equality with values from :func:`make_args_data_from_iir`.
"""
data = ModuleData()
node = pipeline.full()
write_fields = (node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("left").if_isinstance(
gtir.FieldAccess).getattr("name").to_set())
read_fields: Set[str] = set()
for expr in node.iter_tree().if_isinstance(
gtir.ParAssignStmt).getattr("right"):
read_fields |= expr.iter_tree().if_isinstance(
gtir.FieldAccess).getattr("name").to_set()
referenced_field_params = [
param.name for param in node.params
if isinstance(param, gtir.FieldDecl)
]
field_extents = compute_legacy_extents(node, mask_inwards=legacy)
k_boundary = (compute_k_boundary(node) if not legacy else
{v: (0, 0)
for v in referenced_field_params})
for name in sorted(referenced_field_params):
access = AccessKind.NONE
if name in read_fields:
access |= AccessKind.READ
if name in write_fields:
access |= AccessKind.WRITE
boundary = Boundary(*field_extents[name].to_boundary()[0:2],
k_boundary[name])
data.field_info[name] = FieldInfo(
access=access,
boundary=boundary,
axes=tuple(
dimension_flags_to_names(
node.symtable_[name].dimensions).upper()),
data_dims=tuple(node.symtable_[name].data_dims),
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()),
)
referenced_scalar_params = [
param.name for param in node.params
if param.name not in referenced_field_params
]
for name in sorted(referenced_scalar_params):
data.parameter_info[name] = ParameterInfo(
dtype=numpy.dtype(node.symtable_[name].dtype.name.lower()))
unref_params = get_unused_params_from_gtir(pipeline)
for param in sorted(unref_params, key=lambda decl: decl.name):
if isinstance(param, gtir.FieldDecl):
data.field_info[param.name] = None
elif isinstance(param, gtir.ScalarDecl):
data.parameter_info[param.name] = None
data.unreferenced = [*sorted(param.name for param in unref_params)]
return data
