PyPTO Tiling 配置实战用 set_vec_tile_shapes 与 set_cube_tile_shapes 精细化切分算子数据【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pyptoTileShape切分形状是 PyPTO 算子性能调优的核心旋钮它决定了数据如何在 NPU 的不同计算单元向量单元的 Unified Buffer、矩阵单元的 L0/L1 多级缓冲中被切分与调度直接影响数据搬运开销和计算流水效率。本文以 PyPTO 框架Parallel Tensor/Tile Operation 编程范式中的 Tiling 配置为主题完整讲解向量计算与矩阵计算的 TileShape 设置接口、参数语义、使用约束、性能影响以及配套的验证手段帮助你掌握同一算子、不同切分、性能迥异的调优方法论。为什么 TileShape 决定算子性能TileShape 定义了数据在硬件不同计算单元中的切分方式。PyPTO 的核心设计理念是让开发者用接近数学表达式的形式描述张量/分块运算而 TileShape 则是把这种高层描述映射到硬件执行单元的关键参数向量计算通过set_vec_tile_shapes设置向量在各维度上的切分大小。合理的切分可以让数据充分利用统一缓冲区Unified BufferUB使向量计算单元上的取数、计算、写回形成高效流水。矩阵计算矩阵相乘的形状变化记为(m, k) x (k, n) (m, n)通过set_cube_tile_shapes依次设置矩阵在 m、k、n 维度上的切分大小。合理的切分可以充分利用 L0、L1 缓冲区减少数据在 MTE内存搬运引擎上的重复搬运缩短端到端运行时间。从源码实现看这两组接口定义在 python/pypto/_controller.py 中与set_conv_tile_shapes、set_convbp_input_tile_shapes等同属 PyPTO 的 Tiling 配置家族。它们并非独立的全局变量而是通过pypto_impl.SetScope写入当前编译作用域Scopedef set_vec_tile_shapes(*shapes: int): concrete_shapes [it.concrete() if isinstance(it, SymbolicScalar) else it for it in shapes] pypto_impl.SetScope({vec_tile_shapes: concrete_shapes}) def set_cube_tile_shapes(m: List[int], k: List[int], n: List[int], enable_split_k: bool False): cube_tile CubeTile(m, k, n, enable_split_k) pypto_impl.SetScope({cube_tile_shapes: cube_tile.impl()})也就是说TileShape 配置是作用域级的在同一个pypto.frontend.jitkernel 函数内可以在不同计算段前后分别设置不同的 TileShape见下文create_different_tile_shapes_kernel示例框架会按当前作用域内的最新设置进行切分编译。配置项由 python/pypto/config.py 中的get_current_scope()承载读取侧get_vec_tile_shapes()/get_cube_tile_shapes()返回当前作用域中的配置快照。Vector 计算的 Tiling 配置接口与基本用法set_vec_tile_shapes用于设置向量计算中各维度的 TileShape其返回值为空配套的get_vec_tile_shapes用于读取当前配置# 设置向量计算的 TileShape pypto.set_vec_tile_shapes(1, 1, 8, 8) # 获取并打印设置的 TileShape print(pypto.get_vec_tile_shapes()) # 输出: [1, 1, 8, 8]pypto.set_vec_tile_shapes(1, 1, 8, 8)表示该向量有四个维度每个维度分别按照 1、1、8、8 的大小进行切分并按(1, 1, 8, 8)的切分大小将原向量搬移到 UB 上进行运算。维度数与切分参数一一对应这是最容易出错的点切分参数的个数必须与待计算 Tensor 的维度数量一致。对应接口的单元测试位于 python/tests/ut/interface/test_pto_vec_tiles_shape.py分别覆盖了 2D(8, 16)、3D(1, 2, 3)、4D(1, 2, 3, 8)三种维度规模的设置与读取一致性校验def test_tile_shape_set_vec_tiles_shape_3d(): expected (1, 2, 3) pypto.set_vec_tile_shapes(*expected) actual pypto.get_vec_tile_shapes() assert tuple(actual) expected完整用例向量加法以下用例展示了在简单向量加法场景下set_vec_tile_shapes的应用。kernel 通过set_shapes元组参数动态传入切分配置使用pypto.frontend.jit装饰器编译out[:] pypto.add(a, b)完成按位相加并写回输出pypto.frontend.jit def compute_with_vec_tile_shapes_kernel( a: pypto.Tensor((32, 32), pypto.DT_FP32), b: pypto.Tensor((32, 32), pypto.DT_FP32), out: pypto.Tensor((32, 32), pypto.DT_FP32), set_shapes: tuple ): pypto.set_vec_tile_shapes(*set_shapes) out[:] pypto.add(a, b) def compute_with_vec_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, set_shapes: tuple, dynamic: bool False) - torch.Tensor: # 直接传入torch tensor调用 out torch.empty_like(a) compute_with_vec_tile_shapes_kernel(a, b, out, set_shapes) return out def test_set_vec_tile_shapes_basic(): ... a torch.tensor([[[1, 2, 3], [1, 2, 3]]], dtypedtype, devicefnpu:{device_id}) b torch.tensor([[[4, 5, 6], [4, 5, 6]]], dtypedtype, devicefnpu:{device_id}) expected torch.tensor([[[5, 7, 9], [5, 7, 9]]], dtypedtype, devicefnpu:{device_id}) set_shapes (1, 2, 8) out compute_with_vec_tile_shapes_op(a, b, set_shapes) assert_allclose(out.cpu().numpy(), expected.cpu().numpy(), rtol1e-3, atol1e-3)注意这里输入 shape 为(1, 2, 3)切分参数为(1, 2, 8)前两维与输入严格一致尾维 8 大于实际长度 3PyPTO 会按合法范围切分并保证结果正确assert_allclose校验通过。这印证了文档中的约束原则切分大小只要与 Tensor 的维度数量匹配且不越界即可数值不必等于真实 shape。不同 TileShape 只影响性能不影响结果需要说明的是通常设置不同的 TileShape不影响向量的计算结果但会影响向量计算的运行时间。下述用例对 shape 为(4, 32, 64, 256)的 FP32 向量相加分别使用(1, 2, 4, 128)与(2, 4, 8, 256)两组切分并用time.perf_counter粗测运行时间pypto.frontend.jit def compute_with_vec_specific_tile_shapes_kernel( a: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), b: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), out: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), ): pypto.set_vec_tile_shapes(1, 2, 4, 128) out[:] pypto.add(a, b) pypto.frontend.jit def compute_with_vec_another_tile_shapes_kernel( a: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), b: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), out: pypto.Tensor((4, 32, 64, 256), pypto.DT_FP32), ): pypto.set_vec_tile_shapes(2, 4, 8, 256) out[:] pypto.add(a, b) def compute_with_vec_specific_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, dynamic: bool False) - torch.Tensor: out torch.empty_like(a) compute_with_vec_specific_tile_shapes_kernel(a, b, out) return out def compute_with_vec_another_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, dynamic: bool False) - torch.Tensor: out torch.empty_like(a) compute_with_vec_another_tile_shapes_kernel(a, b, out) return out def test_set_vec_different_tile_shapes_runtime(): ... a torch.randn((4, 32, 64, 256), dtypedtype, devicefnpu:{device_id}) b torch.randn((4, 32, 64, 256), dtypedtype, devicefnpu:{device_id}) TEST_TIME 1 start time.perf_counter() for _ in range(TEST_TIME): out1 compute_with_vec_specific_tile_shapes_op(a, b) runtime_1 time.perf_counter() - start start time.perf_counter() for _ in range(TEST_TIME): out2 compute_with_vec_another_tile_shapes_op(a, b) runtime_2 time.perf_counter() - start print(fruntime_1(pypto.set_vec_tile_shapes(1, 2, 4, 128)): {runtime_1}) print(fruntime_2(pypto.set_vec_tile_shapes(2, 4, 8, 256)): {runtime_2})在该示例中set_vec_tile_shapes(1, 2, 4, 128)的运行时间明显比set_vec_tile_shapes(2, 4, 8, 256)要长。原因在于切分过小导致切分次数TensorShape 与 TileShape 各维度比值的乘积急剧增大UB 每次只装载少量数据搬运与计算的流水难以打满而接近硬件容量的切分能减少搬运轮次、提高单轮计算密度。Cube 计算的 Tiling 配置接口与参数语义set_cube_tile_shapes用于设置矩阵计算中各矩阵在 m、k、n 维度上的 TileShape。将矩阵相乘形状变化记为(m, k) x (k, n) (m, n)三个列表分别设置 m、k、n 维度的切分大小每个列表的两个元素中第一个元素是 L0 缓冲区的切分大小第二个元素是 L1 缓冲区的切分大小# 设置Cube计算的TileShape 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], [128, 128], False]其中enable_split_k参数为是否开启多核切 K 功能默认为 False。对于 M、N 较小而 K 轴较大的场景仅在 M、N 轴做分核可能无法用满核导致整体性能较差此时可以设置enable_split_kTrue以使能 K 轴分核。从实现看该接口内部会构造 python/pypto/config.py 中的CubeTile对象构造时对参数做了严格校验m、n 列表长度必须为 2k 列表长度允许为 2 或 3——当只传两个值时内部会补充k[2] k[1]即 A、B 两个矩阵在 L1 上的 K 轴切分默认相同传三个值时[kL0, kAL1, kBL1]可分别指定 A、B 矩阵的 K 轴 L1 切分高级用法详见下文。class CubeTile: def __init__(self, m: List[int], k: List[int], n: List[int], enable_split_k: bool False): if len(m) ! 2: raise FeError(ValueError(fm must have exactly 2 elements, got {len(m)})) if len(n) ! 2: raise FeError(ValueError(fn must have exactly 2 elements, got {len(n)})) if len(k) not in [2, 3]: raise FeError(ValueError(fk must have 2 or 3 elements, got {len(k)})) k_padded list(k) if len(k_padded) 2: k_padded.append(k_padded[1]) # k[2] k[1] self._impl pypto_impl.CubeTile(list(m), k_padded, list(n), enable_split_k)与之对应的单元测试见 python/tests/ut/interface/test_pto_vec_tiles_shape.py 中的test_cube_tile_shapes([16, 16], [256, 512, 512], [128, 128], False)的设置与读取完全一致印证了 k 列表三元素含自动补齐与 enable_split_k 的读写语义。完整用例矩阵乘法pypto.frontend.jit def compute_with_cube_tile_shapes_kernel( a: pypto.Tensor((64, 64), pypto.DT_FP32), b: pypto.Tensor((64, 64), pypto.DT_FP32), out: pypto.Tensor((64, 64), pypto.DT_FP32), set_shapes: list ): pypto.set_cube_tile_shapes(*set_shapes) out[:] pypto.matmul(a, b, a.dtype) def compute_with_cube_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, set_shapes: list, dynamic: bool False) - torch.Tensor: # 直接传入torch tensor调用 out torch.empty((64, 64), dtypea.dtype, devicea.device) compute_with_cube_tile_shapes_kernel(a, b, out, set_shapes) return out def test_set_cube_tile_shapes_basic(): ... a torch.tensor([[1, 2], [3, 4]], dtypedtype, devicefnpu:{device_id}) b torch.tensor([[5, 6], [7, 8]], dtypedtype, devicefnpu:{device_id}) expected torch.tensor([[19, 22], [43, 50]], dtypedtype, devicefnpu:{device_id}) set_shapes [[32, 32], [64, 64], [64, 64]] out compute_with_cube_tile_shapes_op(a, b, set_shapes) assert_allclose(out.cpu().numpy(), expected.cpu().numpy(), rtol1e-3, atol1e-3)上述用例展示了简单矩阵乘法场景下set_cube_tile_shapes的应用输入仅为(2, 2)的矩阵切分配置却可以设置为[32, 32], [64, 64], [64, 64]计算结果的正确性由assert_allclose保证。这说明切分大小是上限意义上的配置——只要不超出硬件缓冲容量即使大于实际 shape 也能正确计算。不同 TileShape 对 Matmul 运行时间的影响与向量场景同理设置不同的 TileShape 通常不影响矩阵的计算结果但会影响矩阵计算的运行时间。以下用例对 shape 分别为(4, 64, 512)与(4, 128, 512)的矩阵做转置乘b_transTrue等价于(4, 64, 512) x (4, 512, 128)输出(4, 64, 128)import pypto import torch import time pypto.frontend.jit def compute_with_cube_specific_tile_shapes_kernel( a: pypto.Tensor((4, 64, 512), pypto.DT_FP32), b: pypto.Tensor((4, 128, 512), pypto.DT_FP32), out: pypto.Tensor((4, 64, 128), pypto.DT_FP32), ): pypto.set_cube_tile_shapes([32, 32], [32, 32], [32, 32]) out[:] pypto.matmul(a, b, a.dtype, b_transTrue) pypto.frontend.jit def compute_with_cube_another_tile_shapes_kernel( a: pypto.Tensor((4, 64, 512), pypto.DT_FP32), b: pypto.Tensor((4, 128, 512), pypto.DT_FP32), out: pypto.Tensor((4, 64, 128), pypto.DT_FP32), ): pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128]) out[:] pypto.matmul(a, b, a.dtype, b_transTrue) def compute_with_cube_specific_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, dynamic: bool False) - torch.Tensor: out torch.empty((4, 64, 128), dtypea.dtype, devicea.device) compute_with_cube_specific_tile_shapes_kernel(a, b, out) return out def compute_with_cube_another_tile_shapes_op(a: torch.Tensor, b: torch.Tensor, dynamic: bool False) - torch.Tensor: out torch.empty((4, 64, 128), dtypea.dtype, devicea.device) compute_with_cube_another_tile_shapes_kernel(a, b, out) return out def test_set_cube_different_tile_shapes_runtime(): ... a torch.randn((4, 64, 512), dtypedtype, devicefnpu:{device_id}) b torch.randn((4, 128, 512), dtypedtype, devicefnpu:{device_id}) TEST_TIME 1 start time.perf_counter() for _ in range(TEST_TIME): out1 compute_with_cube_specific_tile_shapes_op(a, b) runtime_1 time.perf_counter() - start start time.perf_counter() for _ in range(TEST_TIME): out2 compute_with_cube_another_tile_shapes_op(a, b) runtime_2 time.perf_counter() - start print(fruntime_1(pypto.set_cube_tile_shapes([32, 32], [32, 32], [32, 32])): {runtime_1}) print(fruntime_2(pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128])): {runtime_2})在该示例中set_cube_tile_shapes([32, 32], [32, 32], [32, 32])的运行时间明显比set_cube_tile_shapes([64, 64], [128, 128], [128, 128])要长过小的切分导致 M、N 轴切分出的任务块数过多、单块计算量过小L0/L1 多级缓冲的复用率降低MTE2 重复载入量上升整体流水效率下降。在一个 Kernel 内对比多组切分实际调优时更常见的是一次编译、多组切分对照。参考 examples/01_beginner/tiling/tiling_config.py 中的create_different_tile_shapes_kernel可以在同一个 kernel 中依次设置多组 Cube TileShape分别计算并直接对比结果与get_cube_tile_shapes的输出def create_different_tile_shapes_kernel(run_modeglobal_run_mode): b, m_batch, k_batch, n_batch 2, 2, 2, 2 pypto.frontend.jit(runtime_options{run_mode: run_mode}) def compute_with_different_tile_shapes( x: pypto.Tensor((b, m_batch, k_batch), pypto.DT_FP32), y: pypto.Tensor((b, k_batch, n_batch), pypto.DT_FP32), out1: pypto.Tensor((b, m_batch, n_batch), pypto.DT_FP32), out2: pypto.Tensor((b, m_batch, n_batch), pypto.DT_FP32), out3: pypto.Tensor((b, m_batch, n_batch), pypto.DT_FP32), ): pypto.set_cube_tile_shapes([32, 32], [16, 16], [32, 32]) out1[:] pypto.matmul(x, y, x.dtype) pypto.set_cube_tile_shapes([32, 32], [16, 64], [32, 128]) out2[:] pypto.matmul(x, y, x.dtype) pypto.set_cube_tile_shapes([64, 64], [128, 128], [128, 128]) out3[:] pypto.matmul(x, y, x.dtype) return compute_with_different_tile_shapes该样例文件还提供了命令行运行入口支持--run_modenpu/sim与按用例 ID 选择执行例如export TILE_FWK_DEVICE_ID0 python examples/01_beginner/tiling/tiling_config.py --list python examples/01_beginner/tiling/tiling_config.py cube_tile::test_set_cube_tile_shapes_basic python examples/01_beginner/tiling/tiling_config.py vec_tile::test_set_vec_different_tile_shapes_runtime在 NPU 模式运行前需确保 CANN 环境已配置且 NPU 可用在无硬件环境时可用--run_mode sim走仿真模式。用例内部会在 NPU 模式下用assert_allclose校验结果在 sim 模式下仅打印输出。使用约束切分过小与过大的边界设置 TileShape 参数时必须满足约束条件应与需要处理的 Tensor 的 Shape 维度数量和大小相匹配且数值不能过小或过大。切分不能过小过小的 TileShape 会导致切分次数TensorShape/TileShape即 TensorShape 与 TileShape 每个维度比值的乘积过大从而使表达式在线循环展开的次数过大。这可能导致表达式表Expression Table编译失败并增加运行时的头开销。表达式表的大小与在线循环展开次数以及算子输入个数有关建议控制(TensorShape/TileShape) * (1 算子输入个数)的值小于18000。切分不能过大过大的 TileShape 会超出相应硬件缓冲区的存储大小应控制切分后的数据大小数据类型大小与切分后数据各维度大小乘积不大于对应硬件单元存储容量。例如向量切分要落在 UB 容量内矩阵切分要同时满足 L0L0A/L0B/L0C与 L1 的容量约束。对齐与维度约束set_vec_tile_shapes的维度数量不大于 5尾轴切分大小需满足32B 对齐set_cube_tile_shapes要求 kL0、kL1、nL0、nL1 均满足32 字节对齐更详细的配置要求请参见相关接口文档。从原理理解切分算数强度与缓冲复用在 docs/zh/guide/programming_guide/tensor/debug/matmul_performance_guide.md 中PyPTO 给出了 Matmul Tile 配置的量化分析方法可用于解释本文示例中的性能差异算数强度Arithmetic Intensity定义为单位字节数据访问对应的浮点运算次数FLOPs/Byte。切分必然会引入重复搬运切分越多重复搬运量越大、算数强度越低而切分又受片上多级缓存L1、L0容量限制不能无限增大。当算数强度大于平台的算力带宽比时可达到 Compute Bound反之受限于 Memory Bound。推荐起点配置以 Atlas A3/A2 训练与推理系列产品、A、B 矩阵均为 FP16 为例满足 Buffer 空间约束的推荐 Tile 配置包括pypto.set_cube_tile_shapes([128, 128], [64, 256], [256, 256], enable_split_kFalse) pypto.set_cube_tile_shapes([256, 256], [64, 256], [128, 128], enable_split_kFalse) pypto.set_cube_tile_shapes([128, 128], [128, 512], [128, 128], enable_split_kFalse)这类配置在满足 L0 Buffer 约束的前提下达到较大的算数强度mL1、nL1 取 128-256 组合时L0A、L0B 空间占用较小可以同时使能 MTE1/MTE2 double buffer配合kL1 kL0的大包搬运进一步提升 MTE2 带宽利用率。需要注意的是mL1、nL1 为 128-256 组合时 L0C 无法开启 nbuffer因此更适合 K 轴较大搬出次数相对较少的场景。K 轴分核对于 M、N 较小而 K 轴较大的场景仅 M、N 轴分核无法用满核可采用enable_split_kTrue自动使能 K 轴分核默认按kL1切 K 轴单核计算kL1长度部分和后搬出累加适合快速验证深度调优场景则可手动循环切 K 并叠加pypto.add累加部分和。这些优化结论与本文示例相互印证切分大小直接决定任务数量分核数与计算轮次与数据重复载入量是 Matmul 性能的第一决定因素。实践建议如何评估与挑选 TileShape先验证正确性再优化性能切换 TileShape 不应改变计算结果除非设置极端值。可参考 python/tests/st/operation/matmul/test_matmul_basic.py 等 Matmul 用例的组织方式把不同切分配置的结果与 torch 参考结果做assert_allclose比对。性能观察通过性能分析工具如泳道图观察不同 TileShape 下的搬运与计算流水评估 TileShape 设置的合理性获取当前场景下的最优 TileShape。泳道图可以直观看到 MTE2 搬运、CUBE/Vector 计算是否被小切分切碎导致流水空隙。分场景调参训练场景 M、N 足够大优先选择 128-256 组合的大切分以增大算数强度推理等小 M、N 场景优先保证分满核分核数达到总核数 0.8 以上再尽量减少 MTE2 重复载入。注意上板时间与开销权衡设置 TileShape 会影响上板时间。一般来说越能充分利用硬件单元容量一次计算数据量越大运行时间越短但 TileShape 参数设置得越大并不一定意味着上板运行越快还需要考虑数据搬运等环节的开销——这正是过大切分反而变慢的原因。完整可运行样例请参考 examples/01_beginner/tiling/tiling_config.py其覆盖了 Vector 与 Cube 两组接口的基础用法、多组切分结果一致性对比、不同切分运行时间对比共六个用例接口底层实现见 python/pypto/_controller.py 与 python/pypto/config.pyMatmul 场景的完整调优方法论见 Matmul高性能编程指南。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考