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.

Parameters:

  • arch (tilelang.carver.arch.TileDevice)

  • tags (Optional[Dict])

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')¶
Parameters:

  • func (tvm.tir.PrimFunc)

  • arch (tilelang.carver.arch.TileDevice)

  • tags (Optional[Dict])

  • name (str)

classmethod from_output_nodes(nodes, arch, tags=None)¶
Parameters:
emit_config(topk)¶
Parameters:

topk (int)

Return type:

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

dfs_smem_tile(init_tile, rstep_map)¶
Return type:

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

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]

compute_workload_per_item(output_tile)¶
Return type:

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.

Parameters:

n (int) – The block size to score.

Returns:

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

Return type:

Tuple[float, float]

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)¶

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, node)¶

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

compute_node_stride_map(node, td)¶

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, rstep_map)¶

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

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.

Parameters:

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

Return type:

bool

recommend_block_size(td)¶

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]

assign_block_size(td, 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.

plan_rasterization(td)¶

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.