tvm.dlight

tvm.dlight#

DLight package provides efficient schedules out-of-box for deep learning workloads.

class tvm.dlight.ApplyDefaultSchedule(*rules: ScheduleRule)[源代码]#

A IRModule pass that applies a list of ScheduleRules to all PrimFuncs in the module.

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

Information about a TIR block.

dom() → List[int | PrimExpr][源代码]#

The iteration domain of the block.

dom_kind() → str[源代码]#

The iteration domain kind of the block, for example, SSSS, SSSR.

is_elementwise(sch: Schedule) → bool[源代码]#

Whether the block is elementwise, i.e. trivial mapping between read/write region

is_gemm() → bool[源代码]#

Whether the block is a GEMM workload.

is_gemv() → bool[源代码]#

Whether the block is a GEMV workload.

is_injective() → bool[源代码]#

Whether the block is injective, i.e. all its iteration domains are injective.

is_reduction() → bool[源代码]#

Whether the block is a reduction workload.

class tvm.dlight.IterInfo(kind: Literal['S', 'R', 'O'], var: Var, dom: PrimExpr, loop_rv: LoopRV)[源代码]#

Information about a loop/iter var.

property dom: int | PrimExpr#

The iteration domain of the loop.

class tvm.dlight.ScheduleRule[源代码]#

A thin wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc.

Given a PrimFunc, a target, and a tunable flag, the apply method of a ScheduleRule returns either a Schedule, a list of Schedules, or None, where None means that the rule is not applicable to the given PrimFunc. If the tunable flag is True, the ScheduleRule is allowed to return either a Schedule or a list of Schedules, and the Schedules are allowed to contain tunable instructions. If the tunable flag is False, the ScheduleRule is only allowed to return a Schedule, and the Schedule is not allowed to contain tunable instructions.

apply(func: PrimFunc, target: Target, tunable: bool) → None | Schedule | List[Schedule][源代码]#

Apply the ScheduleRule to the given PrimFunc.

参数:

  • func (tir.PrimFunc) -- The PrimFunc to apply the ScheduleRule to.

  • target (Target) -- The compilation target the schedule is supposed to be built for.

  • tunable (bool) -- Whether the schedule is allowed to contain tunable instructions.

返回:

results -- Either a Schedule, a list of Schedules, or None, where None means that the rule is not applicable to the given PrimFunc.

返回类型:

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

static from_callable(name) → Callable[[Callable[[PrimFunc, Target, bool], None | Schedule | List[Schedule]]], ScheduleRule][源代码]#

Create a ScheduleRule from a callable.

参数:

name (str)

返回:

decorator -- A decorator that takes a callable and returns a 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[源代码]#

Check whether the rule is available for the given target.

参数:

target (Target) -- The compilation target the schedule is supposed to be built for.

返回:

available -- Whether the rule is available for the given target.

返回类型:

bool

tvm.dlight.normalize_prim_func(sch: Schedule) → List[BlockInfo] | None[源代码]#

Normalize the primfunc to normal form

tvm.dlight.try_inline(sch: Schedule, blocks: List[BlockInfo]) → List[BlockInfo][源代码]#

Try to inline as many blocks as possible, and return the remaining blocks.

参数:

  • sch (tir.Schedule) -- The TIR schedule used to inline blocks.

  • blocks (List[BlockInfo]) -- The blocks to be inlined.

返回:

remaining -- The remaining blocks that cannot be inlined.

返回类型:

List[BlockInfo]

tvm.dlight.try_inline_contiguous_spatial(sch: Schedule, block_infos: List[BlockInfo]) → List[BlockInfo][源代码]#

Try to inline contiguous spatial blocks in a schedule

参数:

  • sch (tir.Schedule) -- The TIR schedule used to inline blocks.

  • block_infos (List[BlockInfo]) -- The blocks to be try.

返回:

remaining -- The remaining blocks that cannot be inlined.

返回类型:

List[BlockInfo]