mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Lei Wang
85fd8fc2d3
* [Refactor][Language] Make T.Kernel target-neutral with dialect-owned launch annotations T.Kernel is traced before the Target is known, yet it used to emit threadIdx.* thread_binding loops with a hard-coded default of 128 threads. That is a SIMT launch baked into the frontend: CPU lowering silently pinned the thread index to 0 (so per-thread code computed 1/128 of the work), and every backend inherited CUDA's launch vocabulary. Split the launch into the part every target shares and the parts a backend interprets: - The frontend records the grid as blockIdx.* thread_binding loops (the program-index space all targets share) and reserves thread identities as `tx/ty/tz = tl.launch_thread_idx(axis)` Bind placeholders. `threads=` becomes the `tl.launch_threads` annotation on the tilelang_root block; the frontend no longer chooses a default. - MaterializeKernelLaunch(lower_thread_binding, default_threads, unsupported_annotations) is where the target-dependent meaning is decided. SIMT backends rebind each placeholder as a threadIdx.* thread_extent over the same Var (extent from the annotation, else the backend default; the CUDA/ROCm/Metal/WebGPU pipelines pass 128). Backends without SIMT drop the placeholders and reject a body that references a thread index instead of miscomputing. Launch annotations a backend cannot honor (cluster_dims on cpu/rocm/metal/webgpu) are rejected here rather than ignored downstream. - KernelLaunchFrame exposes grid_vars/grid_extents/thread_vars/ thread_extents explicitly; get_thread_extent() raises at trace time when threads= was omitted, since only the backend knows the extent. - Each language dialect owns its `Kernel` with explicit keyword parameters: cuda offers threads/prelude/cluster_dims, rocm/metal threads/prelude, webgpu threads, cpu prelude, and the common `Kernel(*grid)` none. All of them funnel into `tilelang.language.kernel.launch_kernel`. This is what editors show on hover and what type checkers verify; unknown keywords are rejected by Python itself. ClusterKernel keeps its interface. - Cluster identity: `T.get_cluster_id/ids/dims/size/extent(s)` are program-space arithmetic (`bx // cluster_dims`) computed at trace time and therefore target-neutral, while `T.block_rank_in_cluster()` stays the CUDA hardware intrinsic. The two agree under the cluster-aware threadblock swizzle, which permutes the grid at cluster granularity. Tests that applied passes directly to a traced function now materialize the launch first (thread indices otherwise carry no extent), fixtures that wrote per-thread code and compiled it for the `c` target are rewritten at the tile level, and the span-preservation check asserts that no statement loses its span rather than that the spanned count never decreases. * [Fix][Language] Unwind kernel frames after tracing errors Pass the active exception to the pending frame exits when tracing fails, so the Python launch state (KernelLaunchFrame stack) is cleared without finalizing incomplete TIR. Previously the next trace observed the failed kernel's launch frame as KernelLaunchFrame.Current(). Covered for both the prim_func and eager JIT tracing paths. * [Fix][Language] Recognize every dialect's Kernel as a launch in the eager rewriter Eager JIT phase 1 infers the signature with symbolic T.const values and must not execute the kernel body; the AST rewriter guards `with T.Kernel(...)` with `if __tb.skip_kernel_ctx(): return`. It found the launch by object identity against the common `Kernel`, which no longer holds now that each dialect (and the `tilelang.language` facade) exports its own `Kernel`: the guard was not injected, phase 1 ran the body, and e.g. `T.gemm` on a tile whose extent comes from `T.const` failed its static-shape check. Mark every launch factory (`Kernel` of each dialect, `ClusterKernel`) with `kernel_launch_factory` and let the rewriter test that mark instead. It now resolves plain names and attribute chains only, so aliases such as `K = T.Kernel` are recognized and no factory expression is executed at rewrite time. Covered by a rewriter-level test over all dialects and an end-to-end eager gemm whose tile extent depends on a const.