tilelang.tileop.base¶

Classes¶

GemmWarpPolicy

Enumeration for GEMM Warp Partitioning Policies.

Module Contents¶

class tilelang.tileop.base.GemmWarpPolicy¶

Bases: enum.IntEnum

Enumeration for GEMM Warp Partitioning Policies.

Square = 0¶
FullRow = 1¶
FullCol = 2¶
is_square()¶

Check if the policy is a square partitioning.

Returns:

True if the policy is square, False otherwise.

Return type:

bool

is_full_row()¶

Check if the policy is a full row partitioning.

Returns:

True if the policy is full row, False otherwise.

Return type:

bool

is_full_col()¶

Check if the policy is a full column partitioning.

Returns:

True if the policy is full column, False otherwise.

Return type:

bool

static to_prime_factors(num)¶

Compute the prime factorization of a given number.

Parameters:

num (int) – The number to factorize.

Returns:

A list of prime factors of the number.

Return type:

list

compute_warp_partition(M, N, num_warps)¶

Compute the warp partition (m_warp, n_warp) based on the given policy.

Parameters:

  • M (int) – The number of rows in the GEMM workload.

  • N (int) – The number of columns in the GEMM workload.

  • num_warps (int) – The total number of warps available.

Returns:

A tuple (m_warp, n_warp) representing the partitioning of warps.

Return type:

tuple

Raises:

  • ValueError – If the policy is invalid or the partitioning fails.

  • AssertionError – If M or N is not divisible by the required factor for FullRow or FullCol policies.

classmethod from_warp_partition(m_warp, n_warp)¶

Determine the warp policy based on the given warp partitioning.

Parameters:

  • m_warp (int) – Number of warps in the row dimension

  • n_warp (int) – Number of warps in the column dimension

Returns:

The corresponding warp policy

Return type:

GemmWarpPolicy

Examples

>>> GemmWarpPolicy.from_block_row_cols(4, 1)  # All warps in rows
GemmWarpPolicy.FullRow
>>> GemmWarpPolicy.from_block_row_cols(1, 4)  # All warps in columns
GemmWarpPolicy.FullCol
>>> GemmWarpPolicy.from_block_row_cols(2, 2)  # Balanced distribution
GemmWarpPolicy.Square