hw-native-sys/pto-isa/pto-isa-operator-implementation

用户指定算子功能
    ↓
步骤1: 阅读PTOISA_zh.md
    ↓
步骤2: 分析算子需求，列举ISA指令
    ↓
步骤3: 按数据流顺序解释指令功能
    ↓
步骤4: 输出kernel代码实现

算子: 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

数据流向: gm → ub → vector → ub → gm
阶段1: GM → UB (数据加载，使用TLOAD)
阶段2: UB → Vector (计算准备)
阶段3: Vector计算 (核心计算，使用TADD/TMUL/TEXP等)
阶段4: Vector → UB (计算结果)
阶段5: UB → GM (数据存储，使用TSTORE)

数据流向: GM → L1 → L0A/L0B → L0C → GM
阶段1: GM → L1 (矩阵数据加载，使用TLOAD)
阶段2: L1 → L0A/L0B (数据搬运到矩阵计算单元，使用TMOV)
阶段3: Cube计算 (矩阵乘法，使用TMATMUL，结果到L0C)
阶段4: L0C → GM (计算结果存储，使用TSTORE)

指令: [指令名称]
阶段: [GM/UB/Vector阶段]
功能: [具体功能描述]
数据流: [输入 → 输出的数据流向]
示例: [具体使用示例]
同步需求: [是否需要同步，如何同步]

指令: 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返回

逐元素加法: 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);

矩阵乘法: 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);

指令: 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);

/**
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