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