Tile Language (tile-lang)

TileLang implements stub libraries to improve portability and compatibility:

CUDA Stubs

• cuda_stub (libstub_cuda.so): Allows importing TileLang on systems without a GPU (like CI nodes) by lazy-loading libcuda.so only when needed.
• cudart_stub (libstub_cudart.so) and nvrtc_stub (libstub_nvrtc.so): Resolves SONAME versioning mismatches (e.g., libcudart.so.11 vs libcudart.so.12). They attempt to reuse CUDA libraries already loaded by frameworks like PyTorch to ensure a single build works across different CUDA versions.

ROCm Stubs

• hip_stub (libstub_hip.so): Allows importing TileLang on systems without ROCm installed by lazy-loading libamdhip64.so. It uses RTLD_DEFAULT / RTLD_NEXT to interoperate with frameworks that have already loaded HIP symbols.
• hiprtc_stub (libstub_hiprtc.so): Lazily loads libhiprtc.so and exposes the minimal HIPRTC API subset used by TileLang/TVM.

To profile a kernel, wrap the compilation process in an iket.session(...). This ensures the session is active when tilelang.compile(...) generates CUDA source. You can use iket.range as a context manager and iket.mark for instant events within your T.prim_func.

import tilelang
import tilelang.language as T
from tilelang.tools.cuda import iket


def instrumented_add(n: int, threads: int = 128):
    @T.prim_func
    def main(
        A: T.Tensor((n,), T.float32),
        B: T.Tensor((n,), T.float32),
        C: T.Tensor((n,), T.float32),
    ):
        with T.Kernel(T.ceildiv(n, threads), threads=threads) as bx:
            with iket.range("block_total"):
                for tx in T.Parallel(threads):
                    i = bx * threads + tx
                    if i < n:
                        iket.mark("before_store")
                        C[i] = A[i] + B[i]
                        iket.mark("after_store")

    return main


with iket.session(output_dir="/tmp/tilelang_iket"):
    program = instrumented_add(1024)
    kernel = tilelang.compile(
        program,
        out_idx=-1,
        target="cuda",
        execution_backend="cython",
    )

To implement an elementwise operator, define a function that returns a T.prim_func. All kernel logic must reside within a T.Kernel(...) scope.

Key concepts:

• T.Kernel(grid_size, threads=threads): Defines the grid size and threads per block. The returned values (e.g., bx) correspond to CUDA's blockIdx.
• T.Parallel(...): Used inside the kernel to process the data tile assigned to the block.
• Block-level programming: Code inside T.Kernel operates at the block level. TileLang automatically handles the mapping to individual threads and applies optimizations during compilation.

def elementwise_add(N, threads=256, dtype=T.bfloat16):
    @T.prim_func
    def main(A: T.Tensor((N), dtype), B: T.Tensor((N), dtype), C: T.Tensor((N), dtype)):
        with T.Kernel(T.ceildiv(N, threads), threads=threads) as (b_x):
            # vector add.
            for i in T.Parallel(threads):
                C[b_x * threads + i] = A[b_x * threads + i] + B[b_x * threads + i]

    return main