多核并行编程

本文档说明 PTO Tile Lib 中常见的多核编程方式，并重点介绍与当前编程模型相匹配的工作划分模式。

本文档围绕基于 tile 的工作分解、输出归属、负载均衡和局部性展开。

1. 概述

PTO kernel 通常采用类似 SPMD 的执行方式：多个核心运行同一份 kernel 主体代码，不同核心根据 block 或 core 身份处理不同的工作范围。

这种方式与 PTO 文档中一贯采用的 tile 化编程模型相匹配。

入门示例可参考：

2. 当前仓库中最主要的多核模型

2.1 类 SPMD 的工作划分

本仓库中最常见的模式是：

  • 所有核心执行相同的 kernel 代码
  • 每个核心处理输入或输出的不同区域
  • 划分方式通常围绕行、列、tile 或 block 范围展开

这种方式天然适用于：

  • 逐元素算子
  • 基于 tile 的归约算子
  • GEMM 类算子
  • attention 类的分块计算

2.2 为什么更推荐这种方式

类 SPMD 的划分方式更符合 PTO 的编程特点，因为它与以下因素天然一致：

  • 基于 tile 的工作分解
  • 可预测的 GM 访问模式
  • 更直接的负载均衡
  • 更简单的同步结构

在大多数情况下，让每个核心负责一段规则且连续的工作区域，通常比引入不规则的核间协作更容易分析与优化。

3. 实际划分建议

3.1 按输出归属划分

一个常见且稳妥的默认策略，是按输出区域的归属来划分工作。

例如：

  • 对于向量类算子，可以按线性输出区间划分
  • 对于矩阵类算子，可以按 tile 行、tile 列或二维 block 网格划分
  • 对于按行归约的算子，可以给每个核心分配一行或多行输出

这样做的好处是：

  • 负责计算输出 tile 的核心也负责存储该 tile
  • 中间状态尽量保持在本地
  • 可以避免核间写冲突

3.2 尽量保持工作量均衡

为核心分配 tile 时，建议：

  • 尽量让每个核心承担相近的计算量
  • 避免把尾部的小块工作全部集中到单个过载核心上
  • 在可能的情况下同时兼顾规则访问与均衡划分

需要注意的是，数学上平均的划分如果破坏了内存局部性，依然可能导致较差性能，因此负载均衡和局部性应同时考虑。

3.3 保持规则的 tile 循环结构

当每个核心都遵循相同的 tile 循环结构时，多核 kernel 往往更容易验证和优化。

典型结构通常是：

  1. 确定当前核心负责的 tile 范围
  2. 在该范围内迭代
  3. 执行 TLOAD -> transform / compute -> TSTORE
  4. 在需要时通过 valid region 处理边界 tile

这种写法也遵循了 快速开始教程 和 Tile 编程模型 中描述的 tile 化编程模型。

4. PTO 中真正重要的多核问题

4.1 负载均衡

负载均衡之所以重要，是因为 PTO kernel 常常同时包含：

  • GM 数据搬运
  • 布局变换
  • vector 或 cube 计算
  • 显式同步

如果某个核心分到了明显更多的 tile，或者分到的 tile 计算代价更高，整体吞吐最终就会受最慢核心限制。

在实践中，建议重点检查：

  • 输出空间是否划分得足够均匀
  • 边界 tile 是否过度集中在少数核心上
  • 是否只有部分核心承担了额外的变换或归约工作

4.2 内存局部性

良好的多核划分还需要尽量保持 GM 局部性。

更理想的模式通常具备以下特点：

  • 连续的读写访问
  • 对邻近 tensor 区域的重复利用
  • 稳定的 tile shape 与 stride

如果局部性不好，常见表现就是数据搬运开销相对计算开销偏高。

4.3 核间通信

本仓库在 docs/isa/comm/ 下提供了通信指令文档，但对于一般计算 kernel，不应默认把任意的跨核 producer-consumer 调度视为 PTO 的标准日常模型。

对大多数计算 kernel，更稳妥的方式仍然是：

  • 尽量减少跨核依赖
  • 清晰划分输出归属
  • 仅在确有必要时保留真实的 producer-consumer 同步关系

如果某个 kernel 依赖通信指令，应参考 通信 ISA 参考 及相应指令文档。

5. 与流水线优化的关系

多核并行和流水线重叠解决的是两个不同层面的问题：

  • 多核并行：通过把工作分配到多个核心来提升吞吐
  • 流水线重叠：通过重叠 load / transform / compute / store 阶段来提升单核利用率

一个高性能 kernel 往往同时需要：

  1. 合理的每核 tile 划分
  2. 高效的核内流水线组织

关于重叠、缓冲与同步，可参考 流水线并行 和 事件与同步。

6. 编程边界

多核 PTO kernel 通常围绕 tile 归属、规则工作划分和显式依赖来描述。除非在专门的运行时或 backend 文档中另行定义，以下内容不属于本文档的描述范围：

  • 脱离仓库上下文、凭空假设的 get_block_idx() 一类运行时接口契约
  • TCOMPUTE、TFILL 这类并非当前公开 PTO intrinsic 的占位式指令
  • 在 TLOAD / TSTORE 中使用 Python 风格张量切片的伪语法
  • 把 MPMD 直接写成当前仓库普通 PTO kernel 的标准公开编程模型

这类写法在其他场景中可以用于说明思路，但不适合作为严谨的仓库文档。

7. 多核开发流程

在开发多核 PTO kernel 时，一个更实用的流程是：

  1. 先从单 tile 或单核的正确实现开始
  2. 明确每个核心的输出归属
  3. 将工作划分为规则的 tile 范围
  4. 先在 CPU 仿真上验证正确性
  5. 再在目标 backend 上调优 tile 大小、划分方式和重叠策略

这种流程更容易与现有 PTO 文档保持一致，也能避免过早引入不必要的复杂度。

8. 结语

在当前 PTO Tile Lib 仓库中，更可靠的多核编程理解方式是：

  • 默认采用类 SPMD 的工作划分
  • 以输出归属和规则 tile 范围为核心组织工作
  • 保持规则、连续、均衡的访问模式
  • 将多核划分与单核流水线优化结合起来

相比把一些推测性的伪 API 或未经验证的执行模型写成既定接口，这种写法更准确，也更符合本仓库的实际情况。