tvm.s_tir.dlight

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

class tvm.s_tir.dlight.SBlockInfo(name: str, iters: list[IterInfo], block_rv: SBlockRV, 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_injective() → bool: Whether the SBlock is injective, i.e. all its iteration domains are injective.

is_elementwise(sch: Schedule) → bool: Whether the SBlock is elementwise, i.e. trivial mapping between read/write region

is_reduction() → bool: Whether the SBlock is a reduction workload.

is_layout_transform(sch: Schedule) → bool: Whether the SBlock can be considered having a Layout Transform Pattern

is_data_pad(sch: Schedule) → bool: Whether the SBlock can be considered having a data pad pattern

is_convolution() → bool: Whether a SBlock can be considered having Convolution Pattern

is_pool() → bool: Whether a SBlock can be considered having Pooling Pattern

is_gemv() → bool: Whether the SBlock is a GEMV workload.

is_gemm() → bool: Whether the SBlock is a GEMM workload.

class tvm.s_tir.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.

tvm.s_tir.dlight.normalize_prim_func(sch: Schedule) → list[SBlockInfo] | None: Normalize the primfunc to normal form

class tvm.s_tir.dlight.ApplyDefaultSchedule(*rules: ScheduleRule): A IRModule pass that applies a list of ScheduleRules to all PrimFuncs in the module.

class tvm.s_tir.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.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

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

Create a ScheduleRule from a callable.

Parameters:: name (str)
Returns:: decorator – A decorator that takes a callable and returns a ScheduleRule.
Return type:: Callable

Examples

@ScheduleRule.from_callable("MyRule")
def my_rule(func: tirx.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.

Parameters:: target (Target) – The compilation target the schedule is supposed to be built for.
Returns:: available – Whether the rule is available for the given target.
Return type:: bool

tvm.s_tir.dlight.try_inline(sch: Schedule, blocks: list[SBlockInfo]) → list[SBlockInfo]

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

Parameters:

sch (s_tir.Schedule) – The TIR schedule used to inline blocks.
blocks (List[SBlockInfo]) – The blocks to be inlined.

Returns:

remaining – The remaining blocks that cannot be inlined.

Return type:

List[SBlockInfo]

tvm.s_tir.dlight.try_inline_contiguous_spatial(sch: Schedule, block_infos: list[SBlockInfo]) → list[SBlockInfo]

Try to inline contiguous spatial blocks in a schedule

Parameters:

sch (s_tir.Schedule) – The TIR schedule used to inline blocks.
block_infos (List[SBlockInfo]) – The blocks to be try.

Returns:

remaining – The remaining blocks that cannot be inlined.

Return type:

List[SBlockInfo]

tvm.s_tir.dlight.gpu

GPU-generic schedule rules. For CUDA/ROCm/Vulkan/Metal-specific rules, use tvm.s_tir.dlight.cuda/rocm/vulkan/metal instead

class tvm.s_tir.dlight.gpu.GEMV

A rule for GEMV and DecodeGEMV.

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

sch_inner_reduction(sch: Schedule, target: Target, block: SBlockRV, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None): Schedule the inner reduction block.

sch_outer_reduction(sch: Schedule, target: Target, block: SBlockRV, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None): Schedule the outer reduction block.

sch_outer_reduction_fallback(sch: Schedule, target: Target, block: SBlockRV, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None): Schedule the outer reduction block.

class tvm.s_tir.dlight.gpu.LowBatchGEMV(bucket=4)

A rule for low batch GEMM / decode-GEMM.

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

sch_inner_reduction(sch: Schedule, target: Target, block: SBlockRV, dequantize_block: SBlockRV | None, pad_input_block: SBlockRV | None, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None, batch_pad: int): Schedule the inner reduction block.

sch_outer_reduction(sch: Schedule, target: Target, block: SBlockRV, dequantize_block: SBlockRV | None, pad_input_block: SBlockRV | None, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None, batch_pad: int): Schedule the outer reduction block.

class tvm.s_tir.dlight.gpu.Fallback

A fallback schedule rule for all GPU operators. It will try to inline all the blocks first, and then apply a simple block/grid mapping to the spatial loops on top of the remaining blocks.

apply(func: PrimFunc, target: Target, _: bool) → Schedule

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.gpu.Matmul

The schedule rule for matmul-like computation

class Config(block_size_x: int = 8, block_size_y: int = 8, vthread_x: int = 1, vthread_y: int = 1, micro_size_x: int = 4, micro_size_y: int = 4, micro_size_k: int = 8, vector_size: int = 1, unroll: int = 256, use_shared: bool = True, storage_align: bool = False, inner_x: bool = False)

get_configs(target: Target) → Config: Get the schedule config for the target

apply(func: PrimFunc, target: Target, _: bool) → Schedule | None

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

sch_outer_reduction(sch: Schedule, config: Config, reduction_block: SBlockRV, blocks: list[SBlockRV]) → Schedule | None: Get vectorization factor

class tvm.s_tir.dlight.gpu.Reduction

A rule for Reduction.

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.gpu.Transpose

Schedule rule for transpose

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.gpu.GeneralReduction

General Reduction rule for operators including softmax, layer norm, RMS norm, etc

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.gpu.RMSNorm

A rule for RMS norm.

apply(func: PrimFunc, target: Target, _: bool) → Schedule

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

tvm.s_tir.dlight.adreno

Adreno schedule rules.

class tvm.s_tir.dlight.adreno.Conv2d

The schedule rule for convolution computation

apply(func: PrimFunc | Schedule, target: Target, _: bool) → Schedule | None

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.adreno.LayoutTransform(use_op_name=True)

Texture based Layout Transform Dlight Schedule for Adreno

apply(func: PrimFunc | Schedule, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.adreno.Fallback

Texture Based Fallback Schedule(s) for Adreno

static schedule_inline_blocks(sch: Schedule, blocks: list[SBlockRV]) → list[SBlockRV]: Auto Inlines Injective and Element-wise Operations while trying to omit data pad blocks…

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

class tvm.s_tir.dlight.adreno.Pool2D

apply(func: PrimFunc, target: Target, _: bool) → Schedule

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

tvm.s_tir.dlight.cpu

CPU-generic schedule rules.

class tvm.s_tir.dlight.cpu.GEMV

A rule for GEMV and DecodeGEMV.

apply(func: PrimFunc, target: Target, _: bool) → None | Schedule | list[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tirx.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.

Returns:

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

Return type:

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

sch_inner_reduction(sch: Schedule, target: Target, block: SBlockRV, vector_input_buffers: list[Buffer], epilogue_info: SBlockInfo | None): Schedule the inner reduction block.

tvm.s_tir.dlight.analysis

Base infra

class tvm.s_tir.dlight.analysis.SBlockInfo(name: str, iters: list[IterInfo], block_rv: SBlockRV, 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_injective() → bool: Whether the SBlock is injective, i.e. all its iteration domains are injective.

is_elementwise(sch: Schedule) → bool: Whether the SBlock is elementwise, i.e. trivial mapping between read/write region

is_reduction() → bool: Whether the SBlock is a reduction workload.

is_layout_transform(sch: Schedule) → bool: Whether the SBlock can be considered having a Layout Transform Pattern

is_data_pad(sch: Schedule) → bool: Whether the SBlock can be considered having a data pad pattern

is_convolution() → bool: Whether a SBlock can be considered having Convolution Pattern

is_pool() → bool: Whether a SBlock can be considered having Pooling Pattern

is_gemv() → bool: Whether the SBlock is a GEMV workload.

is_gemm() → bool: Whether the SBlock is a GEMM workload.

class tvm.s_tir.dlight.analysis.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.

tvm.s_tir.dlight.analysis.collect_block_iter_vars_used_in_access_region(block: SBlock, region: list[Range]) → set[Var]: Collect the block iter variables used in the access region of a buffer region.

tvm.s_tir.dlight.analysis.collect_vars_used_in_prim_expr(expr: PrimExpr) → set[Var]: Collect the variables used in the PrimExpr.

tvm.s_tir.dlight.analysis.detect_dominant_read(block: SBlock) → PrimExpr: Detect the dominant read indices in the block.

tvm.s_tir.dlight.analysis.is_broadcast_epilogue(sch: Schedule, block: SBlockRV, epilogue: SBlockRV) → bool: Check if the epilogue block is a broadcast pattern

tvm.s_tir.dlight.analysis.normalize_prim_func(sch: Schedule) → list[SBlockInfo] | None: Normalize the primfunc to normal form

tvm.s_tir.dlight.analysis.is_gemv(sch: Schedule, block_info: SBlockInfo) → list[Buffer] | None

Check if the block is a GEMV.

Parameters:

sch (s_tir.Schedule) – The schedule
block_info (SBlockInfo) – The block info to be checked

Returns:

ret – The vector buffers used in the GEMV if it is a GEMV, otherwise None.

Return type:

Optional[List[tirx.Buffer]]

tvm.s_tir.dlight.analysis.normalize(sch: Schedule, block_info: SBlockInfo) → bool | None: Normalize the main block.

tvm.s_tir.dlight.base

Base infra

tvm.s_tir.dlight.benchmark

Benchmarking dynamic shape workloads

Benchmark a PrimFunc or IRModule with dynamic input shapes.

Parameters:

mod_or_func (Union[PrimFunc, IRModule]) – The PrimFunc or IRModule to be benchmarked.
dym_var_sample (Optional[Dict[str, int]]) – The dynamic shape variable sample, e.g., {“n”: 64, “m”: 128}.
args (Optional[List[Union[relax.TensorStructInfo, Tuple[Tuple[Union[int, str], ...], str]]]]) – The input tensor information, including shape and dtype. If none, will use the input information from the PrimFunc or IRModule.
target (Optional[Union[str, tvm.target.Target]]) – The target to be benchmarked on, if none, will get the target from context.
func_name (Optional[str]) – The name of the function to be benchmarked, will use “main” by default.
evaluator_config (Optional["EvaluatorConfig"]) – The evaluator configuration to use. If none, will use default evaluator configuration.
rpc_config (Optional["RPCConfig"]) – The RPC configuration to connect to the remote device. If none, will use local mode.

Returns:

input_infos (List[Tuple[Tuple[int, …], str]]) – The input tensor information, including shape and dtype.
median (float) – The median of the benchmarking results.
std (float) – The standard deviation of the benchmarking results.

tvm.s_tir.dlight.benchmark.benchmark_prim_func(mod_or_func: ~tvm.tirx.function.PrimFunc | ~tvm.ir.module.IRModule, *, dym_var_sample_func: ~collections.abc.Callable[[dict[str, str]], dict[str, int]] = <function default_dym_var_sample_func>, args: list[~tvm.relax.struct_info.TensorStructInfo | tuple[tuple[int | str, ...], str]] | None = None, dym_var_dict: dict[str, str] | None = None, sample_number: int = 5, target: str | ~tvm.target.target.Target | None = None, weight: int | None = 1, relax_func_name: str | None = None, prim_func_name: str | None = None, evaluator_config: EvaluatorConfig | None = None, rpc_config: RPCConfig | None = None, sort_by: str | None = None, desc: bool | None = True)

Benchmark a PrimFunc or IRModule with dynamic input shapes and show results.

Parameters:

mod_or_func (Union[PrimFunc, IRModule]) – The PrimFunc or IRModule to be benchmarked.
dym_var_sample_func (Callable[[Dict[str, str]], Dict[str, int]]) – The function to sample dynamic shape variables.
dym_var_dict (Optional[Dict[str, str]]) – Dynamic shape variable dictionary, e.g., {“n”: “int32”, “m”: “int32”}. If none, will use the input information from the PrimFunc or IRModule.
args (Optional[List[Union[relax.TensorStructInfo, Tuple[Tuple[Union[int, str], ...], str]]]]) – The input tensor information, including shape and dtype. If none, will use the input information from the PrimFunc or IRModule.
sample_number (int) – The number of times to sample dynamic shape variables.
target (Optional[Union[str, tvm.target.Target]]) – The target to be benchmarked on, if none, will get the target from context.
weight (Optional[int]) – The weight of this PrimFunc.
relax_func_name (Optional[str]) – The name of the relax function.
prim_func_name (Optional[str]) – The name of the PrimFunc.
evaluator_config (Optional["EvaluatorConfig"]) – The evaluator configuration to use. If none, will use default evaluator configuration.
rpc_config (Optional["RPCConfig"]) – The RPC configuration to connect to the remote device. If none, will use local mode.
sort_by (Optional[str]) – Sort results by this key, if None, no sorting.
desc (Optional[bool]) – Whether to sort results in descending order.

tvm.s_tir.dlight.benchmark.benchmark_relax_func(mod: ~tvm.ir.module.IRModule, relax_func: ~tvm.ir.expr.GlobalVar | str, sample_number: int = 2, dym_var_sample_func: ~collections.abc.Callable[[dict[str, str]], dict[str, int]] = <function default_dym_var_sample_func>, target: str | dict | ~tvm.target.target.Target = None, evaluator_config: EvaluatorConfig | None = None, rpc_config: RPCConfig | None = None) → None

Benchmark a relax function with dynamic input shapes.

Parameters:

mod (tvm.ir.IRModule) – The IRModule to be benchmarked.
relax_func (Union[tvm.ir.GlobalVar, str]) – The relax function to be benchmarked.
sample_number (int) – The number of times to sample dynamic shape variables.
dym_var_sample_func (Callable[[Dict[str, str]], Dict[str, int]]) – The function to sample dynamic shape variables.
target (Union[str, tvm.target.Target]) – The target to be benchmarked on.
dev (tvm.runtime.Device) – The device to be benchmarked on.
evaluator_config (Optional["EvaluatorConfig"]) – The evaluator configuration to use. If none, will use default evaluator configuration.
rpc_config (Optional["RPCConfig"]) – The RPC configuration to connect to the remote device.

tvm.s_tir.dlight.benchmark.extract_prim_func(model_name: str, relax_func_name: str, prim_func_name: str, func: PrimFunc, *, func_args: list[tuple[tuple[Call | int, ...], str]] | None = None, dym_var_dict: dict[str, str] | None = None, weight: int = 1, sample_number: int = 5, target: str | dict | Target | None = None) → str

Extract a self-contained PrimFunc test file from a Relax module.

Parameters:

model_name (str) – The name of the model.
relax_func_name (str) – The name of the Relax function.
prim_func_name (str) – The name of the prim function.
func (tvm.tirx.PrimFunc) – The PrimFunc to be extracted.
func_args (Optional[List[Tuple[Tuple[Union[tvm.relax.expr.Call, int], ...], str]]]) – The arguments of the prim function, including both static and dynamic shape arguments. Given in format [ …, ((1, n, 128), “float32”), … ]. If not given, the arguments will be extracted from the PrimFunc.
dym_var_dict (Optional[Dict[str, str]]) – The dictionary of dynamic shape variables. Given in format {“n”: “int32”, “m”: “int32”}. If not given, the dictionary will be extracted from the PrimFunc.
weight (int) – The weight of the prim function, by default 1.
sample_number (int) – The number of times to sample dynamic shape variables, by default 5.
target (Optional[Union[str, dict, tvm.target.Target]]) – The target device to run the PrimFunc. If None, will use target from the context.

Returns:

result – The extracted PrimFunc test file content.

Return type:

str

tvm.s_tir.dlight.benchmark.extract_from_relax(mod: IRModule, model_name: str, file_path: str, target: str | dict | Target | None = None) → None

Extract self-contained PrimFunc test files from a Relax module.

Parameters:

mod (tvm.ir.IRModule) – The Relax module to be extracted.
model_name (str) – The name of the model.
file_path (str) – The path to store the extracted files.
target (Optional[Union[str, tvm.target.Target]]) – The target device to run the PrimFunc. If None, will use target from the context.

tvm.s_tir.dlight.benchmark.extract_func_info_from_prim_func(func: PrimFunc) → tuple[list[tuple[tuple[Var | int, ...], str]], dict[str, str]]

Extract function input information from a PrimFunc.

Parameters:

func (tvm.tirx.PrimFunc) – The PrimFunc to be analyzed.

Returns:

result (Tuple[) – List[Tuple[Tuple[Union[tvm.tirx.Var, int], …], str]], Dict[str, str],
] – The function input information and dynamic shape variable dictionary.

tvm.s_tir.dlight.benchmark.extract_all_func_info_from_relax(mod: IRModule) → tuple[dict[GlobalVar, dict[GlobalVar, list[tuple[list, int]]]], dict[GlobalVar, dict[str, str]]]

Extract function input information from a relax module.

Parameters:

mod (tvm.ir.IRModule) – The Relax module to be analyzed.

Returns:

result (Tuple[) – Dict[tvm.ir.GlobalVar, Dict[tvm.ir.GlobalVar, List[Tuple[List, int]]]], Dict[tvm.ir.GlobalVar, Dict[str, str]],
] – The function input information and dynamic shape variable dictionary.