Method call_tir

python/tvm/relax/base_py_module.py:221–254

Call a TIR function with PyTorch tensors.

(self, tir_func, args, out_sinfo)

            register_py_func(func_name, wrapped_func)

    def call_tir(self, tir_func, args, out_sinfo):
        """Call a TIR function with PyTorch tensors."""
        # Try to get function name from different sources
        if isinstance(tir_func, str):
            func_name = tir_func
        elif hasattr(tir_func, "name"):
            func_name = tir_func.name
        elif hasattr(tir_func, "__name__"):
            func_name = tir_func.__name__
        else:
            # Try to find by function object reference
            for name, func in self.compiled_tir_funcs.items():
                if func == tir_func:
                    func_name = name
                    break
            else:
                func_name = None

        if not func_name or func_name not in self.compiled_tir_funcs:
            available_funcs = list(self.compiled_tir_funcs.keys())
            raise ValueError(
                f"Could not resolve or find compiled TIR function: {tir_func}. "
                f"Available functions: {available_funcs}"
            )
        func = self.compiled_tir_funcs[func_name]

        out = self._create_output_tensors(out_sinfo, args)
        tvm_args = self._convert_pytorch_to_tvm(args)
        tvm_out = self._convert_pytorch_to_tvm(out)

        func(*tvm_args, *tvm_out)

        result = self._convert_tvm_to_pytorch(tvm_out)
        return result[0] if len(result) == 1 else result

    def call_dps_packed(self, func_name: str, args, out_sinfo):
        """Call a packed function with PyTorch tensors, converting TVM Tensors via DLPack."""