def parse_definition(
definition_func: Callable[..., None],
*,
name: str,
module: str,
externals: Optional[Dict[str, Any]] = None,
dtypes: Dict[Type, Type] = None,
rebuild=False,
**kwargs,
):
original_annotations = gtscript._set_arg_dtypes(definition_func, dtypes=dtypes or {})
build_options = gt_definitions.BuildOptions(
name=name,
module=module,
rebuild=rebuild,
backend_opts=kwargs,
build_info=None,
)
gt_frontend.GTScriptFrontend.prepare_stencil_definition(
definition_func, externals=externals or {}
)
definition_ir = gt_frontend.GTScriptParser(
definition_func, externals=externals or {}, options=build_options
).run()
setattr(definition_func, "__annotations__", original_annotations)
return definition_ir
def compile_definition(
definition_func,
name: str,
module: str,
*,
externals: dict = None,
dtypes: dict = None,
rebuild=False,
**kwargs,
):
_, original_annotations = gtscript._set_arg_dtypes(definition_func, dtypes=dtypes or {})
build_options = gt_definitions.BuildOptions(
name=name, module=module, rebuild=rebuild, backend_opts=kwargs, build_info=None
)
options_id = gt_utils.shashed_id(build_options)
stencil_id = frontend.get_stencil_id(
build_options.qualified_name, definition_func, externals, options_id
)
definition_ir = gt_frontend.GTScriptParser(
definition_func, externals=externals or {}, options=build_options
).run()
setattr(definition_func, "__annotations__", original_annotations)
return stencil_id, definition_ir
def analyze(name):
module_name = "_test_module." + name
stencil_name = name + "_stencil"
options = gt_definitions.BuildOptions(name=stencil_name, module=module_name, rebuild=True)
definition_ir_factory = def_ir_registry[name]
definition_ir = definition_ir_factory()
iir = gt_analysis.transform(definition_ir, options)
return iir
def compile_definition(
definition_func, name: str, module: str, *, externals: dict, rebuild=False, **kwargs
):
build_options = gt_definitions.BuildOptions(
name=name, module=module, rebuild=rebuild, backend_opts=kwargs, build_info=None
)
options_id = gt_utils.shashed_id(build_options)
_ = frontend.get_stencil_id(
build_options.qualified_name, definition_func, externals, options_id
)
gt_frontend.GTScriptParser(definition_func, externals=externals, options=build_options).run()
def generate_test_module(name, backend, *, id_version, rebuild=True):
module_name = "_test_module." + name
stencil_name = name + "_stencil"
backend_opts = {}
if issubclass(backend, gt_back.BaseGTBackend):
backend_opts["debug_mode"] = False
backend_opts["add_profile_info"] = True
backend_opts["verbose"] = True
options = gt_defs.BuildOptions(
name=stencil_name, module=module_name, rebuild=rebuild, backend_opts=backend_opts
)
return build_def_ir_stencil(name, options, backend, id_version=id_version)
def generate_test_module(name, backend, *, id_version, rebuild=True):
module_name = "_test_module." + name
stencil_name = name
backend_opts = {}
if issubclass(backend, gt_back.BaseGTBackend):
backend_opts["debug_mode"] = False
backend_opts["add_profile_info"] = True
backend_opts["verbose"] = True
options = gt_defs.BuildOptions(
name=stencil_name, module=module_name, rebuild=rebuild, backend_opts=backend_opts
)
decorator = gtscript.stencil(backend=backend.name, externals=EXTERNALS_REGISTRY[stencil_name])
stencil_definition = stencil_registry[name]
return decorator(stencil_definition)
def compile_and_run_vertical_regions(backend, id_version):
import gt4py.definitions as gt_defs
from gt4py import ir as gt_ir
from gt4py.definitions import StencilID
iir = make_test_vertical_regions()
module_name = "_test_module." + "test_vertical_regions"
stencil_name = "test_vertical_regions_stencil"
options = gt_defs.BuildOptions(name=stencil_name,
module=module_name,
rebuild=False)
stencil_id = StencilID("{}.{}".format(options.module, options.name),
id_version)
stencil_class = backend.load(stencil_id, None, options)
if stencil_class is None:
stencil_class = backend.generate(stencil_id, iir, None, options)
stencil_implementation = stencil_class()
field_out = gt_store.from_array(
np.zeros([d + 2 * o for d, o in zip((10, 10, 10), (0, 0, 0))]),
default_origin=(0, 0, 0),
shape=(10, 10, 10),
dtype=np.float_,
mask=[True, True, True],
backend=backend.name,
)
field_out_ref = np.zeros(
[d + 2 * o for d, o in zip((10, 10, 10), (0, 0, 0))])
field_out_ref[:, :, 0] = 1.0
field_out_ref[:, :, 1:-1] = 2.0
field_out_ref[:, :, -1:] = 3.0
args = dict(_origin_=dict(out=(0, 0, 0)),
_domain_=(10, 10, 10),
out=field_out)
if hasattr(field_out, "host_to_device"):
field_out.host_to_device()
stencil_implementation.run(**args, exec_info=None)
if hasattr(field_out, "device_to_host"):
field_out.device_to_host(force=True)
return (field_out, field_out_ref)
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 generate(def_ir, backend, *, id_version):
module_name = "_test_module." + def_ir.name
stencil_name = def_ir.name + "_stencil"
options = gt_definitions.BuildOptions(name=stencil_name, module=module_name, rebuild=True)
stencil_id = StencilID("{}.{}".format(options.module, options.name), id_version)
if options.rebuild:
# Force recompilation
stencil_class = None
else:
# Use cached version (if id_version matches)
stencil_class = backend.load(stencil_id, None, options)
if stencil_class is None:
stencil_class = backend.build(stencil_id, def_ir, None, options)
stencil_implementation = stencil_class()
return stencil_implementation
def stencil(
backend,
definition=None,
*,
build_info=None,
dtypes=None,
externals=None,
format_source=True,
name=None,
rebuild=False,
**kwargs,
):
"""Generate an implementation of the stencil definition with the specified backend.
It can be used as a parametrized function decorator or as a regular function.
Parameters
----------
backend : `str`
Name of the implementation backend.
definition : `None` when used as a decorator, otherwise a `function` or a `:class:`gt4py.StencilObject`
Function object defining the stencil.
build_info : `dict`, optional
Dictionary used to store information about the stencil generation.
(`None` by default).
dtypes: `dict`[`str`, dtype_definition], optional
Specify dtypes for string keys in the argument annotations.
externals: `dict`, optional
Specify values for otherwise unbound symbols.
format_source : `bool`, optional
Format generated sources when possible (`True` by default).
name : `str`, optional
The fully qualified name of the generated :class:`StencilObject`.
If `None`, it will be set to the qualified name of the definition function.
(`None` by default).
rebuild : `bool`, optional
Force rebuild of the :class:`gt4py.StencilObject` even if it is
found in the cache. (`False` by default).
**kwargs: `dict`, optional
Extra backend-specific options. Check the specific backend
documentation for further information.
Returns
-------
:class:`gridtools.StencilObject`
Properly initialized instance of a dynamically-generated
subclass of :class:`gt4py.StencilObject`.
Raises
-------
ValueError
If inconsistent arguments are specified.
Examples
--------
TODO
"""
from gt4py import loader as gt_loader
if build_info is not None and not isinstance(build_info, dict):
raise ValueError(f"Invalid 'build_info' dictionary ('{build_info}')")
if dtypes is not None and not isinstance(dtypes, dict):
raise ValueError(f"Invalid 'dtypes' dictionary ('{dtypes}')")
if externals is not None and not isinstance(externals, dict):
raise ValueError(f"Invalid 'externals' dictionary ('{externals}')")
if not isinstance(format_source, bool):
raise ValueError(f"Invalid 'format_source' bool value ('{name}')")
if name is not None and not isinstance(name, str):
raise ValueError(f"Invalid 'name' string ('{name}')")
if not isinstance(rebuild, bool):
raise ValueError(f"Invalid 'rebuild' bool value ('{rebuild}')")
module = None
if name:
name_components = name.split(".")
name = name_components[-1]
module = ".".join(name_components[:-1])
name = name or ""
module = (
module or inspect.currentframe().f_back.f_globals["__name__"]
) # definition_func.__globals__["__name__"] ??,
# Move hidden "_option" keys to _impl_opts
_impl_opts = {}
for key, value in kwargs.items():
if key.startswith("_"):
_impl_opts[key] = value
for key in _impl_opts:
kwargs.pop(key)
build_options = gt_definitions.BuildOptions(
name=name,
module=module,
format_source=format_source,
rebuild=rebuild,
backend_opts=kwargs,
build_info=build_info,
impl_opts=_impl_opts,
)
def _decorator(definition_func):
if not isinstance(definition_func, types.FunctionType):
if hasattr(definition_func, "definition_func"): # StencilObject
definition_func = definition_func.definition_func
elif callable(definition_func): # General callable
definition_func = definition_func.__call__
_, original_annotations = _set_arg_dtypes(definition_func, dtypes or {})
out = gt_loader.gtscript_loader(
definition_func,
backend=backend,
build_options=build_options,
externals=externals or {},
)
setattr(definition_func, "__annotations__", original_annotations)
return out
if definition is None:
return _decorator
else:
return _decorator(definition)
def stencil(
backend,
definition=None,
*,
build_info=None,
externals=None,
name=None,
rebuild=False,
**kwargs,
):
"""Generate an implementation of the stencil definition with the specified backend.
It can be used as a parametrized function decorator or as a regular function.
Parameters
----------
backend : `str`
Name of the implementation backend.
definition : ``None` when used as a decorator, otherwise a `function` or a `:class:`gt4py.StencilObject`
Function object defining the stencil.
build_info : `dict`, optional
Dictionary used to store information about the stencil generation.
(`None` by default).
externals: `dict`, optional
Specify values for otherwise unbound symbols.
name : `str`, optional
The fully qualified name of the generated :class:`StencilObject`.
If `None`, it will be set to the qualified name of the definition function.
(`None` by default).
rebuild : `bool`, optional
Force rebuild of the :class:`gt4py.StencilObject` even if it is
found in the cache. (`False` by default).
**kwargs: `dict`, optional
Extra backend-specific options. Check the specific backend
documentation for further information.
Returns
-------
:class:`gridtools.StencilObject`
Properly initialized instance of a dynamically-generated
subclass of :class:`gt4py.StencilObject`.
Raises
-------
ValueError
If inconsistent arguments are specified.
Examples
--------
TODO
"""
from gt4py import loader as gt_loader
if build_info is not None and not isinstance(build_info, dict):
raise ValueError(f"Invalid 'build_info' dictionary ('{build_info}')")
if externals is not None and not isinstance(externals, dict):
raise ValueError(f"Invalid 'externals' dictionary ('{externals}')")
if name is not None and not isinstance(name, str):
raise ValueError(f"Invalid 'name' string ('{name}')")
if not isinstance(rebuild, bool):
raise ValueError(f"Invalid 'rebuild' bool value ('{rebuild}')")
module = None
if name:
name_components = name.split(".")
name = name_components[-1]
module = ".".join(name_components[:-1])
name = name or ""
module = (module or inspect.currentframe().f_back.f_globals["__name__"]
) # definition_func.__globals__["__name__"] ??,
build_options = gt_definitions.BuildOptions(name=name,
module=module,
rebuild=rebuild,
backend_opts=kwargs,
build_info=build_info)
if definition is None:
return functools.partial(
gt_loader.gtscript_loader,
backend=backend,
build_options=build_options,
externals=externals or {},
)
else:
return gt_loader.gtscript_loader(definition,
backend=backend,
build_options=build_options,
externals=externals or {})
def compile_and_run_average_stencil(backend, id_version, origin_in, origin_out,
domain):
import gt4py.definitions as gt_defs
from gt4py import ir as gt_ir
from gt4py.definitions import StencilID
def stage(field_in, field_out):
field_out = (field_in[-1, 0, 0] + field_in[1, 0, 0] +
field_in[0, -1, 0] + field_in[0, 1, 0] +
field_in[0, 0, -1] + field_in[0, 0, 1]) / 6
stage_1 = gt_ir.utils.make_stage(
stage,
compute_extent=[(1, 1), (1, 1), (1, 1)],
fields_with_access={
"field_in": (gt_ir.AccessIntent.READ_ONLY, [(1, 1), (1, 1),
(1, 1)]),
"field_out": (gt_ir.AccessIntent.READ_WRITE, [(0, 0), (0, 0),
(0, 0)]),
},
)
multi_stage_1 = gt_ir.utils.make_multi_stage("multi_stage_1",
gt_ir.IterationOrder.PARALLEL,
[stage_1])
iir = gt_ir.utils.make_implementation(
"test_iteration_domain_stencil",
args_list=["field_in", "field_out"],
fields_with_description={
"field_in": dict(is_api=True, extent=[(1, 1), (1, 1), (1, 1)]),
"field_out": dict(is_api=True, extent=[(0, 0), (0, 0), (0, 0)]),
},
parameters_with_type={},
multi_stages=[multi_stage_1],
)
module_name = "_test_module." + "test_iteration_domain"
stencil_name = "test_iteration_domain_stencil"
options = gt_defs.BuildOptions(name=stencil_name,
module=module_name,
rebuild=False,
backend_opts={"verbose": True})
stencil_id = StencilID("{}.{}".format(options.module, options.name),
id_version)
stencil_class = backend.load(stencil_id, None, options)
if stencil_class is None:
stencil_class = backend.generate(stencil_id, iir, None, options)
stencil_implementation = stencil_class()
field_in = gt_store.from_array(
np.ones([d + 2 * o for d, o in zip(domain, origin_in)]),
default_origin=origin_in,
dtype=np.float_,
mask=[True, True, True],
backend=backend.name,
)
field_out = gt_store.from_array(
np.zeros([d + 2 * o for d, o in zip(domain, origin_out)]),
dtype=np.float_,
mask=[True, True, True],
default_origin=origin_out,
backend=backend.name,
)
field_out_ref = np.zeros([d + 2 * o for d, o in zip(domain, origin_out)])
slices_out = tuple(
[slice(o, o + d, None) for o, d in zip(origin_out, domain)])
field_out_ref[slices_out] = 1.0
args = dict(
_origin_=dict(field_in=origin_in, field_out=origin_out),
_domain_=domain,
field_in=field_in,
field_out=field_out,
)
if hasattr(field_in, "host_to_device"):
field_in.host_to_device()
if hasattr(field_out, "host_to_device"):
field_out.host_to_device()
stencil_implementation.run(**args, exec_info=None)
if hasattr(field_in, "device_to_host"):
field_in.device_to_host(force=True)
if hasattr(field_out, "device_to_host"):
field_out.device_to_host(force=True)
return (field_out.data, field_out_ref)
