tilelang-ascend

# TileLang Ascend NPU Programming Use this guide to write, port, tune, and debug kernels for Ascend NPUs. Import the Ascend dialect: ```python import tilelang.ascend.language as T ``` The explicit dialect exposes Ascend-specific APIs such as `T.MixedKernel` and `T.SimdVF`. For installation, follow the repository's current build documentation and [tilelang-build](../tilelang-build/SKILL.md). ## Task-specific references - [Scheduling and synchronization](references/synchronization.md): automatic, T.Stage performance tuning, buffer ownership and version modes, conditional execution, conflict hints, and fully manual protocols. - [Profiling recipes](references/profiling.md): device timing, regression tools, generated-source experiments, and operation cycle measurements with npusim. For a suspected compiler defect, preserve the smallest failing kernel together with its generated source or lowering trace. --- ## 1. Mental Model: Ascend vs CUDA | Concept | CUDA | Ascend NPU | |---|---|---| | Kernel launch | `T.Kernel(grid, threads=128)` | `T.Kernel(N)` -- **no threads arg** | | Launch grid | 1-D, 2-D, or 3-D | 1-D grids only | | Thread access | `T.get_thread_binding()` | No `threadIdx` directly; use `T.SimtVF(threads=N)` | | Compute cores | SM (unified) | **Two distinct cores**: AIC (Cube) + AIV (Vector) | | Matrix multiply | `T.gemm(...)` | Current L1-input path uses `transpose_B=True`; check the selected lowering | | Shared memory | `T.alloc_shared(...)` ("shared") | `T.alloc_shared(...)` maps to UB (Unified Buffer) | | Global memory | Global | GM (Global Memory / HBM) | | L1 cache / data | N/A | `T.alloc_l1(...)` -- Cube Buffer (CBuf) | | Register fragments | `T.alloc_fragment(...)` | `T.alloc_fragment(...)` | | Cube engine buffers | N/A | `T.alloc_l0a/l0b/l0c(...)` -- L0A/L0B/L0C | | Warp-level ops | `__shfl_sync` | `T.alloc_reducer` / `T.warp_reduce_sum` | | Synchronization | `__syncthreads()`, `mbarrier`, `cp.async` | `asc_syncthreads()` (auto-inserted in SimtVF), Flag-based: `T.ascend_set_flag` / `T.ascend_wait_flag` | | Auto-Schedule | `T.Pipelined(...)` | `T.Pipelined(...)` with `num_stages=...` | | Compiler | nvcc | Bisheng (Huawei Ascend compiler) | --- ## 2. Thread Hierarchy ### 2.1 Kernel Launch Ascend supports **only 1-D grids** and does **not** accept a `threads=` parameter on `T.Kernel`. Thread domains are defined separately. ```python # CORRECT - Ascend style: with T.Kernel(NUM_CORES) as bx: ... # WRONG - will raise TypeError: with T.Kernel(NUM_CORES, threads=128) as bx: ... ``` ### 2.2 SIMT VF: Thread-Parallel Compute Use `T.SimtVF(threads=N)` to define a thread-parallel region. `threads` also accepts a tuple/list of up to three dimensions, for example `(32, 4)`; omitted dimensions default to one. Inside this block, `T.Parallel` loops distribute work across the declared thread domain. ```python with T.SimtVF(threads=128): for i in T.Parallel(N): out[i] = a[i] + b[i] ``` Auto-sync: The compiler inserts `asc_syncthreads()` when it detects cross-thread data hazards; regression coverage is in `testing/ascend/transform/test_ascend_thread_sync.py`. Thread binding is accessed via `T.get_thread_binding()`. For random values inside a VF, use `T.rng_init(seed, seq=..., off=...)` with `T.rng_rand()` or `T.rng_rand_float()` (uniform) / `T.rng_rand_float(dist="normal")`. Choose explicit sequences when reproducibility must be independent of launch geometry. ### 2.3 SIMD VF: Register-Level Vector (No Threads) `T.SimdVF()` maps to CCE MicroAPI vector instructions (2048-bit vectors, 64 x float32 or 128 x float16). There is **no thread concept** -- the code runs as SIMD vector instructions directly on the vector unit. ```python with T.SimdVF(): mask = T.simd.pset(32) # All lanes active for float32 elements. for i in range(N // 64): r0 = T.simd.vld(buf[i * 64]) r1 = T.simd.vadd(r0, r0, mask) T.simd.vsts(out[i * 64], r1, mask) ``` ### 2.4 Task Boundaries `T.Task()` groups serial statements into one scheduling task. All statements must use one hardware PIPE and one AIC/AIV core affinity; the marker does not select a core or insert synchronization inside its body. Split work requiring an internal cross-PIPE handoff into separate tasks. Optional `latency=` and `ii=` values override the cost of the complete task: latency is the output-ready delay and II is the minimum issue interval on its PIPE. The final pair must satisfy `latency >= ii > 0`; omitted values are estimated. These are distinct from the enclosing loop's solved II. ```python with T.Task(): # Two serial GM-to-L1 copies on MTE2. T.copy(A_tile, a_l1) T.copy(B_tile, b_l1) ``` ### 2.5 Per-Core Task Contract `T.PerCoreTask()` describes one logical operation whose concrete execution site is selected by the kernel block index. It is not a general-purpose conditional scope. Every dependency-bearing basic operation inside it is an explicit `T.Task()` or a statement that the compiler can infer as one. For every dynamic invocation and every participating `blockIdx.x` value, the candidate guards for that logical operation must be mutually exclusive and collectively exhaustive: **exactly one candidate `T.Task()` must execute on each core**. Conditions may dispatch work by block/core ID, but must not make the operation data-dependent, optional, or executable more than once on one core. All candidates must use the same single hardware pipe and represent the same logical read/write role. Put the complete block-index dispatch inside one `T.PerCoreTask()` so the scheduler can treat all branches as alternative sites of the same operation: ```python # Correct: every block executes exactly one candidate of one logical task. with T.PerCoreTask(): if bx == 0: with T.Task(): T.copy(A, tile) else: with T.Task(): T.copy(B, tile) ``` Do not split the alternatives into independently guarded `T.PerCoreTask` regions, and do not omit a block-index case: ```python # Wrong: these are two unrelated logical tasks, not two candidates of one task. if bx == 0: with T.PerCoreTask(): T.copy(A, tile) if bx != 0: with T.PerCoreTask(): T.copy(B, tile) # Wrong: blocks for which `runtime_condition` is false execute no candidate. with T.PerCoreTask(): if runtime_condition: with T.Task(): T.copy(A, tile) ``` AutoSchedule inserts the logical task's synchronization beside every candidate `T.Task()` while preserving its block-index branch. Correctness therefore depends on the exact-once contract above; the compiler validates structure, pipe compatibility, and guard relationships, but cannot generally prove that arbitrary runtime predicates form an exact partition. --- ## 3. Core Architecture: AIC vs AIV The Ascend NPU has two distinct compute cores on each die: | Core | Engine | Purpose | TileLang Scope | |---|---|---|---| | AIC | Cube unit | Matrix multiply (GEMM, convolutions) | `T.Cube()` | | AIV | Vector unit | Element-wise ops, reductions, data movement | `T.Vector()` | ### 3.1 Pure Kernel (Single-Scope) Use `T.Kernel(N)` without explicit `T.Cube()`/`T.Vector()` scopes. The codegen auto-detects the kernel type based on the operations inside: Cube ops (`T.gemm`) produce a `__cube__` kernel on AIC, while Vector ops produce a `__vector__` kernel on AIV. ```python with T.Kernel(NUM_BLOCKS) as bx: x_l1 = T.alloc_l1((TILE_M, TILE_K), dtype) w_l1 = T.alloc_l1((TILE_N, TILE_K), dtype) res = T.alloc_l0c((TILE_M, TILE_N), accum_dtype) # ... T.copy into L1, T.gemm, T.copy out ... ``` ### 3.2 MixedKernel: Sub-Core Access and Selection Use `T.MixedKernel(NUM_BLOCKS, sids=2)` when the program needs the AIV sub-core ID returned with `bx`, or explicitly controls the active AIV count (`sids=1` or `2`). If neither is needed, use `T.Kernel`; the compiler handles core splitting and the appropriate `T.dual_copy` lowering. ```python with T.MixedKernel(NUM_BLOCKS, sids=2) as (bx, sid): ... # Partition each AIV's accesses according to sid and buffer ownership. ``` See the mixed-core epilogues in `examples/ascend/deepgemm/kernels/epilogue.py`. Follow the actual buffer extents when partitioning output; a local half-sized buffer is not a full-sized buffer to split again. ### 3.3 Fully Manual Mixed Kernels Explicit `T.Cube()` and `T.Vector()` scopes with manually managed buffer slots and intra-/cross-core flags belong to the fully manual path. Disable the scheduling path with `PassConfigKey.TL_ENABLE_AUTO_SCHEDULE: False`. Read the [synchronization reference](references/synchronization.md) for the complete initialization, reuse and final-drain protocol. Small manual mixed-core regressions are in `testing/ascend/language/test_tilelang_ascend_nd2nz_scatter.py`. --- ## 4. Memory Hierarchy and Allocation Ascend memory hierarchy (top to bottom): ``` GM (HBM) -- global memory, large capacity ↕ DMA (MTE pipes) UB (Unified Buffer, "shared") -- T.alloc_shared(...) ↕ DMA L1 (CBuf, Cube Buffer) -- T.alloc_l1(...) ↕ load L0A / L0B (Cube inputs) -- T.alloc_l0a/l0b(...) → Cube MAD unit L0C (accumulator) -- T.alloc_l0c(...) ``` ### 4.1 Allocation API ```python # Unified Buffer (UB) -- general-purpose on-chip memory buf = T.alloc_shared((STAGES, TILE_M, TILE_K), "bfloat16") # Cube Buffer (L1) -- fast buffer for Cube engine input buf = T.alloc_l1((TILE_M, TILE_K), "bfloat16") # Cube L0 buffers -- direct Cube engine inputs/outputs buf_a = T.alloc_l0a((STAGES, TILE_M, TILE_K_SUB), "bfloat16") buf_b = T.alloc_l0b((STAGES, TILE_N, TILE_K_SUB), "bfloat16") buf_c = T.alloc_l0c((TILE_M, TILE_N), "float32") # Fragment (register-level, inside SimtVF/SIMD) frag = T.alloc_fragment((TILE,), "float32") ``` For a deferred reduction inside `T.SimtVF`, use the reducer epoch API. Given float32 UB buffers `x_ub[TILE]` and `sum_ub[1]`: ```python with T.SimtVF(threads=128): acc = T.alloc_reducer((1,), "float32", op="sum") T.reducer_init(acc) for i in T.Parallel(TILE): T.reducer_update(acc[0], x_ub[i]) result = T.alloc_fragment((1,), "float32") T.finalize_reducer(acc, result) T.copy(result, sum_ub) ``` Read the independent destination fragment after finalization. The reducer handle itself cannot be read, assigned directly, cleared/filled with `T.clear` or `T.fill`, or finalized in place. Keep init and finalize in the same enclosing loop/branch scope. Ascend currently supports reducer `sum`, `max`, and `min`; bitwise collectives are unsupported. ### 4.2 Buffer Versioning Use `T.annotate_buffer_versions` to fix a version count, select `auto`, `iteration`, or `counter` indexing, or specify both. For example, `T.annotate_buffer_versions({buf: (2, "counter")})` fixes two counter-indexed versions. See [version semantics](references/synchronization.md#buffer-versions-and-offsets) for ownership, control-flow, and automatic buffer sizing. In particular, `{buf: 1}` opts out of multibuffer eligibility, while `{buf: (1, "auto")}` keeps eligibility with one physical version. Explicit `counter` mode cannot serve accesses to the same buffer at different stages within one owner. Prefer compiler-managed multibuffering. [Manual multibuffering](references/synchronization.md#manual-multibuffering-limited-support) has limited support and is not recommended: complex access or control-flow patterns may produce incorrect dependency analysis and synchronization. ### 4.3 Physical Capacity and Reuse Allocation shape, a copy's access region, and its physical layout are different quantities. Include layout padding, buffer versions, and compiler reservations when checking capacity. Layout normalization does not initialize padded lanes. Different Buffer views can alias, and separate logical buffers can share an allocation when their asynchronous lifetimes permit reuse. Follow the [allocation ownership rules](references/synchronization.md#allocation-ownership) when changing manual kernels. For capacity or launch failures, record the generated allocation sizes, buffer versions, VF mode, and compiler options. Check SDK reservations as well as user buffers before changing tile sizes. `T.annotate_unlimit_memory("shared", "shared.l1")` inside a kernel removes the scheduler's capacity constraint for those pools. Other supported scopes are `shared.l0a`, `shared.l0b`, and `shared.l0c`. This is useful for investigating a capacity-constrained schedule; it does not enlarge hardware memory or make an over-capacity allocation executable. Verify physical allocations before launch. --- ## 5. GEMM Constraints ### 5.1 Select the GEMM Lowering Path The current L1-input path in `tilelang/ascend/op/gemm/gemm_mad.py` requires `transpose_A=False` and `transpose_B=True`. L0-input GEMM has its own layout/transpose handling; do not generalize the L1 restriction to every Ascend GEMM form. When sharing GEMM macros across loop orders, recheck compact L0 layouts, buffer lifetimes, and stage counts; equivalent mathematics alone does not make the storage and scheduling contracts interchangeable. ```python T.gemm(a_l1, b_l1, c_l0c, transpose_B=True, clear_accum=(kt == 0)) ``` ### 5.2 Weight Layout Convention For this L1-input NT convention, weights are typically stored as `[N, K]` in global memory (transposed relative to standard `[K, N]`): ```python # Weight buffer shape: [N, K], not [K, N] W: T.Buffer((N_DIM, K_DIM), dtype) ```