def from_source( input_func: Union[str, Callable], tir_prefix: Optional[List[str]] = None) -> Union[PrimFunc, IRModule]: """Parse function or string into PrimFunc or IRModule. If possible, pass the TVM script in as a function so that line numbers and filename will be accurate. Parameters ---------- input_module : Union[str, Callable] The python function to be parsed. tir_prefix : Optional[List[str]] The tir prefix list. Only works for str input, default by "tir" and "T". Returns ------- output : Union[Function, Module] The Function or Module in IR. """ if isinstance(input_func, str): tir_prefix = ["T", "tir"] if tir_prefix is None else tir_prefix return to_ast(input_func, TVMDiagnosticCtx(), TVMScriptParser(0, tir_prefix)) elif inspect.isfunction(input_func): _, start_line = inspect.getsourcelines(input_func) env: Dict[str, Any] = input_func.__globals__ namespace = [key for key in env.keys() if env[key] == tir] parser = TVMScriptParser(start_line, namespace) result = to_ast(input_func, TVMDiagnosticCtx(), parser) return result else: raise TypeError("Only function definitions are supported.")
def to_ast(program: Any) -> Any: diag_ctx = synr.PrinterDiagnosticContext() transformer = None res = synr.to_ast(program, diag_ctx, transformer) if isinstance(res, str): raise (RuntimeError(res)) return res
def from_source(src): """Parse function or string into TIR. If possible, pass the TVM script in as a function so that line numbers and filename will be accurate. Parameters ---------- src : [str, function, class] Pruned source of original script Returns ------- functions : PrimFunc or IRModule The PrimFunc or IRModule in IR. """ if isinstance(src, str): start_line = 0 else: _, start_line = inspect.getsourcelines(src) parser = TVMScriptParser(start_line) return to_ast(src, TVMDiagnosticCtx(), parser)
def to_ast_err(program: Any) -> Any: diag_ctx = ErrorAccumulator() transformer = None return synr.to_ast(program, diag_ctx, transformer)