Just-in-time compile and link the modules. The Executable returned by tvm.compile may not be directly runnable as they may contain CUDA source files and objects that are yet to be compiled and linked. This function helps to create a runtime.Module for these cases.
(
self,
*,
fcompile: Callable[[str, list[str], dict[str, Any]], None] | None = None,
addons: list[str] | None = None,
force_recompile: bool = False,
**kwargs,
)
| 46 | return self.jit().main(*args, **kwargs) |
| 47 | |
| 48 | def jit( |
| 49 | self, |
| 50 | *, |
| 51 | fcompile: Callable[[str, list[str], dict[str, Any]], None] | None = None, |
| 52 | addons: list[str] | None = None, |
| 53 | force_recompile: bool = False, |
| 54 | **kwargs, |
| 55 | ) -> Module: |
| 56 | """Just-in-time compile and link the modules. |
| 57 | |
| 58 | The Executable returned by tvm.compile may not be directly |
| 59 | runnable as they may contain CUDA source files and objects that |
| 60 | are yet to be compiled and linked. |
| 61 | This function helps to create a runtime.Module for these cases. |
| 62 | |
| 63 | Parameters |
| 64 | ---------- |
| 65 | fcompile : function(target, file_list, kwargs), optional |
| 66 | The compilation function to use create the final library object during |
| 67 | |
| 68 | addons : list of str, optional |
| 69 | Additional object files to link against. |
| 70 | |
| 71 | force_recompile : bool, optional |
| 72 | If True, force a recompile of the module. |
| 73 | |
| 74 | kwargs : dict, optional |
| 75 | Additional arguments passed to fcompile |
| 76 | |
| 77 | Returns |
| 78 | ------- |
| 79 | rt_mod: tvm.runtime.Module |
| 80 | A runnable runtime module that can be passed to VirtualMachine. |
| 81 | |
| 82 | Examples |
| 83 | -------- |
| 84 | .. code:: python |
| 85 | |
| 86 | ex = tvm.compile(mod, target) |
| 87 | rt_mod = ex.jit() |
| 88 | |
| 89 | """ |
| 90 | |
| 91 | # If the module is already jitted and we don't want to force a recompile, |
| 92 | # return the cached module |
| 93 | if self._jitted_mod is not None and not force_recompile: |
| 94 | return self._jitted_mod |
| 95 | |
| 96 | # TODO(tvm-team): Update runtime.Module interface |
| 97 | # to query these properties as bitmask. |
| 98 | def _not_runnable(x): |
| 99 | return x.kind in ("c", "static_library") |
| 100 | |
| 101 | # pylint:disable = protected-access |
| 102 | not_runnable_list = self.mod._collect_from_import_tree(_not_runnable) |
| 103 | |
| 104 | # everything is runnable, directly return mod. |
| 105 | if len(not_runnable_list) == 0: |