def test_setters():
# should set some build data (internal cache)
builder = StencilBuilder(simple_stencil).with_externals({"a": 1.0}).with_backend_data({"b": 2})
assert builder.externals == {"a": 1.0}
assert builder.backend_data == {"b": 2}
version = builder.stencil_id.version
assert version
# should reset build data, stencil_id particularly should be recomputed
builder.with_backend("gtc:numpy")
assert builder.is_build_data_empty
assert builder.externals == {"a": 1.0}
assert builder.backend_data == {}
assert builder.stencil_id.version
assert builder.stencil_id.version == version
# should change the stencil version
builder.with_externals({"a": 2.0})
assert builder.stencil_id.version != version
# should reset build data, stencil_id should be recomputed according to
# new caching strategy
builder.with_backend_data({"b": 2.0}).with_caching("nocaching")
assert builder.is_build_data_empty
assert not builder.stencil_id.version
# should reset build data
builder.with_backend_data({"b": 2}).with_externals({"a": 3.0})
assert builder.is_build_data_empty
assert builder.externals == {"a": 3.0}
def load_stencil(
frontend_name: str,
backend_name: str,
definition_func: StencilFunc,
externals: Dict[str, Any],
build_options: "BuildOptions",
) -> Type["StencilObject"]:
"""Generate a new class object implementing the provided definition."""
# Load components
backend_cls = gt_backend.from_name(backend_name)
if backend_cls is None:
raise ValueError(
"Unknown backend name ({name})".format(name=backend_name))
frontend = gt_frontend.from_name(frontend_name)
if frontend is None:
raise ValueError("Invalid frontend specification ({name})".format(
name=frontend_name))
builder = StencilBuilder(definition_func,
options=build_options,
backend=backend_cls,
frontend=frontend).with_externals(externals)
return builder.build()
def test_lazy_syntax_check(frontend, backend):
"""Test syntax checking."""
lazy_pass = LazyStencil(
StencilBuilder(copy_stencil_definition, frontend=frontend, backend=backend)
)
lazy_fail = LazyStencil(
StencilBuilder(wrong_syntax_stencil_definition, frontend=frontend, backend=backend)
)
lazy_pass.check_syntax()
with pytest.raises(GTScriptDefinitionError):
lazy_fail.check_syntax()
def test_k_bounds(definition, expected_k_bounds):
builder = StencilBuilder(definition, backend=from_name("debug"))
k_boundary = compute_k_boundary(builder.gtir_pipeline.full(skip=[prune_unused_parameters]))[
"field_b"
]
assert expected_k_bounds == k_boundary
def test_generate_computation(backend, tmp_path):
"""Test if the :py:meth:`gt4py.backend.CLIBackendMixin.generate_computation` generates code."""
# note: if a backend is added that doesn't use CliBackendMixin it will
# have to be special cased in the backend fixture
builder = StencilBuilder(init_1, backend=backend).with_caching(
"nocaching", output_path=tmp_path / __name__ / "generate_computation")
result = builder.backend.generate_computation()
# python backends only need to generate one module, either the stencil module...
py_result = backend.languages[
"computation"] == "python" and "init_1.py" in result
# ... or a standalone computation module
py_standalone_result = (backend.languages["computation"] == "python"
and "computation.py" in result)
# c++ / cuda files generated by gt backends.
# computation does not include bindings
gt_result = (backend.languages["computation"] in {"c++", "cuda"}
and "init_1_src" in result
and "computation.hpp" in result["init_1_src"]
and ("computation.cpp" in result["init_1_src"]
or "computation.cu" in result["init_1_src"])
and "bindings.cpp" not in result["init_1_src"])
# TODO(havogt) remove once gtc:gt produces a cpp-file for computation
gtc_result = ((backend.name.startswith("gtc")
and backend.languages["computation"] in ["c++", "cuda"])
and "computation.hpp" in result["init_1_src"]
and "bindings.cpp" not in result["init_1_src"])
assert py_result or py_standalone_result or gt_result or gtc_result
def test_make_args_data_from_iir(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
iir = builder.implementation_ir
args_data = make_args_data_from_iir(iir)
args_list = set(inspect.signature(stencil_def).parameters.keys())
args_found = set()
for key in args_data.field_info:
assert key in args_list
if key in field_info_val[mode]:
assert args_data.field_info[key] is not None
else:
assert args_data.field_info[key] is None
assert key not in args_found
args_found.add(key)
for key in args_data.parameter_info:
assert key in args_list
if key in parameter_info_val[mode]:
assert args_data.parameter_info[key] is not None
else:
assert args_data.parameter_info[key] is None
assert key not in args_found
args_found.add(key)
for key in args_data.unreferenced:
assert key in args_list
assert key not in field_info_val[mode]
assert key not in parameter_info_val[mode]
assert key in unreferenced_val[mode]
assert key in args_found
def test_generate_bindings(backend, tmp_path):
"""Test :py:meth:`gt4py.backend.CLIBackendMixin.generate_bindings`."""
builder = StencilBuilder(init_1, backend=backend).with_caching(
"nocaching", output_path=tmp_path / __name__ / "generate_bindings")
if not backend.languages["bindings"]:
# no bindings supported
with pytest.raises(NotImplementedError):
result = builder.backend.generate_bindings("python")
elif "python" in backend.languages["bindings"] and backend.languages[
"computation"] == "python":
# standalone python computation module imported by stencil module
result = builder.backend.generate_bindings("python")
assert "init_1.py" in result
assert re.search(r"import computation", result["init_1.py"],
re.MULTILINE)
else:
# assumption: only gt backends support python bindings for other languages than python
if backend.name.startswith("gtc:"):
result = builder.backend.generate_bindings(
"python", ir=builder.definition_ir)
else:
result = builder.backend.generate_bindings("python")
assert "init_1_src" in result
srcs = result["init_1_src"]
assert "bindings.cpp" in srcs or "bindings.cu" in srcs
def test_module_data_equivalence():
builder = StencilBuilder(sample_stencil_with_args)
legacy_module_data = make_args_data_from_iir(builder.implementation_ir)
gtc_module_data = make_args_data_from_gtir(builder.gtir_pipeline)
assert legacy_module_data == gtc_module_data
def test_invalid_temporary_access(definition):
builder = StencilBuilder(definition, backend=from_name("gtc:numpy"))
with pytest.raises(
TypeError,
match="Invalid access with offset in k to temporary field tmp."):
k_boundary = compute_k_boundary(
builder.gtir_pipeline.full(skip=[prune_unused_parameters]))
def test_lazy_stencil():
"""Test lazy stencil construction."""
builder = StencilBuilder(copy_stencil_definition).with_options(
name="copy_stencil", module=copy_stencil_definition.__module__)
lazy_s = LazyStencil(builder)
assert lazy_s.backend.name == "debug"
def test_daa(definition, valid):
builder = StencilBuilder(definition, backend=from_name("debug"))
gtir_stencil_expr = builder.gtir_pipeline.full()
invalid_accesses = daa.analyze(gtir_stencil_expr)
if valid:
assert len(invalid_accesses) == 0
else:
assert len(invalid_accesses) == 1 and invalid_accesses[0].name == "tmp"
def test_make_args_data_from_gtir(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals({"MODE": mode})
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
iir_args_data = make_args_data_from_iir(builder.implementation_ir)
assert args_data.field_info == iir_args_data.field_info
assert args_data.parameter_info == iir_args_data.parameter_info
assert args_data.unreferenced == iir_args_data.unreferenced
def _decorator(func):
_set_arg_dtypes(func, dtypes or {})
options = gt_definitions.BuildOptions(
**{
**StencilBuilder.default_options_dict(func),
**StencilBuilder.name_to_options_args(name),
"rebuild": rebuild,
"build_info": build_info,
**StencilBuilder.nest_impl_options(kwargs),
})
stencil = LazyStencil(
StencilBuilder(func, backend=backend,
options=options).with_externals(externals or {}))
if eager:
stencil = stencil.implementation
elif check_syntax:
stencil.check_syntax()
return stencil
def make_builder(definition, backend_name="gtc:numpy", module=None):
"""Make a builder instance a definition and default params."""
return StencilBuilder(
definition,
backend=gt4py.backend.from_name(backend_name),
options=gt4py.definitions.BuildOptions(name="foo",
module=module
or f"{__name__}"),
)
def test_j():
def stencil(field_a: gs.Field[float], field_b: gs.Field[float, gs.J]):
with computation(PARALLEL), interval(...):
field_a = field_b[1] + field_b[-2]
builder = StencilBuilder(stencil, backend=from_name("debug"))
old_ext = builder.implementation_ir.fields_extents
legacy_ext = compute_legacy_extents(prepare_gtir(builder))
for name, ext in old_ext.items():
assert legacy_ext[name] == ext
def test_module_data():
builder = StencilBuilder(sample_stencil_with_args)
module_data = make_args_data_from_gtir(builder.gtir_pipeline)
assert module_data.field_info["used_io_field"].access == AccessKind.WRITE
assert module_data.field_info["used_in_field"].access == AccessKind.READ
assert module_data.field_info["unused_field"].access == AccessKind.NONE
assert module_data.parameter_info["used_scalar"].access == AccessKind.READ
assert module_data.parameter_info[
"unused_scalar"].access == AccessKind.NONE
def test_demote_temporaries_to_variables_pass_forward():
builder = StencilBuilder(double_smul_forward)
iir = builder.implementation_ir
assert iir.temporary_fields == ["tmp_f"]
multi_stage = iir.multi_stages[0]
assert len(multi_stage.groups) == 2
for i in range(len(multi_stage.groups)):
stage = multi_stage.groups[i].stages[0]
assert len(stage.apply_blocks[0].body.stmts) == 2
def test_single_k_offset():
def stencil(field_a: gs.Field[float], field_b: gs.Field[float]):
with computation(PARALLEL), interval(...):
field_a = field_b[0, 0, 1]
builder = StencilBuilder(stencil, backend=from_name("debug"))
old_ext = builder.implementation_ir.fields_extents
legacy_ext = compute_legacy_extents(prepare_gtir(builder),
mask_inwards=True)
for name, ext in old_ext.items():
assert legacy_ext[name] == ext
def test_regression_run_gtir_pipeline_twice(tmp_path):
builder = (
StencilBuilder(assign_bool_float)
.with_backend("numpy")
.with_externals({"a": 1.0})
.with_caching("nocaching", output_path=tmp_path)
.with_options(name="simple_stencil", module="", rebuild=True)
)
# property caching should not reevaluate the analysis pipeline as a side effect.
ir = builder.gtir_pipeline.full()
assert ir is builder.gtir_pipeline.full()
def test_regression_run_analysis_twice(tmp_path):
builder = (
StencilBuilder(assign_bool_float)
.with_backend("gtc:numpy")
.with_externals({"a": 1.0})
.with_caching("nocaching", output_path=tmp_path)
.with_options(name="simple_stencil", module="", rebuild=True)
)
# property caching should not reevaluate the analysis pipeline as a side effect.
ir = builder.implementation_ir
# this raises an error if the analysis pipeline is reevaluated:
assert ir is builder.implementation_ir
def test_generate_bindings(backend, tmp_path):
"""Test :py:meth:`gt4py.backend.CLIBackendMixin.generate_bindings`."""
builder = StencilBuilder(init_1, backend=backend).with_caching(
"nocaching", output_path=tmp_path / __name__ / "generate_bindings")
if not backend.languages["bindings"]:
# no bindings supported
with pytest.raises(NotImplementedError):
result = builder.backend.generate_bindings("python")
else:
# assumption: only gt backends support python bindings
result = builder.backend.generate_bindings("python")
assert "init_1_src" in result
assert "bindings.cpp" in result["init_1_src"]
def test_generate_post_run(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals({"MODE": mode})
iir = builder.implementation_ir
args_data = backend_cls.make_args_data_from_iir(iir)
module_generator = backend_cls.MODULE_GENERATOR_CLASS()
module_generator.args_data = args_data
source = module_generator.generate_post_run()
if backend_name in CPU_BACKENDS:
assert source == ""
else:
assert source == "out._set_device_modified()"
def test_generate_post_run(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
module_generator = backend_cls.MODULE_GENERATOR_CLASS()
module_generator.args_data = args_data
source = module_generator.generate_post_run()
if gt_backend.from_name(backend_name).storage_info["device"] == "cpu":
assert source == ""
else:
assert source == "out._set_device_modified()"
def test_usage_numpy_nocaching(tmp_path):
builder = (
StencilBuilder(simple_stencil)
.with_backend("gtc:numpy")
.with_externals({"a": 1.0})
.with_caching("nocaching", output_path=tmp_path)
.with_options(name="simple_stencil", module="")
)
computation_src = builder.generate_computation()
assert "computation.py" in computation_src
builder.build()
assert tmp_path.joinpath("simple_stencil", "computation.py").exists(), list(tmp_path.iterdir())
def test_offset_chain():
def stencil(field_a: gs.Field[float], field_b: gs.Field[float]):
with computation(PARALLEL), interval(...):
field_a = field_b[1, 0, 1]
with computation(PARALLEL), interval(...):
field_b = field_a[1, 0, 0]
with computation(PARALLEL), interval(...):
tmp = field_b[0, -1, 0] + field_b[0, 1, 0]
field_a = tmp[0, 0, 0] + tmp[0, 0, -1]
builder = StencilBuilder(stencil, backend=from_name("debug"))
old_ext = builder.implementation_ir.fields_extents
legacy_ext = compute_legacy_extents(prepare_gtir(builder))
for name, ext in old_ext.items():
assert legacy_ext[name] == ext
def _decorator(func):
# Move backend options to `backend_opts`
backend_opts: Dict[str, Any] = {}
for backend_opt in ("device_sync", "skip_passes", "verbose"):
if backend_opt in kwargs:
backend_opts[backend_opt] = kwargs.pop(backend_opt)
builder = (StencilBuilder(func).with_backend(backend).with_externals(
externals or {}).with_options(
name=func.__name__,
module=func.__module__,
rebuild=rebuild,
backend_opts=backend_opts,
**kwargs,
))
return FutureStencil(builder, wrapper)
def test_generate_pre_run(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
module_generator = backend_cls.MODULE_GENERATOR_CLASS()
module_generator.args_data = args_data
source = module_generator.generate_pre_run()
if gt_backend.from_name(backend_name).storage_info["device"] == "cpu":
assert source == ""
else:
for key in field_info_val[mode]:
assert f"{key}.host_to_device()" in source
for key in unreferenced_val[mode]:
assert f"{key}.host_to_device()" not in source
def test_lazy_call(frontend, backend):
"""Test that the lazy stencil is callable like the compiled stencil object."""
import numpy
a = gt4py.storage.from_array(
numpy.array([[[1.0]]]), default_origin=(0, 0, 0), backend=backend.name
)
b = gt4py.storage.from_array(
numpy.array([[[0.0]]]), default_origin=(0, 0, 0), backend=backend.name
)
lazy_s = LazyStencil(
StencilBuilder(copy_stencil_definition, frontend=frontend, backend=backend).with_options(
name="copy", module=copy_stencil_definition.__module__, rebuild=True
)
)
lazy_s(b, a)
assert b[0, 0, 0] == 1.0
def test_make_args_data_from_gtir(backend_name, mode):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
assert set(args_data.unreferenced) == set(unreferenced_val[mode])
field_info_from_gtir = {(
p.name,
np.dtype(p.dtype.name.lower()),
utils.dimension_flags_to_names(p.dimensions).upper(),
p.data_dims,
)
for p in builder.gtir.params
if isinstance(p, gtir.FieldDecl)}
field_info_from_args_data = {(name, d.dtype, "".join(d.axes), d.data_dims)
for name, d in args_data.field_info.items()
if name not in args_data.unreferenced}
assert field_info_from_gtir == field_info_from_args_data
param_info_from_gtir = {(p.name, np.dtype(p.dtype.name.lower()))
for p in builder.gtir.params
if isinstance(p, gtir.ScalarDecl)}
param_info_from_args_data = {
(name, d.dtype)
for name, d in args_data.parameter_info.items()
if name not in args_data.unreferenced
}
assert param_info_from_gtir == param_info_from_args_data
for name, field_info in args_data.field_info.items():
if name == "out":
access = AccessKind.WRITE
elif name in field_info_val[mode]:
access = AccessKind.READ
else:
access = AccessKind.NONE
assert field_info.access == access
for name, param_info in args_data.parameter_info.items():
if name in parameter_info_val[mode]:
access = AccessKind.READ
else:
access = AccessKind.NONE
assert param_info.access == access
def test_device_sync_option(backend_name, mode, device_sync):
backend_cls = backend_registry[backend_name]
builder = StencilBuilder(stencil_def, backend=backend_cls).with_externals(
{"MODE": mode})
builder.options.backend_opts["device_sync"] = device_sync
args_data = make_args_data_from_gtir(builder.gtir_pipeline)
module_generator = backend_cls.MODULE_GENERATOR_CLASS()
source = module_generator(
args_data,
builder,
pyext_module_name=builder.module_name,
pyext_file_path=str(builder.module_path),
)
if device_sync:
assert "cupy.cuda.Device(0).synchronize()" in source
else:
assert "cupy.cuda.Device(0).synchronize()" not in source