name: PTO-ISA算子实现指南 description: 使用PTO-ISA实现指定算子功能的完整流程指南,涵盖ISA指令选择、数据流分析、指令功能解释和kernel代码生成 license: CANN Open Software License Agreement Version 2.0
PTO-ISA算子实现指南¶
本指南为使用PTO-ISA实现指定算子功能提供完整的流程指导。
目录¶
概述¶
本skill专门用于帮助开发者从PTO-ISA指令集中选择合适的指令来实现指定的算子功能,并生成完整的kernel代码。
适用场景¶
- 用户指定具体的算子功能需求(如: "实现GELU激活函数"、"实现Batch Normalization")
- 需要分析PTO-ISA指令集并选择合适的指令组合
- 需要理解数据在硬件上的流动过程
- 需要生成可直接使用的kernel代码
关键特性¶
- ISA指令选择: 从PTOISA_zh.md文档中分析并选择合适的指令
- 数据流分析: 按数据处理顺序分析数据流向
- Vector计算: gm → ub → vector → ub → gm
- Cube计算(矩阵乘): GM → L1 → L0A/L0B → L0C → GM
- 指令功能解释: 详细解释每个ISA指令在算子实现中的作用
- 代码生成: 输出完整的kernel代码实现
工作流程¶
当用户指定算子功能后,遵循以下工作流程:
用户指定算子功能
↓
步骤1: 阅读PTOISA_zh.md
↓
步骤2: 分析算子需求,列举ISA指令
↓
步骤3: 按数据流顺序解释指令功能
↓
步骤4: 输出kernel代码实现
步骤详解¶
步骤1: 阅读PTOISA_zh.md文档¶
目标: 全面了解PTO-ISA指令集,识别可能与算子相关的指令类别。
行动:
1. 阅读文档路径: pto-isa/docs/PTOISA_zh.md
2. 重点关注指令索引表,识别以下类别的指令:
- 内存指令: TLOAD, TSTORE, TPREFETCH (GM <-> Tile数据搬运)
- 逐元素计算: TADD, TSUB, TMUL, TDIV, TMAX, TMIN (Tile-Tile操作)
- 标量操作: TADDS, TSUBS, TMULS, TDIVS (Tile-标量操作)
- 数学函数: TLOG, TEXP, TSQRT, TRSQRT, TPOW (数学运算)
- 激活函数: TRELU, TPRELU, TLRELU (激活操作)
- 轴归约/扩展: TROWSUM, TCOLSUM, TROWMAX, TCOLMAX (轴操作)
- 广播操作: TROWEXPANDADD, TCOLEXPANDADD (广播加法)
- 类型转换: TCVT (类型转换)
- 选择操作: TSEL, TSELS (条件选择)
- 矩阵操作: TMATMUL, TGEMV (矩阵计算)
- 记录每个相关指令的:
- 指令名称
- 功能描述
- 所属类别
- 适用场景
输出: 相关ISA指令列表
步骤2: 整理并列举需要使用的ISA指令¶
目标: 根据算子功能需求,确定具体的ISA指令组合。
分析框架:
2.1 算子功能分解¶
将算子功能分解为基本操作:
| 算子类型 | 分解步骤 | 典型指令 |
|---|---|---|
| 激活函数 | 输入加载 + 计算 + 输出存储 | TLOAD + TEXP/TLOG/TRELU + TSTORE |
| 归约操作 | 输入加载 + 归约 + 输出存储 | TLOAD + TROWSUM/TCOLSUM + TSTORE |
| 逐元素运算 | 输入加载 + 运算 + 输出存储 | TLOAD + TADD/TSUB/TMUL/TDIV + TSTORE |
| 广播操作 | 输入加载 + 广播 + 运算 + 存储 | TLOAD + TROWEXPANDADD + TSTORE |
| 矩阵运算(Cube) | 输入加载 + 数据搬运 + 矩阵乘 + 输出存储 | TLOAD + TMOV + TMATMUL + TSTORE (GM→L1→L0A/L0B→L0C→GM) |
| 类型转换 | 输入加载 + 转换 + 输出存储 | TLOAD + TCVT + TSTORE |
| 条件运算 | 输入加载 + 比较 + 选择 + 存储 | TLOAD + TCMP + TSEL + TSTORE |
2.2 ISA指令选择原则¶
最小化原则: 使用最少的指令完成功能,减少数据搬运。
数据流优化: - Vector计算: 遵循 gm → ub → vector → ub → gm 的基本数据流 - Cube计算(矩阵乘): 遵循 GM → L1 → L0A/L0B → L0C → GM 的矩阵数据流
同步考虑: 指令间使用Event同步或手动标志同步。
示例:
算子: GELU激活函数
GELU(x) = x * Φ(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
指令选择:
1. TLOAD: 从GM加载输入x到UB
2. TMUL: 计算 x^3 (x * x * x)
3. TMULS: 计算 0.044715 * x^3 (标量乘法)
4. TADD: 计算 x + 0.044715 * x^3
5. TMULS: 计算 sqrt(2/π) * (x + ...) (标量乘法)
6. TEXP/TLOG: 计算tanh函数 (可选,或使用近似)
7. TADD: 计算 1 + tanh(...)
8. TMULS: 计算 0.5 * (结果) (标量乘法)
9. TMUL: 计算 x * 最终结果
10. TSTORE: 将结果从UB存储到GM
输出: 按执行顺序排列的ISA指令列表
步骤3: 按数据处理顺序详细解释指令功能¶
目标: 详细说明每个指令在数据流中的作用。
数据流框架:
Vector计算数据流¶
数据流向: gm → ub → vector → ub → gm
阶段1: GM → UB (数据加载,使用TLOAD)
阶段2: UB → Vector (计算准备)
阶段3: Vector计算 (核心计算,使用TADD/TMUL/TEXP等)
阶段4: Vector → UB (计算结果)
阶段5: UB → GM (数据存储,使用TSTORE)
Cube计算数据流 (矩阵乘)¶
数据流向: GM → L1 → L0A/L0B → L0C → GM
阶段1: GM → L1 (矩阵数据加载,使用TLOAD)
阶段2: L1 → L0A/L0B (数据搬运到矩阵计算单元,使用TMOV)
阶段3: Cube计算 (矩阵乘法,使用TMATMUL,结果到L0C)
阶段4: L0C → GM (计算结果存储,使用TSTORE)
关键区别: - Vector计算: 使用UB(Unified Buffer)作为中间缓冲区,执行逐元素操作 - Cube计算: 使用L1和L0缓冲区(L0A/L0B/L0C),执行矩阵乘法操作
3.1 指令功能解释模板¶
对每个指令,按以下模板解释:
指令: [指令名称]
阶段: [GM/UB/Vector阶段]
功能: [具体功能描述]
数据流: [输入 → 输出的数据流向]
示例: [具体使用示例]
同步需求: [是否需要同步,如何同步]
3.2 阶段详解¶
阶段1: GM → UB (数据加载)
指令: TLOAD
功能: 从GlobalTensor (GM) 加载数据到Tile (UB)
数据流: GlobalMemory[srcGlobal] → UnifiedBuffer[srcTile]
输入: GlobalTensor对象,描述GM上的数据布局
输出: Tile对象,存储加载到UB的数据
同步需求:
- 推荐使用Event同步: Event<Op::TLOAD, Op::NextOp>
- 或手动同步: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0)
示例:
TLOAD(srcTile, srcGlobal);
event0 = TLOAD(src1Tile, src1Global); // 带Event返回
阶段2/3: Vector计算 (核心计算)
根据具体算子,选择相应的计算指令:
逐元素加法: TADD
功能: 两个Tile的逐元素加法
数据流: UB[src0Tile] + UB[src1Tile] → UB[dstTile]
输入: 两个源Tile
输出: 一个目标Tile
同步需求: Event<Op::TLOAD, Op::TADD>
示例:
event1 = TADD(dstTile, src0Tile, src1Tile, event0);
标量乘法: TMULS
功能: Tile与标量的逐元素乘法
数据流: UB[srcTile] * scalar → UB[dstTile]
输入: 一个源Tile + 一个标量值
输出: 一个目标Tile
同步需求: Event<Op::PreviousOp, Op::TMULS>
示例:
event2 = TMULS(dstTile, srcTile, (T)scalar, event1);
指数运算: TEXP
功能: Tile的逐元素指数运算 (e^x)
数据流: exp(UB[srcTile]) → UB[dstTile]
输入: 一个源Tile
输出: 一个目标Tile
同步需求: Event<Op::PreviousOp, Op::TEXP>
示例:
event3 = TEXP(dstTile, srcTile, event2);
最大值选择: TMAX
功能: 两个Tile的逐元素最大值
数据流: max(UB[src0Tile], UB[src1Tile]) → UB[dstTile]
输入: 两个源Tile
输出: 一个目标Tile
同步需求: Event<Op::PreviousOp, Op::TMAX]
示例:
event4 = TMAX(dstTile, src0Tile, src1Tile, event3);
Cube计算阶段: 矩阵乘法 (GM → L1 → L0A/L0B → L0C → GM)
矩阵乘法: TMATMUL
功能: 矩阵乘法计算 C = A * B
数据流:
- GM → L1: GlobalMemory[矩阵A/B] → L1Buffer[MatTile] (TLOAD)
- L1 → L0A/L0B: L1Buffer[MatTile] → L0Buffer[Left/RightTile] (TMOV)
- L0A/L0B → L0C: 矩阵乘法计算 (TMATMUL)
- L0C → GM: L0Buffer[AccTile] → GlobalMemory[结果] (TSTORE)
输入: 矩阵A和B (通过MatTile加载)
输出: 矩阵C (通过AccTile存储)
同步需求:
- TLOAD完成后: Event<Op::TLOAD, Op::TMOV> 或 set_flag(PIPE_MTE2, PIPE_MTE1)
- TMOV完成后: Event<Op::TMOV, Op::TMATMUL> 或 set_flag(PIPE_MTE1, PIPE_M)
- TMATMUL完成后: Event<Op::TMATMUL, Op::TSTORE_VEC> 或 set_flag(PIPE_M, PIPE_FIX)
示例:
// 1. 加载矩阵到L1
TLOAD(aMatTile, src0Global);
TLOAD(bMatTile, src1Global);
// 2. 搬运数据到L0A/L0B
#ifndef __PTO_AUTO__
set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
#endif
TMOV(aTile, aMatTile);
TMOV(bTile, bMatTile);
// 3. 矩阵乘法计算
#ifndef __PTO_AUTO__
set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
#endif
TMATMUL(cTile, aTile, bTile);
// 4. 存储结果
#ifndef __PTO_AUTO__
set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
#endif
TSTORE(dstGlobal, cTile);
阶段5: UB → GM (数据存储)
指令: TSTORE
功能: 将Tile数据存储到GlobalTensor (GM)
数据流: UnifiedBuffer[dstTile] → GlobalMemory[dstGlobal]
输入: Tile对象,UB上的数据
输出: GlobalTensor对象,GM上的数据
同步需求:
- 推荐使用Event同步: Event<Op::LastOp, Op::TSTORE_VEC>
- 或手动同步: set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0)
示例:
TSTORE(dstGlobal, dstTile, eventLast);
输出: 按数据流顺序的完整指令功能解释文档
步骤4: 输出kernel代码实现¶
目标: 生成完整的、可运行的kernel代码。
代码结构:
/**
Copyright (c) 2025 Huawei Technologies Co., Ltd.
...
*/
#include <pto/pto-inst.hpp>
#include "acl/acl.h"
using namespace pto;
namespace OperatorName {
// ==================== Device函数 ====================
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
__global__ AICORE void runOperator(__gm__ T *out, __gm__ T *src0, ...)
{
// 1. 类型定义
using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>;
using DynStrideDim5 = Stride<1, 1, 1, vCols, 1>;
using GlobalData = GlobalTensor<T, DynShapeDim5, DynStrideDim5>;
using TileData = Tile<TileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1>;
// 2. Tile和GlobalTensor声明
TileData src0Tile(vRows, vCols);
TileData dstTile(vRows, vCols);
TASSIGN(src0Tile, 0x0);
TASSIGN(dstTile, sizeof(T) * TileData::Numel);
GlobalData src0Global(src0);
GlobalData dstGlobal(out);
// 3. Event声明 (推荐使用Event同步)
Event<Op::TLOAD, Op::CALC_OP> event0;
Event<Op::CALC_OP, Op::TSTORE_VEC> event1;
// 4. 数据加载 (gm → ub)
event0 = TLOAD(src0Tile, src0Global);
// 5. 核心计算 (vector计算)
event1 = CALC_OP(dstTile, src0Tile, ..., event0);
// 6. 数据存储 (ub → gm)
TSTORE(dstGlobal, dstTile, event1);
out = dstGlobal.data();
}
// ==================== Host函数 ====================
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
void launchOperator(T *out, T *src0, ..., void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16>) {
runOperator<half, kTRows_, kTCols_, vRows, vCols>
<<<1, nullptr, stream>>>((half *)out, (half *)src0, ...);
} else {
runOperator<T, kTRows_, kTCols_, vRows, vCols><<<1, nullptr, stream>>>(out, src0, ...);
}
}
// ==================== 模板实例化 ====================
template void launchOperator<float, 64, 64, 64, 64>(float *out, float *src0, ..., void *stream);
template void launchOperator<aclFloat16, 16, 256, 16, 256>(aclFloat16 *out, aclFloat16 *src0, ..., void *stream);
} // namespace OperatorName
代码生成要点:
- 命名规范: kernel文件命名为
t<操作指令>_kernel.cpp - 模板参数: T (数据类型), kTRows_, kTCols_ (Tile维度), vRows, vCols (有效数据维度)
- 类型转换: aclFloat16需要转换为half类型
- 缓冲区分配: 使用TASSIGN紧密排布Tile地址
- 同步策略: 推荐使用Event同步,备选手动标志同步
- 模板实例化: 为常用配置提供实例化
输出: 完整的kernel代码文件
ISA指令分类参考¶
内存指令 (GM <-> Tile)¶
| 指令 | 功能 | 数据流 | 适用场景 |
|---|---|---|---|
| TLOAD | GM → UB/L1 | GlobalMemory → UnifiedBuffer/L1Buffer | Vector和Cube计算 |
| TSTORE | UB/L0C → GM | UnifiedBuffer/L0Buffer → GlobalMemory | Vector和Cube计算 |
| TPREFETCH | 预取到UB缓存 | 提示性预取 | Vector计算优化 |
| MGATHER | 索引收集加载 | GM[索引] → UB | Vector计算 |
| MSCATTER | 索引散播存储 | UB → GM[索引] | Vector计算 |
注意: - TLOAD可以加载到UB (Vector计算) 或L1Buffer (Cube计算) - TSTORE可以存储UB (Vector计算) 或L0C (Cube计算) 的数据
逐元素计算指令 (Tile-Tile)¶
| 指令 | 功能 | 表达式 |
|---|---|---|
| TADD | 逐元素加法 | dst = src0 + src1 |
| TSUB | 逐元素减法 | dst = src0 - src1 |
| TMUL | 逐元素乘法 | dst = src0 * src1 |
| TDIV | 逐元素除法 | dst = src0 / src1 |
| TMAX | 逐元素最大值 | dst = max(src0, src1) |
| TMIN | 逐元素最小值 | dst = min(src0, src1) |
| TCMP | 比较(生成掩码) | predicate = cmp(src0, src1) |
| TSHL | 逐元素左移 | dst = src0 << src1 |
| TSHR | 逐元素右移 | dst = src0 >> src1 |
| TAND | 逐元素按位与 | dst = src0 & src1 |
| TOR | 逐元素按位或 | dst = src0 |
| TXOR | 逐元素按位异或 | dst = src0 ^ src1 |
| TNOT | 逐元素按位取反 | dst = ~src |
标量操作指令 (Tile-标量)¶
| 指令 | 功能 | 表达式 |
|---|---|---|
| TADDS | Tile加标量 | dst = src + scalar |
| TSUBS | Tile减标量 | dst = src - scalar |
| TMULS | Tile乘标量 | dst = src * scalar |
| TDIVS | Tile除标量 | dst = src / scalar |
| TMINS | Tile与标量最小值 | dst = min(src, scalar) |
| TMAXS | Tile与标量最大值 | dst = max(src, scalar) |
| TCMPS | Tile与标量比较 | predicate = cmp(src, scalar) |
| TEXPANDS | 标量广播到Tile | dst = broadcast(scalar) |
数学函数指令¶
| 指令 | 功能 | 表达式 |
|---|---|---|
| TLOG | 自然对数 | dst = log(src) |
| TEXP | 指数运算 | dst = exp(src) |
| TSQRT | 平方根 | dst = sqrt(src) |
| TRSQRT | 倒数平方根 | dst = 1/sqrt(src) |
| TPOW | 幂运算 | dst = src0 ^ src1 |
| TRECIP | 倒数 | dst = 1/src |
| TABS | 绝对值 | dst = abs(src) |
| TNEG | 取负 | dst = -src |
激活函数指令¶
| 指令 | 功能 | 表达式 |
|---|---|---|
| TRELU | ReLU | dst = max(0, src) |
| TPRELU | PReLU | dst = max(0, src) + slope * min(0, src) |
| TLRELU | Leaky ReLU (标量斜率) | dst = max(0, src) + scalar * min(0, src) |
轴归约/扩展指令¶
| 指令 | 功能 | 操作 |
|---|---|---|
| TROWSUM | 行求和 | 每行所有列求和 |
| TROWPROD | 行乘积 | 每行所有列乘积 |
| TROWMAX | 行最大值 | 每行所有列最大值 |
| TROWMIN | 行最小值 | 每行所有列最小值 |
| TROWARGMAX | 行argmax | 每行最大值列索引 |
| TROWARGMIN | 行argmin | 每行最小值列索引 |
| TCOLSUM | 列求和 | 每列所有行求和 |
| TCOLPROD | 列乘积 | 每列所有行乘积 |
| TCOLMAX | 列最大值 | 每列所有行最大值 |
| TCOLMIN | 列最小值 | 每列所有行最小值 |
| TCOLARGMAX | 列argmax | 每列最大值行索引 |
| TCOLARGMIN | 列argmin | 每列最小值行索引 |
| TROWEXPAND | 行广播 | 将行首元素广播到整行 |
| TCOLEXPAND | 列广播 | 将列首元素广播到整列 |
广播运算指令¶
| 指令 | 功能 | 操作 |
|---|---|---|
| TROWEXPANDADD | 行广播加法 | 每行 + 广播标量向量 |
| TROWEXPANDSUB | 行广播减法 | 每行 - 广播标量向量 |
| TROWEXPANDMUL | 行广播乘法 | 每行 * 广播标量向量 |
| TROWEXPANDDIV | 行广播除法 | 每行 / 广播标量向量 |
| TROWEXPANDMAX | 行广播最大值 | max(每行, 广播标量向量) |
| TROWEXPANDMIN | 行广播最小值 | min(每行, 广播标量向量) |
| TROWEXPANDEXPDIF | 行指数差 | exp(每行 - 广播标量向量) |
| TCOLEXPANDADD | 列广播加法 | 每列 + 广播标量向量 |
| TCOLEXPANDSUB | 列广播减法 | 每列 - 广播标量向量 |
| TCOLEXPANDMUL | 列广播乘法 | 每列 * 广播标量向量 |
| TCOLEXPANDDIV | 列广播除法 | 每列 / 广播标量向量 |
| TCOLEXPANDMAX | 列广播最大值 | max(每列, 广播标量向量) |
| TCOLEXPANDMIN | 列广播最小值 | min(每列, 广播标量向量) |
| TCOLEXPANDEXPDIF | 列指数差 | exp(每列 - 广播标量向量) |
数据搬运指令 (缓冲区数据移动)¶
| 指令 | 功能 | 数据流 | 适用场景 |
|---|---|---|---|
| TMOV | L1 → L0A/L0B | MatTile → LeftTile/RightTile | Cube计算 (矩阵乘法) |
| TMOV | Tile之间移动 | srcTile → dstTile | 数据格式转换 |
| TMOV_FP | 带缩放的移动 | srcTile * scale → dstTile | 量化操作 |
| TRESHAPE | Tile重解释 | 保持字节,改变类型/形状 | 类型转换 |
| TTRANS | Tile转置 | srcTile^T → dstTile | 矩阵转置 |
TMOV在矩阵乘法中的关键作用: - 将L1Buffer的MatTile数据搬运到L0Buffer - 准备LeftTile和RightTile供Cube计算单元使用 - 数据流: L1 → L0A/L0B → TMATMUL → L0C
类型转换与选择指令¶
| 指令 | 功能 | 操作 |
|---|---|---|
| TCVT | 类型转换 | src_type → dst_type |
| TSEL | 条件选择(Tile) | mask ? src0 : src1 |
| TSELS | 条件选择(Tile-标量) | mask ? src : scalar |
矩阵运算指令 (使用Cube核心)¶
重要: 矩阵运算使用Cube核心,数据流为 GM → L1 → L0A/L0B → L0C → GM
| 指令 | 功能 | 数据流 | 表达式 |
|---|---|---|---|
| TLOAD | GM → L1 | GlobalMemory → L1Buffer (加载矩阵数据) | MatTile加载 |
| TMOV | L1 → L0A/L0B | L1Buffer → L0Buffer (搬运到计算单元) | MatTile → LeftTile/RightTile |
| TMATMUL | L0A/L0B → L0C | Cube矩阵乘法计算 | C = A * B |
| TSTORE | L0C → GM | L0Buffer → GlobalMemory (存储结果) | AccTile → GlobalMemory |
| TMATMUL_ACC | 矩阵乘法(累加) | L0A/L0B → L0C (带累加) | C = A * B + C |
| TMATMUL_BIAS | 矩阵乘法(加偏置) | L0A/L0B → L0C + bias | C = A * B + bias |
| TGEMV | 矩阵向量乘 | L0A/L0B → L0C | y = A * x |
| TGEMV_ACC | 矩阵向量乘(累加) | L0A/L0B → L0C (带累加) | y = A * x + y |
| TGEMV_BIAS | 矩阵向量乘(加偏置) | L0A/L0B → L0C + bias | y = A * x + bias |
矩阵乘法完整数据流示例:
GM → L1 (TLOAD) → L0A/L0B (TMOV) → L0C (TMATMUL) → GM (TSTORE)
详细步骤:
1. TLOAD: 加载矩阵A和B从GM到L1Buffer (MatTile)
2. TMOV: 将MatTile数据搬运到L0Buffer (LeftTile和RightTile)
3. TMATMUL: 在Cube核心执行矩阵乘法,结果存储到L0C (AccTile)
4. TSTORE: 将AccTile结果存储到GM
三元运算指令¶
| 指令 | 功能 | 表达式 |
|---|---|---|
| TADDC | 三元加法 | dst = src0 + src1 + src2 |
| TSUBC | 三元减法 | dst = src0 - src1 + src2 |
| TADDSC | Tile+标量+Tile加法 | dst = src0 + scalar + src1 |
| TSUBSC | Tile-标量+Tile运算 | dst = src0 - scalar + src1 |
数据流分析框架¶
基本数据流模式¶
Vector计算数据流 (逐元素操作)¶
┌─────────────────────────────────────────────────────────────┐
│ Vector计算数据流 (gm → ub → vector → ub → gm) │
├─────────────────────────────────────────────────────────────┤
│ │
│ GlobalMemory (GM) │
│ │ │
│ │ TLOAD │
│ ↓ │
│ UnifiedBuffer (UB) │
│ │ │
│ │ 计算指令 (TADD/TMUL/TEXP等) │
│ ↓ │
│ Vector计算单元 │
│ │ │
│ │ 计算结果 │
│ ↓ │
│ UnifiedBuffer (UB) │
│ │ │
│ │ TSTORE │
│ ↓ │
│ GlobalMemory (GM) │
│ │
└─────────────────────────────────────────────────────────────┘
Cube计算数据流 (矩阵乘法)¶
┌─────────────────────────────────────────────────────────────┐
│ Cube计算数据流 (GM → L1 → L0 → GM) │
├─────────────────────────────────────────────────────────────┤
│ │
│ GlobalMemory (GM) │
│ │ │
│ │ TLOAD │
│ ↓ │
│ L1Buffer (L1) │
│ │ │
│ │ TMOV │
│ ↓ │
│ L0Buffer (L0A/L0B) │
│ │ │
│ │ TMATMUL │
│ ↓ │
│ L0Buffer (L0C) │
│ │ │
│ │ TSTORE │
│ ↓ │
│ GlobalMemory (GM) │
│ │
└─────────────────────────────────────────────────────────────┘
两种数据流的区别:
| 特性 | Vector计算 | Cube计算 |
|---|---|---|
| 计算单元 | Vector Unit (PIPE_V) | Matrix Unit (PIPE_M) |
| 中间缓冲 | UnifiedBuffer (UB) | L1Buffer + L0Buffer (L0A/L0B/L0C) |
| 适用场景 | 逐元素操作 (TADD/TMUL/TEXP等) | 矩阵乘法 (TMATMUL) |
| 数据流路径 | GM → UB → V → UB → GM | GM → L1 → L0A/L0B → L0C → GM |
| 同步流水线 | MTE2 → V → MTE3 | MTE2 → MTE1 → M → FIX → MTE3 |
同步机制¶
Event同步(推荐):
Event<Op::TLOAD, Op::TADD> event0;
Event<Op::TADD, Op::TSTORE_VEC> event1;
event0 = TLOAD(srcTile, srcGlobal); // TLOAD完成时event0触发
event1 = TADD(dstTile, src0Tile, src1Tile, event0); // 等待event0,完成后触发event1
TSTORE(dstGlobal, dstTile, event1); // 等待event1
手动标志同步:
TLOAD(src0Tile, src0Global);
TLOAD(src1Tile, src1Global);
#ifndef __PTO_AUTO__
set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); // MTE2(内存加载) → V(向量计算)
wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); // 等待内存加载完成
#endif
TADD(dstTile, src0Tile, src1Tile);
#ifndef __PTO_AUTO__
set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // V(向量计算) → MTE3(内存存储)
wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // 等待计算完成
#endif
TSTORE(dstGlobal, dstTile);
流水线阶段¶
| 流水线 | 缩写 | 功能 | 适用场景 |
|---|---|---|---|
| Memory Transfer Engine 1 | PIPE_MTE1 | 矩阵数据搬运 (L1 → L0) | Cube计算 (矩阵乘法) |
| Memory Transfer Engine 2 | PIPE_MTE2 | 向量/矩阵数据加载 (GM → UB/L1) | Vector和Cube计算 |
| Memory Transfer Engine 3 | PIPE_MTE3 | 数据存储 (UB/L0C → GM) | Vector和Cube计算 |
| Vector Unit | PIPE_V | 向量计算 (逐元素操作) | Vector计算 (TADD/TMUL等) |
| Matrix Unit | PIPE_M | 矩阵计算 (矩阵乘法) | Cube计算 (TMATMUL) |
| Fix Unit | PIPE_FIX | 格式转换 | Cube计算结果格式化 |
| Scalar Unit | PIPE_S | 标量计算 | 标量操作 |
| All Pipelines | PIPE_ALL | 所有流水线 | 全局同步 |
流水线同步策略: - Vector计算: MTE2 (GM → UB) → V (计算) → MTE3 (UB → GM) - Cube计算: MTE2 (GM → L1) → MTE1 (L1 → L0) → M (计算) → FIX (格式转换) → MTE3 (L0 → GM)
核间同步机制 (TPUSH/TPOP)¶
重要: Cube核与Vector核之间的数据传输必须使用TPUSH/TPOP,不能使用简单的TMOV。
核间同步架构¶
┌─────────────────────────────────────────────┐
│ A5 AI Core Architecture │
├─────────────────────────────────────────────┤
│ │
│ ┌──────────────┐ ┌──────────────────┐ │
│ │ Cube Core │ │ Vector Core 0 │ │
│ │ (PIPE_M) │ │ (PIPE_V) │ │
│ │ L1/L0 │ │ UB │ │
│ └ │ │ │ │
│ │ TPUSH (C2V) │ │ TPOP (C2V) │ │
│ │ TPOP (V2C) │ │ TPUSH (V2C) │ │
│ └ │ │ │ │
│ └──────────────┘ └──────────────────┘ │
│ │
│ ┌──────────────────┐ │
│ │ Vector Core 1 │ │
│ │ (PIPE_V) │ │
│ │ UB │ │
│ │ │ │
│ │ TPOP (C2V) │ │
│ │ TPUSH (V2C) │ │
│ │ │ │
│ └──────────────────┘ │
│ │
│ 核间同步:TPUSH → TPOP → TFREE │
│ │
└─────────────────────────────────────────────┘
核间同步方向类型¶
| 方向类型 | 定义 | 数据流 | 生产者流水线 | 消费者流水线 |
|---|---|---|---|---|
| DIR_C2V | Cube → Vector | L0C → UB | PIPE_FIX | PIPE_V |
| DIR_V2C | Vector → Cube | UB → L1 | PIPE_MTE3 | PIPE_MTE1 |
| DIR_BOTH | 双向 | L0C ↔ UB | PIPE_FIX + PIPE_MTE3 | PIPE_V + PIPE_MTE1 |
Vector/Cube核代码区分宏定义¶
重要: PTO-ISA使用编译器宏来区分Vector核和Cube核的执行路径。同一份kernel代码会被编译两次,分别生成Vector核和Cube核的可执行文件。
宏定义模式:
// 编译器在编译不同核时自动定义以下宏:
// - 编译Vector核时:定义 __DAV_VEC__
// - 编译Cube核时:定义 __DAV_CUBE__
#ifdef __DAV_CUBE__
constexpr bool DAV_CUBE = true;
#else
constexpr bool DAV_CUBE = false;
#endif
#ifdef __DAV_VEC__
constexpr bool DAV_VEC = true;
#else
constexpr bool DAV_VEC = false;
#endif
使用示例:
template <typename T, int M, int K, int N>
__global__ AICORE void runOperator(__gm__ T *out, __gm__ T *srcA, __gm__ T *srcB)
{
// Vector核执行路径
if constexpr (DAV_VEC) {
// Vector计算:逐元素操作、激活函数、归约等
TLOAD(vecTile, srcGlobal);
TADD(dstTile, src0Tile, src1Tile);
// V2C: TPUSH数据到Cube核
TPUSH<V2CPipe, VecTileNZ, TileSplitAxis::TILE_NO_SPLIT>(pipe, vecTile);
// C2V: TPOP从Cube核接收数据
TPOP<C2VPipe, VecTile, TileSplitAxis::TILE_NO_SPLIT>(pipe, recvTile);
TFREE<C2VPipe, TileSplitAxis::TILE_NO_SPLIT>(pipe);
TSTORE(dstGlobal, dstTile);
}
// Cube核执行路径
if constexpr (DAV_CUBE) {
// Cube计算:矩阵乘法
TLOAD(matTileA, srcAGlobal);
TLOAD(matTileB, srcBGlobal);
// V2C: TPOP从Vector核接收数据
TPOP<V2CPipe, MatTile, TileSplitAxis::TILE_NO_SPLIT>(pipe, matTileB);
TFREE<V2CPipe, TileSplitAxis::TILE_NO_SPLIT>(pipe);
TMOV(leftTile, matTileA);
TMOV(rightTile, matTileB);
TMATMUL(accTile, leftTile, rightTile);
// C2V: TPUSH数据到Vector核
TPUSH<C2VPipe, AccTile, TileSplitAxis::TILE_NO_SPLIT>(pipe, accTile);
TSTORE(dstGlobal, accTile);
}
}
宏定义规则:
| 宏名称 | 定义时机 | 适用场景 |
|---|---|---|
__DAV_VEC__ |
编译Vector核时 | Vector计算、UB操作、PIPE_V流水线 |
__DAV_CUBE__ |
编译Cube核时 | Cube计算、L1/L0操作、PIPE_M流水线 |
注意事项:
1. 使用 if constexpr (DAV_VEC) 和 if constexpr (DAV_CUBE) 进行分支判断
2. 不要在同一个核的执行路径中混用另一核的Tile类型
3. TPUSH/TPOP必须成对使用:Vector核TPUSH对应Cube核TPOP
4. 每个TPOP后必须调用TFREE释放缓冲区
TPUSH三步流程¶
TPUSH用于生产者核推送数据到消费者核:
步骤1: Alloc (分配空间)
- 生产者核等待消费者核释放空间
- C2V: wait_intra_block(PIPE_FIX, FlagID+1)
- V2C: wait_intra_block(PIPE_MTE3, FlagID+1)
步骤2: Store (写入数据)
- 根据TileType和FIFO类型选择搬运方式
- AccTile → VecFIFO: pushAcc2VecFiFo (L0C → UB)
- VecTile → MatFIFO: pushVec2MatFiFo (UB → L1)
步骤3: Commit (信号通知)
- 通知消费者核数据已就绪
- C2V: set_intra_block(PIPE_FIX, FlagID)
- V2C: set_intra_block(PIPE_MTE3, FlagID)
TPOP三步流程¶
TPOP用于消费者核从生产者核读取数据:
步骤1: Wait (等待数据)
- 消费者核等待生产者核数据就绪
- C2V: wait_intra_block(PIPE_V, FlagID)
- V2C: wait_intra_block(PIPE_MTE1, FlagID)
步骤2: Pop (读取数据)
- 根据TileType和FIFO类型选择读取方式
- VecFIFO → VecTile: popTileFromVecFiFo
- MatFIFO → MatTile: popTileFromMatFiFo
步骤3: Free (释放空间) - 通知生产者核空间已释放 - 使用TFREE指令
TPipe结构定义¶
template <uint8_t FlagID, uint8_t DirType, uint32_t SlotSize, uint32_t SlotNum>
using TPipe = TPipe<FlagID, DirType, SlotSize, SlotNum>;
// 参数说明:
// FlagID: 核间同步标志ID (0-7)
// DirType: 通信方向 (DIR_C2V=1, DIR_V2C=2, DIR_BOTH=3)
// SlotSize: FIFO槽大小(字节)
// SlotNum: FIFO槽数量(建议2)
// TPipe初始化:
// GM_SLOT_BUFFER: GM FIFO基地址
// C2V_CONSUMER_BUF: Cube→Vec消费者UB地址
// V2C_CONSUMER_BUF: Vec→Cube消费者L1地址
using MatPipe = TPipe<FLAG_ID, Direction::DIR_C2V, sizeof(T) * M * N, 2>;
MatPipe mPipe((__gm__ void *)(uint64_t)0x0, (uint32_t)0x0, (uint32_t)0x20000);
TileSplitAxis分块模式¶
| SplitAxis | 说明 | Vector核分配 |
|---|---|---|
| TILE_UP_DOWN | 沿行分块 | Vec0处理上半部分,Vec1处理下半部分 |
| TILE_LEFT_RIGHT | 沿列分块 | Vec0处理左半部分,Vec1处理右半部分 |
| TILE_NO_SPLIT | 不分块 | 单Vector核处理全部 |
FlagID分配策略¶
A5架构提供8个FlagID(0-7),用于核间同步:
| FlagID | 用途 | 说明 |
|---|---|---|
| FlagID | 数据就绪信号 | 生产者设置,消费者等待 |
| FlagID+1 | 空间释放信号 | 消费者设置,生产者等待 |
| FlagID+16 | Vec核1信号 | 双Vector核时使用 |
双Vector核时的FlagID分配:
Vec0: FlagID (主核)
Vec1: FlagID+16 (从核)
Cube核需要等待双核:
wait_intra_block(PIPE_FIX, FlagID); // Vec0信号
wait_intra_block(PIPE_FIX, FlagID+16); // Vec1信号
核间同步最佳实践¶
1. FlagID管理: 为每个TPipe分配独立的FlagID,避免冲突
2. FIFO深度设置: 推荐使用深度=2
3. 同步顺序匹配: 一个TPUSH必须对应一个TPOP + TFREE
4. 错误示例:
// 错误:连续两次TPUSH,没有对应的TPOP
TPUSH(pipe, tile1);
TPUSH(pipe, tile2); // ERROR
// 正确:
TPUSH(pipe, tile1);
// ... 消费者核 ...
TPOP(pipe, vecTile1);
TFREE(pipe);
// 然后才能进行下一次TPUSH
融合算子中的核间同步¶
当算子涉及Vector计算和Cube计算的交替使用时,需要在切换点使用TPUSH/TPOP:
| 切换场景 | 数据流 | 方向 | 核间同步 |
|---|---|---|---|
| Vector → Cube | UB → L1 | V2C | TPUSH (Vec) + TPOP (Cube) |
| Cube → Vector | L0C → UB | C2V | TPUSH (Cube) + TPOP (Vec) |
Flash Attention核间同步示例: - Phase 1 (Vector): K转置 → TPUSH K^T到Cube核 (V2C) - Phase 2 (Cube): QK矩阵乘 → TPUSH Score到Vector核 (C2V) - Phase 3 (Vector): Softmax归一化 → TPUSH P到Cube核 (V2C) - Phase 4 (Cube): PV矩阵乘 → TSTORE输出
详细参考: /home/developer/.agents/skills/pto-isa-operator-implementation/TPUSH_TPOP_GUIDE.md
示例¶
示例1: ReLU激活函数¶
算子功能: ReLU(x) = max(0, x)
ISA指令: TLOAD → TRELU → TSTORE
Kernel代码:
namespace ReLU {
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
__global__ AICORE void runReLU(__gm__ T *out, __gm__ T *src)
{
using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>;
using DynStrideDim5 = Stride<1, 1, 1, vCols, 1>;
using GlobalData = GlobalTensor<T, DynShapeDim5, DynStrideDim5>;
using TileData = Tile<TileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1>;
TileData srcTile(vRows, vCols);
TileData dstTile(vRows, vCols);
TASSIGN(srcTile, 0x0);
TASSIGN(dstTile, sizeof(T) * TileData::Numel);
GlobalData srcGlobal(src);
GlobalData dstGlobal(out);
Event<Op::TLOAD, Op::TRELU> event0;
Event<Op::TRELU, Op::TSTORE_VEC> event1;
event0 = TLOAD(srcTile, srcGlobal);
event1 = TRELU(dstTile, srcTile, event0);
TSTORE(dstGlobal, dstTile, event1);
out = dstGlobal.data();
}
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
void launchReLU(T *out, T *src, void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16>) {
runReLU<half, kTRows_, kTCols_, vRows, vCols>
<<<1, nullptr, stream>>>((half *)out, (half *)src);
} else {
runReLU<T, kTRows_, kTCols_, vRows, vCols><<<1, nullptr, stream>>>(out, src);
}
}
template void launchReLU<float, 64, 64, 64, 64>(float *out, float *src, void *stream);
template void launchReLU<aclFloat16, 16, 256, 16, 256>(aclFloat16 *out, aclFloat16 *src, void *stream);
} // namespace ReLU
示例2: 逐元素加法 (TADD)¶
算子功能: dst = src0 + src1
ISA指令: TLOAD → TLOAD → TADD → TSTORE
Kernel代码:
namespace TAdd {
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
__global__ AICORE void runTAdd(__gm__ T *out, __gm__ T *src0, __gm__ T *src1)
{
using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>;
using DynStrideDim5 = Stride<1, 1, 1, vCols, 1>;
using GlobalData = GlobalTensor<T, DynShapeDim5, DynStrideDim5>;
using TileData = Tile<TileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1>;
TileData src0Tile(vRows, vCols);
TileData src1Tile(vRows, vCols);
TileData dstTile(vRows, vCols);
TASSIGN(src0Tile, 0x0);
TASSIGN(src1Tile, sizeof(T) * TileData::Numel);
TASSIGN(dstTile, 2 * sizeof(T) * TileData::Numel);
GlobalData src0Global(src0);
GlobalData src1Global(src1);
GlobalData dstGlobal(out);
Event<Op::TLOAD, Op::TADD> event0;
Event<Op::TADD, Op::TSTORE_VEC> event1;
TLOAD(src0Tile, src0Global);
event0 = TLOAD(src1Tile, src1Global);
event1 = TADD(dstTile, src0Tile, src1Tile, event0);
TSTORE(dstGlobal, dstTile, event1);
out = dstGlobal.data();
}
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
void launchTAdd(T *out, T *src0, T *src1, void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16>) {
runTAdd<half, kTRows_, kTCols_, vRows, vCols>
<<<1, nullptr, stream>>>((half *)out, (half *)src0, (half *)src1);
} else {
runTAdd<T, kTRows_, kTCols_, vRows, vCols><<<1, nullptr, stream>>>(out, src0, src1);
}
}
template void launchTAdd<float, 64, 64, 64, 64>(float *out, float *src0, float *src1, void *stream);
template void launchTAdd<aclFloat16, 16, 256, 16, 256>(aclFloat16 *out, aclFloat16 *src0, aclFloat16 *src1, void *stream);
} // namespace TAdd
示例3: Softmax归一化¶
算子功能: Softmax(x) = exp(x) / sum(exp(x))
ISA指令: TLOAD → TEXP → TCOLSUM → TCOLEXPANDDIV → TSTORE
Kernel代码:
namespace Softmax {
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
__global__ AICORE void runSoftmax(__gm__ T *out, __gm__ T *src)
{
using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>;
using DynStrideDim5 = Stride<1, 1, 1, vCols, 1>;
using GlobalData = GlobalTensor<T, DynShapeDim5, DynStrideDim5>;
using TileData = Tile<TileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1>;
using SumTileData = Tile<TileType::Vec, T, vRows, 1, BLayout::ColMajor, -1, -1>;
TileData srcTile(vRows, vCols);
TileData expTile(vRows, vCols);
TileData dstTile(vRows, vCols);
SumTileData sumTile(vRows, 1);
TASSIGN(srcTile, 0x0);
TASSIGN(expTile, sizeof(T) * TileData::Numel);
TASSIGN(sumTile, 2 * sizeof(T) * TileData::Numel);
TASSIGN(dstTile, 3 * sizeof(T) * TileData::Numel);
GlobalData srcGlobal(src);
GlobalData dstGlobal(out);
Event<Op::TLOAD, Op::TEXP> event0;
Event<Op::TEXP, Op::TCOLSUM> event1;
Event<Op::TCOLSUM, Op::TCOLEXPANDDIV> event2;
Event<Op::TCOLEXPANDDIV, Op::TSTORE_VEC> event3;
event0 = TLOAD(srcTile, srcGlobal);
event1 = TEXP(expTile, srcTile, event0);
event2 = TCOLSUM(sumTile, expTile, event1);
event3 = TCOLEXPANDDIV(dstTile, expTile, sumTile, event2);
TSTORE(dstGlobal, dstTile, event3);
out = dstGlobal.data();
}
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
void launchSoftmax(T *out, T *src, void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16>) {
runSoftmax<half, kTRows_, kTCols_, vRows, vCols>
<<<1, nullptr, stream>>>((half *)out, (half *)src);
} else {
runSoftmax<T, kTRows_, kTCols_, vRows, vCols><<<1, nullptr, stream>>>(out, src);
}
}
template void launchSoftmax<float, 64, 64, 64, 64>(float *out, float *src, void *stream);
template void launchSoftmax<aclFloat16, 16, 256, 16, 256>(aclFloat16 *out, aclFloat16 *src, void *stream);
} // namespace Softmax
示例4: Batch Normalization¶
算子功能: BN(x) = (x - mean) / sqrt(var + eps) * gamma + beta
ISA指令: TLOAD → TSUBS → TDIVS → TMULS → TADDS → TSTORE
Kernel代码:
namespace BatchNorm {
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
__global__ AICORE void runBatchNorm(__gm__ T *out, __gm__ T *src,
float mean, float var, float eps, float gamma, float beta)
{
using DynShapeDim5 = Shape<1, 1, 1, vRows, vCols>;
using DynStrideDim5 = Stride<1, 1, 1, vCols, 1>;
using GlobalData = GlobalTensor<T, DynShapeDim5, DynStrideDim5>;
using TileData = Tile<TileType::Vec, T, kTRows_, kTCols_, BLayout::RowMajor, -1, -1>;
TileData srcTile(vRows, vCols);
TileData normTile(vRows, vCols);
TileData dstTile(vRows, vCols);
TASSIGN(srcTile, 0x0);
TASSIGN(normTile, sizeof(T) * TileData::Numel);
TASSIGN(dstTile, 2 * sizeof(T) * TileData::Numel);
GlobalData srcGlobal(src);
GlobalData dstGlobal(out);
Event<Op::TLOAD, Op::TSUBS> event0;
Event<Op::TSUBS, Op::TDIVS> event1;
Event<Op::TDIVS, Op::TMULS> event2;
Event<Op::TMULS, Op::TADDS> event3;
Event<Op::TADDS, Op::TSTORE_VEC> event4;
T std_val = (T)sqrt(var + eps);
event0 = TLOAD(srcTile, srcGlobal);
event1 = TSUBS(normTile, srcTile, (T)mean, event0);
event2 = TDIVS(normTile, normTile, std_val, event1);
event3 = TMULS(dstTile, normTile, (T)gamma, event2);
event4 = TADDS(dstTile, dstTile, (T)beta, event3);
TSTORE(dstGlobal, dstTile, event4);
out = dstGlobal.data();
}
template <typename T, int kTRows_, int kTCols_, int vRows, int vCols>
void launchBatchNorm(T *out, T *src, float mean, float var, float eps, float gamma, float beta, void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16>) {
runBatchNorm<half, kTRows_, kTCols_, vRows, vCols>
<<<1, nullptr, stream>>>((half *)out, (half *)src, mean, var, eps, gamma, beta);
} else {
runBatchNorm<T, kTRows_, kTCols_, vRows, vCols><<<1, nullptr, stream>>>(out, src, mean, var, eps, gamma, beta);
}
}
template void launchBatchNorm<float, 64, 64, 64, 64>(float *out, float *src, float mean, float var, float eps, float gamma, float beta, void *stream);
template void launchBatchNorm<aclFloat16, 16, 256, 16, 256>(aclFloat16 *out, aclFloat16 *src, float mean, float var, float eps, float gamma, float beta, void *stream);
} // namespace BatchNorm
示例5: 矩阵乘法 (TMATMUL) - Cube核心¶
算子功能: C = A * B (矩阵乘法)
ISA指令: TLOAD → TMOV → TMATMUL → TSTORE (GM → L1 → L0 → L0C → GM)
重要: 矩阵乘法使用Cube核心,需要使用DAV_CUBE宏判断执行路径。
Kernel代码:
namespace MatMul {
#ifdef __DAV_CUBE__
constexpr bool DAV_CUBE = true;
#else
constexpr bool DAV_CUBE = false;
#endif
template <typename T, typename U, typename S, int validM, int validK, int validN>
__global__ AICORE void runMatMul(__gm__ T *out, __gm__ U *src0, __gm__ S *src1)
{
if constexpr (DAV_CUBE) {
constexpr int blockAlign = C0_SIZE_BYTE / sizeof(U);
constexpr int M = CeilAlign<int>(validM, 16);
constexpr int N = CeilAlign<int>(validN, blockAlign);
constexpr int K = CeilAlign<int>(validK, blockAlign);
using GlobalDataSrc0 = GlobalTensor<U, pto::Shape<1, 1, 1, validM, validK>,
pto::Stride<1 * validM * validK, 1 * validM * validK, validM * validK, validK, 1>>;
using GlobalDataSrc1 = GlobalTensor<S, pto::Shape<1, 1, 1, validK, validN>,
pto::Stride<1 * validK * validN, 1 * validK * validN, validK * validN, validN, 1>>;
using GlobalDataOut = GlobalTensor<T, pto::Shape<1, 1, 1, validM, validN>,
pto::Stride<1 * validM * validN, 1 * validM * validN, validM * validN, validN, 1>>;
GlobalDataSrc0 src0Global(src0);
GlobalDataSrc1 src1Global(src1);
GlobalDataOut dstGlobal(out);
using TileMatAData = Tile<TileType::Mat, U, M, K, BLayout::ColMajor, validM, validK, SLayout::RowMajor, 512>;
using TileMatBData = Tile<TileType::Mat, S, K, N, BLayout::ColMajor, validK, validN, SLayout::RowMajor, 512>;
using LeftTile = TileLeft<U, M, K, validM, validK>;
using RightTile = TileRight<S, K, N, validK, validN>;
using AccTile = TileAcc<T, M, N, validM, validN>;
TileMatAData aMatTile;
TileMatBData bMatTile;
LeftTile aTile;
RightTile bTile;
AccTile cTile;
TASSIGN(aMatTile, 0x0);
TASSIGN(bMatTile, 0x20000);
TLOAD(aMatTile, src0Global);
TLOAD(bMatTile, src1Global);
#ifndef __PTO_AUTO__
set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0);
#endif
TMOV(aTile, aMatTile);
TMOV(bTile, bMatTile);
#ifndef __PTO_AUTO__
set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0);
#endif
TMATMUL(cTile, aTile, bTile);
#ifndef __PTO_AUTO__
set_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0);
#endif
TSTORE(dstGlobal, cTile);
out = dstGlobal.data();
}
}
template <typename T, typename U, typename S, int validM, int validK, int validN>
void launchMatMul(T *out, U *src0, S *src1, void *stream)
{
if constexpr (std::is_same_v<T, aclFloat16> || std::is_same_v<U, aclFloat16> || std::is_same_v<S, aclFloat16>) {
runMatMul<half, half, half, validM, validK, validN>
<<<1, nullptr, stream>>>((half *)out, (half *)src0, (half *)src1);
} else {
runMatMul<T, U, S, validM, validK, validN><<<1, nullptr, stream>>>(out, src0, src1);
}
}
template void launchMatMul<float, float, float, 16, 16, 16>(float *out, float *src0, float *src1, void *stream);
template void launchMatMul<half, half, half, 16, 16, 16>(half *out, half *src0, half *src1, void *stream);
} // namespace MatMul
最佳实践¶
1. ISA指令选择原则¶
最小化指令数量: 使用最少指令完成功能,减少数据搬运开销。
优先使用融合指令: 选择融合指令减少中间步骤: - TAXPY (融合乘加) - Vector计算 - TMATMUL_ACC (融合累加矩阵乘) - Cube计算 - TMATMUL_BIAS (融合加偏置矩阵乘) - Cube计算 - TROWEXPANDADD (融合广播加法) - Vector计算 - TADDC/TSUBC (融合三元运算) - Vector计算
选择合适的数据流: 根据算子特性选择最优数据流路径。 - Vector计算: 使用GM → UB → V → UB → GM数据流,适用于逐元素操作 - Cube计算: 使用GM → L1 → L0A/L0B → L0C → GM数据流,适用于矩阵乘法
2. 数据流优化¶
Vector计算优化: - 尽量在UB上进行多次计算 - 使用原地操作减少中间Tile - 合理使用缓冲区重用 - 避免不必要的TMOV操作
Cube计算优化: - 合理规划L1和L0缓冲区大小 - 使用TileType::Mat和TileType::Vec的转换 - 优化矩阵分块策略 - 减少GM访问次数
对齐和布局: - RowMajor: cols需要32字节对齐 - ColMajor: rows需要32字节对齐 - 使用constexpr计算对齐维度
3. 同步策略选择¶
推荐Event同步: - 自动依赖跟踪 - 编译器优化友好 - 代码简洁易维护 - 支持手动和自动模式
备选手动同步: - 复杂流水线控制 - 需要细粒度同步 - 性能关键路径优化
4. Tile维度选择¶
常见Tile维度配置: | 数据类型 | 推荐Tile维度 | |---------|-------------| | float | 64x64, 32x32, 16x16 | | aclFloat16 | 16x256, 8x768, 4x1024 | | int32 | 64x64, 32x32 | | int16 | 64x128, 32x256 |
5. 类型处理¶
aclFloat16转换: - API类型: aclFloat16 - 硬件类型: half - launch函数中进行转换
混合精度支持: - 使用模板参数支持多类型 - 标量参数统一使用float
6. 代码组织¶
命名规范:
- Kernel文件: t<操作指令>_kernel.cpp
- 命名空间: OperatorName
- 函数名: runOperator, launchOperator
模板实例化: - 为常用配置提供显式实例化 - 减少编译时间 - 确保代码可链接
常见问题¶
Q1: 如何选择合适的ISA指令组合?¶
回答: 1. 分析算子功能,分解为基本操作 2. 在PTOISA_zh.md中查找对应指令 3. 按数据流顺序排列指令 4. 检查指令间的依赖关系 5. 考虑融合指令减少步骤
Q2: 数据流顺序是什么?¶
回答: PTO有两种主要的数据流模式:
1. Vector计算数据流 (逐元素操作): gm → ub → vector → ub → gm - 阶段1: GM → UB (TLOAD) - 阶段2/3: Vector计算 (TADD/TMUL/TEXP等) - 阶段5: UB → GM (TSTORE)
2. Cube计算数据流 (矩阵乘法): GM → L1 → L0A/L0B → L0C → GM - 阶段1: GM → L1 (TLOAD) - 阶段2: L1 → L0A/L0B (TMOV) - 阶段3: Cube计算 (TMATMUL) - 阶段4: L0C → GM (TSTORE)
关键区别: - Vector计算使用UB缓冲区,适用于逐元素操作 - Cube计算使用L1和L0缓冲区,适用于矩阵乘法
Q3: 何时使用Event同步,何时使用手动同步?¶
回答: - Event同步(推荐): 简单融合、清晰依赖关系、自动模式支持 - 手动同步: 复杂流水线、细粒度控制、性能优化
Q4: 如何处理复杂算子(如GELU、LayerNorm)?¶
回答: 1. 将复杂算子分解为多个基本操作 2. 为每个基本操作选择对应ISA指令 3. 合理安排中间Tile缓冲区 4. 优化数据流减少搬运 5. 考虑使用近似计算简化实现
Q5: Tile维度如何选择?¶
回答: - 根据数据类型选择对齐维度 - 考虑片上存储容量限制 - 平衡计算效率和存储开销 - 使用constexpr计算对齐维度
Q6: 如何验证ISA指令选择是否正确?¶
回答: 1. 检查指令功能是否匹配算子需求 2. 验证数据流完整性(GM → UB → GM) 3. 确认同步机制正确设置 4. 在CPU模拟器上测试 5. 与golden结果对比验证
Q7: 标量参数如何处理?¶
回答:
- 标量参数统一使用float类型
- 在指令调用前转换为Tile数据类型: (T)scalar
- 使用标量指令(TADDS/TMULS等)而不是Tile指令
Q8: 何时使用Vector数据流,何时使用Cube数据流?¶
回答: 根据算子类型选择合适的数据流:
使用Vector数据流 (GM → UB → V → UB → GM): - 逐元素操作: TADD, TSUB, TMUL, TDIV, TMAX, TMIN - 数学函数: TEXP, TLOG, TSQRT, TPOW - 激活函数: TRELU, TPRELU, TLRELU - 标量操作: TADDS, TMULS, TDIVS - 轴归约: TROWSUM, TCOLSUM, TROWMAX - 广播操作: TROWEXPANDADD, TCOLEXPANDADD - 类型转换: TCVT
使用Cube数据流 (GM → L1 → L0A/L0B → L0C → GM): - 矩阵乘法: TMATMUL, TMATMUL_ACC, TMATMUL_BIAS - 矩阵向量乘: TGEMV, TGEMV_ACC, TGEMV_BIAS - 需要使用TileType::Mat的矩阵操作
判断方法: - 如果使用TileType::Vec → Vector数据流 - 如果使用TileType::Mat → Cube数据流
Q9: 矩阵乘法中TMOV的作用是什么?¶
回答: TMOV在矩阵乘法中用于数据搬运:
数据流: L1 → L0A/L0B
具体作用: - TLOAD将矩阵数据加载到L1Buffer (MatTile) - TMOV将MatTile数据搬运到L0Buffer (LeftTile和RightTile) - TMATMUL在L0Buffer执行计算
为什么需要TMOV: - L1Buffer和L0Buffer是不同的物理存储区域 - L1Buffer用于存储加载的原始数据 - L0Buffer是Cube计算单元的专用缓冲区 (L0A/L0B) - TMOV将数据从L1搬运到L0,准备矩阵乘法计算
Q10: 核间同步(TPUSH/TPOP)何时使用?¶
回答: 当算子涉及Vector核和Cube核之间的数据传输时,必须使用TPUSH/TPOP:
使用场景: - Vector核计算结果需要传给Cube核进行矩阵乘法 - Cube核矩阵乘法结果需要传给Vector核进行逐元素操作 - 融合算子中Vector/Cube交替使用
不使用TPUSH/TPOP的场景: - 同一核内部的数据搬运使用TMOV - 纯Vector计算或纯Cube计算不需要核间同步
参考资料¶
- ISA参考:
pto-isa/docs/PTOISA_zh.md- PTO指令索引 - ISA详细文档:
pto-isa/docs/isa/- 各指令详细说明 docs/isa/TMATMUL_zh.md- 矩阵乘法指令docs/isa/TLOAD_zh.md- 数据加载指令docs/isa/TMOV_zh.md- 数据搬运指令- C++ API:
include/pto/pto-inst.hpp- PTO指令C++接口 - 常量定义:
include/pto/common/constants.hpp- 流水线、事件ID等常量 - 测试示例:
tests/npu/a2a3/src/st/testcase/- 算子实现示例 - 融合算子指南:
vector-fusion-operator-generateskill - 融合算子开发完整流程
总结¶
本skill提供了使用PTO-ISA实现指定算子功能的完整流程:
- 步骤1: 阅读PTOISA_zh.md,了解指令集
- 步骤2: 分析算子需求,列举ISA指令
- 步骤3: 按数据流顺序解释指令功能
- Vector计算: GM → UB → V → UB → GM
- Cube计算: GM → L1 → L0A/L0B → L0C → GM
- 步骤4: 输出完整kernel代码
通过遵循本指南,开发者可以系统性地选择ISA指令、理解两种数据流模式(Vector和Cube)、生成高质量kernel代码。
关键要点: - Vector计算: 使用UB缓冲区,适用于逐元素操作,流水线 MTE2 → V → MTE3 - Cube计算: 使用L1和L0缓冲区,适用于矩阵乘法,流水线 MTE2 → MTE1 → M → FIX → MTE3 - TMOV关键作用: 在矩阵乘法中将L1数据搬运到L0,准备Cube计算