[Requirement|需求建议]: LpNormV2算子AscendC实现贡献 -ops-nn-AtomGit

已关闭

[Requirement|需求建议]: LpNormV2算子AscendC实现贡献 #218

石向阳创建于 2025年12月8日关闭于 5月12日

石向阳成员

2025年12月8日创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud（背景信息）

使用AscendC对LpNormV2算子进行实现，实现了AscendC实现的LpNormV2算子对Atlas 200/500 A2推理产品和Atlas 800I/T A2硬件的适配。

Origin（信息来源）

哈工大算子团队

Benefit / Necessity （价值/作用）

LpNormV2 算子的主要功能是对输入张量在指定维度（dim）上计算 Lp 次范数并归一化，支持配置范数阶数 p（如 p=1、p=2、p=∞ 等），支持配置轴（默认为-1即求全局范数并规约，支持axis=0，=1，分别按照行和列归约）常用于深度学习模型的特征归一化（如防止梯度爆炸 / 消失、统一特征尺度）、距离度量（如 L2 范数用于欧氏距离计算）、注意力机制的权重规范化、生成模型与分类模型的正则化约束、计算机视觉的特征提取后处理、自然语言处理的词向量标准化等场景。
目前该算子支持 float / float16 数据类型，支持维度为 1–2D Tensor

Design（设计方案）

一、总体设计目标

LpNormV3 算子用于在 Ascend C 平台上实现通用 Lp 范数归一化计算，支持 axis = -1 / 0 / 1 以及 p 为有限值、正无穷、负无穷三类情形。整体设计目标如下：

功能完备：完整覆盖 Lp-Norm 的数学定义，包括 ±∞ 范数的特殊分支。
高并行度：通过 Host 侧 tiling 与 Kernel 侧多核协同，实现数据级并行。
数值稳定性：通过 epsilon、原子累加与多级同步保证数值正确性。

二、Host 端设计方案（LpNormV3TilingFunc）

2.x Host 端核心流程伪代码

// Host side tiling entry
Status LpNormV3TilingFunc(Context &ctx) {
    // 1. 获取平台信息
    PlatformInfo plat = GetPlatformInfo(ctx);
    uint32_t coreNum = plat.aivCoreNum;
    uint32_t ubSize  = plat.ubSize;

    // 2. 解析输入 shape / 属性
    Shape xShape = ctx.GetInputShape(0);   // [rows, cols]
    int axis = ctx.GetAttr("axis");       // -1 / 0 / 1
    float p  = ctx.GetAttr("p");

    // 3. 计算总数据量
    int64_t totalDataNum = xShape.rows * xShape.cols;

    // 4. core 任务均分
    CoreSplit split = SplitTask(totalDataNum, coreNum);
    // split.bigCoreNum / split.smallCoreNum

    // 5. tile 切分（基于 UB）
    TileInfo tile;
    tile.tileDataNum = CalcTileDataNum(ubSize, sizeof(dtype));
    tile.tileNum     = split.coreDataNum / tile.tileDataNum;
    tile.tailDataNum = split.coreDataNum % tile.tileDataNum;

    // 6. workspace 规划（按 axis 分配）
    int64_t workspaceSize = CalcWorkspaceSize(axis, xShape);
    ctx.SetWorkspaceSize(workspaceSize);

    // 7. 设置 tiling 数据
    TilingData tiling;
    tiling.axis          = axis;
    tiling.p             = p;
    tiling.coreDataNum   = split.coreDataNum;
    tiling.tileDataNum   = tile.tileDataNum;
    tiling.tileNum       = tile.tileNum;
    tiling.tailDataNum   = tile.tailDataNum;

    ctx.SetTilingData(tiling);

    // 8. tilingKey（仅区分 dtype / 模板）
    ctx.SetTilingKey(CalcTilingKey(dtype));
    return SUCCESS;
}

Host 端主要负责 平台感知、数据切分、tiling 参数生成以及 kernel 分支选择，整体流程如下：

LpNormV3TilingFunc()
├─ 获取平台信息
├─ 获取输入 shape
├─ 计算 coreNum
├─ 计算每 core 数据块数
├─ 设置 tiling 数据
├─ 计算 workspace
└─ 设置 tilingKey

2.1 平台信息获取

通过 GetPlatformInfo 获取：
- UB 内存大小（用于 tileDataNum 计算）
- AIV 核心数量（用于多核任务均分）
若核心数或 UB 大小非法，直接返回失败，保证后续计算安全。

2.2 输入 Shape 与属性解析

从 context 中获取：
- 输入张量 shape（rows × cols）
- axis 参数（-1 / 0 / 1）
- p 参数（float）
统一将输入视为 二维张量：
- axis = -1：整体范数
- axis = 0：按行归约
- axis = 1：按列归约

2.3 CoreNum 与任务均分策略

设计原则：优先满核，其次均分
根据 totalDataNum 与 coreNum 计算：
- bigCoreDataNum / smallCoreDataNum
- tailBlockNum（需要多处理一个 tile 的核心数）
若数据量无法整除：
- 将多余的数据块分配给前 tailBlockNum 个核心
- 其余核心按 smallCoreDataNum 处理

2.4 Tile 切分策略

根据 UB 大小、BUFFER_NUM、dtype size 计算：
- tileDataNum：单次搬运的数据量
- tileNum：每个 core 的 tile 次数
- tailDataNum：最后一个 tile 的实际数据量
Tile 设计目标：
- 尽量填满 UB
- 尾块可自然融入主流程，避免额外分支

2.5 Workspace 规划

Workspace 用于 多核范数归约结果的全局同步：
- 按 axis 维度划分槽位
- 每个范数值占用一个 64B 对齐槽（SLOT_STRIDE）
设计原因：
- 避免多核在同一 cache line 内原子写导致的数据覆写
- 方便 Cache Clean & Invalid 操作

2.6 TilingKey 规划策略

不基于 shape / axis 划分 tilingKey
tilingKey 主要用于：
- 区分不同 dtype 或实现模板
范数类型（普通 / ±∞）不通过 tilingKey 区分：
- 由 kernel 侧在运行时根据 p 值判断

该策略降低了 kernel 数量，提升代码复用率。

三、Kernel 端设计方案

3.x Kernel 端整体控制流伪代码

// Kernel entry for each core
__aicore__ void LpNormV3Kernel(TilingData tiling) {
    Init(tiling);

    if (!IsInfNorm(tiling.p)) {
        NormalProcess(tiling);
    } else {
        InfProcess(tiling);
    }
}

Kernel 端采用 Init + Process 的两阶段模型。

Init()
├─ 绑定 global buffer (x, z, workspace)
├─ 分配 local buffer (pipe)
├─ 初始化计算 buffer (poweredBuf, tmp_norm, tmp_base)
└─ 初始化 localSum（axis-dependent）

Process()
├─ if P not ±inf: NormalProcess()
│ ├─ Reduce loop
│ ├─ SumAndSyncAll
│ └─ Normalization loop
└─ else: InfProcess()

四、Init 阶段详细设计

4.x Init 阶段伪代码

__aicore__ void Init(TilingData tiling) {
    // 1. 计算当前 core 的数据范围
    coreId = GetCoreId();
    globalOffset = CalcGlobalOffset(coreId, tiling);

    // 2. 绑定 global memory
    xGm   = GlobalTensor(xAddr + globalOffset);
    yGm   = GlobalTensor(yAddr + globalOffset);
    workGm= GlobalTensor(workAddr);

    // 3. 初始化 pipe / queue（双缓冲）
    pipe.Init();
    inQueue.Init();
    outQueue.Init();

    // 4. 初始化 localSum（axis 相关）
    if (axis == -1) {
        localSum = 0;
    } else if (axis == 0) {
        localSum[rows] = 0;
    } else { // axis == 1
        localSum[cols] = 0;
    }
}

4.1 Global Buffer 绑定

根据 tilingData 计算当前核心：
- globalBufferIndex
- coreDataNum
绑定：
- xGm：输入数据
- yGm：输出结果
- workGm：跨核共享 workspace（严格按 axis 分配）

4.2 Local Buffer 与 Pipe 初始化

使用 双缓冲（BUFFER_NUM = 2）：
- inQueue / outQueue
计算缓冲区：
- poweredBuf：存储 |x|^p 或 |x|
- tmp_norm：归一化所需范数
- tmp_base：原子累加与中间计算

4.3 localSum 初始化（axis 相关）

axis = -1：
- localSum 为标量
axis = 0：
- localSum 大小为 rows
axis = 1：
- localSum 大小为 cols

初始化策略：

普通 p：初始化为 0
+∞ / -∞：初始化为 ±p 对应的极值，用于 min/max 归约

五、Process 阶段设计

5.1 NormalProcess（p 为有限值）

NormalProcess 伪代码

__aicore__ void NormalProcess(TilingData tiling) {
    // Reduce 阶段
    for (int t = 0; t < tiling.tileNum; ++t) {
        CopyIn(xGm, inQueue, t);
        ComputeReduce(inQueue, localSum, tiling.axis, tiling.p);
    }
    if (tiling.tailDataNum > 0) {
        CopyInTail(xGm, inQueue);
        ComputeReduce(inQueue, localSum, tiling.axis, tiling.p);
    }

    // 多核归约与同步
    AtomicAdd(workGm, localSum);
    SyncAll();

    // 核 0 计算最终范数
    if (GetCoreId() == 0) {
        norm = Pow(workGm + epsilon, 1.0f / tiling.p);
        WriteBack(norm);
        CacheCleanInvalid(workGm);
    }
    SyncAll();

    // Normalization 阶段
    for (int t = 0; t < tiling.tileNum; ++t) {
        CopyIn(xGm, inQueue, t);
        CopyIn(norm, tmp_norm);
        ComputeDiv(inQueue, tmp_norm, outQueue);
        CopyOut(outQueue, yGm, t);
    }
}

5.1.1 Reduce 阶段

CopyIn：tile 数据搬入 UB
Reduce：
- Cast → Abs → Power
- 按 axis 将结果累加到 localSum

5.1.2 SumAndSyncAll 阶段

各核心通过 AtomicAdd 将 localSum 写入 workGm
使用：
- SyncAll + Barrier 确保所有核写入完成
核 0：
- 计算 (sum + epsilon)^(1/p)
- 回写并执行 Cache Clean & Invalid

5.1.3 Normalization 阶段

再次 CopyIn 原始数据
从 workGm 读取范数
执行逐元素除法并 CopyOut

5.2 InfProcess（p = ±∞）

InfProcess 伪代码

__aicore__ void InfProcess(TilingData tiling) {
    // ReduceInf 阶段
    for (int t = 0; t < tiling.tileNum; ++t) {
        CopyIn(xGm, inQueue, t);
        ComputeReduceInf(inQueue, localSum, tiling.axis, tiling.p);
    }

    // 多核 min / max 归约
    if (tiling.p > 0) {
        AtomicMax(workGm, localSum);
    } else {
        AtomicMin(workGm, localSum);
    }
    SyncAll();

    // 核 0 写回最终范数
    if (GetCoreId() == 0) {
        WriteBack(workGm);
        CacheCleanInvalid(workGm);
    }
    SyncAll();

    // Normalization
    for (int t = 0; t < tiling.tileNum; ++t) {
        CopyIn(xGm, inQueue, t);
        CopyIn(workGm, tmp_norm);
        ComputeDiv(inQueue, tmp_norm, outQueue);
        CopyOut(outQueue, yGm, t);
    }
}