tilelang.jit.adapter.utils¶

Functions¶

match_global_kernel(source[, annotation])

match_declare_kernel(source[, annotation])

match_declare_kernel_cpu(source[, annotation])

is_cuda_target(target)

is_hip_target(target)

is_cpu_target(target)

get_annotated_mod(func_or_mod[, target, target_host, ...])

pythonic_expr(expr[, dtype_map])

Converts a TVM PrimExpr into a Python-style string, correctly handling operator precedence.

Module Contents¶

tilelang.jit.adapter.utils.match_global_kernel(source, annotation='__global__')¶
Parameters:

  • source (str)

  • annotation (str)

Return type:

int

tilelang.jit.adapter.utils.match_declare_kernel(source, annotation='__global__')¶
Parameters:

  • source (str)

  • annotation (str)

Return type:

int

tilelang.jit.adapter.utils.match_declare_kernel_cpu(source, annotation='int32_t')¶
Parameters:

  • source (str)

  • annotation (str)

Return type:

int

tilelang.jit.adapter.utils.is_cuda_target(target)¶
Parameters:

target (tvm.target.Target)

Return type:

bool

tilelang.jit.adapter.utils.is_hip_target(target)¶
Parameters:

target (tvm.target.Target)

Return type:

bool

tilelang.jit.adapter.utils.is_cpu_target(target)¶
Parameters:

target (tvm.target.Target)

Return type:

bool

tilelang.jit.adapter.utils.get_annotated_mod(func_or_mod, target='auto', target_host=None, model_type='all')¶
Parameters:

  • func_or_mod (Union[tvm.tir.PrimFunc, tilelang.tvm.IRModule])

  • target (Union[str, tvm.target.Target])

  • target_host (Optional[Union[str, tvm.target.Target]])

  • model_type (Literal['device', 'host', 'all'])

Return type:

Union[tvm.IRModule, tuple[tvm.IRModule, tvm.IRModule]]

tilelang.jit.adapter.utils.pythonic_expr(expr, dtype_map=None)¶

Converts a TVM PrimExpr into a Python-style string, correctly handling operator precedence.

Parameters:

  • expr (tilelang.tvm.tir.PrimExpr) – The TVM PrimExpr to convert.

  • dtype_map (Optional[Dict[str, str]])

Returns:

A string representation of the expression.

Return type:

str