cann/cannbot-skills Kernel定义内存分配

Kernel 定义、内存分配与数据搬运【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills1. Kernel 定义与启动T.prim_func定义一个 TileLang kernel 函数。参数类型为T.Tensor或T.Buffer。T.prim_func def add_kernel( A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype), C: T.Tensor((M, N), dtype), ): ...支持的 dtypefloat16, float32, bfloat16, int8, int16, int32, int64, uint8, uint16, uint32, uint64动态 shape 符号T.dyn[...]通过 buffer 的 shape 属性获取动态维度K T.dyn[K] T.prim_func def foo(A: T.Tensor((K,), float32)): N A.shape[0] for i in T.serial(N): ...T.dynamic(name, dtype)创建可直接使用的 tir.VarK T.dynamic(K, int32) T.prim_func def bar(A: T.Tensor((K,), float32)): for i in T.serial(K): ...T.Kernel定义 kernel 运行上下文创建 tile block 与逻辑核的绑定。with T.Kernel(m_num * n_num, is_npuTrue) as (cid, vid): bx cid // n_num by cid % n_num ...cid计算任务 ID范围 [0, block_num)vidVector 单元索引0 或 1A2/A3 架构 CV 核配比可为 1:2 或 1:1VEC_NUM通常设为 2表示每个 AI Core 有 2 个 Vector 计算单元jit 装饰器触发即时编译将 kernel 编译为 NPU 可执行代码。jit(out_idx[-1], pass_configspass_configs) def tile_add(M, N, block_M, block_N, dtypefloat): T.prim_func def main(...): ... return main参数out_idx指定输出参数索引如[-1]表示最后一个参数为输出workspace_idx工作空间参数索引如 Flash Attention 中workspace_idx[4,5,6]pass_configs编译配置选项常用 pass_configspass_configs { tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True, # 自动同步插入 tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True, # 自动内存规划 tilelang.PassConfigKey.TL_ASCEND_AUTO_CV_COMBINE: True, # 自动CV分离核间流水线需要 }查看生成的 AscendC 代码func tile_add(M, N, block_M, block_N) print(f{func.get_kernel_source()})2. 内存分配原语Developer 模式TileLang 对存储层级进行了抽象分为 Global、shared 和 fragment 三个级别。在 Ascend 平台中shared 层级对应 L1 Buffer 和 Unified Bufferfragment 层级对应 L0A/L0B/L0C Buffer。用户无需指定具体硬件存储TileLang 编译器会根据程序上下文自动识别。T.alloc_shared(shape, dtype)分配 shared 层级的存储空间。A_L1 T.alloc_shared((block_M, block_K), dtype)T.alloc_fragment(shape, dtype)分配 fragment 层级的存储空间。C_L0 T.alloc_fragment((block_M, block_N), accum_dtype)T.alloc_var(dtype, init, scopelocal.var)分配标量变量支持初始化。适用于标志位、计数器、临时标量。flag T.alloc_var(bool, initFalse) counter T.alloc_var(int32, init1) b T.alloc_var(int32, inita) # 用另一个变量的值初始化Expert 模式显式指定存储位置适用于需要精确控制内存分配的场景。API存储层级说明T.alloc_ub(shape, dtype)Unified Buffer (UB)Vector 计算T.alloc_L1(shape, dtype)L1 Buffer片上缓存T.alloc_L0A(shape, dtype)L0A BufferCube 左矩阵T.alloc_L0B(shape, dtype)L0B BufferCube 右矩阵T.alloc_L0C(shape, dtype)L0C BufferCube 输出/累加实际使用示例来自examples/gemm/example_gemm.pyA_L1 T.alloc_L1([block_M, block_K], dtype) B_L1 T.alloc_L1([block_K, block_N], dtype) C_L0 T.alloc_L0C([block_M, block_N], accum_dtype)3. 数据搬运原语T.copy(src, dst)在不同内存层级之间搬运 tile 数据块。支持 tir.Buffer、BufferLoad、BufferRegion 类型。支持的搬运路径srcdst说明GML1Global Memory → L1 BufferL1L0AL1 Buffer → L0A BufferCube 左矩阵L1L0BL1 Buffer → L0B BufferCube 右矩阵L0CGML0C Buffer → Global MemoryGMUBGlobal Memory → Unified BufferUBGMUnified Buffer → Global MemoryUBUBUnified Buffer → Unified BufferUBL1Unified Buffer → L1 Buffer使用示例# GM → L1 T.copy(A[bx * block_M, k * block_K], A_L1) # GM → UBvid 切分 T.copy(A[bx * block_M vid * block_M // VEC_NUM, by * block_N], a_ub) # UB → GM T.copy(c_ub, C[bx * block_M vid * block_M // VEC_NUM, by * block_N]) # L0C → GM T.copy(C_L0, C[bx * block_M, by * block_N]) # BufferRegion 切片搬运 T.copy(K[bz, by, k * block_N:(k 1) * block_N, :], k_l1)4. 完整示例来自docs/TileLang-Ascend Programming Guide.md§2.2import tilelang import tilelang.language as T from tilelang import jit import torch M, N 1024, 1024 block_M, block_N 128, 128 VEC_NUM 2 pass_configs { tilelang.PassConfigKey.TL_ASCEND_AUTO_SYNC: True, tilelang.PassConfigKey.TL_ASCEND_MEMORY_PLANNING: True, } jit(out_idx[-1], pass_configspass_configs) def tile_add(M: int, N: int, block_M: int, block_N: int, dtype: str float): m_num M // block_M n_num N // block_N T.prim_func def add_kernel( A: T.Tensor((M, N), dtype), B: T.Tensor((M, N), dtype), C: T.Tensor((M, N), dtype), ): with T.Kernel(m_num * n_num, is_npuTrue) as (cid, vid): bx cid // n_num by cid % n_num a_ub T.alloc_shared((block_M // VEC_NUM, block_N), dtype) b_ub T.alloc_shared((block_M // VEC_NUM, block_N), dtype) c_ub T.alloc_shared((block_M // VEC_NUM, block_N), dtype) T.copy(A[bx * block_M vid * block_M // VEC_NUM, by * block_N], a_ub) T.copy(B[bx * block_M vid * block_M // VEC_NUM, by * block_N], b_ub) for i, j in T.Parallel(block_M // VEC_NUM, block_N): c_ub[i, j] a_ub[i, j] b_ub[i, j] T.copy(c_ub, C[bx * block_M vid * block_M // VEC_NUM, by * block_N]) return add_kernel func tile_add(M, N, block_M, block_N) a torch.randn(M, N).npu() b torch.randn(M, N).npu() c func(a, b)【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考