Abstract:

Tilus is a novel domain-specific language (DSL) for GPU programming, designed to simplify the development of high-performance kernels with fine-grained control over hardware resources. It operates at the thread-block level, treating tensors as primary data types, and offers explicit management of shared memory and registers, alongside support for arbitrary low-precision integer and floating-point types (1-8 bits). The language features automatic tuning, caching of compiled kernels, and a Pythonic interface for ease of use. Built upon the Hidet IR and runtime system, Tilus incorporates advanced optimization techniques like asynchronous memory copies, software pipelining, and split-K aggregation, demonstrating performance comparable to vendor libraries for tasks like matrix multiplication and flash attention. Contributions to the project require signed commits adhering to the Developer Certificate of Origin.

Summarizing Tilus: A GPU Kernel Programming Language

• Introduction to Tilus (README.md): Tilus is a powerful domain-specific language (DSL) for GPU programming, emphasizing thread-block-level granularity with tensors. It offers explicit control over shared memory and registers, supports low-precision types (1-8 bits), and includes automatic tuning, caching, and a Pythonic interface.
• Installation (docs/source/getting-started/install.rst): Tilus can be easily installed via pip install tilus or from source for development purposes.
• Programming Model - Tilus Script (docs/source/programming-guides/tilus-script.rst): Kernels are defined by subclassing tilus.Script, implementing __init__ for compilation-time setup (e.g., hyperparameters) and __call__ for the core kernel logic. Key attributes like self.attrs.blocks (number of thread blocks) and self.attrs.warps (warps per block) must be specified.
• Just-In-Time (JIT) Compilation (docs/source/programming-guides/tilus-script.rst): Tilus scripts are JIT-compiled on their first call. Parameters annotated with Python native types (int, float, bool) are treated as compile-time constants, triggering new compilations for different values. Parameters with Tilus-specific types (int32, float16) are runtime-known.
• Type System (docs/source/programming-guides/type-system/init.rst): Tilus supports Scalar, Pointer, and Tensor types (Global, Shared, Register). It includes a wide range of data types, from standard float32, int32 to arbitrary bit-width integers (e.g., int1b to int7b) and sub-byte floating-point types (e.g., float3_e1m1 to float8_e5m2).
• Memory Management & Tensors (docs/source/programming-guides/type-system/*.rst):
  • Global Tensors: Stored in GPU global memory, defined with global_view from pointers or global_tensor for kernel-managed allocation.
  • Shared Tensors: Stored in GPU shared memory, explicitly allocated with shared_tensor and freed with free_shared. All threads in a block can access shared memory, with performance considerations for access patterns (e.g., bank conflicts).
  • Register Tensors: Stored in GPU registers, distributed among threads within a block. Defined with register_tensor.
• Layout System (docs/source/programming-guides/layout-system/init.rst): Tilus provides flexible layout definitions for global, shared, and register tensors, controlling how elements are stored in memory. It supports automatic layout inference and explicit control using functions like spatial, local, compose, and reduce. Register layouts are "distributed layouts" mapping (thread_id, local_id) to tensor indices.
• Instructions (docs/source/programming-guides/instructions.rst): Tilus offers block-level instructions for tensor operations, including:
  • Creation & Free: global_view, register_tensor, shared_tensor, free_shared.
  • Load & Store: load_global, store_global, load_shared, store_shared.
  • Asynchronous Copy: copy_async, copy_async_commit_group, copy_async_wait_group, copy_async_wait_all (for overlapping computation and memory transfers).
  • Linear Algebra: dot (for matrix multiplication).
  • Element-wise: add, exp, maximum, where, etc.
  • Reduction: max, min, sum.
  • Transform: assign, cast, view, transpose, squeeze, unsqueeze, repeat, repeat_interleave.
  • Atomic & Semaphore: lock_semaphore, release_semaphore.
  • Synchronization: sync.
  • Miscellaneous: assume, annotate_layout, print_tensor, printf.
• Control Flow (docs/source/programming-guides/control-flow.rst): Standard control flow statements are supported: if/elif/else, for (with optional unroll parameter for loop unrolling), while, break, and continue.
• Autotuning (docs/source/programming-guides/autotuning.rst): The @tilus.autotune decorator allows defining search spaces for hyperparameters (e.g., warps, tile sizes, optimization knobs). Tilus automatically explores these configurations to find the best performing kernel for a given input.
• Caching (docs/source/programming-guides/cache.rst): Tilus caches compilation results in ~/.cache/tilus (or a custom directory) to speed up subsequent calls. Users can also dump intermediate representations (IRs) for debugging.
• Optimization Examples (examples/matmul/, examples/attention/, examples/quantization/): The codebase includes practical examples demonstrating various GPU optimization techniques:
  • Shared Memory (matmul_v1.py): Utilizing fast shared memory for data exchange within thread blocks.
  • Asynchronous Copy (matmul_v3.py): Overlapping memory transfers and computation using hardware-accelerated asynchronous copies.
  • Software Pipelining (matmul_v4.py): Structuring loops to hide memory latency by overlapping fetching data for the next iteration with processing data from the current one.
  • Split-K (matmul_v5.py): Distributing the K-dimension of matrix multiplication across multiple thread blocks and using semaphores for aggregation.
  • Flash Attention (attention/flash_attention_v.py):* Implementing an efficient attention mechanism.
  • Quantized Matrix Multiplication (quantization/matmul_a16wx.py): Support for quantized weights using low-precision data types and specialized kernels.
• Build Process (python/tilus/drivers.py): The build_program function orchestrates program verification, optimization (Tilus IR and Hidet IR), CUDA C code generation, and final compilation into a shared library (.so).
• Contribution Guidelines (CONTRIBUTING.md): All contributions require signed commits, certifying compliance with the Developer Certificate of Origin (DCO) to ensure integrity and authenticity.
• Licensing (LICENSE): Tilus is released under the Apache License, Version 2.0.