人工智能编译器模型编译深度学习高性能计算CANNAscend【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载PyPTOPro 前端中算子 Kernel 需要在 Device 侧被多个 AI Core 并行执行而每个核的片上 Buffer 容量有限因此必须在 Host 侧规划好数据切分并把切分参数传给 Kernel。本文基于仓库中 Tiling 实现指南Tiling 实现 及其三篇子文档系统讲解为什么需要 Tiling、核间/核内切分的两个层次、多核切分的通用写法与 MIX Kernel 的切分方法以及TilingData与TilingKey两类运行时/编译期参数的声明、传递与校验规则并结合开源仓库源码印证其底层实现机制。读完后你能够独立完成一个算子的多核切分方案并正确地把 shape、stride、循环边界等参数从 Host 传入 Kernel。一、Tiling 概述两个硬件约束决定了切分的必要性TilingTiling 实现指南 Tiling概述指把算子的输入数据划分为基本任务块Tile将一组任务块分配给各个核再为核内处理配置 Buffer 的过程。它是算子开发的第一步切分方式决定了各核负载是否均衡、片上 Buffer 能否放下、数据搬运是否连续直接影响算子性能。Tiling 的必要性来自两个硬件约束Device 上有多个 AI Core。算子数据若只交给一个核处理其余核处于空闲多核算力无法发挥。必须把数据划分成基本任务块并将一组任务块分配给每个核并行处理才能把算力用起来。单个核的片上 Buffer 容量有限。计算只能在片上 Buffer 中进行而算子的输入输出通常远大于片上容量无法一次完整装入。基本任务块的大小需要适配片上容量核内可为已分配的任务划分双缓冲或多个 Buffer 槽位轮转执行搬入—计算—搬出让不同任务的搬运与计算尽量并行。因此 Tiling 包含两个层次层次含义详细文档核间切分把算子的数据划分为基本任务块将一组任务块分配给每个 AI Core 并行处理多核Tiling切分核内切分基于已分配的任务划分核内 Buffer例如使用双缓冲或多个 Buffer 槽位轮转让搬运与计算并行Tile创建和操作、Tile计算AI Core 与执行域Device 上有多个 AI Core每个 AI Core 内含 Cube 核AIC和 Vector 核AIV。在 AIC 与 AIV 为 1:2 的芯片上一个 AI Core 内有 1 个 Cube 核和 2 个 Vector 核Cube 核承担矩阵计算Vector 核承担矢量计算两者各有自己的片上 Buffer 和指令流水。Kernel 通过执行域声明代码运行在哪类核上pypto_pro.language.section_cube()中的代码由 Cube 核执行pypto_pro.language.section_vector()中的代码由 Vector 核执行。Kernel 包含哪些执行域决定了它使用哪类核也决定了核间切分时可用的核数。片上 Buffer 的层级、容量和数据通路参见 SIMD 抽象硬件架构。设计 Tiling 时的考虑因素Tiling 概述文档给出五条设计准则实际开发中可逐条对照检查负载均衡各个核分到的任务数和计算量应尽量接近避免个别核成为长尾拖慢整体。片上容量单次搬入的 Tile 必须能放进对应的片上 Buffer并为双缓冲等优化预留空间。接口约束Tile 的形状需要满足数据搬运和计算接口在对齐、layout 和分形上的要求。数据复用相邻任务如果复用同一块数据切分顺序应尽量让它们落在同一个核上。尾块处理数据量通常不能被 Tile 大小整除任务数需要向上取整最后一块用有效形状表达实际范围。二、多核 Tiling 切分任务编号与连续分核基本写法以矩阵乘为例输出被划分为m_tiles * n_tiles个 Tile每个 Tile 的计算相互独立。Tile 数按向上取整计算保证最后不足一个 Tile 的数据也被覆盖m_tiles (m tile_m - 1) // tile_m n_tiles (n tile_n - 1) // tile_n total_tiles m_tiles * n_tiles把这些 Tile 按行优先统一编号成total_tiles个任务每个核负责其中连续的一段再由任务号换算回二维下标import pypto_pro.language as pl core_id pl.get_block_idx() // pl.get_subblock_num() core_num pl.get_block_num() tiles_per_core (total_tiles core_num - 1) // core_num start core_id * tiles_per_core end pl.min(start tiles_per_core, total_tiles) for task_idx in pl.range(start, end, 1): tile_m_idx task_idx // n_tiles tile_n_idx task_idx % n_tiles # 计算第 (tile_m_idx, tile_n_idx) 个输出Tile ...这里有几个关键点来自 多核Tiling切分pypto_pro.language.get_block_num()API 文档返回本次实际生效的核数计算每个核的任务范围时应使用它的返回值而不是 Host 侧请求的block_dim。block_dim的默认值、调用形式和核数限制参见 Kernel核函数。core_id由pypto_pro.language.get_block_idx()除以pypto_pro.language.get_subblock_num()得到。纯 Vector 和纯 Cube Kernel 中get_subblock_num()返回 1该除法不改变结果MIX Kernel 中 Vector 侧的get_block_idx()是 Vector 核的全局索引除以每个 AI Core 内的 Vector 核数后与 Cube 侧得到同一个 AI Core 编号两侧才能按同一套切分处理同一批数据。这种分配不要求任务数被核数整除tiles_per_core向上取整末尾的核用pypto_pro.language.min()把上界限制在total_tiles剩余任务自然落在最后一个进入循环的核上不需要为除不尽单独编写分支。任务总数少于核数时编号超出范围的核得到空区间不会进入循环。算子的执行时间取决于任务最多的那个核为ceil(total_tiles / core_num)个任务的耗时。上面的编号方式是一种通用选择。算子完全可以按自身的数据复用和依赖关系定义别的切分方式例如输出 Tile 之间存在跨迭代累加时被累加的那一维必须完整留在同一个核内。MIX Kernel 的切分Device 上有多个 AI Core每个 AI Core 内含 Cube 核AIC和 Vector 核AIV在 AIC 与 AIV 为 1:2 的芯片上一个 AI Core 内有 1 个 Cube 核和 2 个 Vector 核。同时包含 Cube 和 Vector 执行域 的 Kernel 称为MIX Kernel。MIX Kernel 以 AI Core 为启动单位核间切分也只在 AI Core 这一层做一次get_block_num()在两侧都返回 AI Core 数core_id在两侧都取get_block_idx() // get_subblock_num()AIC 上get_subblock_num()返回 1AIV 上返回 2两侧因此得到同一个 AI Core 编号。上一节的任务划分直接沿用Cube 和 Vector 在同一个任务循环里按同一个task_idx推进import pypto_pro.language as pl core_id pl.get_block_idx() // pl.get_subblock_num() core_num pl.get_block_num() tiles_per_core (total_tiles core_num - 1) // core_num start core_id * tiles_per_core end pl.min(start tiles_per_core, total_tiles) for task_idx in pl.range(start, end, 1): with pl.section_cube(): # 计算本任务的整块输出 ... with pl.section_vector(): # 处理Cube刚算出的这块数据 ...在 1:2 的芯片上一个 AI Core 内有 2 个 Vector 核两者要分摊同一个任务里 Cube 产出的这份数据因此每个 Vector 核处理的数据量是 Cube 的一半。这一步用pypto_pro.language.get_subblock_idx()在某一维上把数据切成两半属于核内切分不参与核间任务的分配import pypto_pro.language as pl sub_id pl.get_subblock_idx() first_half (rows 1) // 2 half_rows first_half if sub_id 1: half_rows rows - first_half for task_idx in pl.range(start, end, 1): with pl.section_cube(): # 计算本任务的rows行输出 ... with pl.section_vector(): if half_rows 0: # 起始行 sub_id * first_half # 行数 half_rows ...切分的维度由算子自身决定行、列或 batch 方向都可以前提是两个 Vector 核的工作互不重叠且合起来覆盖 Cube 的全部输出。负载均衡与切分粒度上面的分配方式保证了各核Tile 块数接近但没有保证各核计算量接近。每个 Tile 块的计算量相近时两者等价计算量差异较大时例如因果掩码下不同行块的有效长度不同需要额外调整调整 Tile 块的编号顺序使计算量大的 Tile 块分散到不同的核而不是集中在少数核上。拆分计算量过大的 Tile 块减少单个 Tile 块造成的尾部等待。切分粒度的选择需要同时考虑两端Tile 块总数要足以覆盖核数。Tile 块数少于核数时部分核不会进入循环Tile 块数接近核数时零头会明显放大为负载差异。每核分配的 Tile 块不宜过多。在数据规模一定时每核处理过多 Tile 块通常意味着使用核数偏少串行迭代过长难以充分发挥多核算力。单个 Tile 块不宜过小。Tile 块越小循环控制、数据搬运启动和同步的开销占比越高且计算指令占比越小算力不能充分发挥。设计切分方式时重点检查三点Tile 块总数是否足以覆盖计划使用的核各核分到的 Tile 块数和计算量是否接近单核是否承担了过多 Tile 块切分后的 Tile 形状是否满足数据搬运和计算接口的约束。尾块处理尾块场景需要向上取整计算 Tile 数以覆盖全部任务并根据数据搬运和计算接口的对齐要求选择 Tile 尺寸有效形状、填充和计算方法参考 Tile 计算中的尾块处理。三、Tiling 参数定义与传递TilingKey 与 TilingData切分需要的参数每个核处理多少任务、每个 Tile 多大、循环多少次等通常依赖运行时的 shape无法在编写 Kernel 时写死需要在 Host 侧计算后传给 Kernel。这些参数由 TilingKey 和 TilingData 共同承担Tiling参数定义与传递类型生效阶段承担的部分TilingKeyKernel 编译时选择本次使用哪一套实现例如对齐与非对齐、不同的 Tile 模板或数据排布。候选数量有限每个取值组合生成一个专用 Kernel 实例TilingDataKernel 运行时携带这一套实现所需的具体数值例如 shape、stride、循环边界和缩放系数。数值变化时复用同一份已编译的 Kernel即 TilingKey 决定走哪个分支TilingData 提供该分支运行所需的数据。TilingDatadataclass 声明与运行时传值TilingData 使用 Pythondataclass声明。框架根据字段标注生成设备侧结构体并在启动 Kernel 时序列化字段值。一个 Kernel 最多只能声明一个 TilingData 参数且该参数必须位于形参列表末尾。from __future__ import annotations from dataclasses import dataclass dataclass class AddTiling: rows: int cols: int scale: float strides: int[2]字段类型。TilingData 的字段支持标量int、float、bool以及对应的定长数组int[N]、float[N]、bool[N]。其中int对应DT_INT64float对应DT_FP32bool对应DT_BOOL。字段声明需满足以下要求类中至少包含一个字段每个字段只能使用上述类型。定长数组的长度N必须直接写成 12048 的整数值运行时传入的列表长度必须与声明一致。使用int[N]、float[N]或bool[N]时需要启用from __future__ import annotations。字段可以设置dataclass默认值数组默认值使用dataclasses.field(default_factory...)。字段顺序决定设备侧结构体的字段顺序、偏移和对齐不能在 Host 侧和设备侧使用不同顺序解释同一份数据。Host 侧传入。Host 侧构造 TilingData 实例并将其放在 Kernel 启动实参列表末尾下面的add_kernel是一个接收AddTiling的 Kernel定义见下文tiling AddTiling( rows513, cols511, scale1.0, strides[511, 1], ) block_dim 8 add_kernelNone, block_dim使用默认值时未显式传入的字段采用dataclass默认值from dataclasses import dataclass, field dataclass class CopyTiling: length: int scale: float 1.0 offsets: int[4] field(default_factorylambda: [0, 0, 0, 0])带默认值的字段必须位于无默认值字段之后。Kernel 中读取。在 Kernel 函数签名中声明 TilingData 参数后即可通过字段名读取运行时值import pypto_pro.language as pl pl.jit(auto_mutexTrue) def add_kernel( x: pl.Ptr[pl.DT_FP16], y: pl.Ptr[pl.DT_FP16], z: pl.Ptr[pl.DT_FP16], tiling: AddTiling, ): strides [tiling.strides[0], tiling.strides[1]] tensor_x pl.make_tensor(x, [tiling.rows, tiling.cols], strides) tensor_y pl.make_tensor(y, [tiling.rows, tiling.cols], strides) tensor_z pl.make_tensor(z, [tiling.rows, tiling.cols], strides) tile_rows (tiling.rows 127) // 128 tile_cols (tiling.cols 127) // 128 for tile_idx in pl.range( pl.get_block_idx(), tile_rows * tile_cols, pl.get_block_num(), ): ...TilingData 字段是运行时值可以用于pypto_pro.language.make_tensor的 shape 和 stride。pypto_pro.language.range的循环边界。Tile 数量、偏移和有效形状等标量计算。运行时if条件和算术表达式。TilingData 字段不能用于要求编译期 Python 常量的参数例如TileType的静态shape。这类有限候选项应通过 TilingKey 选择。源码级实现字段如何变成设备侧结构体从源码结构看TilingData 的校验与序列化逻辑集中在 python/pypto_pro/language/typing/_tiling.pyis_tiling_class()判定一个类是否为合法 Tiling 类必须是带至少一个字段的dataclass且所有字段标注都能被_parse_field_annotation()解析为标量类型或T[N]定长数组。其中N必须是正整数字面量并通过check_in_range(size, 1, _MAX_ARRAY_SIZE, ...)限制在 12048_MAX_ARRAY_SIZE 2048与文档描述的约束一一对应。_PYTHON_TYPE_TO_DTYPE明确给出int → DataType.INT64、float → DataType.FP32、bool → DataType.BOOL的映射即文档中int对应 DT_INT64的出处。get_tiling_tuple_type()把各字段降级为 IR 的TupleType定长数组T[N]降级为 N 个标量组成的元组用于 Kernel 参数类型检查。get_tiling_ctype_struct()依据字段顺序生成ctypes.Structure其原生对齐规则与 CCE codegen 输出的 Cstruct Tiling保持一致从而保证 Host 侧tiling_instance_to_bytes()序列化出的字节布局与设备侧结构体的sizeof、字段偏移、对齐完全相同——这就是字段顺序决定设备侧结构体布局这句话的底层依据。额外的防御性检查_check_duplicate_annotations()通过解析类源码 AST 检测重复字段名Python 运行时的__annotations__会去重、看不到重复项tiling_instance_to_bytes()会在传入的列表长度与T[N]声明不一致时抛出InvalidType。单测 python/tests/ut/pypto_pro/language/tiling_class/test_array_type.py 覆盖了int[4]、float[2]、bool[3]的合法解析以及str[4]、int[0]等非法标注的拒绝行为可作为字段类型约束的验证参考。TilingKey编译期常量与专用 Kernel 实例TilingKey 描述一组有限的编译期配置。同一份 Kernel 源码可以针对不同 Key 生成多个专用实例并在启动时选择其中一个。Key 字段在 Kernel 中是编译期常量对应的条件分支可以在编译阶段确定。声明 TilingKey。TilingKey 使用普通 Python 类声明类属性使用TilingKeyField定义候选值from pypto_pro.runtime.tilingkey import TilingKeyField class AddKey: UseScale TilingKeyField(bits1, values[0, 1]) BlockM TilingKeyField(bits2, values[64, 128]) def is_valid(self, key): use_scale, block_m key return use_scale 0 or block_m 128字段按照类中的定义顺序收集。该顺序决定is_valid()参数元组的顺序和 64-bit Key 中各字段的位置。各字段候选值的下标打包为 64-bit Key是为了与离线二进制编译场景使用的uint64_tTilingKey 配合。TilingKeyField需满足以下约束项目约束bits必须大于 0values候选值只能是整数不能使用字符串、浮点数或bool集合必须非空且值互不重复编码容量候选值数量不能超过2**bits总位宽所有字段的bits之和不能超过 64is_valid可选用于拒绝不支持的字段组合字段编码保存候选值在values中的下标而不是候选值本身。例如class BlockKey: BlockM TilingKeyField(bits2, values[32, 64, 128])BlockM32、64、128分别编码为0、1、2。Host 侧启动 Kernel 时仍传入实际值不需要计算候选下标。将 TilingKey 绑定到 Kernel。通过pypto_pro.language.jit(tiling_key...)绑定 TilingKey。字段名可以直接在 Kernel 函数体中使用无需加入形参列表import pypto_pro.language as pl pl.jit(auto_mutexTrue, tiling_keyAddKey) def add_kernel( x: pl.Ptr[pl.DT_FP16], y: pl.Ptr[pl.DT_FP16], z: pl.Ptr[pl.DT_FP16], tiling: AddTiling, ): tile_type pl.TileType( shape[BlockM, 128], dtypepl.DT_FP16, target_memorypl.MemorySpace.Vec, ) if UseScale 1: # 此分支在编译时确定 ... else: ...TilingKey 字段名不能与 Kernel 形参或模块级普通变量重名。字段应只描述会影响代码生成的有限模式不要为任意运行时 shape 分别定义 Key。选择 TilingKey 并启动 Kernel。Key 字典位于启动参数中的stream和block_dim之后key {UseScale: 1, BlockM: 128} add_kernelNone, block_dim, keyKey 字典必须满足以下要求包含 Schema 中的全部字段且不能包含额外字段。每个值都属于对应字段的values。字段组合能够通过is_valid()校验。传入字典中的值是实际候选值不是编码后的下标。同时使用 datatype 特化时datatype 字典位于 TilingKey 之后add_kernelNone, block_dim, key, datatype二进制交付时TilingKey Schema 用于生成字段声明和合法的 Key 组合具体流程参考 离线二进制编译。源码级实现64-bit 打包与启动时校验上述约束在 python/pypto_pro/runtime/tilingkey.py 的TilingKeySchema中逐条落地_validate_and_assign_offsets()第 156179 行依次检查bits 0values非空且元素为int显式排除bool即isinstance(v, bool)会被拒绝且互不重复候选值数量不超过1 bits所有字段位宽之和不超过_MAX_TOTAL_BITS 64。同时按定义顺序为每个字段分配位偏移offset。pack()第 130136 行实现候选值下标打包为 64-bit Keypacked | value_index field.offset与文档描述一致。模块 docstring 还说明该编码方式与 AscendC 模板 Tiling Key 的位选择语义对齐。validate_concrete()第 138154 行对应启动时对 Key 字典的校验必须是 dict、键集合与 Schema 字段完全一致、每个值都在候选集合内、is_valid返回真否则分别抛出InvalidType/InvalidArgument/InvalidVal。enumerate_valid()对全部字段组合做笛卡尔积并用is_valid过滤服务于离线二进制编译场景下枚举合法 Key 组合、逐一生成专用实例的流程。单测 python/tests/ut/pypto_pro/runtime/tiling_key/test_validation.py 给出了完整的校验行为示例{OpType: 99}候选值之外、{}缺字段、{OpType: 0, ExtraField: 1}多字段、[0]非 dict、以及被is_valid拒绝的组合都会在kernel[None, 1, key]阶段抛出异常。系统级用例 python/tests/st/pypto_pro/frontend/matmul/test_matmul_tilingkey.py 则演示了 TilingKey 的典型价值Dtype0: FP16 / 1: BF16和TileBlock0: 64 / 1: 128两个 1-bit 字段让同一份 matmul 源码在TileType(shape...)处按 Key 常量折叠出 4 个专用实例Kernel 体内的if Dtype 0: io_dtype pl.DT_FP16 ...分支在编译时确定。测试通过matmul_tilingkeyNone, 1, {Dtype: dtype_key, TileBlock: tile_key}启动并与torch.matmul对照数值覆盖了全部 4 个 Key 组合。TilingData 与 TilingKey 配合使用同一个 Kernel 可以使用 TilingKey 选择编译期实现再通过 TilingData 接收本次调用的运行时参数key {UseScale: 1, BlockM: 128} tiling AddTiling(rowsm, colsn, scale0.5, strides[n, 1]) add_kernelNone, block_dim, key选择参数传输方式时可以按以下规则判断参数特点传输方式每次调用都可能变化且不要求重新生成代码TilingData候选数量有限并影响代码分支、Tile 模板或排布TilingKey只决定启动的核数block_dim参考 blockDim 的含义与设置其中block_dim表示请求启动的逻辑核数Vector Kernel 中是逻辑 VectorAIV核数Cube Kernel 中是逻辑 CubeAIC核数MIX Kernel 中是逻辑执行组数1:2 芯片上每组含 2 个 Vector 核和 1 个 Cube 核实际生效数量一律以get_block_num()为准见 Kernel核函数。使用限制与建议不要把任意 shape 值放入 TilingKey否则候选组合会快速增加并产生大量专用 Kernel。不要把影响静态 Tile 类型的参数放入 TilingData这类参数在 Kernel 运行时才可用。修改 TilingData 字段顺序或类型会改变设备侧结构体布局Host 侧与 Kernel 必须使用同一定义。数组字段长度必须与声明完全一致运行时可变长度数据应拆成长度字段和固定容量数组或通过 Tensor 传递。同时使用 TilingKey、datatype 特化和 TilingData 时分别检查启动方括号参数与函数调用参数的顺序。四、延伸阅读多核Tiling切分核间切分的具体写法包括任务编号与连续分核、MIX Kernel 的切分和负载均衡。Tiling参数定义与传递TilingData 和 TilingKey 的声明、传递与配合使用。Tile创建和操作Tile 的规格描述、片上地址绑定和有效形状设置。SIMD抽象硬件架构AIC 与 AIV 的计算单元、片上存储层级和数据通路。实现与测试TilingKey Schema、Tiling 类类型与序列化、TilingKey 校验单测、matmul TilingKey 系统测试、Tiling 数组字段单测。赞分享人工智能编译器模型编译深度学习高性能计算CANNAscend【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址https://gitcode.com/cann/pypto点击查看免费下载相关推荐PyPTO Tiling参数定义与传递TilingKey 与 TilingData 的声明、绑定与启动实践PyPTO Tiling参数定义与传递TilingKey 与 TilingData 的声明、绑定与启动实践 导读 在 PyPTOParallel Tenso人工智能编译器模型编译深度学习高性能计算CANNAscendPyPTO Tiling 机制全解析多核切分、TilingKey/TilingData 与片上 Tile 数据流PyPTO Tiling 机制全解析多核切分、TilingKey/TilingData 与片上 Tile 数据流 本文是 CANN pypto 项目PyPT人工智能编译器模型编译深度学习高性能计算CANNAscendPyPTO Pro 多核 Tiling 切分AI Core 并行任务划分与负载均衡实战PyPTO Pro 多核 Tiling 切分AI Core 并行任务划分与负载均衡实战 本篇技术指南围绕 PyPTO Pro 编程范式中的 核间 Tiling人工智能编译器模型编译深度学习高性能计算CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考