tvm.dlight

DLight 包为深度学习工作负载提供了开箱即用的高效调度方案。

class tvm.dlight.BlockInfo(name: str, iters: 列表[IterInfo], block_rv: BlockRV, reduction_block: bool = False)

关于TIR块的信息。

dom() → 列表[int | PrimExpr]: 块的迭代域。

dom_kind() → str: 块的迭代域类型，例如SSSS、SSSR。

is_injective() → bool: 该块是否是单射的，即其所有迭代域都是单射的。

is_elementwise(sch: 日程安排) → bool: 该块是否为元素级的，即在读写区域之间存在简单的映射关系

is_reduction() → bool: 该块是否为归约工作负载。

is_gemv() → bool: 该块是否为GEMV工作负载。

is_gemm() → bool: 该块是否是GEMM工作负载。

class tvm.dlight.IterInfo(kind: typing_extensions.Literal[S, R, O], var: 变量, dom: PrimExpr, loop_rv: LoopRV)

关于循环/迭代变量的信息。

property dom: int | PrimExpr: 循环的迭代域。

tvm.dlight.normalize_prim_func(sch: 日程安排) → 列表[BlockInfo] | None: 将primfunc规范化为标准形式

class tvm.dlight.ApplyDefaultSchedule(*rules: ScheduleRule): 一个IRModule传递过程，将一系列ScheduleRules应用于模块中的所有PrimFuncs。

class tvm.dlight.ScheduleRule

一个对任意函数的轻量级封装，可用于调度TIR PrimFunc。

给定一个PrimFunc、目标和一个可调优标志，ScheduleRule的apply方法会返回一个Schedule、一个Schedule列表或None，其中None表示该规则不适用于给定的PrimFunc。如果可调优标志为True，则ScheduleRule可以返回一个Schedule或一个Schedule列表，并且这些Schedule可以包含可调优指令。如果可调优标志为False，则ScheduleRule只能返回一个Schedule，并且该Schedule不能包含可调优指令。

apply(func: PrimFunc, target: 目标, tunable: bool) → None | 日程安排 | 列表[日程安排]

将ScheduleRule应用于给定的PrimFunc。

Parameters:

func (tir.PrimFunc) – 要应用ScheduleRule的PrimFunc函数。
target (Target) – 调度计划预期构建的编译目标。
tunable (bool) – 该调度是否允许包含可调优指令。

Returns:

results – 可以是Schedule、Schedule列表或None，其中None表示该规则不适用于给定的PrimFunc。

Return type:

Union[None, tir.Schedule, List[tir.Schedule]]

static from_callable(name) → 可调用[[可调用[[PrimFunc, 目标, bool], None | 日程安排 | 列表[日程安排]]], ScheduleRule]

从可调用对象创建ScheduleRule。

Parameters:: 名称 (str)
Returns:: decorator – 一个装饰器，接收可调用对象并返回一个ScheduleRule。
Return type:: 可调用

示例

@ScheduleRule.from_callable("MyRule")
def my_rule(func: tir.PrimFunc, target: Target, tunable: bool) -> Union[None, Schedule]
    # Do something with func and target

is_target_available(target: 目标) → bool

检查该规则是否适用于给定的目标。

Parameters:: 目标 (Target) – 该调度计划预期构建的编译目标。
Returns:: available – 该规则是否适用于给定的目标。
Return type:: bool

tvm.dlight.try_inline(sch: 日程安排, blocks: 列表[BlockInfo]) → 列表[BlockInfo]

尝试内联尽可能多的块，并返回剩余的块。

Parameters:

sch (tir.Schedule) – 用于内联块的TIR调度器。
blocks (List[BlockInfo]) – 需要内联的块。

Returns:

remaining – 无法内联的剩余代码块。

Return type:

列表[BlockInfo]

tvm.dlight.try_inline_contiguous_spatial(sch: 日程安排, block_infos: 列表[BlockInfo]) → 列表[BlockInfo]

尝试在调度中内联连续的空间块

Parameters:

sch (tir.Schedule) – 用于内联块的TIR调度表。
block_infos (List[BlockInfo]) – 待尝试的区块。

Returns:

remaining – 无法内联的剩余块。

Return type:

列表[BlockInfo]