def build_def_ir_stencil(name,
options,
backend="gtx86",
*,
id_version="xxxxxx"):
if isinstance(backend, str):
backend = gt_backend.from_name(backend)
if not issubclass(backend, gt_backend.Backend):
raise TypeError(
"Backend must be a backend identifier string or a gt4py Backend class."
)
def_ir_factory = REGISTRY[name]
def_ir = def_ir_factory()
iir = gt_analysis.transform(def_ir, options)
import tests.test_iir.iir_stencil_definitions as iir_stencil_definitions
ref_iir = iir_stencil_definitions.make_vertical_advection_dycore()
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.generate(stencil_id, iir, None, options)
stencil_implementation = stencil_class()
return stencil_implementation
def load_stencil(frontend_name, backend_name, definition_func, externals, options):
"""Generate a new class object implementing the provided definition.
"""
# Load components
backend = gt_backend.from_name(backend_name)
if backend 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))
# Create ID
options_id = backend.get_options_id(options)
stencil_id = frontend.get_stencil_id(
options.qualified_name, definition_func, externals, options_id
)
# Load or generate class
stencil_class = None if options.rebuild else backend.load(stencil_id, definition_func, options)
if stencil_class is None:
definition_ir = frontend.generate(definition_func, externals, options)
implementation_ir = gt_analysis.transform(definition_ir, options)
stencil_class = backend.build(stencil_id, implementation_ir, definition_func, options)
return stencil_class
def generate(
cls,
stencil_id: gt_definitions.StencilID,
definition_ir: gt_ir.StencilDefinition,
definition_func: types.FunctionType,
options: gt_definitions.BuildOptions,
):
# TODO: move this import to the top and find a better way to avoid circular imports
from gt4py import gt_src_manager
cls._check_options(options)
# Generate the Python binary extension (checking if GridTools sources are installed)
if not gt_src_manager.has_gt_sources(
) and not gt_src_manager.install_gt_sources():
raise RuntimeError("Missing GridTools sources.")
implementation_ir = gt_analysis.transform(definition_ir, options)
pyext_module_name, pyext_file_path = cls.generate_extension(
stencil_id,
definition_ir,
options,
implementation_ir=implementation_ir)
# Generate and return the Python wrapper class
return cls._generate_module(
stencil_id,
definition_ir,
definition_func,
options,
extra_cache_info={"pyext_file_path": pyext_file_path},
implementation_ir=implementation_ir,
pyext_module_name=pyext_module_name,
pyext_file_path=pyext_file_path,
)
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 generate(cls, stencil_id, definition_ir, definition_func, options):
from gt4py import gt_src_manager
cls._check_options(options)
implementation_ir = gt_analysis.transform(definition_ir, options)
# Generate the Python binary extension (checking if GridTools sources are installed)
if not gt_src_manager.has_gt_sources() and not gt_src_manager.install_gt_sources():
raise RuntimeError("Missing GridTools sources.")
pyext_module_name, pyext_file_path = cls.generate_extension(
stencil_id, implementation_ir, options
)
# Generate and return the Python wrapper class
generator_options = options.as_dict()
generator_options["pyext_module_name"] = pyext_module_name
generator_options["pyext_file_path"] = pyext_file_path
extra_cache_info = {"pyext_file_path": pyext_file_path}
return super(BaseGTBackend, cls)._generate_module(
stencil_id, implementation_ir, definition_func, generator_options, extra_cache_info
)
def _generate_implementation_ir(self) -> gt_ir.StencilImplementation:
implementation_ir = gt_analysis.transform(self.definition_ir,
self.options)
return implementation_ir
def generate(cls, stencil_id, definition_ir, definition_func, options):
cls._check_options(options)
implementation_ir = gt_analysis.transform(definition_ir, options)
return cls._generate_module(
stencil_id, implementation_ir, definition_func, copy.deepcopy(options.as_dict()), {}
)
