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 sample_args_data():
dtype = np.dtype(np.float_)
yield ModuleData(
field_info={
"in_field":
FieldInfo(
access=AccessKind.READ_WRITE,
boundary=Boundary.zeros(ndims=3),
axes=("I", "J", "K"),
data_dims=tuple([]),
dtype=dtype,
)
},
parameter_info={
"param": ParameterInfo(access=AccessKind.READ, dtype=dtype)
},
)
def _run_test_implementation(cls, parameters_dict, implementation): # noqa: C901 # too complex
input_data, exec_info = parameters_dict
origin = cls.origin
max_boundary = Boundary(cls.max_boundary)
field_params = cls.field_params
field_masks = {}
for name, value in input_data.items():
if isinstance(value, np.ndarray):
field_masks[name] = tuple(
ax in field_params[name][0] for ax in CartesianSpace.names
)
data_shape = Shape((sys.maxsize,) * 3)
for name, data in input_data.items():
if isinstance(data, np.ndarray):
data_shape &= Shape(
interpolate_mask(data.shape, field_masks[name], default=sys.maxsize)
)
domain = data_shape - (
Index(max_boundary.lower_indices) + Index(max_boundary.upper_indices)
)
referenced_inputs = {
name: info for name, info in implementation.field_info.items() if info is not None
}
referenced_inputs.update(
{name: info for name, info in implementation.parameter_info.items() if info is not None}
)
# set externals for validation method
for k, v in implementation.constants.items():
sys.modules[cls.__module__].__dict__[k] = v
# copy input data
test_values = {}
validation_values = {}
for name, data in input_data.items():
data = input_data[name]
if name in referenced_inputs:
info = referenced_inputs[name]
if isinstance(info, FieldInfo):
data_dims = field_params[name][1]
if data_dims:
dtype = (data.dtype, data_dims)
shape = data.shape[: -len(data_dims)]
else:
dtype = data.dtype
shape = data.shape
test_values[name] = gt_storage.from_array(
data,
dtype=dtype,
shape=shape,
mask=field_masks[name],
default_origin=origin,
backend=implementation.backend,
)
validation_values[name] = np.array(data)
else:
test_values[name] = data
validation_values[name] = data
else:
test_values[name] = None
validation_values[name] = None
# call implementation
implementation(**test_values, origin=origin, exec_info=exec_info)
assert domain == exec_info["domain"]
# for validation data, data is cropped to actually touched domain, so that origin offseting
# does not have to be implemented for every test suite. This is done based on info
# specified in test suite
cropped_validation_values = {}
for name, data in validation_values.items():
sym = cls.symbols[name]
if data is not None and sym.kind == SymbolKind.FIELD:
field_extent_low = tuple(b[0] for b in sym.boundary)
offset_low = tuple(b[0] - e for b, e in zip(max_boundary, field_extent_low))
field_extent_high = tuple(b[1] for b in sym.boundary)
offset_high = tuple(b[1] - e for b, e in zip(max_boundary, field_extent_high))
validation_slice = filter_mask(
tuple(slice(o, s - h) for o, s, h in zip(offset_low, data_shape, offset_high)),
field_masks[name],
)
data_dims = field_params[name][1]
if data_dims:
validation_slice = tuple([*validation_slice] + [slice(None)] * len(data_dims))
cropped_validation_values[name] = data[validation_slice]
else:
cropped_validation_values[name] = data
cls.validation(
**cropped_validation_values,
domain=domain,
origin={
name: info.boundary.lower_indices
for name, info in implementation.field_info.items()
if info is not None
},
)
# Test values
for name, value in test_values.items():
if isinstance(value, np.ndarray):
expected_value = validation_values[name]
if gt_backend.from_name(value.backend).storage_info["device"] == "gpu":
value.synchronize()
value = value.data.get()
else:
value = value.data
np.testing.assert_allclose(
value,
expected_value,
rtol=RTOL,
atol=ATOL,
equal_nan=EQUAL_NAN,
err_msg="Wrong data in output field '{name}'".format(name=name),
)
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
def _run_test_implementation(self, parameters_dict, implementation):
cls = type(self)
input_data, exec_info = parameters_dict
origin = cls.origin
max_boundary = Boundary(cls.max_boundary)
shape_iter = (Shape(v.shape) for v in input_data.values()
if isinstance(v, np.ndarray))
shape = next(shape_iter)
assert all(shape == sh for sh in shape_iter)
domain = shape - (Index(max_boundary.lower_indices) +
Index(max_boundary.upper_indices))
referenced_inputs = {
name: info
for name, info in implementation.field_info.items()
if info is not None
}
referenced_inputs.update({
name: info
for name, info in implementation.parameter_info.items()
if info is not None
})
# set externals for validation method
for k, v in implementation.constants.items():
sys.modules[self.__module__].__dict__[k] = v
# copy input data
inputs = {}
validation_inputs = {}
for name, data in input_data.items():
data = input_data[name]
if name in referenced_inputs:
info = referenced_inputs[name]
if isinstance(info, FieldInfo):
inputs[name] = gt_storage.from_array(
data,
dtype=data.dtype,
shape=shape,
default_origin=origin,
backend=implementation.backend,
)
validation_inputs[name] = np.array(data)
else:
inputs[name] = data
validation_inputs[name] = data
else:
inputs[name] = None
validation_inputs[name] = None
# call implementation
implementation(**inputs, origin=origin, exec_info=exec_info)
assert domain == exec_info["domain"]
# for validation data, data is cropped to actually touched domain, so that origin offseting
# does not have to be implemented for every test suite. This is done based on info
# specified in test suite
cropped_validation_inputs = {}
for name, data in validation_inputs.items():
sym = cls.symbols[name]
if data is not None and sym.kind == SymbolKind.FIELD:
field_extent_low = tuple(b[0] for b in sym.boundary)
offset_low = tuple(
b[0] - e for b, e in zip(max_boundary, field_extent_low))
field_extent_high = tuple(b[1] for b in sym.boundary)
offset_high = tuple(
b[1] - e for b, e in zip(max_boundary, field_extent_high))
validation_slice = tuple(
slice(o, s - h)
for o, s, h in zip(offset_low, shape, offset_high))
cropped_validation_inputs[name] = data[validation_slice]
else:
cropped_validation_inputs[name] = data
cls.validation(
**cropped_validation_inputs,
domain=domain,
origin={
name: info.boundary.lower_indices
for name, info in implementation.field_info.items()
if info is not None
},
)
# Test values
for name, value in inputs.items():
if isinstance(value, np.ndarray):
expected_value = validation_inputs[name]
if gt_backend.from_name(
value.backend).storage_info["device"] == "gpu":
value.synchronize()
value = value.data.get()
else:
value = value.data
np.testing.assert_allclose(
value,
expected_value,
rtol=RTOL,
atol=ATOL,
equal_nan=EQUAL_NAN,
err_msg="Wrong data in output field '{name}'".format(
name=name),
)