def build_transform(self):
definition = self.build()
return TransformData(
definition_ir=definition,
implementation_ir=init_implementation_from_definition(definition),
options=BuildOptions(name=self.name, module=__name__),
)
def stencil_def_to_oir(stencil_def, externals):
build_options = BuildOptions(
name=stencil_def.__name__, module=__name__, rebuild=True, backend_opts={}, build_info=None
)
definition_ir = GTScriptFrontend.generate(
stencil_def, externals=externals, options=build_options
)
gtir = GtirPipeline(DefIRToGTIR.apply(definition_ir)).full()
return GTIRToOIR().visit(gtir)
def with_options(self: "StencilBuilder", *, name: str, module: str,
**kwargs: Any) -> "StencilBuilder":
"""
Fluidly set the build options.
Parameters
----------
name:
Name of the stencil
module:
qualified module name of the stencil's module
kwargs:
passed through to the options instance
"""
self._build_data = {}
# mypy attribkwclass bug
self.options = BuildOptions(name=name, module=module,
**kwargs) # type: ignore
return self
def make_transform_data(
*,
name: str,
domain: Domain,
fields: List[str],
body: BodyType,
iteration_order: IterationOrder,
) -> TransformData:
definition = make_definition(name, fields, domain, body, iteration_order)
return TransformData(
definition_ir=definition,
implementation_ir=init_implementation_from_definition(definition),
options=BuildOptions(name=name, module=__name__),
)
def __init__(
self,
definition_func: Union[StencilFunc, AnnotatedStencilFunc],
*,
backend: Optional[Type["BackendType"]] = None,
options: Optional[BuildOptions] = None,
frontend: Optional["FrontendType"] = None,
):
self._definition = definition_func
# type ignore explanation: Attribclass generated init not recognized by mypy
self.options = options or BuildOptions( # type: ignore
**self.default_options_dict(definition_func))
self.backend: "BackendType" = (backend(self) if backend else
gt4py.backend.from_name("debug")(self))
self.frontend: "FrontendType" = frontend or gt4py.frontend.from_name(
"gtscript")
self.caching = gt4py.caching.strategy_factory("jit", self)
self._build_data: Dict[str, Any] = {}
self._externals: Dict[str, Any] = {}
class StencilBuilder:
"""
Orchestrates code generation and compilation.
Parameters
----------
definition_func:
Stencil definition function, before or after annotation.
backend:
Backend class to be instanciated for this build
options:
Build options, defaults to name and module of the :code:`definition_func`
frontend:
Frontend class to be used
Notes
-----
Backends can use the :py:meth:`with_backend_data` method the :py:meth:`backend_data` property
to set and retrieve backend specific information for use in later steps.
"""
def __init__(
self,
definition_func: Union[StencilFunc, AnnotatedStencilFunc],
*,
backend: Optional[Type["BackendType"]] = None,
options: Optional[BuildOptions] = None,
frontend: Optional["FrontendType"] = None,
):
self._definition = definition_func
# type ignore explanation: Attribclass generated init not recognized by mypy
self.options = options or BuildOptions( # type: ignore
**self.default_options_dict(definition_func))
self.backend: "BackendType" = (backend(self) if backend else
gt4py.backend.from_name("debug")(self))
self.frontend: "FrontendType" = frontend or gt4py.frontend.from_name(
"gtscript")
self.caching = gt4py.caching.strategy_factory("jit", self)
self._build_data: Dict[str, Any] = {}
self._externals: Dict[str, Any] = {}
def build(self) -> Type["StencilObject"]:
"""Generate, compile and/or load everything necessary to provide a usable stencil class."""
# load or generate
stencil_class = None if self.options.rebuild else self.backend.load()
if stencil_class is None:
stencil_class = self.backend.generate()
return stencil_class
def generate_computation(self) -> Dict[str, Union[str, Dict]]:
"""Generate the stencil source code, fail if backend does not support CLI."""
return self.cli_backend.generate_computation()
def generate_bindings(self,
targe_language: str) -> Dict[str, Union[str, Dict]]:
"""Generate ``target_language`` bindings source, fail if backend does not support CLI."""
return self.cli_backend.generate_bindings(targe_language)
def with_caching(self: "StencilBuilder", caching_strategy_name: str,
*args: Any, **kwargs: Any) -> "StencilBuilder":
"""
Fluidly set the caching strategy from the name.
Parameters
----------
caching_strategy_name:
Name of the caching strategy to be passed to the factory.
args:
Passed through to the caching strategy factory
kwargs:
Passed through to the caching strategy factory
Notes
-----
Resets all cached build data.
"""
self._build_data = {}
self.caching = gt4py.caching.strategy_factory(caching_strategy_name,
self, *args, **kwargs)
return self
def with_options(self: "StencilBuilder", *, name: str, module: str,
**kwargs: Any) -> "StencilBuilder":
"""
Fluidly set the build options.
Parameters
----------
name:
Name of the stencil
module:
qualified module name of the stencil's module
kwargs:
passed through to the options instance
"""
self._build_data = {}
# mypy attribkwclass bug
self.options = BuildOptions(name=name, module=module,
**kwargs) # type: ignore
return self
def with_changed_options(self: "StencilBuilder",
**kwargs: Dict[str, Any]) -> "StencilBuilder":
old_options = self.options.as_dict()
# BuildOptions constructor expects ``impl_opts`` keyword
# but BuildOptions.as_dict outputs ``_impl_opts`` key
old_options["impl_opts"] = old_options.pop("_impl_opts")
return self.with_options(**{**old_options, **kwargs})
def with_backend(self: "StencilBuilder",
backend_name: str) -> "StencilBuilder":
"""
Fluidly set the backend type from backend name.
Parameters
----------
backend_name:
Name of the backend type.
Notes
-----
Resets all cached build data.
"""
self._build_data = {}
self.backend = gt4py.backend.from_name(backend_name)(self)
return self
@classmethod
def default_options_dict(
cls, definition_func: Union[StencilFunc, AnnotatedStencilFunc]
) -> Dict[str, Any]:
return {
"name": definition_func.__name__,
"module": definition_func.__module__
}
@classmethod
def name_to_options_args(cls, name: Optional[str]) -> Dict[str, str]:
"""Check for qualified name, extract also module option in that case."""
if not name:
return {}
components = name.rsplit(".")
data = {"name": name}
if len(components) > 1:
data = {"module": components[0], "name": components[1]}
return data
@classmethod
def nest_impl_options(cls, options_dict: Dict[str, Any]) -> Dict[str, Any]:
impl_opts = options_dict.setdefault("impl_opts", {})
# The following is not a dict comprehension because:
# The backend-specific options (starting with ``_``) are nested under
# options_dict["impl_opts"] *and at the same time removed* from
# options_dict itself.
for impl_key in (k for k in options_dict.keys() if k.startswith("_")):
impl_opts[impl_key] = options_dict.pop(impl_key)
return options_dict
@property
def definition(self) -> AnnotatedStencilFunc:
return self._build_data.get(
"prepared_def") or self._build_data.setdefault(
"prepared_def",
self.frontend.prepare_stencil_definition(
self._definition, self.externals),
)
@property
def externals(self) -> Dict[str, Any]:
return self._build_data.get(
"externals") or self._build_data.setdefault(
"externals", self._externals.copy())
def with_externals(self: "StencilBuilder",
externals: Dict[str, Any]) -> "StencilBuilder":
"""
Fluidly set externals for this build.
Resets all cached build data.
"""
self._build_data = {}
self._externals = externals
return self
@property
def backend_data(self) -> Dict[str, Any]:
return self._build_data.get("backend_data", {}).copy()
def with_backend_data(self: "StencilBuilder",
data: Dict[str, Any]) -> "StencilBuilder":
self._build_data["backend_data"] = {**self.backend_data, **data}
return self
@property
def stencil_id(self) -> StencilID:
return self._build_data.setdefault("id", self.caching.stencil_id)
@property
def root_pkg_name(self) -> str:
return self._build_data.setdefault(
"root_pkg_name", gt4py.config.code_settings["root_package_name"])
def with_root_pkg_name(self: "StencilBuilder",
name: str) -> "StencilBuilder":
self._build_data["root_pkg_name"] = name
return self
@property
def pkg_name(self) -> str:
return self.options.name
@property
def pkg_qualname(self) -> str:
return self.root_pkg_name + "." + self.options.qualified_name
@property
def pkg_path(self) -> pathlib.Path:
return self.caching.backend_root_path.joinpath(
*self.options.qualified_name.split("."))
@property
def definition_ir(self) -> "StencilDefinition":
return self._build_data.get("ir") or self._build_data.setdefault(
"ir",
self.frontend.generate(self.definition, self.externals,
self.options))
@property
def implementation_ir(self) -> "StencilImplementation":
return self._build_data.get("iir") or self._build_data.setdefault(
"iir", gt4py.analysis.transform(self.definition_ir, self.options))
@property
def module_name(self) -> str:
return self.caching.module_prefix + self.options.name + self.caching.module_postfix
@property
def module_qualname(self) -> str:
return f"{self.pkg_qualname}.{self.module_name}"
@property
def module_path(self) -> pathlib.Path:
file_name = self.module_name + ".py"
return self.pkg_path / file_name
@property
def stencil_source(self) -> str:
"""
Read the stencil source from the written module file.
Raises
------
:py:class:`FileNotFoundError`
if the file has not been written yet
"""
if not self.module_path.exists():
return ""
return self.module_path.read_text()
@property
def class_name(self) -> str:
return self.caching.class_name
@property
def is_build_data_empty(self) -> bool:
return not bool(self._build_data)
@property
def cli_backend(self) -> "CLIBackendMixin":
from gt4py.backend.base import CLIBackendMixin
if not isinstance(self.backend, CLIBackendMixin):
raise RuntimeError(
"backend of StencilBuilder instance is not CLI enabled.")
return self.backend