tilelang.rocm.language.kernel¶

ROCm dialect of T.Kernel: the common launch plus ROCm launch annotations.

Functions¶

Kernel(*blocks[, threads, prelude])

Construct a kernel launch frame for ROCm: a grid of workgroups.

Module Contents¶

tilelang.rocm.language.kernel.Kernel(*blocks, threads=None, prelude=None)¶

Construct a kernel launch frame for ROCm: a grid of workgroups.

Code inside the launch operates at the workgroup level: T.Parallel, T.copy and friends are mapped onto threads by the compiler. T.get_thread_binding() exposes the thread index for thread-level code. The keyword arguments are recorded at trace time and materialized by the ROCm pipeline once the target is known.

Parameters:

  • *blocks (int | PrimExpr) – Grid extent along each axis (1-3 dimensions). The launch yields one workgroup index per axis.

  • threads (int | list[int] | tuple[int, ...], optional) – Threads per workgroup: a count or up to three per-dimension extents. Defaults to 128 when omitted.

  • prelude (str, optional) – Source injected before the generated kernel, e.g. #include lines or helper functions.

Return type:

tilelang.language.kernel.KernelLaunchFrame

Examples

with T.Kernel(T.ceildiv(N, 128), threads=128) as bx:
    ...