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:
  • transform (Pass) – 要应用的从IRModule到IRModule的转换过程。

  • func_name_regex (str) – 用于选择要更新的函数的正则表达式。该过程将应用于所有名称匹配该正则表达式的函数。

  • error_if_no_function_matches_regex (bool) – 指定当IRModule中不包含任何匹配提供正则表达式的函数时的行为。如果为true,将引发错误。如果为false(默认值),将返回未经修改的IRModule。

Returns:

new_transform – 修改后的IRModule到IRModule转换过程。

Return type:

Pass

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()

返回当前的传递上下文。

static list_configs()

列出所有已注册的PassContext配置名称和元数据。

Returns:

配置

Return type:

字典[str, 字典[str, str]]

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()

获取堆栈顶部的追踪信息。

set_num_evals(num: int)

设置流水线中执行的评估次数。

inc_num_evals(num: int)

增加管道中执行的评估次数。

get_tuning_api_database()

获取调优API数据库。

class tvm.transform.PassInfo(opt_level, name, required=None, traceable=False)

该类包含一个pass所需的元数据。它是运行优化或分析时所需信息的容器。当需要更多元数据时,可以通过添加新成员来扩展此类。

Parameters:
  • opt_level (int) – 该优化通道的优化级别。

  • name (str) – 该pass的名称。

  • required (List[str]) – 某个特定pass所依赖的pass列表。

tvm.transform.PrintIR(header='', show_meta_data=False)

一个特殊的跟踪过程,用于打印头部和IR。

Parameters:
  • header (str) – 要与转储内容一起显示的标题。

  • show_meta_data (bool) – 一个布尔标志,用于指示是否应打印元数据。

Return type:

该通道

class tvm.transform.Sequential(passes=None, opt_level=0, name='sequential', required=None, traceable=False)

一个作用于一系列pass对象的处理过程。通过这个类可以顺序执行多个pass。

请注意,用户在运行顺序pass时也可以提供一系列他们不希望应用的pass。pass依赖关系也将在后端得到解决。

Parameters:
  • passes (可选[列表[Pass]]) – 用于优化的候选pass序列。

  • opt_level (可选[int]) – 该顺序优化通道的优化级别。 默认顺序通道的opt_level设置为0。 请注意,如果Sequential中的某些通道的opt_level高于提供的opt_level,这些通道可能仍不会执行。

  • name (可选参数[str]) – 该顺序处理步骤的名称。

  • 必选 (可选[列表[str]]) – 顺序执行所依赖的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.