Compile TIR and Relax functions using JIT compilation.
(self)
| 129 | self.relax_func_names.append(global_var.name_hint) |
| 130 | |
| 131 | def _compile_functions(self): |
| 132 | """Compile TIR and Relax functions using JIT compilation.""" |
| 133 | # Compile TIR functions first |
| 134 | tir_mod = tvm.IRModule( |
| 135 | { |
| 136 | gv: func |
| 137 | for gv, func in self.ir_mod.functions_items() |
| 138 | if isinstance(func, tirx.PrimFunc) |
| 139 | } |
| 140 | ) |
| 141 | if tir_mod: |
| 142 | try: |
| 143 | tir_exec_mod = tvm.compile(tir_mod, target=self.target) |
| 144 | for func_name in self.tir_func_names: |
| 145 | self.compiled_tir_funcs[func_name] = tir_exec_mod[func_name] |
| 146 | # pylint: disable=broad-exception-caught |
| 147 | except Exception as error: |
| 148 | print(f"Warning: Failed to compile one or more TIR functions: {error}") |
| 149 | |
| 150 | if self.relax_func_names: |
| 151 | try: |
| 152 | exec_mod = tvm.compile(self.ir_mod, target=self.target) |
| 153 | self.relax_vm = relax.VirtualMachine(exec_mod, self.device) |
| 154 | # pylint: disable=broad-exception-caught |
| 155 | except Exception as error: |
| 156 | print(f"Warning: Failed to compile Relax VM: {error}") |
| 157 | self.relax_vm = None |
| 158 | |
| 159 | def _wrap_tir_functions(self): |
| 160 | """Wrap TIR functions to make them accessible as instance attributes.""" |
no test coverage detected