状态（Status）： Reviewing
作者（Authors）： @samulenzz
创建日期（Created）： 2026-07-26
更新日期（Updated）： 2026-07-26
相关 Issue/PR： #227

1. 概述

1.1 简介

本提案针对 DrivingSDK 仓中 GraphSoftmax 自定义算子进行 Ascend C 内核级性能优化。GraphSoftmax 是自动驾驶场景中图神经网络（GNN）的核心算子，用于对图结构中按节点分组的边特征执行 softmax 归一化。优化后在保持 FP32 精度的前提下，通过消除冗余 Global Memory 访问、引入 run 级聚合和唯一索引快速路径等手段，实现算子端到端性能提升 20% 以上。

1.2 动机

使用场景：GraphSoftmax 广泛应用于 QCNet 等自动驾驶感知模型中，对图中每条边的 8 维特征按目标节点分组做 softmax。在 QCNet 推理和训练流程中，该算子被高频调用，其性能直接影响模型整体推理延迟。

当前痛点：

• 原始实现采用逐 edge 的 scatter 操作（ScatterMax + ScatterSum），每条边单独发起一次 MTE3 原子搬运，导致 Global Memory 原子操作次数与边数成正比，带宽利用率极低。
• 在循环中分块加载数据（singleLoopTaskCount 硬编码为 2360），多次 CopyIn/CopyOut 带来重复的 MTE2/MTE3 同步开销。
• 对于 index 严格递增的场景（每条边指向不同节点），softmax 结果恒为 1.0，原实现仍执行完整的 scatter-reduce 流程，浪费算力。
• ComputeSum 和 ComputeSoftmax 阶段各需一轮完整的 CopyIn → Gather → 计算 → CopyOut 流水，数据多次在 GM 和 UB 之间往返。

不做此提案的影响：在大规模图（边数 > 10000）场景下，GraphSoftmax 将成为 QCNet 推理的性能瓶颈；随着自动驾驶模型对图结构利用的加深，该瓶颈会愈发突出。

1.3 目标

目标：

1. GraphSoftmax 算子在全覆盖测试用例下端到端性能提升 ≥ 20%。
2. 优化具有泛化性，不局限于单点 case 的提升，适用于 E ∈ [1, 50000]、index ∈ [0, 1500000] 的全范围数据。
3. FP32 精度符合 DrivingSDK 官方精度对比工具标准（err_threshold = 1e-4, pass_rate ≥ 99.99%）。
4. 适配 Ascend 800T A2 (ascend910b) 和 Ascend 800T A3 (ascend910_93) 两款硬件。

非目标：

• 不优化 GraphSoftmaxGrad（反向算子）。
• 不改变算子对外 API 接口（输入 / 输出 / 属性保持一致）。
• 不引入新的数据类型支持（仅 FP32）。
• 不修改 PyTorch 侧的 mx_driving.graph_softmax 封装逻辑。

---

2. 用例分析

2.1 功能用例

场景	描述	关键约束
QCNet 推理	图中 E 条边、8 维特征按 N 个节点分组做 softmax	E ∈ [1, 50000], N ∈ [1, 1500000], feature = 8, dtype = FP32
QCNet 训练前向	同推理，但需在 forward 中缓存输出供 backward 使用	同上
单边/单节点极端 case	E = 1 或所有边指向同一节点	softmax 结果退化为全 1.0 或标准单组 softmax
稀疏图	index 严格递增，每条边指向不同节点	softmax 结果恒为 1.0

2.2 性能指标

指标	要求
性能提升	≥ 20%（Device Time 对比 baseline）
泛化性	官方 UT 全部 case 均有提升，不依赖单点优化
精度	FP32 err_threshold = 1e-4，pass_rate ≥ 99.99%

2.3 DFX 要求

属性	要求
兼容性	不改变算子对外接口，现有上层调用代码无需修改
可维护性	代码结构清晰，各优化策略模块化封装
可测试性	复用现有 UT 框架，精度/性能可通过官方工具自动化验证
可靠性	新增完整的空指针守护和输入校验，防止越界访问

2.4 使用限制

• 输入 src 的 shape 必须为二维 [E, 8]，feature 维度固定为 8。
• index 元素取值范围 [0, 1500000)。
• 边数 E 不超过 50000。
• 仅支持 FP32 数据类型。

---

3. 方案设计

3.1 总体方案

3.1.1 GraphSoftmax 算法概述

GraphSoftmax 对图中的边特征 src[E, 8] 按 index[E] 分组执行数值稳定的 softmax：

对每条边 e，其所属组 g = index[e]：
  max_g    = max_{e' ∈ group(g)} src[e']         // 按组取最大值
  exp_e    = exp(src[e] - max_g)                 // 减最大值后取指数
  sum_g    = Σ_{e' ∈ group(g)} exp[e']           // 按组求和
  output_e = exp_e / (sum_g + ε)                 // 归一化

原始实现和优化后实现都遵循三阶段流程：ScatterMax → ComputeExp + ScatterSum → ComputeOutput。核心区别在于如何执行 scatter 聚合以及如何减少 GM 访问。

3.1.2 原始实现分析

原始 kernel 的主要性能瓶颈如下图所示：

flowchart TD
    subgraph INIT["初始化阶段"]
        A1[InitGlobalMemory<br/>初始化 scatterMax = -INF<br/>初始化 scatterSum = 0]
    end

    subgraph SCATTER_MAX["Phase 1: ScatterMax（分块循环）"]
        B1[CopyIn: src + index<br/>每次加载 ≤ 2360 条边]
        B2[逐 edge Atomic ScatterMax<br/>每条边 1 次 MTE3 原子写]
        B1 --> B2
        B2 -->|下一块| B1
    end

    subgraph SCATTER_SUM["Phase 2: ComputeExp + ScatterSum（分块循环）"]
        C1[CopyIn: src + index]
        C2["Gather scatterMax[index[e]]<br/>逐 edge 读 GM"]
        C3["Sub + Exp（向量计算）"]
        C4[逐 edge Atomic ScatterAdd<br/>每条边 1 次 MTE3 原子写]
        C5[CopyOut: exp 中间结果]
        C1 --> C2 --> C3 --> C4 --> C5
        C5 -->|下一块| C1
    end

    subgraph SOFTMAX["Phase 3: ComputeSoftmax（分块循环）"]
        D1[CopyIn: exp 中间结果 + index]
        D2["Gather scatterSum[index[e]]<br/>逐 edge 读 GM"]
        D3["Adds(ε) + Div（向量计算）"]
        D4[CopyOut: 最终结果]
        D1 --> D2 --> D3 --> D4
        D4 -->|下一块| D1
    end

    INIT --> SCATTER_MAX --> SCATTER_SUM --> SOFTMAX

    style B2 fill:#ff6b6b,color:#fff
    style C2 fill:#ff6b6b,color:#fff
    style C4 fill:#ff6b6b,color:#fff
    style D2 fill:#ff6b6b,color:#fff

红色节点标注了主要性能瓶颈：逐 edge 的 GM 原子操作和逐 edge 的 GM gather。

瓶颈量化：对于 E = 30000 的典型用例：

• ScatterMax 阶段：30000 次 MTE3 原子写
• ComputeExp 阶段：30000 次 MTE2 gather + 30000 次 MTE3 原子写
• ComputeSoftmax 阶段：30000 次 MTE2 gather
• 总计：90000 次 GM 单行访问

3.1.3 优化方案总览

优化后的实现引入三条互斥的执行路径，根据运行时数据特征自适应选择最优路径：

flowchart TD
    START([Process 入口])
    CHECK_LARGE{"totalTask ≥ 5120?"}
    COPY_IDX[CopyIn: index 数组]
    PUBLISH[向量化检测 index 有序性<br/>PublishIndexOrder]
    SYNC1[SyncAll：全核同步]
    LOAD_FLAGS[CopyIndexOrderFlagsIn<br/>加载所有核的有序标记]
    ALL_UNIQUE{"AllIndexes<br/>StrictlyIncreasing?"}

    FAST_PATH["<b>快速路径</b><br/>Duplicate(1.0) → CopyOut<br/>直接输出全 1.0"]

    INIT_WS["InitWorkspace<br/>（延迟初始化 scatter 空间）"]
    LOAD_SRC[CopySrcIn: 一次加载全部 src]
    SHOULD_RUN{"ShouldBuildRuns?<br/>（采样前 64 行估算）"}

    BUILD_RUNS["BuildRuns<br/>连续相同 index 合并为 run"]
    RUN_CHECK{"runCount × 3 ≤ coreTask?"}

    subgraph RUNWISE["Runwise 路径"]
        R_MAX["ReduceMaxRun (UB 内规约)"]
        R_SMAX[ScatterMax per-run<br/>仅 runCount 次 MTE3]
        R_EXP["GatherByRun + RunwiseSub + Exp"]
        R_SUM_CHECK{"RunGroups<br/>StrictlyIncreasing?"}
        R_OSUM["ReduceSumRun + Ordered ScatterAdd<br/>（无原子冲突）"]
        R_ASUM["ScatterAdd per-edge<br/>（原子聚合兜底）"]
        R_DIV["GatherByRun + RunwiseDiv"]
    end

    subgraph EDGEWISE["Edgewise 路径（兜底）"]
        E_MAX[ScatterMax per-edge<br/>逐 edge MTE3 原子写]
        E_EXP["GatherByEdge + Sub + Exp"]
        E_SUM[ScatterAdd per-edge]
        E_DIV["GatherByEdge + Adds(ε) + Div"]
    end

    COPYOUT[CopyOut: 写回最终结果]
    DONE([返回])

    START --> CHECK_LARGE
    CHECK_LARGE -->|是| COPY_IDX
    CHECK_LARGE -->|否| INIT_WS
    COPY_IDX --> PUBLISH --> SYNC1 --> LOAD_FLAGS --> ALL_UNIQUE

    ALL_UNIQUE -->|是| FAST_PATH --> DONE
    ALL_UNIQUE -->|否| INIT_WS

    INIT_WS --> LOAD_SRC --> SHOULD_RUN
    SHOULD_RUN -->|是| BUILD_RUNS --> RUN_CHECK
    SHOULD_RUN -->|否| E_MAX

    RUN_CHECK -->|是：进入 Runwise 路径| R_MAX
    RUN_CHECK -->|否：回退 Edgewise| E_MAX
    R_MAX --> R_SMAX --> R_EXP --> R_SUM_CHECK
    R_SUM_CHECK -->|是| R_OSUM --> R_DIV
    R_SUM_CHECK -->|否| R_ASUM --> R_DIV
    R_DIV --> COPYOUT

    E_MAX --> E_EXP --> E_SUM --> E_DIV --> COPYOUT
    COPYOUT --> DONE

    style FAST_PATH fill:#4caf50,color:#fff
    style R_MAX fill:#2196f3,color:#fff
    style R_SMAX fill:#2196f3,color:#fff
    style R_EXP fill:#2196f3,color:#fff
    style R_OSUM fill:#2196f3,color:#fff
    style R_DIV fill:#2196f3,color:#fff

绿色 = 唯一索引快速路径，蓝色 = Run 级聚合路径。

3.1.4 优化策略详解

优化 1：唯一索引快速路径（Unique-Index Fast Path）

当 index 数组严格递增（每条边指向不同节点）时，每组仅含 1 条边，softmax 结果恒为 1.0。优化后直接跳过全部 scatter 操作，用 Duplicate(1.0) 填充输出。

检测过程使用向量化 Compare 指令批量比较相邻 index 值，并通过全核协同（每核发布有序标记到 GM，再批量加载判定）实现全局判定：

sequenceDiagram
    participant Core0 as AIV Core 0
    participant Core1 as AIV Core 1
    participant CoreN as AIV Core N
    participant GM as Global Memory

    Note over Core0,CoreN: Phase 1: 各核独立检测本地 index 有序性
    Core0->>Core0: CopyIndexIn + CopyShiftedIndexIn
    Core1->>Core1: CopyIndexIn + CopyShiftedIndexIn
    CoreN->>CoreN: CopyIndexIn + CopyShiftedIndexIn

    Core0->>Core0: Compare(shifted, original, LE)<br/>向量化比较 64 元素/批
    Core1->>Core1: Compare(shifted, original, LE)
    CoreN->>CoreN: Compare(shifted, original, LE)

    Note over Core0,CoreN: Phase 2: 发布有序标记
    Core0->>GM: indexOrderFlagGm[0] = 0 (有序)
    Core1->>GM: indexOrderFlagGm[1] = 1 (无序)
    CoreN->>GM: indexOrderFlagGm[N] = 0 (有序)

    Note over Core0,CoreN: SyncAll 全核同步

    Note over Core0,CoreN: Phase 3: 加载并汇总
    GM->>Core0: CopyIndexOrderFlagsIn (所有核的标记)
    Core0->>Core0: AllIndexesStrictlyIncreasing?
    
    alt 全部有序
        Core0->>Core0: Duplicate(1.0) + CopyOut
        Note over Core0: 提前返回，跳过所有 scatter
    else 存在无序
        Note over Core0: 继续标准流程
    end

优化 2：Run 级聚合（Run-Based Aggregation）

对于连续相同 index 的边序列（"run"），先在 UB 内做 ReduceMax / ReduceSum 将每个 run 归约为一行，再做 scatter 操作。这将 scatter 次数从 E（边数）降低到 runCount（run 数），显著减少 GM 原子操作：

graph LR
    subgraph 原始["原始方案: 逐 edge scatter"]
        E0["edge 0 → group 2"] -->|"Atomic MTE3"| G2a["GM[2]"]
        E1["edge 1 → group 2"] -->|"Atomic MTE3"| G2b["GM[2]"]
        E2["edge 2 → group 2"] -->|"Atomic MTE3"| G2c["GM[2]"]
        E3["edge 3 → group 5"] -->|"Atomic MTE3"| G5a["GM[5]"]
        E4["edge 4 → group 5"] -->|"Atomic MTE3"| G5b["GM[5]"]
        E5["edge 5 → group 7"] -->|"Atomic MTE3"| G7a["GM[7]"]
    end

    subgraph 优化["优化方案: Run 级 scatter"]
        R0["Run 0: edge 0-2<br/>UB ReduceMax → 1 行"] -->|"1 次 MTE3"| G2["GM[2]"]
        R1["Run 1: edge 3-4<br/>UB ReduceMax → 1 行"] -->|"1 次 MTE3"| G5["GM[5]"]
        R2["Run 2: edge 5<br/>直接使用 src"] -->|"1 次 MTE3"| G7["GM[7]"]
    end

    style G2a fill:#ff6b6b,color:#fff
    style G2b fill:#ff6b6b,color:#fff
    style G2c fill:#ff6b6b,color:#fff
    style G5a fill:#ff6b6b,color:#fff
    style G5b fill:#ff6b6b,color:#fff
    style G7a fill:#ff6b6b,color:#fff
    style G2 fill:#4caf50,color:#fff
    style G5 fill:#4caf50,color:#fff
    style G7 fill:#4caf50,color:#fff

Run 构建采用启发式决策：先对前 64 行采样估算 run 密度（ShouldBuildRuns），仅当 sampleRunCount × 3 ≤ sampleRows 时才构建 run 表，避免在极端稀疏（几乎每条边一个 run）场景下引入额外开销。

优化 3：Runwise 广播计算

在 ComputeExp 和 ComputeOutput 阶段，原始实现对每条边都需从 GM gather 对应的 max/sum 值。优化后改为 per-run gather（仅 runCount 次 MTE2），再用 BinaryRepeatParams 对 run 内所有边做广播 Sub/Div：

graph TB
    subgraph 原始["原始: 逐 edge Gather + 计算"]
        direction LR
        G0["GM → statLocal[0]"] --> S0["Sub(src[0], stat[0])"]
        G1["GM → statLocal[1]"] --> S1["Sub(src[1], stat[1])"]
        G2["GM → statLocal[2]"] --> S2["Sub(src[2], stat[2])"]
        G3["GM → statLocal[3]"] --> S3["Sub(src[3], stat[3])"]
    end

    subgraph 优化["优化: per-run Gather + 广播计算"]
        direction LR
        GR["GM → statLocal[runStart]<br/>仅 1 次 MTE2 per run"]
        GR --> BR["Sub(src[0..3], stat[runStart])<br/>repeat=4, 广播同一行"]
    end

    style G0 fill:#ff6b6b,color:#fff
    style G1 fill:#ff6b6b,color:#fff
    style G2 fill:#ff6b6b,color:#fff
    style G3 fill:#ff6b6b,color:#fff
    style GR fill:#4caf50,color:#fff
    style BR fill:#4caf50,color:#fff

优化 4：有序 Run ScatterSum

当所有 run 的 group index 严格递增（RunGroupsStrictlyIncreasing），ScatterSum 阶段各 run 不会写到同一个 GM 位置，原子操作不存在实际竞争。此时先在 UB 内 ReduceSum，再直接原子写结果，避免了逐 edge 原子加的开销。

优化 5：延迟 Workspace 初始化

原始实现在 InitGM 中无条件初始化 scatter workspace（InitGlobalMemory 将 scatterMax 填充为 -INF、scatterSum 填充为 0）。优化后将初始化延迟到唯一索引快速路径判定之后——若快速路径命中则完全跳过初始化，节省 2 × coreWorkspace × 8 × sizeof(float) 字节的 GM 写操作。

优化 6：单遍数据加载

移除原始实现中的 singleLoopTaskCount 分块循环。优化后 Tiling 阶段将 singleLoopTaskCount 设为 coreTask（即每核全部任务量），所有 src 和 index 数据一次性加载到 UB，消除多轮 CopyIn/CopyOut 的同步开销。UB 容量经过精确计算（见 3.3 节），确保在 80% UB 可用空间内放下全部 buffer。

3.1.5 优化前后对比

维度	原始实现	优化后实现
数据加载	分块循环，每块 ≤ 2360 行	单次加载全部 coreTask 行
ScatterMax	逐 edge 原子写 (E 次)	Run 级 UB ReduceMax + per-run 原子写 (runCount 次)
Gather Max/Sum	逐 edge 从 GM 读 (E 次 × 2)	per-run 从 GM 读 (runCount 次 × 2)
ScatterSum	逐 edge 原子加 (E 次)	有序时 per-run 原子加 (runCount 次)
Sub/Div 计算	标准向量操作	Runwise 广播 (repeat 模式)
唯一索引场景	完整 scatter 流程	直接输出 1.0，跳过全部 scatter
Workspace 初始化	无条件初始化	延迟到快速路径判定之后
中间结果存储	exp 结果写回 GM 再读	全程在 UB 内保持，不写回

3.2 技术选型

方案	描述	优点	缺点	选择
Run 级聚合 + 快速路径	本方案	GM 操作从 O(E) 降至 O(R)，快速路径 O(1)；自适应选择	增加 run 构建开销	采用
多核并行归约	在归约阶段使用多核做 allreduce	理论并行度高	跨核通信开销大，全局 allreduce 不适合稀疏分组	未采用
基于排序的分组	先按 index 排序边，再做连续归约	排序后可完全避免原子操作	O(E log E) 排序开销在小 E 时显著，且排序后需额外置换恢复	未采用
Tiling 分块优化	仅优化分块大小，保持原有逐 edge scatter	改动最小	无法从根本上减少 GM 原子操作次数	未采用

选择理由：在 QCNet 的实际数据分布中，graph 结构常具有较高的 run 密度（连续边指向相同节点），Run 级聚合能显著减少 GM 操作。同时，稀疏图的唯一索引快速路径可以实现数量级的加速。启发式 run 构建决策确保在不利数据分布下不引入性能退化。

3.3 功能与性能设计

3.3.1 Tiling 设计

优化后的 Tiling 策略做了以下关键改进：

UB Buffer 布局（所有地址对齐到 32 字节）：

Buffer	用途	大小
srcBuf_	源数据 / 最终输出	coreTask × 8 × 4 字节
statBuf_	中间统计量 (max/sum/gather 结果)	coreTask × 8 × 4 字节
indexBuf_	index 数组	coreTask × 4 字节
runStartBuf_	Run 起始位置	coreTask × 4 字节
runLengthBuf_	Run 长度	coreTask × 4 字节
runGroupBuf_	Run 对应的 group index	coreTask × 4 字节
reduceTmpBuf_	ReduceMax/ReduceSum 临时空间	reduceTmpBytes 字节
reducePartialBuf_	分块 Reduce 部分结果	32 字节
indexOrderFlagBuf_	各核有序标记	aivNum × 32 字节

UB 容量校验：Tiling 阶段按以下公式计算所需 UB 空间，若超出 80% 可用 UB 则拒绝任务：

requiredUB = 2 × Align32(coreTask × 8 × 4)    // srcBuf + statBuf
           + Align32(coreTask × 4)              // indexBuf
           + 3 × Align32(coreTask × 4)          // runStart/Length/Group
           + Align32(reduceTmpBytes)             // reduce 临时
           + 2 × 32                             // reducePartial + padding
           + aivNum × 32                        // indexOrderFlags

Reduce 临时空间计算：通过 AscendC::GetReduceMaxMaxMinTmpSize 和 AscendC::GetReduceSumMaxMinTmpSize 在 host 端动态获取，取二者最大值作为 reduceTmpBytes 传入 kernel。

Workspace 布局：

UserWorkspace:
┌─────────────────────────────────────────────┐
│ scatterMaxGm [totalWorkspace × 8] float     │  ← 各核分区，按 coreWorkspace 偏移
├─────────────────────────────────────────────┤
│ scatterSumGm [totalWorkspace × 8] float     │
├─────────────────────────────────────────────┤
│ indexOrderFlagGm [aivNum × 8] uint32        │  ← 各核有序标记（新增）
└─────────────────────────────────────────────┘

3.3.2 核心流程

优化后的 Process() 主流程：

stateDiagram-v2
    [*] --> CheckLargeTask

    CheckLargeTask --> UniqueDetection: totalTask ≥ 5120
    CheckLargeTask --> InitWorkspace: totalTask < 5120

    state UniqueDetection {
        [*] --> CopyIndex
        CopyIndex --> VectorCompare: 向量化比较相邻 index
        VectorCompare --> PublishFlag: 写有序标记到 GM
        PublishFlag --> GlobalSync: SyncAll
        GlobalSync --> LoadAllFlags: 加载全核标记
        LoadAllFlags --> CheckResult
        CheckResult --> FastExit: 全部有序
        CheckResult --> [*]: 存在无序
    }

    FastExit --> OutputOnes: Duplicate(1.0) + CopyOut
    OutputOnes --> [*]

    InitWorkspace --> CopyInAll: InitGlobalMemory + CopyIn src/index
    UniqueDetection --> InitWorkspace: 存在无序

    CopyInAll --> BuildRunDecision

    state BuildRunDecision {
        [*] --> Sample64: 采样前 64 行
        Sample64 --> EstimateDensity
        EstimateDensity --> BuildRuns: sampleRuns × 3 ≤ sampleRows
        EstimateDensity --> SkipRuns: 否则
        BuildRuns --> RunwisePath
        SkipRuns --> EdgewisePath
    }

    state RunwisePath {
        [*] --> ReduceMaxPerRun
        ReduceMaxPerRun --> ScatterMaxPerRun
        ScatterMaxPerRun --> Sync1
        Sync1 --> GatherMaxByRun
        GatherMaxByRun --> RunwiseSub
        RunwiseSub --> Exp
        Exp --> CheckOrdered
        CheckOrdered --> OrderedScatterSum: runs 有序
        CheckOrdered --> EdgewiseScatterSum: runs 无序
        OrderedScatterSum --> Sync2
        EdgewiseScatterSum --> Sync2
        Sync2 --> GatherSumByRun
        GatherSumByRun --> RunwiseDiv
        RunwiseDiv --> [*]
    }

    state EdgewisePath {
        [*] --> ScatterMaxPerEdge
        ScatterMaxPerEdge --> Sync3
        Sync3 --> GatherMaxPerEdge
        GatherMaxPerEdge --> SubExp
        SubExp --> ScatterSumPerEdge
        ScatterSumPerEdge --> Sync4
        Sync4 --> GatherSumPerEdge
        GatherSumPerEdge --> AddsDiv
        AddsDiv --> [*]
    }

    RunwisePath --> CopyOut
    EdgewisePath --> CopyOut
    CopyOut --> [*]

3.3.3 性能提升分析

以 E = 30000（典型 QCNet 用例）为例，假设平均 run 密度为 3 边/run：

操作	原始次数	优化后次数	减少比例
MTE3 原子写 (ScatterMax)	30000	10000 (runCount)	66.7%
MTE2 gather (ComputeExp)	30000	10000	66.7%
MTE3 原子写 (ScatterSum)	30000	10000	66.7%
MTE2 gather (ComputeOutput)	30000	10000	66.7%
CopyIn/CopyOut 同步次数	13 轮 × 3 阶段	1 次	97.4%

对于唯一索引场景（稀疏图），性能提升接近 100%（跳过全部 scatter）。

3.4 安全隐私与 DFX 设计

3.4.1 安全性增强

本次优化新增了完善的输入校验和空指针守护：

检查项	位置	说明
context == nullptr	TilingForGraphSoftmax 入口	防止空 TilingContext
platformInfoPtr == nullptr	Tiling	防止平台信息缺失
aivNum == 0 || ubSize == 0	Tiling	防止无效硬件配置
storageShape.GetDimNum() != 2	Tiling + InferShape	防止非二维输入
*NPtr <= 0 || *NPtr > MAX_GROUP_COUNT	Tiling	防止非法 N 值
numEdgeValue <= 0 || numEdgeValue > MAX_EDGE_COUNT	Tiling	防止边数越界
numFeatureValue != NUM_FEATURE	Tiling	强制 feature = 8
requiredUbBytes > availableUbBytes	Tiling	防止 UB 溢出
reduceTmpBytes == 0	Tiling	防止 Reduce API 返回异常
userWorkspace == nullptr	Kernel 入口	防止空 workspace
context == nullptr	InferShape / InferDataType	防止空推断上下文

3.4.2 兼容性设计

• 接口兼容：算子注册信息不变，Input("src")、Input("index")、Output("softmax_result")、Attr("N") 保持一致。
• 芯片兼容：保留 AddConfig("ascend910b") 和 AddConfig("ascend910_93")，新增 AddConfig("ascend950") 条件编译。
• Tiling 数据兼容：新增 reduceTmpBytes 字段，为前向扩展保留空间。
• 功能兼容：Edgewise 路径作为兜底保证在任何数据分布下的正确性。

3.4.3 可测试性设计

• 复用现有 test_graph_softmax.py UT，覆盖 E ∈ {1, 1112, 2223, ..., 50000} 共 46 个 case。
• 特殊 case：E = 1112 时所有边指向同一节点（index 全 0），测试单组极端场景。
• 精度验证使用 DrivingSDK 官方 precision_compare 工具。
• 性能验证使用 DrivingSDK 官方 ops_performance 工具。

3.4.4 可靠性设计

• 多核同步使用 SyncAll() 确保 scatter 结果全局可见后再进入下一阶段。
• 原子操作使用 SetAtomicMax / SetAtomicAdd / DisableDmaAtomic 严格配对。
• 向量化有序检测使用 Compare + 位掩码检查，不依赖浮点比较。
• Run 构建使用启发式采样 + 完整扫描双重保障。

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境：

项目	要求
硬件平台	Ascend 800T A2 / Ascend 800T A3 (Ascend 910B / 910_93)
开发语言	Ascend C (C++ 扩展)
编程框架	AscendC Kernel 开发框架
工具链	CANN Toolkit (包含 ccec 编译器)
依赖库	kernel_operator.h、kernel_tiling/kernel_tiling.h、tiling/tiling_api.h
上层框架	PyTorch + torch_npu + mx_driving

开发约束：

• 仅支持 FP32 输入/输出。
• Feature 维度固定为 8（NUM_FEATURE = 8）。
• 边数 E ≤ 50000，节点数 N ≤ 1500000。
• UB 使用率不超过 80%。

可验收设计：

验收项	工具	标准
精度	mx_driving.tools.precision_compare	FP32: err_threshold = 1e-4, pass_rate ≥ 99.99%
性能	mx_driving.tools.ops_perf	Device Time 提升 ≥ 20% (对比 baseline)
UT 通过	test_graph_softmax.py	全部 46 个 case 通过

3.5.2 接口定义与设计

3.5.2.1 GraphSoftmax（AscendC Kernel API）

• 接口描述：对图结构中按节点分组的边特征执行数值稳定的 softmax 归一化。
• 接口原型：

extern "C" __global__ __aicore__ void graph_softmax(
    GM_ADDR src, GM_ADDR index, GM_ADDR softmaxResult,
    GM_ADDR workspace, GM_ADDR tilingData);

• 输入/输出参数：

参数名称	输入/输出	类型	描述	取值范围
src	输入	GM_ADDR (float)	边特征张量，shape = [E, 8]	E ∈ [1, 50000]，值 ∈ [-500, 500]
index	输入	GM_ADDR (int32)	边到节点的映射索引，shape = [E]	值 ∈ [0, 1500000)
softmaxResult	输出	GM_ADDR (float)	softmax 结果，shape = [E, 8]	值 ∈ [0, 1]
workspace	输入	GM_ADDR	系统 + 用户工作空间	由 Tiling 计算
tilingData	输入	GM_ADDR	Tiling 参数	由 Host 端生成

• 返回参数：无（void kernel）。

• 异常处理：userWorkspace == nullptr 时直接返回，不执行计算。

• 约束说明：

  • src 的第二维必须为 8。
  • index 的元素值必须在 [0, N) 范围内，N 由 Attr("N") 指定。
  • 仅支持 FP32 数据类型。

• 变更说明：

  • 新增 reduceTmpBytes Tiling 字段。
  • Workspace 末尾新增 indexOrderFlagGm 区域。
  • 内部实现重构，外部接口不变。

• 调用参考代码：

import torch
import mx_driving

src = torch.randn(10000, 8, dtype=torch.float32).npu()
index = torch.randint(0, 5000, (10000,), dtype=torch.int32).npu()

output = mx_driving.graph_softmax(src, index)

3.5.2.2 GraphSoftmaxTilingData（Tiling 数据结构）

• 接口描述：Host 端到 Device 端的 Tiling 参数传递结构。

• 结构定义：

BEGIN_TILING_DATA_DEF(GraphSoftmaxTilingData)
    TILING_DATA_FIELD_DEF(uint32_t, coreTask);
    TILING_DATA_FIELD_DEF(uint32_t, coreWorkspace);
    TILING_DATA_FIELD_DEF(uint32_t, totalTask);
    TILING_DATA_FIELD_DEF(uint32_t, totalWorkspace);
    TILING_DATA_FIELD_DEF(uint32_t, bigCoreCount);
    TILING_DATA_FIELD_DEF(uint32_t, singleLoopTaskCount);
    TILING_DATA_FIELD_DEF(uint32_t, reduceTmpBytes);      // 新增
END_TILING_DATA_DEF;

• 输入/输出参数：

参数名称	输入/输出	类型	描述	取值范围
coreTask	输出	uint32_t	大核每核任务数 (边数)	[1, ⌈E/aivNum⌉]
coreWorkspace	输出	uint32_t	每核 workspace 节点数	[1, ⌈N/aivNum⌉]
totalTask	输出	uint32_t	总任务数 (边数 E)	[1, 50000]
totalWorkspace	输出	uint32_t	workspace 总节点数	coreWorkspace × aivNum
bigCoreCount	输出	uint32_t	大核数量	[1, aivNum]
singleLoopTaskCount	输出	uint32_t	单轮处理任务数	等于 coreTask
reduceTmpBytes	输出	uint32_t	Reduce 操作所需临时空间	> 0

• 变更说明：
  • 新增 reduceTmpBytes 字段，用于传递 ReduceMax/ReduceSum 的临时 buffer 大小。
  • singleLoopTaskCount 语义变更：从固定值 2360 变为等于 coreTask（消除分块循环）。

3.5.3 编程手册设计

本次优化不改变外部 API，无需更新用户侧编程手册。内部开发者参考：

1. Kernel 开发指南：参考 kernels/graph_softmax/op_kernel/graph_softmax.cpp 中的模块化设计，各优化策略通过独立方法实现。
2. Tiling 开发指南：参考 kernels/graph_softmax/op_host/graph_softmax.cpp 中的 UB 容量计算和 Reduce 临时空间动态获取方式。
3. 精度验证：使用 mx_driving.tools.precision_compare 工具，参考 docs/zh/tools/precision_compare.md。
4. 性能验证：使用 mx_driving.tools.ops_perf 工具，参考 docs/zh/tools/ops_perf.md。

---

4. 测试设计

4.1 单元测试 (UT)

复用现有 tests/torch/test_graph_softmax.py，覆盖以下测试维度：

Case	边数 E	Index 特征	覆盖路径
E = 1	1	单边	Edgewise 兜底
E = 1112	1112	全部指向同一节点 (index 全 0)	Run 级聚合（单 run）
E = 2223~48894	每隔 1111 递增	均匀分布 (stride = 1500000/E)	Run 级聚合
E = 50000	50000	均匀分布至 1500000	Run 级聚合（最大规模）
E ≥ 5120 且 index 严格递增	≥ 5120	每条边唯一节点	唯一索引快速路径

精度验证方法：

from mx_driving.tools.precision_compare import precision_compare

result = precision_compare(npu_output, golden_output)
assert result.passed  # FP32: err_threshold=1e-4, pass_rate≥99.99%

性能验证方法：

from mx_driving.tools.ops_perf import ops_performance

# Step 1: 采集 baseline（使用 master 分支编译的算子）
output = ops_performance(
    mx_driving.graph_softmax, src_npu, index_npu,
    _perf={"mode": "baseline", "case_id": f"graph_softmax_E{E}"}
)

# Step 2: 采集 candidate（使用优化分支编译的算子）
output = ops_performance(
    mx_driving.graph_softmax, src_npu, index_npu,
    _perf={
        "mode": "performance",
        "case_id": f"graph_softmax_E{E}",
        "baseline_files": [baseline_csv]
    }
)

4.2 集成测试

• 通过 mx_driving.graph_softmax PyTorch 接口调用，验证 forward + backward 联合正确性。
• 验证 GraphSoftmaxGrad（未修改）与优化后 GraphSoftmax 的配合正确性。

4.3 端到端测试

• 在 QCNet 模型推理流程中验证 GraphSoftmax 优化不影响模型输出。
• 使用 ops_performance 工具在模型流水线中采集算子级性能数据。

---

附录

参考资料链接

资料	路径/链接
DrivingSDK 性能工具文档	docs/zh/tools/ops_perf.md
DrivingSDK 精度对比工具文档	docs/zh/tools/precision_compare.md
算子测试数据约束	算子测试数据约束.md