跳转至

TileLang 昇腾版（tilelang-ascend）

约 713 个字 56 行代码 预计阅读时间 3 分钟

TileLang 是面向高性能算子开发的 Python DSL（基于 TVM 编译栈），用接近数学表达的方式描述计算，同时保留对底层内存层级与流水线的精细控制。tilelang-ascend 是华为深度参与的昇腾适配版，编译后端走 AscendNPU IR / Ascend C 路线，可以视为“介于 Triton 与 Ascend C 之间”的又一选择。

主项目 tile-ai/tilelang 支持 CUDA、HIP 等后端；昇腾版本独立发布在 tilelang-ascend 仓库，接口与主项目保持同源，因此学一次 TileLang，GPU/NPU 两边通用。

环境要求与安装

前置条件：

  • CANN ≥ 8.3.RC1，torch_npu ≥ 2.6.0（需先激活 set_env.sh）
  • 已在 A2/A3 系列硬件上完成验证

三种安装方式任选其一：

# 方式一：wheel 包安装（推荐）
# 从 Releases 页下载与自己的 CANN 版本、Python 版本、CPU 架构匹配的 whl：
# https://github.com/tile-ai/tilelang-ascend/releases
export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest
pip install tilelang-*.whl

# 方式二：从源码构建 wheel
git clone --recursive https://github.com/tile-ai/tilelang-ascend.git
cd tilelang-ascend
export ASCEND_HOME_PATH=/usr/local/Ascend/ascend-toolkit/latest
./build_wheel_ascend.sh
pip install dist/tilelang-*.whl

# 方式三：源码编译安装
bash install_ascend.sh && source set_env.sh

跑一下官方示例验证环境：

cd examples/gemm && python example_gemm.py   # 输出 "Kernel Output Match!" 即成功

小试牛刀：读懂一个 GEMM kernel

TileLang 版本的矩阵乘长这样（节选自官方示例 examples/gemm/example_gemm.py，仅添加了中文注释）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import tilelang
import tilelang.language as T
import torch


@tilelang.jit(out_idx=[-1])
def matmul(M, N, K, block_M, block_N, K_L1, dtype="float16", accum_dtype="float"):
    m_num = M // block_M          # M 方向的 block 数
    n_num = N // block_N          # N 方向的 block 数

    @T.prim_func
    def main(A: T.Tensor((M, K), dtype),
             B: T.Tensor((K, N), dtype),
             C: T.Tensor((M, N), dtype)):
        with T.Kernel(m_num * n_num, is_npu=True) as (cid, _):
            bx = cid // n_num     # 本核负责的 block 行号
            by = cid % n_num      # 本核负责的 block 列号
            A_L1 = T.alloc_L1((block_M, K_L1), dtype)        # L1 片上缓冲
            B_L1 = T.alloc_L1((K_L1, block_N), dtype)
            C_L0 = T.alloc_L0C((block_M, block_N), accum_dtype)
            with T.Scope("C"):                               # Cube 核作用域
                loop_k = T.ceildiv(K, K_L1)
                for k in T.serial(loop_k):
                    T.copy(A[bx * block_M, k * K_L1], A_L1)  # GM -> L1
                    T.copy(B[k * K_L1, by * block_N], B_L1)
                    T.barrier_all()
                    T.gemm_v0(A_L1, B_L1, C_L0, init=(k == 0))  # Cube 矩阵乘累加
                    T.barrier_all()
                T.copy(C_L0, C[bx * block_M, by * block_N])  # L0C -> GM
    return main


torch.manual_seed(0)
a = torch.randn(1024, 1024).half().npu()
b = torch.randn(1024, 1024).half().npu()

kernel = matmul(1024, 1024, 1024, block_M=128, block_N=256, K_L1=64)
c = kernel(a, b)    # JIT 编译并执行
torch.testing.assert_close(c, a @ b, rtol=1e-2, atol=1e-2)
print("Kernel Output Match!")

对照 Ascend C 一章，你会发现这里把概念都“内建”了：

TileLang 原语 对应的昇腾概念
T.alloc_L1 / T.alloc_UB / T.alloc_L0C L1 / UB / L0C 片上内存（GEMM 示例只用到 L1 与 L0C）
T.Scope("C") / T.Scope("V") 显式调度到 Cube / Vector 单元
T.copy DataCopy 搬运
T.gemm_v0 Cube 矩阵乘累加指令
T.Kernel(..., is_npu=True) 多 AICore SPMD 划分

对照代码可以注意到：gemm_v0 的操作数直接位于 L1，T.barrier_all() 负责搬运与计算之间的核内同步。新手可以先用“简化模式”（只写计算、内存层级全部交给编译器自动分配），等理解 tiling 后再切到上面的 Expert 模式手动管理 T.Scope 与缓冲区。编译器还提供软件流水 T.Pipelined、自动同步插入、自动内存复用等优化开关。

三个 DSL 怎么选

Ascend C Triton（triton-ascend） TileLang（昇腾版）
语言 C/C++ Python Python
控制粒度 最细，流水线手写 较粗，编译器自动 中等，可选 Expert 模式
适用场景 官方算子库、极致优化 融合算子快速开发 GEMM 类高性能 kernel 研究
学习曲线 陡 平缓 中等

入门顺序建议：Triton → TileLang → Ascend C，先会用编译器替你干活的工具，再往下抠硬件细节。

学习&拓展：

  1. 跑通 examples/gemm 的 GEMM 示例，然后修改 block_M/block_N/K_L1 做一轮手调，记录最优配置。
  2. 对比同一尺寸矩阵乘在 torch_npu、triton-ascend、tilelang 三种实现下的性能（这是超算队竞赛的日常）。
  3. 阅读昇腾社区关于 AscendNPU IR 的文章，理解 tilelang-ascend 的编译后端为什么选择引入一层中间表示。
  4. 尝试在 Expert 模式下把一段 T.Scope("C") 中的计算改到 T.Scope("V")（Vector 单元）完成，体会双核异构的调度方式。

一些链接

tilelang-ascend 官方仓库

TileLang-Ascend 编程指南

TileLang 主项目

昇腾社区：TileLang AscendNPU IR 知识点