tvm.transform
跨IR变体的通用传递基础设施。
- tvm.transform.ApplyPassToFunction(transform: 通过, func_name_regex: str, error_if_no_function_matches_regex: bool = False) 通过
用于将通道应用于IRModule中特定函数的实用工具
TVM在所有降阶阶段都使用IRModule到IRModule的转换。这些转换在手动编写优化模型时可能很有用,或者用于对IRModule中的特定内核执行优化。该实用程序允许将优化过程应用于指定函数,而不会修改模块中的其他函数。
- Parameters:
- Returns:
new_transform – 修改后的IRModule到IRModule转换过程。
- Return type:
- class tvm.transform.ModulePass
一个作用于tvm.IRModule的转换过程。用户无需直接与该类交互,而应通过module_pass创建模块级转换,因为module_pass API的设计足够灵活,能以不同方式处理模块级转换的创建。此外,模块级转换的所有成员都可以从基类访问。该规则同样适用于FunctionPass。
- class tvm.transform.Pass
所有Pass的基类。这里的所有方法都只是后端实现的简单封装,定义这些方法是为了方便用户与基类进行交互。
- property info
获取传递的元数据。
- class tvm.transform.PassContext(opt_level=2, required_pass=None, disabled_pass=None, instruments=None, config=None, trace=None, trace_stack=None, make_traceable=None, num_evals=0, tuning_api_database=None)
TVM优化/分析运行的基础。 每个处理上下文包含一些辅助信息,用于帮助优化处理。这些信息包括用于记录优化过程中错误的错误报告器等。
- opt_levelOptional[int]
此优化过程的优化级别。
- required_passOptional[Union[List[str], Set[str], Tuple[str]]]
某个特定pass所需的pass列表。
- disabled_passOptional[Union[List[str], Set[str], Tuple[str]]]
被禁用的pass列表。
- instrumentsOptional[Sequence[PassInstrument]]
通行工具实现列表。
- configOptional[Dict[str, Object]]
特定通道的额外配置。
- trace: Optional[relax.tuning.Trace]
跟踪模式的初始跟踪。
- trace_stack: Optional[List[relax.tuning_api.Trace]]
追踪模式的初始追踪堆栈。
- make_traceable: Optional[List[str]]
使可追踪的pass列表。
- num_evals: int
在流程中初始进行的评估次数。
tuning_api_database: 可选[relax.tuning_api.JSONDatabase]
- override_instruments(instruments)
覆盖此PassContext中的工具。
如果存在现有工具,则会调用它们的
exit_pass_ctx回调函数。 然后切换到新工具并调用新的enter_pass_ctx回调函数。- instrumentsSequence[PassInstrument]
通行工具实现列表。
- static current()
返回当前的传递上下文。
- push_trace(trace)
将追踪信息压入堆栈。
- pop_trace(return_current=True)
从堆栈中弹出最顶层的追踪记录。 :returns: Trace :rtype: Optional[relax.tuning.Trace]
- get_trace_stack()
获取当前跟踪堆栈。
- get_trace_stack_size()
获取当前堆栈的大小。
- get_current_trace()
获取堆栈顶部的追踪信息。
- get_tuning_api_database()
获取调优API数据库。
- class tvm.transform.PassInfo(opt_level, name, required=None, traceable=False)
该类包含一个pass所需的元数据。它是运行优化或分析时所需信息的容器。当需要更多元数据时,可以通过添加新成员来扩展此类。
- tvm.transform.PrintIR(header='', show_meta_data=False)
一个特殊的跟踪过程,用于打印头部和IR。
- class tvm.transform.Sequential(passes=None, opt_level=0, name='sequential', required=None, traceable=False)
一个作用于一系列pass对象的处理过程。通过这个类可以顺序执行多个pass。
请注意,用户在运行顺序pass时也可以提供一系列他们不希望应用的pass。pass依赖关系也将在后端得到解决。
- tvm.transform.module_pass(pass_func=None, opt_level=None, name=None, required=None, traceable=False)
装饰一个模块传递。
当提供pass_func时,此函数返回一个回调函数。否则,它充当装饰器函数。
pass_func 也可以是一个带有 transform_module 方法的类类型。 该函数将使用 transform_module 作为传递函数创建一个装饰后的 ModulePass。
- Parameters:
pass_func (可选[Callable[(Module, PassContext) ->Module]]) – 转换函数或类。
opt_level (int) – 该模块传递的优化级别。
name (Optional[str]) – 模块传递的名称。名称可以为空。在这种情况下,将使用优化函数的名称作为传递名称。
required (Optional[List[str]]) – 模块转换所依赖的转换过程列表。
traceable (布尔值) - 布尔变量,表示模块传递是否可追踪
- Returns:
create_module_pass – 如果未提供pass_func参数,将返回一个装饰器, 否则返回装饰后的结果。 返回的装饰器根据输入有两种行为: 当我们装饰一个pass函数时,将返回一个新的ModulePass。 当我们装饰一个类类型时,将返回一个新的ModulePass类。
- Return type:
Union[Callable, ModulePass]
示例
以下代码块装饰了一个模块传递类。
@tvm.ir.transform.module_pass class CustomPipeline: def __init__(self, enable_fold): self.enable_fold = enable_fold self.const_fold = relax.transform.FoldConstant() def transform_module(self, mod, ctx): if self.enable_fold: mod = self.const_fold(mod, ctx) return mod # create an instance of customized pipeline pipeline = CustomPipeline(enable_fold=False) assert isinstance(pipeline, transform.ModulePass) # run the pipeline. output_module = pipeline(input_module)
以下代码通过装饰用户自定义的转换函数来创建一个模块pass。
@tvm.ir.transform.module_pass(opt_level=2) def transform(mod, ctx): return relax.transform.FoldConstant(mod) module_pass = transform assert isinstance(module_pass, transform.ModulePass) assert module_pass.info.opt_level == 2 # Given a module m, the optimization could be invoked as the follwoing: updated_mod = module_pass(m) # Now a function abs should be added to the module m.