tilelang.carver.roller.policy.default module#

Policy for cuda core schedule

class tilelang.carver.roller.policy.default.DefaultPolicy(arch: TileDevice, tags: Optional[Dict] = None)#

Bases: object

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

arch: TileDevice#
assign_block_size(td: TileDict, topk=1)#

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

Parameters:

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

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

Yields:

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

check_tile_shape_isvalid(td: TileDict) → bool#

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.

compute_node_stride_map(node: PrimFuncNode, td: TileDict)#

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

Parameters:

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

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

Returns:

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

Return type:

Tuple[Dict, Dict]

compute_tile_dict(output_tile: List[int], rstep_map) → TileDict#

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

Parameters:

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

  • rstep_map (Dict) – The reduction step map.

Returns:

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

Return type:

TileDict

compute_workload_per_item(output_tile) → float#
dfs_smem_tile(init_tile, rstep_map) → Iterable[TileDict]#
emit_config(topk: int) → List[Hint]#
classmethod from_output_nodes(nodes: List[OutputNode], arch: TileDevice, tags: Optional[Dict] = None)#
classmethod from_prim_func(func: PrimFunc, arch: TileDevice, tags: Optional[Dict] = None, name: str = 'PrimFuncNode')#
func: PrimFunc#
get_base_tile()#

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

Returns:

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

Return type:

List[int]

get_block_size(n)#

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

Parameters:

n (int) – The constraint size.

Returns:

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

Return type:

int

get_node_reduce_step_candidates(node: PrimFuncNode)#

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.

Parameters:

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

Returns:

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.

Return type:

Dict[str, List[int]]

infer_node_smem_usage(td: TileDict, node: PrimFuncNode)#

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

Parameters:

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

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

Returns:

The estimated amount of shared memory used by the node.

Return type:

int

nodes: List[PrimFuncNode] = []#
plan_rasterization(td: TileDict)#

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

Parameters:

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

Raises:

RasterRationPlan – This function is not implemented yet.

recommend_block_size(td: TileDict) → List[int]#

Recommends optimal block sizes based on the TileDict configuration.

Parameters:

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

Returns:

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

Return type:

List[int]

score_block_size(n)#

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

Parameters:

n (int) – The block size to score.

Returns:

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

Return type:

Tuple[float, float]

tags: Dict#