PyPTO 的 get_cube_tile_shapes 使用指南读取 Cube 计算 TileShape 与多核切 K 开关【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto本篇技术指南聚焦 CANN PyPTO 的配置查询接口pypto.get_cube_tile_shapes()说明其函数原型、返回值语义、在 m/k/n 三维切分与多核切 KSplit-K场景中的实际用法并结合仓库源码与官方示例深入讲解其底层实现与配套的set_cube_tile_shapes设置接口。读完本文你将能在矩阵乘法Matmul性能调优中准确读取与校验 Cube 计算的 Tile 切分配置。一、功能概述pypto.get_cube_tile_shapes是 PyPTO Tensor 编程范式下的配置类查询接口用于获取 Cube矩阵计算中设置的 TileShape 大小以及多核切 K 功能的开关使能状态。在 PyPTO 中矩阵乘类运算如pypto.matmul由 Cube 计算单元执行其计算过程按 m、k、n 三个维度进行分块切分Tile 化每个维度的切分大小直接决定数据如何在 L0、L1 两级 Cache 之间搬运与复用进而影响算子的执行性能。get_cube_tile_shapes提供的就是这一配置的读取能力常与设置接口set_cube_tile_shapes配对使用用于校验用户配置是否生效设置后回读比对在 JIT 编译与多 Scope 切换场景中确认当前生效的 Tile 配置结合运行时间对比如官方示例中的 32×32 与 64×64 Tile 对比辅助性能调优。二、产品支持情况依据 pypto-get_cube_tile_shapes.md 文档说明该接口在以下产品形态上均受支持产品系列支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持注文档中A3对应的 NPU 架构标识为 A3其接口能力与上述产品保持一致。三、函数原型与参数说明get_cube_tile_shapes() - Tuple[List[int], List[int], List[int], bool]参数无参数。返回值返回一个包含 4 个元素的元组依次为 m、k、n 方向上的 TileShape 大小列表以及是否开启多核切 K 功能的布尔值。具体地返回值与设置接口的参数一一对应语义如下返回位置类型含义第 1 个元素List[int]m 维度上的 TileShape 大小列表第 2 个元素List[int]k 维度上的 TileShape 大小列表第 3 个元素List[int]n 维度上的 TileShape 大小列表第 4 个元素bool是否开启多核切 KSplit-K功能其中每个维度列表的长度通常为 2两个元素分别表示L0 Cache 层与 L1 Cache 层的切分大小k 维度在设置时也可传 3 个元素见下文源码分析但通过该接口回读时得到的仍是规范的 3 元素表示。约束说明按 pypto-get_cube_tile_shapes.md 文档描述该接口无额外约束。需要注意的是无约束仅表示接口调用层面无限制实际 TileShape 取值仍需满足硬件规格与算子实现要求如对齐约束非法取值会在编译或运行时由框架报错。四、调用示例最简单的调用方式如下import pypto # 直接调用返回当前 Scope 下 Cube 计算的 TileShape 与切 K 开关 result pypto.get_cube_tile_shapes() print(result)与set_cube_tile_shapes配对使用的完整读写流程import pypto # 设置 Cube 计算的 TileShapem[16,16], k[256,512], n[128,128]不开启多核切 K pypto.set_cube_tile_shapes([16, 16], [256, 512], [128, 128], enable_split_kFalse) # 回读并打印设置的 TileShape print(pypto.get_cube_tile_shapes()) # 输出: ([16, 16], [256, 512, 512], [128, 128], False)运行上述代码后返回值中的三个列表分别表示m方向L0 切分为 16L1 切分为 16k方向L0 切分为 256L1 切分为 512回读时按内部规范补全为 3 元素第三元素与 L1 层取值一致n方向L0 切分为 128L1 切分为 128最后一个布尔值False表示未开启多核切 K。五、配套设置接口set_cube_tile_shapesget_cube_tile_shapes读取的配置由 set_cube_tile_shapes 写入。其函数原型为set_cube_tile_shapes(m: List[int], k: List[int], n: List[int], enable_split_k: bool False)mm 维度 TileShape列表长度必须为 2kk 维度 TileShape列表长度可为 2 或 3nn 维度 TileShape列表长度必须为 2enable_split_k是否开启多核切 K默认False即不进行 GM 累加式的 K 轴分核。从源码实现看set_cube_tile_shapes会先构造一个CubeTile对象再通过pypto_impl.SetScope({cube_tile_shapes: cube_tile.impl()})将配置写入当前编译 Scopepython/pypto/_controller.py。因此get_cube_tile_shapes读取到的正是当前 Scope 中生效的配置二者天然配对。CubeTile 的数据模型CubeTile定义于 python/pypto/config.py其语义为m[0]、k[0]、n[0]对应 L0 Cache 的切分大小m[1]、k[1]、n[1]对应 L1 Cache 的切分大小。构造时存在如下校验逻辑python/pypto/config.pym必须恰好包含 2 个元素否则抛出ValueErrorn必须恰好包含 2 个元素k允许 2 或 3 个元素当传入 2 个元素时框架会自动补全k[2] k[1]即 L1 层的 k 切分被复制为第三元素——这正是get_cube_tile_shapes回读时 k 列表出现 3 个元素的原因参见 python/pypto/config.py。enable_split_k参数对应是否开启多核切 K其本质是 matmul 结果是否采用 GM 累加GM ACC模式默认False即不开启。其他设置入口除set_cube_tile_shapes外cube_tile_shapes还可以通过以下两种方式设置三者殊途同归最终都会写入 Scope 的cube_tile_shapes配置键pypto.options装饰器 / 上下文管理器python/pypto/config.py# 作为上下文管理器 with pypto.options(nametest, cube_tile_shapes[[16, 16], [256, 512, 128], [128, 128], True]): pass # 作为装饰器 pypto.options(cube_tile_shapes[[16, 16], [256, 512, 128], [128, 128], True]) def func(): passpypto.set_options全局设置python/pypto/config.pypypto.set_options(cube_tile_shapes[[16, 16], [256, 512, 128], [128, 128], True])无论从哪个入口进入最终都会在_Options.prepare_options中被统一转换为CubeTile对象并写入配置python/pypto/config.py保证读取接口看到的行为一致。六、底层实现从 Python 到 C 的调用链get_cube_tile_shapes的完整实现位于 python/pypto/_controller.py核心逻辑为scope get_current_scope() cube_tile scope.get_cube_tile_shapes() return tuple([cube_tile.m, cube_tile.k, cube_tile.n, cube_tile.enableSplitK])即先获取当前配置 Scopepypto_impl.CurrentScope()再取出其中的CubeTile对象将其m、k、n三个列表与enableSplitK布尔值组装为元组返回。在 C 绑定层对应的接口实现在 python/src/bindings/controller.cppm.def(GetCubeTile, []() { auto cubeTile TileShape::Current().GetCubeTile(); return std::tuple(cubeTile.m, cubeTile.k, cubeTile.n, cubeTile.enableSplitK); });可见 Python 侧返回的元组结构m、k、n、enableSplitK与 C 侧TileShape::Current().GetCubeTile()的返回完全对齐TileShape是框架内部维护当前 Tile 配置的全局单例。而CubeTile类在绑定层通过py::class_CubeTile导出暴露了m、k、n、enableSplitK四个可读写字段及valid、ToString等成员方法python/src/bindings/controller.cpp。也就是说一次简单的get_cube_tile_shapes()调用实际经过了Python 封装 → C 绑定 → TileShape 单例查询的完整链路读取的是框架编译流水线中真实生效的配置状态。七、实战场景结合 Matmul 性能调优get_cube_tile_shapes最常见的实战场景是配合set_cube_tile_shapes进行矩阵乘法的 Tile 切分调优。7.1 为什么需要关注 Cube TileShape在 Tiling 编程指南 中明确说明将矩阵相乘形状记为(m, k) × (k, n) (m, n)set_cube_tile_shapes的三个列表分别设置 m、k、n 维度的切分大小每个列表的第一个元素是 L0 切分大小第二个元素是 L1 切分大小。通常设置不同的 TileShape 不影响矩阵的计算结果但会影响矩阵计算的运行时间——这是性能调优的核心切入点。7.2 多核切 K 的使用场景enable_split_k返回值第 4 个元素对应多核切 K 功能。对于M、N 较小而 K 轴较大的场景仅在 M、N 轴做分核可能无法用满硬件核数导致整体性能较差此时可将enable_split_k设为True以使能 K 轴分核参见 tiling.md。7.3 完整调优示例仓库提供的官方示例位于 examples/01_beginner/tiling/tiling_config.py其中cube_tile相关的用例完整演示了设置—回读—结果校验—耗时对比的调优闭环import time from numpy.testing import assert_allclose import torch import pypto # 在 JIT kernel 内部设置 TileShape并回读校验 pypto.frontend.jit(runtime_options{run_mode: pypto.RunMode.NPU}) def compute_with_different_tile_shapes(x, y, out1, out2, out3): pypto.set_cube_tile_shapes([32, 32], [16, 16], [32, 32]) print(fpypto.get_cube_tile_shapes(): {pypto.get_cube_tile_shapes()}) out1[:] pypto.matmul(x, y, x.dtype) pypto.set_cube_tile_shapes([32, 32], [16, 64], [32, 128]) print(fpypto.get_cube_tile_shapes(): {pypto.get_cube_tile_shapes()}) out2[:] pypto.matmul(x, y, x.dtype) pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128]) print(fpypto.get_cube_tile_shapes(): {pypto.get_cube_tile_shapes()}) out3[:] pypto.matmul(x, y, x.dtype)运行方式需先配置 CANN 环境与 NPU或使用--run_mode sim进入仿真模式export TILE_FWK_DEVICE_ID0 python tiling_config.py cube_tile::test_set_cube_tile_shapes_basic python tiling_config.py cube_tile::test_set_different_tile_shapes_result python tiling_config.py cube_tile::test_set_different_tile_shapes_runtime上述用例验证了两点关键结论结果一致性同一矩阵乘法在不同 TileShape 下32/16/32、32/16/64/128、64/128/128计算结果一致assert_allclose校验通过说明 TileShape 不影响数值正确性性能差异性test_set_different_tile_shapes_runtime对比了[32,32],[32,32],[32,32]与[64,64],[128,128],[128,128]两组配置的耗时证明 TileShape 显著影响运行时间。7.4 单元测试中的读写一致性验证仓库的单测 python/tests/ut/interface/test_pto_vec_tiles_shape.py 直接对读写一致性做了断言可作为接口行为的最简验证范例def test_cube_tile_shapes(): expected ([16, 16], [256, 512, 512], [128, 128], False) pypto.set_cube_tile_shapes(*expected[:3], expected[3]) actual pypto.get_cube_tile_shapes() assert actual expected注意测试中传入的 k 列表为[256, 512]2 个元素回读得到的却是[256, 512, 512]3 个元素——这正是前面提到的k[2] k[1]自动补全规则在读取侧的体现是使用本接口时最值得注意的语义细节。八、小结与相关文档pypto.get_cube_tile_shapes()是一个零参数、语义简单但调优价值明确的配置读取接口返回值(m 列表, k 列表, n 列表, enable_split_k)其中维度列表长度 2 表示 L0/L1 两级切分k 维度回读为 3 元素第三元素等于 L1 层取值适用产品Ascend 950PR/DT、Atlas A3、Atlas A2 系列典型用法与set_cube_tile_shapes/pypto.options/pypto.set_options配对在 Matmul 调优中完成设置—回读—校验—耗时对比的完整流程底层原理通过 C 绑定查询TileShape::Current()单例中的CubeTile配置Python 侧返回结构与 C 侧完全对齐。进一步阅读接口定义文档pypto-get_cube_tile_shapes.md同类配置接口索引见 config/index.md编程指南含完整调优用例与背景说明tiling.md可运行官方示例tiling_config.pyPython 侧实现python/pypto/_controller.py、python/pypto/config.pyC 绑定实现python/src/bindings/controller.cpp。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考