tilelang.language package#

Subpackages#

Submodules#

Module contents#

The language interface for tl programs.

tilelang.language.annotate_l2_hit_ratio(l2_hit_ratio_map: Dict)#

Annotate the L2 hit ratio of the buffer, detailed explanation please refer to: https://docs.nvidia.com/cuda/cuda-c-programming-guide/#l2-policy-for-persisting-accesses

Parameters:

l2_hit_ratio_map (dict) – a dictionary of buffer to L2 hit ratio value

Example

# 0.5 is the hit ratio T.annotate_l2_hit_ratio({A: 0.5})

tilelang.language.annotate_layout(layout_map: Dict)#

Annotate the layout of the buffer

Parameters:

layout_map (Dict) – a dictionary of buffer to layout

Returns:

a block attribute

Return type:

block_attr

Example

@T.prim_func def main(

A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype),

):

# Initialize Kernel Context with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):

A_shared = T.alloc_shared((block_M, block_N), dtype)

T.annotate_layout({A_shared: layout}) for i, j in T.Parallel(block_M, block_N):

A_shared[i, j] = A[by * block_M + i, bx * block_N + j]

for i, j in T.Parallel(block_M, block_N):

B[by * block_M + i, bx * block_N + j] = A_shared[i, j]

return main

tilelang.language.annotate_padding(padding_map: Dict)#

Annotate the padding of the buffer

Parameters:

padding_map (dict) – a dictionary of buffer to padding value

Returns:

a block attribute

Return type:

block_attr

Example

@T.prim_func def main(

A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype),

):

# Initialize Kernel Context with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):

A_shared = T.alloc_shared((block_M, block_N), dtype)

T.annotate_padding({A_shared: pad_value}) for i, j in T.Parallel(block_M, block_N):

A_shared[i, j] = A[by * block_M + i - 10, bx * block_N + j]

for i, j in T.Parallel(block_M, block_N):

B[by * block_M + i, bx * block_N + j] = A_shared[i, j]

return main

tilelang.language.import_source(source: Optional[str] = None)#
tilelang.language.symbolic(name: str, dtype: str = 'int32')#
tilelang.language.use_swizzle(panel_size: int, order: str = 'row', enable: boolean = True)#