CANN ops-math 的 Roll 算子实战指南：aclnnRoll 接口调用、循环位移语义与多核实现原理

• 算子库
• 人工智能
• CANN

Roll 是 CANN ops-math 数学算子库中负责「沿给定维度对张量做循环位移」的基础算子，常用于数据增强、滑动窗口与序列对齐等场景。本文以 experimental/math/roll/README.md 为主线，结合算子目录下的接口文档、Host 侧 Tiling、Kernel 实现与单元测试源码，完整讲解 Roll 的功能语义、参数约束、编译部署、aclnnRoll 两段式调用方式，以及其从参数校验、展平合并、多核切分到 AICore 执行的全链路实现原理，帮助开发者快速上手并在 NPU 上正确、高效地使用该算子。

Roll 算子功能与语义

Roll 沿给定维度对输入张量执行循环位移（circular shift）：当某维度上位移量为正时，元素向该维度增大方向滚动，越界的元素循环回该维度起始位置；位移量为负时方向相反。

• 当dims非空时，shifts[i]作用于dims[i]指定的维度；
• 当dims为空时，先按逻辑视图将输入展平（flatten）成一维，执行一维循环位移，最后按原始形状输出；
• 若shifts与dims中指定了重复维度，Host 侧会将多次位移归一化合并为一个位移量（见下文 Tiling 实现）。

典型示例：对形状为[2, 3]、元素为[0,1,2,3,4,5]的输入执行shifts=[1], dims=[1]（沿最后一维右移 1 位），输出为[[2,0,1],[5,3,4]]，即每行内部循环滚动，行间顺序不变。

算子原型与支持范围

原型信息

该原型在 roll_def.cpp 中通过算子定义注册实现：输入x与输出y均声明了 9 种支持的DataType（DT_BOOL/DT_UINT8/DT_INT8/DT_BF16/DT_FLOAT16/DT_FLOAT/DT_INT32/DT_UINT32/DT_COMPLEX64）且格式限定为ND；属性shifts为REQUIRED的ListInt，dims为OPTIONAL的ListInt，默认值为空列表。同时，该文件为算子在ascend910b、ascend910_93、ascend950三种 AICore 配置上注册了编译入口。

支持的产品型号

• Atlas A2 训练系列产品
• Atlas A3 训练系列产品
• Atlas A5 训练系列产品

约束与限制

使用 Roll 算子前需确认以下约束（来自 README，且与 aclnnRoll.md 及 aclnn_roll.cpp 中的参数校验逻辑一一对应）：

• 仅支持ND数据格式；
• 支持 0 维到 8 维输入（源码中MAX_SUPPORT_DIMS_NUMS = 8，见 aclnn_roll.cpp）；
• dims为空时，shifts长度必须为 1；
• dims非空时，shifts与dims长度必须一致；
• dims取值范围为[-rank, rank)（支持负数索引，-1表示最后一维）；
• 0 维输入时，shifts长度必须为 1，且dims必须为空；
• 输入输出 dtype 必须一致、shape 必须一致；
• 非连续输入会先整理为连续视图后执行，非连续输出会在算子结果生成后做回写（contiguous与view_copy处理见下文调用链路）。

环境准备与编译部署

使用该算子前，请先参考社区版 CANN 开发套件包安装文档完成开发运行环境部署，然后执行编译打包并安装算子包：

cd ${git_clone_path}/ops-math bash build.sh --pkg --experimental --soc=ascend910b --ops=roll ./build_out/cann-ops-<vendor_name>-linux.<arch>.run

命令说明：

• --pkg：生成可安装的算子软件包；
• --experimental：开启实验性算子目录（Roll 位于experimental/math/roll）的构建；
• --soc=ascend910b：指定目标 SoC（本算子同样支持ascend910_93、ascend950）；
• --ops=roll：仅构建 Roll 算子，缩短编译时间；
• 第二个命令将生成的.run包安装到 CANN 运行环境。

算子目录的构建接入方式见 roll/CMakeLists.txt 中的add_all_modules_sources(OPTYPE roll ACLNNTYPE aclnn_exclude)。

aclnnRoll 接口详解

Roll 算子通过标准的aclnn 两段式接口调用，函数原型定义在 aclnn_roll.h：

aclnnStatus aclnnRollGetWorkspaceSize( const aclTensor* x, const aclIntArray* shifts, const aclIntArray* dims, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor); aclnnStatus aclnnRoll( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream);

aclnnRollGetWorkspaceSize 参数

返回值与参数校验

第一段接口aclnnRollGetWorkspaceSize完成全部参数校验（实现见 aclnn_roll.cpp 中的CheckParams及若干子检查函数），返回aclnnStatus，出错场景包括：

• ACLNN_ERR_PARAM_NULLPTR：x、shifts、out、workspaceSize、executor为空；
• ACLNN_ERR_PARAM_INVALID：
  • 输入或输出 dtype 不在支持范围内（CheckDtypeValid比对DTYPE_SUPPORT_LIST）；
  • 输入与输出 dtype 不一致；
  • 输入或输出为私有存储格式（非 ND，CheckFormatValid使用IsPrivateFormat判定）；
  • 输入与输出 shape 不一致（CheckShapeValid）；
  • rank 大于 8（x->GetViewShape().GetDimNum() > MAX_SUPPORT_DIMS_NUMS）；
  • 0 维输入时，shifts长度不为 1 或dims非空；
  • dims为空但shifts长度不为 1；
  • dims非空但shifts与dims长度不一致；
  • dims元素越界（dim < -tensorDim || dim >= tensorDim，见CheckDimsRange）。

特别地，当输入张量为空（x->IsEmpty()）时，接口直接返回ACLNN_SUCCESS且workspaceSize = 0，不会下发 Kernel 任务。测试用例 test_aclnn_roll.cpp 覆盖了上述校验路径，例如case_invalid_dtype（ACL_DOUBLE 返回ACLNN_ERR_PARAM_INVALID）、case_invalid_dims_range（dims=[2]越界）。

aclnnRoll 参数

第二段接口aclnnRoll通过CommonOpExecutorRun把第一段生成的 executor 在指定 stream 上异步执行，不参与参数校验。

完整调用示例

test_aclnn_roll.cpp 提供了一个可直接运行的完整示例：以[2, 3]的 complex64 输入、shifts=[1]、dims=[1]调用 Roll，并对输出做位级精确校验。核心流程如下：

// 1. 初始化 ACL 环境并创建 stream auto ret = aclInit(nullptr); ret = aclrtSetDevice(deviceId); ret = aclrtCreateStream(&stream); // 2. 构造输入/输出 aclTensor（分配 device 内存、H2D 拷贝、计算 strides） ret = aclrtMalloc(&xDeviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); ret = aclrtMemcpy(xDeviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); x = aclCreateTensor(shape.data(), shape.size(), ACL_COMPLEX64, strides.data(), 0, ACL_FORMAT_ND, shape.data(), shape.size(), xDeviceAddr); // 3. 构造 shifts / dims 属性数组 std::vector<int64_t> shiftsData = {1}; std::vector<int64_t> dimsData = {1}; aclIntArray* shifts = aclCreateIntArray(shiftsData.data(), shiftsData.size()); aclIntArray* dims = aclCreateIntArray(dimsData.data(), dimsData.size()); // 4. 第一段接口：参数校验并获取 workspace 大小与 executor uint64_t workspaceSize = 0; aclOpExecutor* executor = nullptr; ret = aclnnRollGetWorkspaceSize(x, shifts, dims, y, &workspaceSize, &executor); // 5. 按需分配 workspace void* workspaceAddr = nullptr; if (workspaceSize > 0) { ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 6. 第二段接口：异步执行，随后同步等待结果 ret = aclnnRoll(workspaceAddr, workspaceSize, executor, stream); ret = aclrtSynchronizeStream(stream); // 7. D2H 拷贝结果并校验 ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), yDeviceAddr, resultData.size() * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); // 8. 释放资源：aclDestroyIntArray / aclDestroyTensor / aclrtFree / aclrtDestroyStream / aclrtResetDevice / aclFinalize

该示例中对形状[2,3]、输入{(0,10),(1,11),(2,12),(3,13),(4,14),(5,15)}、shifts=[1]、dims=[1]的期望输出为{(2,12),(0,10),(1,11),(5,15),(3,13),(4,14)}，与「每行内部右移一位」的语义完全一致，示例最后通过memcmp做 complex64 位级精确比对并打印PASS/FAIL。

调用链路与 Host 侧处理逻辑

从 aclnn 入口到 Kernel 执行，Roll 的调用链路为：

1. aclnnRollGetWorkspaceSize完成参数校验后，根据输入输出张量的布局决定执行路径（aclnn_roll.cpp）：
   • 若输入非连续（!HasDenseViewLayout(x)），先调用l0op::Contiguous将其整理为连续视图；
   • 若输出可以直接写入（CanWriteOutDirectly：密集布局且 storage shape 与 view shape 一致），则直接以out为目标执行；
   • 否则先计算到临时张量，再通过l0op::ViewCopy回写到非连续输出；
2. L0 层l0op::Roll（roll.cpp）通过ADD_TO_LAUNCHER_LIST_AICORE(Roll, OP_INPUT(x), OP_OUTPUT(rollOut), OP_ATTR(shifts, dims))下发 AICore 任务；
3. Host 侧 Tiling（roll_tiling.cpp）计算分核参数并填充RollTilingData；
4. Kernel 端roll核函数（roll.cpp）在 AIV 上执行循环位移。

Host 侧语义归一化

Tiling 阶段对属性做了关键归一化（roll_tiling.cpp）：

• 负数修正与取模：PositiveMod将负的shifts转正（value % mod后若为负则加mod），负数dims加rank转正；
• dims 为空时的展平：将dimNum置 1、shapes[0] = totalNum、strides[0] = 1、shifts[0] = PositiveMod(shift, totalNum)，与 README 描述的「按逻辑视图展平后一维 roll」对应；
• 重复维度合并：对每个dims[i]，执行shifts[dim] = PositiveMod(shifts[dim] + PositiveMod(shifts[i], shapes[dim]), shapes[dim])，使同一维度上的多次位移累加取模，与接口文档「重复维度会在 Host 侧归一化合并」的语义说明一致。

RollTilingData结构体（roll_tiling_data.h）中ROLL_MAX_DIM_NUM = 8，记录了totalNum、各维shapes/strides/shifts、活跃位移维度数activeDimCount与唯一活跃维度activeDim（存在位移的维度）、innerSize/dimSize/outerSize、每核元素数、UB 元素数等关键信息。

多核切分与 Tiling 策略

Tiling 通过Ops::Base::GetAivCoreNum获取可用 AIV 核数，并综合考虑数据类型字节数、GM 块大小（GetUbBlockSize）、GM 带宽对齐（512 字节）与 UB 容量（64 KB）计算blockDim与每核元素数：

• 小数据量（总字节数 ≤ 4096）默认单核执行；
• 多核场景下按对齐后的perCoreElements均分，末核元素数由lastCoreElements记录；
• 针对 bf16、uint8、float16 等数据类型以及「最后一维 roll」「首维大 stride roll」「多活跃维度」等形态，包含多组专门的粒度对齐与拆分分支（如splitHugeLeadingDimRoll、splitHugeTwoWayRoll、splitLargeByInner），以贴合 GM 带宽与 UB 搬移限制；
• workspace 恒为 0（FillWorkspace写入WORKSPACE_SIZE = 0），算子不需要额外 workspace；
• 最终通过context->SetBlockDim(blockDim)与SetTilingKey(GET_TPL_TILING_KEY(ROLL_TPL_SCH_MODE_0))输出调度信息，tiling key 的可选模板参数在 roll_tiling_key.h 中声明。

对应的 Tiling 单测（test_roll_tiling.cpp）验证了典型场景，例如basic_last_dim_roll（float、[2,3]、shifts=[1], dims=[1]）断言totalNum=6、dimNum=2、activeDim=1、dimSize=3、innerSize=1、activeShift=1；flatten_roll_when_dims_empty（float16、[2,3,4]、dims为空）断言dimNum=1、totalNum=24。

Kernel 端实现：多路径数据搬移

Kernel 核函数入口位于 roll.cpp，模板类型做了特殊映射：complex64 按 8 字节的uint64_t处理、bool 按uint8_t处理，其余类型按原 dtype；任务类型固定为KERNEL_TYPE_AIV_ONLY（纯向量核执行）。

核心类RollKernel::Roll<T>（roll.h）在Init中根据GetBlockIdx()计算本核负责的起始下标startIndex_与元素数elementCount_，并在 UB 上申请inQueue_/outQueue_双缓冲。Process按位移形态选择执行路径，包括：

• 无位移（CopyIdentity）：直接按源对齐拷贝；
• 展平一维 roll（CopyFlattenRoll/CopyFlattenRollBySource）：把一维位移拆成两段连续区间交换拷贝；
• 首维位移（CopyLeadingDimRollBySource）：按「块内偏移是否越过位移分界」拆分连续段；
• 末维位移（CopyLastDimRoll系列）：逐行在 UB 内完成行内滚动（CopyRowsRollInUb），并利用CopyStridedSourceSegments以DataCopyExtParams的多 block 跨步搬移一次处理多行；
• 多活跃维度的非末维位移（CopyMultiDimNonLastRollByBlocks）：先按块（blockSize = dimSize * inner）计算源块索引ComputeSourceBlockIndex与连续块运行长度ComputeContiguousSourceBlockRun，再批量搬移；
• 通用路径（CopySegmentedRoll/CopySegment）：通过ComputeInputIndex逐段计算源下标，按「源对齐搬移 / 标量搬移 / flat patch 合并搬移」等策略拷贝。

其中CopyBlockRollByFlatPatch针对小块（块大小 × 类型字节 ≤ 768、总量 ≥ 4096 字节、单活跃维度）在 UB 内一次性完成多块的「前半段 + 后半段」换位后整体回写，减少 GM 访问次数；CopyStridedSingleElementByRowGather则针对 1 字节类型的小块多行场景做行收集优化。这些分支共同保证了不同 shape、dtype、位移维度组合下的搬移效率。

Kernel 侧同样存在测试：test_roll.cpp 与 roll_tiling.h 用于验证 Kernel 数值正确性，构建配置见 tests/ut/op_kernel/CMakeLists.txt。

总结

• 语义：Roll 沿指定维度循环位移；dims为空时按展平视图执行一维 roll，位移量支持负数，重复维度在 Host 侧归一化合并。
• 调用方式：采用aclnnRollGetWorkspaceSize+aclnnRoll两段式接口，完整可运行样例见 test_aclnn_roll.cpp。
• 约束：仅 ND 格式、0-8 维、9 种 dtype、shifts/dims长度规则与dims取值范围需满足，接口层会在第一段调用中返回明确的ACLNN_ERR_PARAM_INVALID等错误码。
• 实现：从 roll_def.cpp 的算子注册，到 roll_tiling.cpp 的多核切分，再到 roll.h 的多路径数据搬移，形成完整的「定义—调度—执行」链路，支持 Atlas A2/A3/A5 训练系列产品。

• 算子库
• 人工智能
• CANN