def test_generation(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend, stencil_definition, externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=(v.dtype, v.data_dims) if v.data_dims else v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
# vertical domain size >= 16 required for test_large_k_interval
stencil(**args, origin=(10, 10, 5), domain=(3, 3, 16))
def test_generation_gpu(name, backend):
stencil_definition = stencil_definitions[name]
externals = externals_registry[name]
stencil = gtscript.stencil(backend,
stencil_definition,
externals=externals)
args = {}
for k, v in stencil_definition.__annotations__.items():
if isinstance(v, gtscript._FieldDescriptor):
args[k] = gt_storage.ones(
dtype=v.dtype,
mask=gtscript.mask_from_axes(v.axes),
backend=backend,
shape=(23, 23, 23),
default_origin=(10, 10, 10),
)
else:
args[k] = v(1.5)
stencil(**args, origin=(10, 10, 10), domain=(3, 3, 3))
def make_definition(
stencil_name: str,
definition_func: types.FunctionType,
args_list: list,
fields_with_storage_descriptor: dict,
temp_fields_with_type: dict,
parameters_with_type: dict,
domain=None,
externals=None,
sources=None,
):
api_signature = make_api_signature(args_list)
domain = domain or Domain.LatLonGrid()
externals = externals or {}
sources = sources or {}
fields_decls = []
for name, descriptor in fields_with_storage_descriptor.items():
fields_decls.append(
make_field_decl(
name=name,
dtype=descriptor.dtype,
masked_axes=[
i
for i, masked in enumerate(gtscript.mask_from_axes(descriptor.axes))
if masked
],
is_api=True,
layout_id=name,
)
)
temp_fields_decls = {
name: make_field_decl(name=name, dtype=dtype, is_api=False)
for name, dtype in temp_fields_with_type.items()
}
parameter_decls = []
for key, value in parameters_with_type.items():
if isinstance(value, tuple):
assert len(value) == 2
data_type = value[0]
length = value[1]
else:
data_type = value
length = 0
parameter_decls.append(
VarDecl(name=key, data_type=DataType.from_dtype(data_type), length=length, is_api=True)
)
computations = make_computations(
definition_func,
fields={decl.name: decl for decl in fields_decls},
parameters={decl.name: decl for decl in parameter_decls},
local_symbols=None,
externals=externals,
domain=domain,
extra_temp_decls=temp_fields_decls,
)
definition = StencilDefinition(
name=stencil_name,
domain=domain,
api_signature=api_signature,
api_fields=fields_decls,
parameters=parameter_decls,
computations=computations,
externals=externals,
sources=sources,
)
return definition
