tilelang.carver.roller.policy.default

Policy for cuda core schedule

Classes

DefaultPolicy

Default Policy for fastdlight, a heuristic plan that tries to

Module Contents

class tilelang.carver.roller.policy.default.DefaultPolicy(arch, tags=None)

Default Policy for fastdlight, a heuristic plan that tries to minimize memory traffic and maximize parallelism.for BitBLAS Schedule.

参数:

  • arch (tilelang.carver.arch.TileDevice)

  • tags (dict | None)

func: tvm.tir.PrimFunc
nodes: list[tilelang.carver.roller.node.PrimFuncNode] = []
arch: tilelang.carver.arch.TileDevice
tags: dict
rasterization
classmethod from_prim_func(func, arch, tags=None, name='PrimFuncNode')
参数:

  • func (tvm.tir.PrimFunc)

  • arch (tilelang.carver.arch.TileDevice)

  • tags (dict | None)

  • name (str)

classmethod from_output_nodes(nodes, arch, tags=None)
参数:
emit_config(topk)
参数:

topk (int)

返回类型:

list[tilelang.carver.roller.hint.Hint]

dfs_smem_tile(init_tile, rstep_map)
返回类型:

collections.abc.Iterable[tilelang.carver.roller.hint.TileDict]

get_base_tile()

Gets the minimum tile configuration that satisfies no redundancy in computation.

返回:

The base tile configuration, which is a list of 1s equal in length to the space dimensions of the primary function node.

返回类型:

List[int]

compute_workload_per_item(output_tile)
返回类型:

float

score_block_size(n)

Scores a block size based on its efficiency and fit relative to the architecture's warp size and SM partition.

参数:

n (int) -- The block size to score.

返回:

A tuple containing two scores representing efficiency and fit, respectively.

返回类型:

Tuple[float, float]

get_block_size(n)

Determines the optimal block size for a given constraint, based on scoring various factors.

参数:

n (int) -- The constraint size.

返回:

The optimal block size chosen from the factors of n, constrained by a maximum of 1024 and scored by the score_block_size method.

返回类型:

int

get_node_reduce_step_candidates(node)

Calculates reduction step candidates for each reduction axis in a PrimFuncNode. General idea : use factor first, since it does not require extra boundary check. for large prime number, which is rare case, use power of 2.

参数:

node (PrimFuncNode) -- The node for which to calculate reduction step candidates. It contains reduction axes (raxis) with their domains (dom.extent).

返回:

A dictionary mapping axis variable names to lists of step candidates. For each axis in the node, this function calculates possible step sizes. For axes with a large prime domain, it uses powers of 2 as step candidates; for others, it uses all factors of the domain.

返回类型:

Dict[str, List[int]]

infer_node_smem_usage(td, node)

Infers the shared memory usage of a node given a TileDict configuration.

参数:

  • td (TileDict) -- The TileDict object containing the tile configuration.

  • node (PrimFuncNode) -- The node for which to infer the shared memory usage.

返回:

The estimated amount of shared memory used by the node.

返回类型:

int

compute_node_stride_map(node, td)

Computes the stride map for a given node based on the TileDict configuration.

参数:

  • node (PrimFuncNode) -- The node for which to compute the stride map.

  • td (TileDict) -- The TileDict object containing the tile configuration.

返回:

A tuple of dictionaries containing the output strides and tensor strides.

返回类型:

Tuple[Dict, Dict]

compute_tile_dict(output_tile, rstep_map)

Computes and returns a TileDict object for a given output tile configuration and reduction step map.

参数:

  • output_tile (List[int]) -- The output tile configuration.

  • rstep_map (Dict) -- The reduction step map.

返回:

A TileDict object containing the computed tile configuration, memory traffic, shared memory cost, grid size, and other related parameters.

返回类型:

TileDict

check_tile_shape_isvalid(td)

Checks if the tile shapes in the TileDict are valid for the nodes in this context.

Parameters: - td (TileDict): The TileDict object containing tile shapes and other configurations.

Returns: - bool: True if all tile shapes are valid, False otherwise.

参数:

td (tilelang.carver.roller.hint.TileDict)

返回类型:

bool

recommend_block_size(td)

Recommends optimal block sizes based on the TileDict configuration.

参数:

td (TileDict) -- The TileDict object containing the tile configuration.

返回:

A list of recommended block sizes sorted based on their score.

返回类型:

List[int]

assign_block_size(td, topk=1)

Assigns block sizes to the TileDict based on the recommended block sizes.

参数:

  • td (TileDict) -- The TileDict object to assign block sizes to.

  • topk (int, optional) -- The number of top block sizes to consider.

生成器:

Dict -- The block size assignment for the primary function node.

plan_rasterization(td)

Plans the rasterization for the given TileDict. This function is not implemented yet.

参数:

td (TileDict) -- The TileDict object to plan rasterization for.

抛出:

RasterRationPlan -- This function is not implemented yet.