已关闭

[Requirement|需求建议]: 【社区任务】Roll 算子需求 #1971

boxw987创建于 6月15日关闭于 7月2日

boxw987

6月15日创建

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

Roll 算子需求与设计说明

一、需求背景

Roll 算子用于沿指定维度对输入张量执行循环位移，常见于序列处理、窗口重排和特征对齐等场景。

本说明聚焦当前最终提交版本的真实能力，目标是明确：

功能边界
设计方案
关键优化思路
测试与验收标准
后续维护约束

二、目标与范围

2.1 功能目标

提供 aclnnRollGetWorkspaceSize 与 aclnnRoll 两阶段 ACLNN 接口。
支持沿单维或多维执行循环位移。
当 dims 为空时，行为等价于先展平再执行一维 roll，最后恢复原始形状。

2.2 支持范围

类别	当前支持范围
数据类型	`uint8`、`int8`、`bfloat16`、`float16`、`float32`、`int32`、`uint32`
数据格式	`ND`
维度范围	0-D 到 8-D
属性语义	`shifts` 为滚动量，`dims` 为滚动维度
特性	负维度归一化、重复维度合并、`dims` 为空的 flatten roll

2.3 非目标范围

当前版本不纳入以下支持范围：

bool
int64
非 ND 数据格式

三、输入约束与数学语义

3.1 输入约束

输入输出 dtype 必须一致。
输入输出 shape 必须一致。
输入张量 rank 不超过 8。
当 dims 为空时，shifts 长度必须为 1。
当 dims 非空时，shifts 与 dims 长度必须一致。
dims 的取值范围为 [-rank, rank)。
对重复维度上的多个 shift，需要先做归一化合并。
0 维输入时，shifts 长度必须为 1，且 dims 为空。

3.2 数学语义

对任意参与滚动的维度，输出元素满足：

y[i] = x[(i - s) mod n]

其中：

s 为该维度归一化后的滚动量
n 为该维度长度

多维滚动时，各维度上的归一化结果共同作用于输出索引映射。

四、问题分析

Roll 的定义简单，但高性能实现存在以下难点：

dims 为空时需要退化为 flatten roll，不能仍按原始多维索引路径执行。
单维滚动、多维滚动、最后一维滚动和非最后一维滚动，对应的访存模式差异很大。
小宽度、非对齐、小块场景若统一走通用路径，性能容易明显回退。
在控制实现复杂度的同时，仍需要保证正式支持范围内的性能优于参考 TBE。

因此，当前实现采取的主线是：

Host 侧负责参数规范化、归一化和切分
Kernel 侧根据归一化后的形态选择搬运路径

五、总体方案

5.1 方案概述

整体实现分为两部分：

Host 侧负责参数校验、维度归一化、重复维度合并和 tiling 数据生成。
Kernel 侧根据 active dim 形态执行不同的 roll/copy 路径。

flowchart TD
    A[aclnnRollGetWorkspaceSize] --> B[参数校验]
    B --> C[dtype format shape 约束校验]
    C --> D[dims shifts 归一化]
    D --> E[重复维度合并]
    E --> F[识别 activeDim 与 activeDimCount]
    F --> G[生成 Tiling 信息]
    G --> H[aclnnRoll]
    H --> I[Kernel Process]
    I --> J{路径选择}
    J --> K[Identity]
    J --> L[Flatten Roll]
    J --> M[Leading Dim Roll]
    J --> N[Last Dim Row Roll]
    J --> O[Single Dim Block Roll]
    J --> P[Multi-Dim Last-Dim Roll]
    J --> Q[Multi-Dim Non-Last Roll]
    J --> R[Segmented Fallback]

5.2 Host 侧职责

Host 侧主要负责：

空指针校验
dtype 支持校验
ND 格式校验
输入输出形状一致性校验
dims/shifts 规则校验
负维度归一化
重复维度 shift 合并
生成 kernel 所需 tiling 数据

5.3 Kernel 侧职责

Kernel 侧不做复杂策略搜索，只消费 Host 给出的 tiling 信息并执行对应搬运路径。

核心原则：

优先把 roll 转为连续段 copy。
最后一维优先按行处理。
非最后一维优先按 block 重排。
保留 segmented fallback 作为通用正确性兜底路径。

六、详细设计

6.1 Host 侧设计

6.1.1 归一化逻辑

Host 侧将输入整理成更适合 Kernel 消费的结构：

dims 为空时，退化为一维 flatten roll。
dims 非空时，将负维度转换为正维度。
对重复维度上的多个 shift 先累加，再按该维长度取模。
统计 activeDimCount。
若仅存在一个 active dim，则进一步提取：
- outerSize
- dimSize
- innerSize
- activeShift

6.1.2 Tiling 数据

Host 侧向 Kernel 传递的核心信息包括：

totalNum
dimNum
shapes[]
strides[]
shifts[]
activeDim
activeDimCount
outerSize
dimSize
innerSize
activeShift
usedCoreNum
perCoreElements
lastCoreElements
ubElements

6.1.3 切分策略

当前版本的 tiling 重点是稳定切分，而非复杂搜索：

小数据量优先单核，降低调度开销。
大数据量按 perCoreElements 切分，保证多核负载基本均衡。
对最后一维滚动、非最后维 block roll、BF16 小宽度和部分 uint8 小块场景，调整对齐与切分粒度。
对大 inner stride 场景避免过粗切分。

6.2 Kernel 侧设计

6.2.1 路径选择树

flowchart TD
    A[Kernel Process] --> B{activeDimCount == 0?}
    B -->|Yes| C[CopyIdentity]
    B -->|No| D{dimNum == 1?}
    D -->|Yes| E[CopyFlattenRollBySource]
    D -->|No| F{single active dim?}
    F -->|Yes| G{activeDim == 0?}
    G -->|Yes| H[CopyLeadingDimRollBySource]
    G -->|No| I{innerSize == 1 and last dim?}
    I -->|Yes| J[CopyLastDimRollByRows]
    I -->|No| K[CopySingleDimRollByBlocks]
    F -->|No| L{last active dim is final dim?}
    L -->|Yes| M[CopyMultiDimLastDimRollByRows]
    L -->|No| N[CopyMultiDimNonLastRollByBlocks]
    N --> O[CopySegmentedRoll]