tvm.dlight

tvm.dlight#

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

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

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

参数:

rules (ScheduleRule)

__init__(*rules)[源代码]#

Construct a new ApplyDefaultSchedule pass.

Parameters#

*rulesScheduleRule

The ScheduleRules to apply to all PrimFuncs in the module.

参数:

rules (ScheduleRule)

class tvm.dlight.BlockInfo(name, iters, block_rv, reduction_block=False)[源代码]#

Information about a TIR block.

参数:
__init__(name, iters, block_rv, reduction_block=False)[源代码]#

Construct a BlockInfo object.

参数:
dom()[源代码]#

The iteration domain of the block.

返回类型:

List[int | PrimExpr]

dom_kind()[源代码]#

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

返回类型:

str

is_elementwise(sch)[源代码]#

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

参数:

sch (Schedule)

返回类型:

bool

is_gemm()[源代码]#

Whether the block is a GEMM workload.

返回类型:

bool

is_gemv()[源代码]#

Whether the block is a GEMV workload.

返回类型:

bool

is_injective()[源代码]#

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

返回类型:

bool

is_reduction()[源代码]#

Whether the block is a reduction workload.

返回类型:

bool

class tvm.dlight.IterInfo(kind, var, dom, loop_rv)[源代码]#

Information about a loop/iter var.

参数:
__init__(kind, var, dom, loop_rv)[源代码]#

Construct an IterInfo object.

参数:
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, target, tunable)[源代码]#

Apply the ScheduleRule to the given PrimFunc.

Parameters#

functir.PrimFunc

The PrimFunc to apply the ScheduleRule to.

targetTarget

The compilation target the schedule is supposed to be built for.

tunablebool

Whether the schedule is allowed to contain tunable instructions.

Returns#

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

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

参数:
返回类型:

None | Schedule | List[Schedule]

static from_callable(name)[源代码]#

Create a ScheduleRule from a callable.

Parameters#

name : str

Returns#

decoratorCallable

A decorator that takes a callable and returns a ScheduleRule.

Examples#

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

Callable[[Callable[[PrimFunc, Target, bool], None | Schedule | List[Schedule]]], ScheduleRule]

is_target_available(target)[源代码]#

Check whether the rule is available for the given target.

Parameters#

targetTarget

The compilation target the schedule is supposed to be built for.

Returns#

availablebool

Whether the rule is available for the given target.

参数:

target (Target)

返回类型:

bool

tvm.dlight.normalize_prim_func(sch)[源代码]#

Normalize the primfunc to normal form

参数:

sch (Schedule)

返回类型:

List[BlockInfo] | None

tvm.dlight.try_inline(sch, blocks)[源代码]#

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

Parameters#

schtir.Schedule

The TIR schedule used to inline blocks.

blocksList[BlockInfo]

The blocks to be inlined.

Returns#

remainingList[BlockInfo]

The remaining blocks that cannot be inlined.

参数:
返回类型:

List[BlockInfo]

tvm.dlight.try_inline_contiguous_spatial(sch, block_infos)[源代码]#

Try to inline contiguous spatial blocks in a schedule

Parameters#

schtir.Schedule

The TIR schedule used to inline blocks.

block_infosList[BlockInfo]

The blocks to be try.

Returns#

remainingList[BlockInfo]

The remaining blocks that cannot be inlined.

参数:
返回类型:

List[BlockInfo]