[Requirement|需求建议]: Index 算子支持 ascend960dt/960pr 平台(DV121 DMA 路径)

Background(背景信息)

Index 算子当前仅支持 ascend950/ascend350 等平台,在 ascend960dt/960pr(A6/DV121)上未注册,无法在 960 环境运行。

960 平台引入 GatherDataCopy(DMA gather)能力,对于大输入的 full-load 场景,可将 input 留 GM 直接按偏移 gather,避免整表载入 UB,显著降低 UB 占用并提升大 shape 场景性能。需要为 Index 新增 DMA 路径,使 full-load 模板在大输入单维索引(1D 及 2D MASK_MODE_0)场景下走 GatherDataCopy。

同时,现有实现直接复用 950 路径在 960 上存在两个缺陷需要一并解决:

1. host 侧 outputQueSize 按"乘积末尾一次对齐"分配,而 kernel 按行 32B 对齐步进写入,1D 场景(lastDimBytes = sizeof(T) < 32)分配量仅为写入量的 1/4~1/32,越界写 UB;
2. 新增 DMA 模板的 SEL 条目缺少 arch 门禁,与实现/注册/tiling 三处的 9201/9202 宏隔离不同源。

Origin(信息来源)

算子组 960 平台适配专项(ops-nn 仓 960_op 分支)。

Benefit / Necessity(价值/作用)

• 使 Index 算子在 ascend960dt/960pr 平台可用,补齐 960 平台算子支持清单;
• 大输入 full-load 场景走 DMA gather 路径,input 不再占用 UB,可支持更大的单维索引规模,降低 UB 压力;
• 修复 UB 越界写隐患,消除静默数据损坏风险;
• SEL 门禁对齐,避免非 9201/9202 平台产生死模板登记面,消除后续演进漂移风险。

Design(设计方案)

总体思路

复用 arch35 full-load tiling 框架,新增 DMA 模板类型与独立 kernel,960 专有逻辑全部用 __NPU_ARCH__ == 9201/9202 宏隔离,其他平台零影响。

tiling 侧(op_host/arch35/index_tiling_full_load.cpp)

IsCapable():
    ...
    isDv121_ = (GetCurNpuArch() == DAV_9201 || GetCurNpuArch() == DAV_9202)
    useDmaPath_ = isDv121_
                  && (dimSize == 1 || maskMode == MASK_MODE_0)
                  && (inputBytes >= 64KB)
                  && (dimSize == 1 || lastDimBytes >= 64B)
    // 非 DMA 路径行为完全不变(!useDmaPath_ 守护原有校验)

CalcUBBuffer():
    if (useDmaPath_):
        // input 留 GM, UB 只存 output + indices + gather offset
        // 按行 32B 对齐预算, 与 kernel 写入布局一致
        outPerIdx = CeilAlign(lastDimBytes, 32) * DOUBLE
        ubIndices_ = (ubSize - fixedAlign) / (outPerIdx + idxPerIdx + offsetPerIdx)

DoOpTiling():
    if (useDmaPath_):
        outputQueSize = ubIndices_ * CeilAlign(lastDimBytes, 32)  // 按行对齐分配
    else:
        outputQueSize = CeilAlign(ubIndices_ * lastDimBytes, 32) // 原公式不变

GetTilingKey():
    useDmaPath_ 时新增 INDEX_FULL_LOAD_DMA_1D / INDEX_FULL_LOAD_DMA_2D_MASK_0 模板类型

kernel 侧(op_kernel/arch92/index_full_load_dma.h,新增)

整文件包在 #if __NPU_ARCH__ == 9201 || 9202 内,仅在 960 编译:

KernelIndexFullLoadDma<T, U, SrcU, MASK_MODE, DIM_NUM>:
    CopyInIndices:  indices GM -> UB(int64 先 cast 到 int32)
    ComputeDma:
        GenOffsetVF:  VF 生成 gather 偏移(负索引回绕 + dim0Shape, offset = idx * lastDimBytes)
        Duplicate:    按行 32B 对齐步长清零输出区
        GatherDataCopy: inputX 留 GM, 按 offset gather 到 UB(burstLen=lastDimBytes, dstStride=CeilAlign(lastDimBytes,32))
    CopyOut:          DataCopyPad UB -> GM(blockLen=lastDimBytes, srcStride=0, UB 侧按 32B slot 寻址与写入布局天然一致)

模板注册与门禁

• tiling_key.h: FULL_LOAD_TYPE 枚举末尾追加 DMA_1D=4/DMA_2D_MASK_0=5(旧值 0-3 编码不变),新增 SEL 条目用 #if (!defined(__CCE_AICORE__) || __NPU_ARCH__ == 9201 || 9202) 包裹(参照 qbmmv3 先例);
• index_apt.cpp: DMA 分发分支同样宏隔离;
• index_def.cpp: AddConfig("ascend960dt"/"ascend960pr");
• CMakeLists: SUPPORT_COMPUTE_UNIT 增加 960dt/960pr(tiling 复用 arch35);
• ascendc_config.json: Index 注册 960dt/960pr,保持 auto_sync=false 与 --cce-no-dcache-flush(kernel 使用手动 SetFlag/WaitFlag 同步);
• 新增 config/ascend960dt、ascend960pr 二进制配置目录与 960 tiling UT 用例。

验证

• 960DT 实机用例覆盖:1D fp32/fp16 大输入、2D MODE_0 非对齐(fp32 68B/int64 72B)、2D 对齐对照组,精度全过;
• tiling key 验证:DMA_1D=4/DMA_2D_MASK_0=5 模板命中正确;
• 非 960 平台逻辑零变化(所有 960 逻辑宏隔离/平台门禁)。