Initialize BasePyModule with JIT compilation and DLPack conversion.
(
self,
ir_mod: IRModule,
device: Device,
target: Target | None = None,
)
| 64 | pass |
| 65 | |
| 66 | def __init__( |
| 67 | self, |
| 68 | ir_mod: IRModule, |
| 69 | device: Device, |
| 70 | target: Target | None = None, |
| 71 | ): |
| 72 | """Initialize BasePyModule with JIT compilation and DLPack conversion.""" |
| 73 | self.device = device |
| 74 | self.ir_mod = ir_mod |
| 75 | |
| 76 | # Delegate IRModule operations |
| 77 | self.functions = ir_mod.functions |
| 78 | self.attrs = ir_mod.attrs |
| 79 | self.global_infos = ir_mod.global_infos |
| 80 | self.__getitem__ = ir_mod.__getitem__ |
| 81 | self.__setitem__ = ir_mod.__setitem__ |
| 82 | self.functions_items = ir_mod.functions_items |
| 83 | self.with_attr = ir_mod.with_attr |
| 84 | self.get_attr = ir_mod.get_attr |
| 85 | self.update_global_info = ir_mod.update_global_info |
| 86 | |
| 87 | def _getattr_python_function(name: str) -> Any: |
| 88 | """Support direct attribute access to funcs and IRModule methods.""" |
| 89 | if name in self.pyfuncs: |
| 90 | return self.pyfuncs[name] |
| 91 | if name in self.compiled_tir_funcs: |
| 92 | return self.compiled_tir_funcs[name] |
| 93 | if self.relax_vm and name in self.relax_func_names: |
| 94 | try: |
| 95 | return self.relax_vm[name] |
| 96 | except AttributeError: # More specific exception |
| 97 | return None |
| 98 | if hasattr(self.ir_mod, name): |
| 99 | return getattr(self.ir_mod, name) |
| 100 | raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") |
| 101 | |
| 102 | self.__getattr__ = _getattr_python_function |
| 103 | |
| 104 | self.compiled_tir_funcs: dict[str, Function] = {} |
| 105 | self.extern_funcs: dict[str, Function] = {} |
| 106 | self.tir_func_names: list[str] = [] |
| 107 | self.relax_func_names: list[str] = [] |
| 108 | self.relax_vm: relax.VirtualMachine | None = None |
| 109 | self.pyfuncs: dict[str, Any] = {} |
| 110 | |
| 111 | if target is None: |
| 112 | target = Target.from_device(device) |
| 113 | elif isinstance(target, str): |
| 114 | target = Target(target) |
| 115 | self.target = target |
| 116 | |
| 117 | self._collect_function_names() |
| 118 | self._compile_functions() |
| 119 | self._wrap_tir_functions() |
| 120 | self._wrap_relax_functions() |
| 121 | self._register_python_functions() |
| 122 | |
| 123 | def _collect_function_names(self): |
nothing calls this directly
no test coverage detected