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__
_set_arg_dtypes(definition_func, dtypes or {})
return gt_loader.gtscript_loader(
definition_func,
backend=backend,
build_options=build_options,
externals=externals or {},
)
def _decorator(def_func):
_set_arg_dtypes(def_func, dtypes or {})
return gt_loader.gtscript_loader(def_func,
backend=backend,
build_options=build_options,
externals=externals or {})
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 {})