tvm.dlight

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

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

关于 TIR 块的信息。

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

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

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

is_elementwise(sch: Schedule) → bool: 块是否是元素级的，即读/写区域之间的简单映射

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

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

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

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

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

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

tvm.dlight.normalize_prim_func(sch: Schedule) → List[BlockInfo] | None: 将 primfunc 规范化为标准形式

class tvm.dlight.ApplyDefaultSchedule(*rules: ScheduleRule): 一个 IRModule pass，它将 ScheduleRule 列表应用于模块中的所有 PrimFunc。

class tvm.dlight.ScheduleRule

一个任意函数的薄封装，可用于调度 TIR PrimFunc。

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

apply(func: PrimFunc, target: Target, tunable: bool) → None | Schedule | List[Schedule]

将 ScheduleRule 应用于给定的 PrimFunc。

参数:

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

返回:

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

返回类型:

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

static from_callable(name) → Callable[[Callable[[PrimFunc, Target, bool], None | Schedule | List[Schedule]]], ScheduleRule]

从可调用对象创建 ScheduleRule。

参数:: name (str)
返回:: decorator – 一个装饰器，它接受一个可调用对象并返回一个 ScheduleRule。
返回类型:: Callable

示例

@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: Target) → bool

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

参数:: target (Target) – 计划为其构建调度的编译目标。
返回:: available – 规则是否适用于给定目标。
返回类型:: bool

tvm.dlight.try_inline(sch: Schedule, blocks: List[BlockInfo]) → List[BlockInfo]

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

参数:

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

返回:

remaining – 无法内联的剩余块。

返回类型:

List[BlockInfo]

tvm.dlight.try_inline_contiguous_spatial(sch: Schedule, block_infos: List[BlockInfo]) → List[BlockInfo]

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

参数:

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

返回:

remaining – 无法内联的剩余块。

返回类型:

List[BlockInfo]