def _get_module_jit(backend_name: str = None,
depth_frame: int = 2,
frame=None):
"""Get the ModuleJIT instance corresponding to the calling module
Parameters
----------
depth_frame : int
Depth of the frame to be selected
"""
if frame is None:
frame = get_frame(depth_frame)
module_name = get_module_name(frame)
if backend_name is None:
backend_name = get_backend_name_module(module_name)
modules = modules_backends[backend_name]
if module_name in modules:
return modules[module_name]
else:
return ModuleJIT(backend_name=backend_name, frame=frame)
def _get_transonic_calling_module(backend_name: str = None):
"""Get the Transonic instance corresponding to the calling module
Parameters
----------
backend_name: str
"""
frame = get_frame(2)
module_name = get_module_name(frame)
if backend_name is None:
backend_name = get_backend_name_module(module_name)
modules = modules_backends[backend_name]
if module_name in modules:
ts = modules[module_name]
if (
ts.is_transpiling != is_transpiling
or ts._compile_at_import_at_creation != has_to_compile_at_import()
or (
hasattr(ts, "path_mod")
and ts.path_backend.exists()
and mpi.has_to_build(ts.path_backend, ts.path_mod)
)
):
ts = Transonic(frame=frame, reuse=False, backend=backend_name)
else:
ts = Transonic(frame=frame, reuse=False, backend=backend_name)
return ts
def use_block(self, name):
"""Use the pythranized version of a code block
Parameters
----------
name : str
The name of the block.
"""
if not self.is_transpiled:
raise ValueError(
"`use_block` has to be used protected by `if ts.is_transpiled`"
)
if self.is_compiling and not self.process.is_alive(raise_if_error=True):
self.is_compiling = False
time.sleep(0.1)
self.module_backend = import_from_path(
self.path_extension, self.module_backend.__name__
)
assert self.backend.check_if_compiled(self.module_backend)
self.is_compiled = True
func = getattr(self.module_backend, name)
argument_names = self.arguments_blocks[name]
locals_caller = get_frame(1).f_locals
arguments = [locals_caller[name] for name in argument_names]
return func(*arguments)
def set_backend_for_this_module(backend_name):
"""Programmatically set a backend for a module"""
backend_name = backend_name.lower()
frame = get_frame(1)
module_name = get_module_name(frame)
if backend_name not in backends.keys():
raise ValueError(f"Bad backend value ({backend_name})")
backend_default_modules[module_name] = backend_name
def __init__(self, backend_name: str, frame=None):
self.backend_name = backend_name
if frame is None:
frame = get_frame(1)
self.filename = _get_filename_from_frame(frame)
if any(
self.filename.startswith(start)
for start in ("<ipython-", "/tmp/ipykernel_")):
self.is_dummy_file = True
self._ipython_src, self.filename = get_info_from_ipython()
self.module_name = self.filename
else:
self.is_dummy_file = False
self.module_name = get_module_name(frame)
modules_backends[backend_name][self.module_name] = self
self.used_functions = {}
self.jit_functions = {}
(
jitted_dicts,
codes_dependance,
codes_dependance_classes,
code_ext,
special,
) = analysis_jit(self.get_source(), self.filename, backend_name)
self.info_analysis = {
"jitted_dicts": jitted_dicts,
"codes_dependance": codes_dependance,
"codes_dependance_classes": codes_dependance_classes,
"special": special,
}
self.backend = backend = backends[backend_name]
path_jit = mpi.Path(backend.jit.path_base)
path_jit_class = mpi.Path(backend.jit.path_class)
# TODO: check if these files have to be written here...
# Write exterior code for functions
for file_name, code in code_ext["function"].items():
path_ext = path_jit / self.module_name.replace(".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info)
# Write exterior code for classes
for file_name, code in code_ext["class"].items():
path_ext = path_jit_class / self.module_name.replace(
".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info)
def jit(func=None,
backend: str = None,
native=False,
xsimd=False,
openmp=False):
"""Decorator to record that the function has to be jit compiled
"""
frame = get_frame(1)
decor = JIT(frame,
backend=backend,
native=native,
xsimd=xsimd,
openmp=openmp)
if callable(func):
return decor(func)
else:
return decor
def __init__(
self, use_transonified=True, frame=None, reuse=True, backend=None
):
if frame is None:
frame = get_frame(1)
self.module_name = module_name = get_module_name(frame)
if backend is None:
backend = get_backend_name_module(module_name)
if isinstance(backend, str):
backend = backends[backend]
self.backend = backend
modules = modules_backends[backend.name]
self._compile_at_import_at_creation = has_to_compile_at_import()
if reuse and module_name in modules:
ts = modules[module_name]
for key, value in ts.__dict__.items():
self.__dict__[key] = value
return
self.is_transpiling = is_transpiling
self.has_to_replace = has_to_replace
if is_transpiling:
self.functions = {}
self.classes = {}
self.signatures_func = {}
self.is_transpiled = False
self.is_compiled = False
return
self.is_compiling = False
if not use_transonified or not has_to_replace:
self.is_transpiled = False
self.is_compiled = False
return
if "." in module_name:
package, module_short_name = module_name.rsplit(".", 1)
module_backend_name = package + "."
else:
module_short_name = module_name
module_backend_name = ""
module_backend_name += f"__{backend.name}__." + module_short_name
self.path_mod = path_mod = Path(_get_filename_from_frame(frame))
suffix = ".py"
self.path_backend = path_backend = (
path_mod.parent / f"__{backend.name}__" / (module_short_name + suffix)
)
path_ext = None
if has_to_compile_at_import() and path_mod.exists():
if mpi.has_to_build(path_backend, path_mod):
if path_backend.exists():
time_backend = mpi.modification_date(path_backend)
else:
time_backend = 0
returncode = None
if mpi.rank == 0:
print(f"Running transonic on file {path_mod}... ", end="")
# better to do this in another process because the file is already run...
os.environ["TRANSONIC_NO_MPI"] = "1"
returncode = subprocess.call(
[
sys.executable,
"-m",
"transonic.run",
"-nc",
str(path_mod),
]
)
del os.environ["TRANSONIC_NO_MPI"]
returncode = mpi.bcast(returncode)
if returncode != 0:
raise RuntimeError(
f"transonic does not manage to produce the {backend.name_capitalized} "
f"file for {path_mod}"
)
if mpi.rank == 0:
print("Done!")
path_ext = path_backend.with_name(
backend.name_ext_from_path_backend(path_backend)
)
time_backend_after = mpi.modification_date(path_backend)
# We have to touch the files to signal that they are up-to-date
if time_backend_after == time_backend and mpi.rank == 0:
if not has_to_build(path_ext, path_backend):
path_backend.touch()
if path_ext.exists():
path_ext.touch()
else:
path_backend.touch()
path_ext = path_ext or path_backend.with_name(
backend.name_ext_from_path_backend(path_backend)
)
self.path_extension = path_ext
if (
has_to_compile_at_import()
and path_mod.exists()
and not self.path_extension.exists()
):
if mpi.rank == 0:
print(
f"Launching {backend.name_capitalized} to compile a new extension..."
)
self.is_compiling, self.process = backend.compile_extension(
path_backend, name_ext_file=self.path_extension.name
)
self.is_compiled = not self.is_compiling
self.is_transpiled = True
if not path_ext.exists() and not self.is_compiling:
path_ext_alt = path_backend.with_suffix(backend.suffix_extension)
if path_ext_alt.exists():
self.path_extension = path_ext = path_ext_alt
self.reload_module_backend(module_backend_name)
if not self.is_transpiled:
logger.warning(
f"Module {path_mod} has not been compiled for "
f"Transonic-{backend.name_capitalized}"
)
else:
self.is_compiled = backend.check_if_compiled(self.module_backend)
if self.is_compiled:
module = inspect.getmodule(frame)
# module can be None if (at least) it has been run with runpy
if module is not None:
if backend.name == "pythran":
module.__pythran__ = self.module_backend.__pythran__
module.__transonic__ = self.module_backend.__transonic__
if hasattr(self.module_backend, "arguments_blocks"):
self.arguments_blocks = getattr(
self.module_backend, "arguments_blocks"
)
modules[module_name] = self