[RFC]: RoiawarePool3d 算子优化设计文档-DrivingSDK-AtomGit

已关闭

[RFC]: RoiawarePool3d 算子优化设计文档 #340

gcw_GF5DMRgv创建于 8月18日关闭于 9月18日

8月18日创建

状态（Status）： reviewing
作者（Authors）： @gcw_GF5DMRgv
创建日期（Created）： 2026-08-18
更新日期（Updated）： 2026-08-18
**相关 Issue/PR：**https://gitcode.com/Ascend/DrivingSDK/issues/230

1. 概述

1.1 简介

本设计文档描述 DrivingSDK 中 RoiawarePool3d 自定义算子在 Ascend NPU 上的性能与健壮性优化。RoiawarePool3d 是 3D 目标检测（PointPillars、CenterPoint 等）中 ROI 特征提取的关键算子：对每个 3D 检测框内的点云，先经旋转局部坐标变换判定框内点，再将其划入 outx × outy × outz 体素网格，最后对每个体素内的点特征执行 max 或 avg 池化，输出池化特征、argmax（最大特征对应点索引）与体素点索引表。

1.2 动机

使用场景：RoiawarePool3d 在 3D 检测模型的 ROI 特征提取阶段被高频调用。模型对每个候选框执行"框内点判定 → 体素划分 → 池化"，框数量级为数百至数千、框内点数万级、通道数 64~512，且每个框的体素网格固定为 outx×outy×outz。算子耗时主要由逐点几何判定与逐体素池化构成。

当前痛点（来自 master 基线与 roi 分支的代码对比分析）：

角度计算逐点重复且为标量：check_pt_in_box3d 对每个点执行一次 Cos/Sin（每次仅 1 元素向量调用 + V→S 同步），同一 box 的 ptsNum 个点重复计算同一角度，冗余且串行化严重。
几何判定逐点标量执行：原实现逐点 GetValue 读取坐标、标量分支判定框内、标量计算体素编号，循环内还伴随多次 PipeBarrier<PIPE_ALL>；点云全部加载到 UB 后，Vector 单元几乎闲置。
体素索引表每点读改写：collect_inside_pts_for_box3d 对每个框内点都做一次"读体素表 → 计数 → 写回"的完整 GM 往返（DataCopy + PipeBarrier），同一体素被多次重复读写。
池化阶段 channel 标量循环：compute_featrue 以 channelIdx 为外层循环，逐 channel 用标量 GetValue 比较更新 argmax；Max/Add 向量指令只在 channelIdx == 0 执行一次，其余 channelNum - 1 个 channel 全部退化为标量读写。
空体素无快速路径：体素内点数为 0 时仍进入完整 channel 双层循环，仅靠末尾兜底赋值。
多核分核缺陷：原 Tiling 直接 SetBlockDim(coreNum) 并用 coreNum 分核，当 boxNum < coreNum 时大量 Core 空转；且无 UB 容量校验，大输入存在 UB 溢出风险。
GM 偏移计算溢出隐患：原 Kernel 中 boxIdx * outx * outy * outz * ... 等乘法在 32 位域完成，大 box 数下可能溢出。
Python 侧重复分配：forward 内联输出分配、校验与数据类型转换逻辑，Host 开销与维护成本高。

不做此提案的影响：在框数多、通道数大的场景下，逐点标量判定与逐 channel 标量池化会放大标量指令占比与 GM 往返次数，使 RoiawarePool3d 成为 ROI 特征提取阶段的可观热点；同时分核缺陷与 UB 无校验会带来空转与稳定性风险。

1.3 目标

目标：

角度计算批量向量化：每个 Core 对全部负责 ROI 一次性 Cos/Sin（替代逐点标量计算）。
几何判定整体向量化：局部坐标变换、框内判定、体素编号全部以整点集 Vector 指令完成，消除逐点标量循环。
体素索引构建去重：先 GatherMask 压紧有效点，再按体素一次性写出索引，消除逐点读改写。
池化整体向量化：Max/Add/Muls 按 channelNum 整行执行，argmax 用 Compare+Select 向量化更新（保持 first-point-wins）。
空体素快速路径：validCount == 0 直接 Duplicate 默认值写回。
健壮性：Tiling 新增 64 位 UB 容量校验；分核改为 min(coreNum, boxNum)；GM 偏移全部 64 位计算。
接口收敛：输出 Tensor 分配与参数校验下沉到 C++ 层，Python 层精简；公共调用方式不变。

非目标：

不改变算子语义：max/avg 池化结果、argmax 的 first-point-wins 行为、体素点索引截断规则均与 master 一致。
不改变任务并行粒度：仍按 ROI 分核，单 box 不拆核。
不新增数据类型：沿用 FP32/FP16。
不修改 roiaware_pool3d_grad 的对外行为。

2. 用例分析

2.1 功能用例

场景	描述	关键约束
3D 检测 ROI 特征提取	每个框内点经旋转局部坐标变换划入 `outx×outy×outz` 体素，逐体素池化	boxNum ≥ 1，outx/outy/outz ∈ [1, 16]
max 池化	`mode=0`：体素内特征逐 channel 取最大，argmax 记录最大特征对应的点索引（首个最大者优先）	空体素 argmax = -1，特征 = 0
avg 池化	`mode=1`：体素内特征逐 channel 取平均，argmax 统一为 0	空体素 argmax = 0，特征 = 0
体素点数截断	落入体素的点数超过 `max_pts_per_voxel - 1` 时截断	`pts_idx_of_voxels[..., 0]` 为计数，未用 slot 保持 0
空体素 / 框外点	体素内无点、或点完全落在所有框外	快速路径，不进入池化循环
多通道特征	channelNum 较大（64~512）	池化向量化收益明显
FP16 输入	特征 FP16，输出转回 FP16	C++ 层统一在 FP32 域计算后转回

2.2 性能指标

指标	要求
角度计算	每 Core 从 `ptsNum × coreRoiNum` 次标量 Cos/Sin 降为 2 次整行向量调用（Cos、Sin 各 1）
几何判定	逐点标量分支 → 整点集 Vector 指令（Adds/Muls/CompareScalar/Select/Cast）
体素索引写	逐点读改写 → 压紧后按体素一次写出
池化	`Max`/`Add`/`Muls` 按整行 channelNum 执行，argmax 向量化
空体素	零长度体素不进入池化循环
分核	实际用核 `U = min(coreNum, boxNum)`，消除空转核
稳定性	UB 容量 64 位校验；GM 偏移全程 64 位

注：roi 分支提交说明未附实测数据，具体 Device Time 提升率需在统一环境对 baseline（master）与 candidate（roi）采集后补录至 3.3.3 节。

2.3 DFX 要求

属性	要求
兼容性	公共 Python API 与返回值不变；`mode`、`max_pts_per_voxel`、`out_size` 语义不变
可靠性	空体素、体素点数截断、框外点、非法参数（shape/dtype/范围）均有明确行为
可维护性	Kernel 按"加载 → 角度 → 几何判定 → 压紧 → 池化"拆分独立方法；C++ 层承担参数校验
可测试性	参数化用例 + 固定随机种子（20260702）+ 模块级 golden，可复现
安全与隐私	仅处理调用方 Tensor，无外部 I/O 与持久化

2.4 使用限制

rois 为 FP32/FP16 [R, 7]（中心 xyz + 尺寸 xyz + 角度 rz）；pts 为同 dtype [P, 3]；pts_feature 为同 dtype [P, C]。
C++ 层新增约束：R > 0、0 < P ≤ 1024、0 < C ≤ 1024、out_size[i] ∈ (0, 16]、max_pts_per_voxel ∈ [2, 256]、mode ∈ {0, 1}、三输入同 dtype 同设备。
单次调用要求 UB 容量满足 GetRequiredUbSize 校验。
体素内最多保留 max_pts_per_voxel - 1 个点（首元素为计数，未用 slot 保持 0）。
输出 shape：pooled_features/argmax 为 [R, outx, outy, outz, C]，pts_idx_of_voxels 为 [R, outx, outy, outz, max_pts_per_voxel]。

3. 方案设计

3.1 总体方案

3.1.1 RoiawarePool3d 算法概述

对第 b 个框，中心 (cx, cy, cz)、尺寸 (sx, sy, sz)、绕 z 轴旋转角 rz。点 p = (x, y, z) 变换到框局部坐标系：

shift_x = x - cx,  shift_y = y - cy
local_x = shift_x·cos(-rz) + shift_y·(-sin(-rz))
local_y = shift_x·sin(-rz) + shift_y·cos(-rz)
local_z = z - cz

框内判定与体素编号：

在框内 ⇔ -sx/2 < local_x < sx/2 且 -sy/2 < local_y < sy/2 且 0 ≤ local_z ≤ sz
x_idx = clamp(floor((local_x + sx/2)·outx/sx), 0, outx-1)
y_idx = clamp(floor((local_y + sy/2)·outy/sy), 0, outy-1)
z_idx = clamp(floor(local_z·outz/sz), 0, outz-1)
voxel_id = x_idx·(outy·outz) + y_idx·outz + z_idx

池化（体素 v 内点集合 S）：

mode=0 (max):  pooled[c] = max_{i∈S} feature[i][c]；argmax[c] = 首个达到最大值的点索引
mode=1 (avg):  pooled[c] = (Σ_{i∈S} feature[i][c]) / |S|；argmax[c] = 0

体素索引表 pts_idx_of_voxels[..., 0] 为计数，[..., 1..count] 为点索引，超 max_pts_per_voxel - 1 截断。

调用链（优化后）：

flowchart LR
    A["mx_driving.roiaware_pool3d<br/>(Python)"] --> B["npu_roiaware_pool3d_forward<br/>(C++, 参数校验 + 输出分配)"]
    B --> C["rois → rois_coord + rois_angle<br/>pts → pts.transpose(0,1)"]
    C --> D["aclnnRoiawarePool3d"]
    D --> E["Host Tiling<br/>分核 + UB 校验"]
    E --> F["Ascend C Kernel<br/>向量化判定 + 池化"]
    F --> G["返回 (pooled_features, argmax, pts_idx_of_voxels)"]

3.1.2 原始实现分析

master 基线 Kernel 主流程与瓶颈：

flowchart TD
    A["Process: 逐 box 循环"] --> B["collect_inside_pts_for_box3d<br/>逐点循环"]
    B --> C["DataCopy pts + roi<br/>(每点搬入, 含 PIPE_ALL)"]
    C --> D["Cos/Sin 逐点标量计算<br/>(每次 1 元素 + V→S 同步)"]
    D --> E["check_pt_in_box3d 标量判定"]
    E --> F["框内: 读体素表 → 计数+1 → 写回<br/>(每点一次 GM 往返)"]
    B --> G["compute_featrue<br/>逐体素 (outx·outy·outz)"]
    G --> H["channel 外层循环"]
    H --> I["逐 channel 标量 GetValue 比较<br/>Max/Add 仅 channel 0 执行一次"]
    style C fill:#ff6b6b,color:#fff
    style D fill:#ff6b6b,color:#fff
    style F fill:#ff6b6b,color:#fff
    style H fill:#ff6b6b,color:#fff
    style I fill:#ff6b6b,color:#fff

瓶颈归纳：

编号	问题	影响
P1	Cos/Sin 逐点标量计算	每 box 重复 `ptsNum` 次 1 元素向量调用 + V→S 同步
P2	几何判定逐点标量	Vector 单元闲置，标量分支 + 大量 `PIPE_ALL` 串行化
P3	体素索引表每点读改写	同一体素多次 GM 往返，`DataCopy` 长度含未用容量
P4	池化 channel 外层标量循环	`Max`/`Add` 仅执行一次，其余 channel 退化标量
P5	空体素无快速路径	零长度体素仍进入双层循环
P6	`SetBlockDim(coreNum)` 固定全核	`boxNum < coreNum` 时空转；无 UB 校验
P7	GM 偏移 32 位乘法	大输入下溢出风险
P8	Python 内联分配/校验	Host 开销与维护成本高

3.1.3 优化方案总览

优化后的 Kernel 将计算分为五个阶段，全部围绕"整点集向量化"组织：

flowchart TD
    S([Process 入口]) --> L["LoadPoints<br/>按维搬入全部点坐标 X/Y/Z<br/>(C++ 已转置为 [3,P])"]
    L --> LA["LoadAndComputeAngles<br/>搬入角度 → Muls(-1) → Cos/Sin 整行"]
    LA --> CI["CreateVecIndex<br/>生成点全局索引 0..P-1"]
    CI --> BLOOP{"逐 box 循环"}
    BLOOP --> LROI["LoadRoi<br/>搬 6 维 ROI + 预计算 halfSize/scale"]
    LROI --> CVI["ComputePointVoxelIds<br/>向量化局部坐标 + 框内判定 + 体素编号"]
    CVI --> GV["GatherValidPoints<br/>GatherMask 压紧 (voxelId, pointId)"]
    GV --> VLOOP{"逐 voxel 循环"}
    VLOOP -->|validCount==0| FAST["Duplicate 默认值直接写回"]
    VLOOP -->|非空| WIDX["写 pts_idx_of_voxels 计数与索引"]
    WIDX --> M{"mode"}
    M -->|0| MP["MaxPool<br/>整行 Max + Compare/Select argmax"]
    M -->|1| AP["AvgPool<br/>整行 Add + Muls(1/count)"]
    FAST --> VLOOP
    MP --> VLOOP
    AP --> VLOOP
    VLOOP -->|完成| BLOOP
    BLOOP -->|完成| E([结束])

3.1.4 优化策略详解

优化 1：角度批量向量化（对应 P1）

新增独立输入 rois_angle（C++ 层从 rois[:, 6] 拆分），每个 Core 一次性搬入自己负责的所有 ROI 角度，Muls(-1) 取负后 Cos/Sin 整行计算：

LoadAndComputeAngles() {
    DataCopyPad(angleLocal_, angleGm_[roiStart_], ...);   // 搬入 coreRoiNum 个角度
    Muls(angleLocal_, angleLocal_, -1, coreRoiNum_);      // 取负
    Cos(cosLocal_, angleLocal_, coreRoiNum_);             // 整行 cos(-angle)
    Sin(sinLocal_, angleLocal_, coreRoiNum_);             // 整行 sin(-angle)
    SetFlag/WaitFlag<V_S>;                                 // 一次同步供后续 Scalar 读取
}

每个 box 处理时只需 cosLocal_.GetValue(localRoiIdx) 一次标量读取，替代 master 版本每点一次的 Cos/Sin + 同步。

优化 2：几何判定整点集向量化（对应 P2）

ComputePointVoxelIds 对全部 ptsNum 个点并行执行：

Adds/Muls/Add  → 局部坐标 (localX/localY/localZ)
CompareScalar + Select  → 框内条件逐条 AND 累加（ApplyCondition × 6）
Adds/Muls + Maxs/Mins + Cast(floor)  → 体素编号 xIdx/yIdx/zIdx
Muls + Add  → 扁平化 voxelId = x·(outy·outz) + y·outz + z

六个框内条件（x/y 的上下界、z 的上下界）通过 CompareScalar 得到 mask 后用 Select 与 validFlag 累乘，一次遍历完成全部点的有效性标记，无任何逐点分支。

优化 3：GatherMask 压紧，消除体素表逐点读改写（对应 P3）

GatherValidPoints：用 CompareScalar(validFlag == 1) 生成有效 mask，GatherMask 一次压紧得到"有效点 voxelId 数组"与"有效点全局索引数组"，复用已废弃的 xIndexLocal_/yIndexLocal_ buffer；
GatherPointsAndPool：对每个体素用 CompareScalar(voxelId == voxelIdx) 生成 mask，GatherMask 收集该体素内的点索引，一次性写计数与索引表：

DataCopyPad(ptsIdxGm_[voxelOffset], countLocal_, countCopyParams_);       // 计数（标量 S_MTE3）
DataCopyPad(ptsIdxGm_[voxelOffset + 1], gatheredIndexLocal_, ...);        // 索引（向量 V_MTE3）

计数与索引分别以 S_MTE3 / V_MTE3 建立依赖，避免 PIPE_ALL 过度串行。

优化 4：池化整体向量化（对应 P4）

MaxPool：Duplicate(poolFeature, -inf) 初始化，循环内 CopyPointFeature 搬整行特征，Compare(pool ≥ tmp) + Select 更新 argmax（GE 保留旧值 → first-point-wins），Max 整行更新池化值，结束 Cast(CAST_RINT) 写回 argmax；
AvgPool：Duplicate(0) 初始化，循环内 Add 整行累加，结束后 Muls(1/validCount) 一次完成平均（倒数乘法替代 Div）。

// MaxPool 核心（替代 channel 外层标量循环）
for (pointOffset ...) {
    CopyPointFeature(pointIdx);                                  // 搬整行 channelNum 特征
    Compare(featureMask, poolFeature, tmpFeature, GE, channelNum_);
    Select(argmaxFloat, featureMask, argmaxFloat, pointIdx, ...); // first-point-wins
    Max(poolFeature, poolFeature, tmpFeature, channelNum_);       // 整行 max
}
Cast(argmaxIntLocal_, argmaxFloatLocal_, CAST_RINT, channelNum_);

优化 5：空体素快速路径（对应 P5）

if (validCount == 0) {
    Duplicate(poolFeatureLocal_, 0, channelNum_);
    Duplicate(argmaxIntLocal_, mode_ == 0 ? -1 : 0, channelNum_);
    DataCopyPad(pooledFeatureGm_[...], poolFeatureLocal_, ...);
    DataCopyPad(argmaxGm_[...], argmaxIntLocal_, ...);
    return;
}

优化 6：分核修正与 64 位 UB 校验（对应 P6）

Tiling 改为 usedCoreNum = min(coreNum, boxNum)，SetBlockDim(usedCoreNum)，coreBoxNums = boxNum / usedCoreNum、coreBoxTail = boxNum % usedCoreNum；
新增 GetRequiredUbSize(ptsNum, channelNum, maxCoreRoiNum)，按 Kernel 实际 buffer 布局（点坐标 10×、掩码、点索引 6×、ROI/角度、特征 3×、特征掩码）以 64 位累加，ub_total_size == 0 || required > ub_total_size 时拒绝任务。

优化 7：GM 偏移 64 位化（对应 P7）

Kernel 中所有输出偏移（voxelOffset、channelOffset、ptsFeatureGm_[pointIdx * channelNum]）均以 uint64_t 计算，避免 boxNum × voxelNum × channelNum 在 32 位域溢出。

优化 8：参数校验与输出分配下沉到 C++ 层（对应 P8）

npu_roiaware_pool3d_forward 重构为"校验 + 分配 + 执行 + 返回"四段：

校验：rois [R,7]、pts [P,3]、pts_feature [P,C]、R>0、P∈(0,1024]、C∈(0,1024]、out_size∈(0,16]³、max_pts_per_voxel∈[2,256]、mode∈{0,1}、三输入同 dtype 同设备；
整理：rois_coord = rois[:, :6]、rois_angle = rois[:, 6]（独立输入供批量向量化）；pts_coord = pts.T（转置为 [3, P] 供按维搬运）；
分配：argmax = at::empty(...)、pts_idx_of_voxels = at::zeros(...)（未用 slot 必须为 0 以保持公共语义）、pooled_features_cast = at::empty(...)（FP32 计算域）；
返回：std::tuple(pooled_features, argmax, pts_idx_of_voxels)，Python 层不再预分配。

Python 层相应精简：normalize_out_size 统一 int/三元组；删除内联校验与输出分配；ctx 直接保存反向所需上下文。

3.1.5 优化前后对比

维度	master 原实现	roi 优化实现
角度计算	逐点标量 Cos/Sin + 同步	每 Core 一次整行 Cos/Sin，逐 box 标量读取
几何判定	逐点标量分支	整点集 Adds/Muls/CompareScalar/Select/Cast
体素索引	每点读改写 GM	GatherMask 压紧后按体素一次写出
池化循环次序	channel 外层 × pts 内层	pts 外层 × 向量整行
Max/Add/Div	仅 channel 0 一次向量指令	每点一次整行向量指令
argmax 更新	逐 channel 标量 GetValue/SetValue	Compare + Select（first-point-wins）
空体素	进入双层循环后兜底	`validCount==0` 快速路径
平均计算	标量累加 + Div	向量 Add + 一次 Muls(1/count)
分核	`SetBlockDim(coreNum)`（可能空转）	`min(coreNum, boxNum)`
UB 校验	无	64 位 `GetRequiredUbSize`
GM 偏移	部分 32 位乘法	全程 64 位
参数校验	Python 内联、不完整	C++ 层完整校验 + 输出分配
输入整理	无	rois 拆 coord/angle，pts 转置

3.2 技术选型

方案	描述	优点	缺点	选择
整点集向量化 + GatherMask 压紧	本方案：几何判定与池化全部向量化，压紧后按体素输出	消除逐点标量循环与 GM 往返	需新增 rois_angle 输入与 pts 转置（C++ 层适配）	采用
保持逐点标量，仅微调	保留 master 结构小幅改动	改动最小	无法消除标量瓶颈与重复角度计算	未采用
按点拆核	单 box 内按点并行	大 box 并行度高	需跨核合并体素计数与截断，破坏稳定顺序	未采用
空间哈希索引	预建点→体素哈希	稀疏场景理论加速	建表、额外 workspace、顺序恢复成本	未采用

选择理由：优化针对的是 master 中最确定的冗余——重复角度计算、逐点标量判定与逐 channel 标量池化。整点集向量化让 Vector 单元承担几何判定主负载，GatherMask 压紧把体素索引写从"每点读改写"降为"每体素一次写"。新增 rois_angle 输入与 pts 转置的成本由 C++ 层一次性承担，换来 Kernel 侧按维搬运与整行角度计算的简化。按 box 分核保持既有并行模型，避免跨核合并复杂度。

3.3 功能与性能设计

3.3.1 Tiling 设计

多核划分：usedCoreNum = min(coreNum, boxNum)，按 box 连续分配：

coreBoxNums = boxNum / usedCoreNum
coreBoxTail = boxNum % usedCoreNum
maxCoreBoxNums = coreBoxNums + (coreBoxTail > 0 ? 1 : 0)

Core c 的 ROI 起始：

c < coreBoxTail : roiStart = (coreBoxNums + 1) * c
c ≥ coreBoxTail : roiStart = coreBoxNums * c + coreBoxTail

TilingData（RoiawarePool3dTilingData，与 master 字段一致）：

BEGIN_TILING_DATA_DEF(RoiawarePool3dTilingData)
    TILING_DATA_FIELD_DEF(uint32_t, coreBoxNums);
    TILING_DATA_FIELD_DEF(uint32_t, coreBoxTail);
    TILING_DATA_FIELD_DEF(uint32_t, boxNum);
    TILING_DATA_FIELD_DEF(uint32_t, ptsNum);
    TILING_DATA_FIELD_DEF(uint32_t, channelNum);
    TILING_DATA_FIELD_DEF(uint32_t, maxPtsPerVoxel);
    TILING_DATA_FIELD_DEF(uint32_t, outx);
    TILING_DATA_FIELD_DEF(uint32_t, outy);
    TILING_DATA_FIELD_DEF(uint32_t, outz);
    TILING_DATA_FIELD_DEF(uint32_t, coreNum);
    TILING_DATA_FIELD_DEF(uint32_t, mode);
END_TILING_DATA_DEF;

UB 布局（GetRequiredUbSize 64 位累加，与 Kernel InitBuffer 一一对应）：

Buffer 分组	内容	大小（对齐后）
点坐标 float	`pointX/Y/Z`、`shiftX/Y`、`localX/Y/Z`、`tmpPts`、`validFlag`	`10 × alignedPtsNum × 4`
点掩码	`pointMask`	`AlignUp((ptsNum+7)/8, 32)`
点索引 int	`xIndex/yIndex/zIndex/voxelId/pointIndex/gatheredIndex`	`6 × alignedPtsNum × 4`
ROI/角度	`roi`、`angle`、`cos`、`sin`	`(8 + 3 × alignedCoreRoiNum) × 4`
计数/argmax int	`count`、`argmaxInt`	`(8 + alignedChannelNum) × 4`
特征 float	`poolFeature`、`tmpFeature`、`argmaxFloat`	`3 × alignedChannelNum × 4`
特征掩码	`featureMask`	`AlignUp((channelNum+7)/8, 32)`

其中 alignedPtsNum = AlignUp(ptsNum, 8)、alignedChannelNum = AlignUp(channelNum, 8)、alignedCoreRoiNum = AlignUp(maxCoreBoxNums, 8)。

输出 shape（InferShape）：

argmax            = [boxNum, outx, outy, outz, channelNum]      # INT32
pts_idx_of_voxels = [boxNum, outx, outy, outz, maxPtsPerVoxel]  # INT32
pooled_features   = [boxNum, outx, outy, outz, channelNum]      # FP32（FP16 输入时 C++ 转回）

3.3.2 核心流程

优化后的 Process() 主流程：

flowchart TD
    START([Process 入口]) --> INIT["Init<br/>Tiling 参数 + GM/UB 绑定<br/>按 box 计算 roiStart"]
    INIT --> LP["LoadPoints<br/>按维搬入全部点坐标 X/Y/Z"]
    LP --> LA["LoadAndComputeAngles<br/>角度取负 + 整行 Cos/Sin"]
    LA --> CI["CreateVecIndex<br/>生成点全局索引 0..P-1"]
    CI --> BLOOP{"逐 box 循环"}

    BLOOP --> LROI["LoadRoi<br/>搬 6 维 ROI + 预计算 halfSize/scale"]
    LROI --> CVI["ComputePointVoxelIds<br/>向量化局部坐标 + 框内判定 + 体素编号"]
    CVI --> GV["GatherValidPoints<br/>GatherMask 压紧 (voxelId, pointId)"]
    GV --> VLOOP{"逐 voxel 循环"}

    VLOOP -->|validCount==0| FAST["Duplicate 默认值<br/>pool=0, argmax=-1/0"]
    VLOOP -->|非空| WIDX["写 pts_idx_of_voxels 计数与索引"]
    WIDX --> M{"mode"}
    M -->|0| MP["MaxPool<br/>整行 Max + Compare/Select argmax"]
    M -->|1| AP["AvgPool<br/>整行 Add + Muls(1/count)"]
    FAST --> VLOOP
    MP --> VLOOP
    AP --> VLOOP
    VLOOP -->|完成| BLOOP
    BLOOP -->|完成| DONE([结束])

池化内层（每体素）：

flowchart LR
    A["CopyPointFeature<br/>DataCopy 整行 channelNum 特征"] --> B{"mode"}
    B -->|0| C["Compare(pool ≥ tmp) GE<br/>Select(argmax, 保留旧值 first-wins)"]
    C --> D["Max(pool, pool, tmp, channelNum)"]
    B -->|1| E["Add(pool, pool, tmp, channelNum)"]
    D --> F["循环完: Cast argmax → INT32"]
    E --> F
    F --> G["DataCopyPad 写 pooled_features / argmax"]

同步设计：点坐标搬入后 MTE2_V；角度向量化后 V_S；ROI 搬入后 MTE2_S；体素索引写计数用 S_MTE3、写索引用 V_MTE3；特征搬移用 V_MTE2/MTE2_V；体素间复用 Vector/标量结果时仅等待对应 MTE3_V/MTE3_S，避免 PIPE_ALL 过度串行。

3.3.3 性能提升分析

理论收益（以每 box ptsNum 个点、voxelNum 个体素、channelNum 个通道计）：

优化项	master 开销	roi 开销	理论变化
Cos/Sin	`ptsNum × 2` 次 1 元素调用/box	每 Core 2 次整行调用	减少 `(ptsNum-1) × 2` 次向量调用与 V→S 同步
几何判定	逐点标量 + 每点 `PIPE_ALL`	整点集 Vector 指令	标量分支与全流水屏障消除
体素索引写	每框内点一次 GM 读改写	每体素一次写	GM 往返降为 `voxelNum` 次/box
Max/Add	每体素 1 次向量 + `(channelNum-1)` 轮标量	每体素 `validCount` 次整行向量	标量逐 channel 循环消除
argmax	逐 channel 标量读写	Compare+Select + Cast 一次	向量化更新
空体素	进入双层循环	2 次 Duplicate + 2 次写回	空体素循环开销归零
avg 平均	标量累加 + Div	1 次 Muls(1/count)	除法降为乘法
分核	固定全核（可能空转）	`min(coreNum, boxNum)`	消除空转核

预计收益来源：角度计算与体素索引写是纯消除项；几何判定与池化的向量化将大量标量读写替换为整行 Vector 指令，对 channelNum ∈ [64, 512]、outx·outy·outz 较大且体素稀疏的典型配置收益最大。

实测数据待补录：需在 A2/A3 同环境对 baseline（master 分支编译）与 candidate（roi 分支编译）采集 Device Time，覆盖多框（R ∈ [1, 1024]）、多通道（C ∈ [64, 512]）、稠密/稀疏体素、FP32/FP16 等矩阵，按 (baseline - candidate) / baseline 计算提升率后填入下表。

平台	Case	Baseline(us)	Candidate(us)	提升率	状态
A2	待补录	—	—	—	待录入
A3	待补录	—	—	—	待录入

3.4 安全隐私与 DFX 设计

3.4.1 安全性增强

检查项	位置	说明
`context` / 输入 Tensor / `platformInfoPtr` / `attrs` 为空	Tiling	空指针守护
`coreNum == 0` / `usedCoreNum == 0`	Tiling	无效硬件/输入配置
`ub_total_size == 0 \|\| GetRequiredUbSize > ub_total_size`	Tiling	防止 UB 溢出（64 位累加）
`GetBlockNum() == 0`	Kernel Init	Kernel 侧守护
rois/pts/feature shape 与 dtype	C++ 层	完整参数校验
`out_size`/`max_pts_per_voxel`/`mode` 范围	C++ 层	越界显式报错
三输入同 dtype 同设备	C++ 层	防止混用

3.4.2 兼容性设计

接口：公共 Python API mx_driving.roiaware_pool3d(rois, pts, pts_feature, out_size, max_pts_per_voxel, mode) 不变，返回值语义不变。
语义：max 池化 argmax 的 first-point-wins 行为通过 Compare(GE) + Select 保留旧值保持；空体素默认值（max=-1、avg=0）与 master 一致；pts_idx_of_voxels 未用 slot 保持 0（C++ 层 at::zeros 分配）。
内部 ABI：Kernel 新增 roisAngle 输入、pts 布局改为转置 [3, P]，Host OpDef、ACLNN 生成物、Kernel 签名与 C++ 调用必须来自同一构建版本。
反向兼容：roiaware_pool3d_grad 依赖的 pts_idx_of_voxels/argmax 布局不变。

3.4.3 可测试性设计

复用 tests/torch/test_roiaware_pool3d.py，参数化用例固定种子 20260702。
RANDOM_CASES 覆盖：max/avg × FP32/FP16 × 通道数 64/256/512 × 体素网格 (4,4,4)/(2,3,4)/(1,1,1) × 点数 96~256。
定向用例：体素点数截断（max_pts_per_voxel=3 保留前 2 点）、空体素（点全在框外）。
CPU golden 用 numpy 实现局部坐标变换、框内判定、体素编号与池化，与 NPU 输出 assertRtolEqual 对比。
C++ 层新增校验通过 TORCH_CHECK 显式报错，UT 覆盖非法参数路径。

3.4.4 可靠性设计

池化循环内 CopyPointFeature 使用 V_MTE2/MTE2_V 建立"索引（Vector）→ 特征搬移（MTE2）→ 向量计算"的依赖，避免 PIPE_ALL 过度串行。
体素间复用 Vector 写结果（gathered index、pooled feature、argmax）与 Scalar 写计数（count），仅等待对应 MTE3_V/MTE3_S 依赖。
计数与索引写入使用 S_MTE3/V_MTE3 精确同步，防止标量计数与向量索引乱序。
argmax 浮点中间值最后 Cast(CAST_RINT) 转 INT32，避免直接浮点→整型截断误差。
所有 GM 偏移与 UB 容量以 uint64_t 计算，防止溢出。

3.5 编程与调用设计

3.5.1 编程模型基本设计

开发环境：

项目	要求
硬件	Atlas A2 / A3（Ascend 910B / 910_93）；Ascend 950 条件编译
语言/框架	Ascend C + AscendC Kernel 开发框架
工具链	CANN Toolkit（ccec 编译器）、自定义 OPP 构建与安装流程
依赖	`kernel_operator.h`、`kernel_tiling/kernel_tiling.h`、`tiling/tiling_api.h`
上层	PyTorch + torch_npu + mx_driving

开发约束：rois/pts 为 FP32/FP16，pts_feature 同 dtype；P ≤ 1024、C ≤ 1024、out_size ≤ 16、max_pts_per_voxel ∈ [2, 256]；UB 受 GetRequiredUbSize 校验。

可验收设计：

验收项	工具/方法	标准
精度	CPU golden 对比	FP32/FP16 `assertRtolEqual` 通过；argmax/索引语义一致
功能	参数化 UT	全部 `RANDOM_CASES` + 定向用例通过
性能	`ops_performance`（待补录）	Device Time 相对 baseline 提升（数值待填 3.3.3）
集成	`mx_driving.roiaware_pool3d` 全链路	forward + backward 联合正确

3.5.2 接口定义与设计

3.5.2.1 RoiawarePool3d Kernel（roi 分支）

接口描述：对每个 3D 框内点执行体素划分与 max/avg 池化，输出池化特征、argmax 与体素点索引表。
接口原型：

extern "C" __global__ __aicore__ void roiaware_pool3d(
    GM_ADDR rois, GM_ADDR roisAngle, GM_ADDR pts, GM_ADDR ptsFeature,
    GM_ADDR argmax, GM_ADDR ptsIdxOfVoxels, GM_ADDR pooledFeatures,
    GM_ADDR workspace, GM_ADDR tiling);

与 master 的差异：新增 roisAngle（角度独立输入，支持批量向量化）；pts 布局为 C++ 转置后的 [3, P]。
参数：

参数	方向	类型/布局	说明
`rois`	输入	FP32 `[R, 6]`	框中心 xyz + 尺寸 xyz（角度已拆出）
`roisAngle`	输入	FP32 `[R]`	框绕 z 轴角度（新增，供批量向量化）
`pts`	输入	FP32 `[3, P]`	转置后的点云坐标
`ptsFeature`	输入	FP32/FP16 `[P, C]`	点特征
`argmax`	输出	INT32 `[R, ox, oy, oz, C]`	最大特征对应点索引
`ptsIdxOfVoxels`	输出	INT32 `[R, ox, oy, oz, maxPtsPerVoxel]`	体素点索引表（[...,0] 为计数）
`pooledFeatures`	输出	FP32	池化特征（FP16 输入时 C++ 转回）
`workspace`	临时	系统 workspace	`SetSysWorkspace`
`tiling`	输入	`RoiawarePool3dTilingData`	Host 生成

3.5.2.2 C++ 封装（`npu_roiaware_pool3d_forward`）

std::tuple<at::Tensor, at::Tensor, at::Tensor> npu_roiaware_pool3d_forward(
    const at::Tensor &rois, const at::Tensor &pts, const at::Tensor &pts_feature,
    at::IntArrayRef out_size, int32_t max_pts_per_voxel, int32_t mode);

完整参数校验（shape/dtype/设备/范围）；
rois 拆分为 rois_coord（前 6 列）与 rois_angle（第 6 列），pts 转置为 [3, P]；
输出分配：argmax = at::empty、pts_idx_of_voxels = at::zeros（未用 slot 保持 0）、pooled_features = at::empty（FP32 域）；
FP16 输入在 FP32 域计算后转回 FP16 返回。

3.5.2.3 Python 封装（`mx_driving/ops/roiaware_pool3d.py`）

normalize_out_size(out_size)：int/三元组统一为三元组；
forward：调用 _C.npu_roiaware_pool3d_forward 返回三元组，ctx 保存反向上下文；
backward：委托 _C.roiaware_pool3d_grad 计算 grad_in，返回与 forward 参数一一对应的梯度元组。

调用参考：

import torch
import mx_driving

rois = torch.randn(10, 7, dtype=torch.float32).npu()
pts = torch.randn(128, 3, dtype=torch.float32).npu()
pts_feature = torch.randn(128, 256, dtype=torch.float32).npu()

# mode=0 为 max 池化，返回 pooled_features [10, 4, 4, 4, 256]
out = mx_driving.roiaware_pool3d(rois, pts, pts_feature, out_size=4, max_pts_per_voxel=128, mode=0)

3.5.3 编程手册设计

公共 API 无变化，无需更新用户侧调用方式。内部开发者参考：

Kernel：kernels/roiaware_pool3d/op_kernel/roiaware_pool3d.cpp——加载 / 角度 / 几何判定 / 压紧 / 池化五段式方法结构。
Tiling：kernels/roiaware_pool3d/op_host/roiaware_pool3d.cpp、roiaware_pool3d_tiling.h——按 box 分核（min(coreNum, boxNum)）与 64 位 UB 容量校验。
C++ 封装：mx_driving/csrc/RoiawarePool3d.cpp——参数校验、输入整理（rois 拆分、pts 转置）、输出分配。
测试：tests/torch/test_roiaware_pool3d.py——参数化用例与固定种子。

4. 测试设计

4.1 单元测试（UT）

复用 tests/torch/test_roiaware_pool3d.py（种子 20260702）：

Case	boxes	out_size	channels	npoints	max_pts	mode	dtype	覆盖点
max_float32_common	10	(4,4,4)	256	128	128	0	FP32	max 池化常规规模
avg_float32_common	10	(4,4,4)	256	128	128	1	FP32	avg 池化常规规模
max_float16_channels_512	20	(4,4,4)	512	256	128	0	FP16	高通道 + FP16
avg_float16_channels_512	20	(4,4,4)	512	256	128	1	FP16	高通道 + FP16
small_voxel_grid	6	(2,3,4)	64	96	32	0	FP32	非立方体素网格
体素点数截断	1	(1,1,1)	2	4	3	0/1	FP32	截断至 max-1
空体素	1	(2,2,2)	2	2	4	0/1	FP32	点全在框外

精度验证：模块级 roiaware_pool3d_golden（numpy 实现局部坐标变换、框内判定、体素编号、max/avg 池化），与 NPU 输出 assertRtolEqual 对比；FP16 用例在 FP32 域计算 golden 后转回 FP16。

性能验证（待补录）：同环境 baseline（master 分支）vs candidate（roi 分支），预热 ≥ 50 次、采样 ≥ 100 次、torch.npu.synchronize() 后统计 Device Time，矩阵覆盖多框（R ∈ [1, 1024]）、多通道（64~512）、稠密/稀疏体素、FP32/FP16。

4.2 集成测试

从公共入口 mx_driving.roiaware_pool3d 验证完整链路：Python → _C.npu_roiaware_pool3d_forward → aclnnRoiawarePool3d → Host Tiling → Kernel → 输出返回。
forward + backward 联合验证：数值梯度对比 roiaware_pool3d_grad，确认 pts_idx_of_voxels/argmax 布局未被破坏。
A2/A3（及 A5 条件编译路径）冒烟：同输入输出一致。
非法参数路径：shape/dtype/范围越界均显式报错。

4.3 端到端测试

在代表性 3D 检测流程（ROI 生成 → RoiawarePool3d → 体素特征 → head）中验证算子级输出一致性与耗时；若仓内模型无直接调用，则以可复现的最小 ROI 池化流程作为端到端基线。
记录统一环境清单（硬件、驱动/固件、CANN、PyTorch、torch_npu、DrivingSDK 提交、编译选项），性能原始数据与分支关联。

附录

关键代码位置

内容	路径（roi 分支）
Kernel	`kernels/roiaware_pool3d/op_kernel/roiaware_pool3d.cpp`
Host 与 Tiling	`kernels/roiaware_pool3d/op_host/roiaware_pool3d.cpp`
TilingData	`kernels/roiaware_pool3d/op_host/roiaware_pool3d_tiling.h`
C++ 封装	`mx_driving/csrc/RoiawarePool3d.cpp`
Python 封装	`mx_driving/ops/roiaware_pool3d.py`
UT	`tests/torch/test_roiaware_pool3d.py`