
CANN ops-math 中 Roll 算子的 aclnnRoll 接口使用指南与实现原理【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math本文以 CANN ops-math 开源仓库中 experimental/math/roll/docs/aclnnRoll.md 为核心主体系统讲解 Roll循环位移算子的 aclnn 两级接口aclnnRollGetWorkspaceSize与aclnnRoll的语义、参数、返回值、约束与调用方式并结合仓库内 示例代码、op_api 实现、Tiling 实现 与 Kernel 实现 进行源码级剖析。读完本文你将能够在 Atlas 训练系列产品上独立完成 Roll 算子的 aclnn 接口调用、参数校验与结果验证并理解其在 Host 侧归一化、多核并行切分与非连续 Tensor 处理上的底层设计。一、功能说明与语义1.1 接口功能aclnnRoll沿给定维度对输入 Tensor 执行循环位移cyclic shift / roll每个元素在其所处维度上整体平移若干个位置移出边界的元素从另一端补回Tensor 的形状与数据类型保持不变。在数学意义上这等价于 PyTorch/TensorFlow 中的roll操作常用于特征对齐、序列平移等场景。该算子归属于 CANN ops-math 仓库的 experimental/math 目录核心实现位于 experimental/math/roll 目录下。1.2 语义说明按官方文档位移语义分三种情况处理dims非空shifts[i]作用于dims[i]对应的维度二者长度必须一致一一对应。dims为空可省略或传空数组先将输入按逻辑视图展平flatten为一维执行一维 roll再按原始形状输出。重复维度若dims中出现重复维度位移量会在 Host 侧Tiling 阶段归一化合并——同一维度上的多次位移先求和再取模避免重复搬运数据。这一“Host 侧归一化合并”的语义在 roll_tiling.cpp 中有直接印证Tiling 阶段对每个dim做正数取模PositiveMod后累加到tilingData-shifts[dim]上再整体取模从而把多次位移折叠为单次有效位移// roll_tiling.cpp 中重复维度归一化的核心逻辑 int64_t PositiveMod(int64_t value, int64_t mod) { if (mod 0) { return 0; } int64_t result value % mod; return result 0 ? result mod : result; } ... tilingData-shifts[dim] PositiveMod( tilingData-shifts[dim] PositiveMod(shifts[i], tilingData-shapes[dim]), tilingData-shapes[dim]);注意PositiveMod对负位移量同样做了处理例如位移 -1 等价于正向位移dimSize - 1负数取模结果会加上模数转为正数因此shifts允许传入负整数。1.3 产品支持情况产品是否支持Atlas A2 训练系列产品支持Atlas A3 训练系列产品支持Atlas A5 训练系列产品支持在算子定义文件 roll_def.cpp 中可以看到对应平台的 AICore 配置为ascend910b、ascend910_93、ascend950。二、函数原型与参数说明aclnnRoll采用 CANN aclnn 标准的两级接口调用模式先用aclnnRollGetWorkspaceSize完成参数校验、算子图构建并查询所需 workspace 大小再调用aclnnRoll在指定 stream 上真正执行计算。2.1 aclnnRollGetWorkspaceSizeaclnnStatus aclnnRollGetWorkspaceSize( const aclTensor* x, const aclIntArray* shifts, const aclIntArray* dims, aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor)参数名输入/输出描述数据类型数据格式维度x输入输入张量bool, uint8, int8, bfloat16, float16, float32, int32, uint32, complex64ND0-8 维shifts输入每个目标维度上的循环位移量aclIntArray*--dims输入循环位移维度可省略或传空数组aclIntArray*--out输出输出张量shape 和 dtype 与 x 一致与 x 相同ND0-8 维workspaceSize输出需要申请的 workspace 大小uint64_t*--executor输出执行器aclOpExecutor**--关于各参数的实际使用补充说明如下x/out均为aclTensor句柄由aclCreateTensor创建参见下文示例。out的 shape、dtype 必须与x完全一致因为 Roll 只做元素重排不改变形状与类型。shifts/dims均为aclIntArrayint64_t数组通过aclCreateIntArray创建。dims可以传nullptr或长度为 0 的空数组此时按展平后的一维 roll 语义执行。dims取值范围[-rank, rank)支持负索引负值表示从末尾倒数例如dims -1表示最后一维Tiling 阶段会通过dim originalDimNum归一化为正索引。2.2 aclnnRollaclnnStatus aclnnRoll( void* workspace, uint64_t workspaceSize, aclOpExecutor* executor, aclrtStream stream)参数名输入/输出描述workspace输入Device 侧 workspace 地址workspaceSize输入Device 侧 workspace 大小executor输入执行器由第一段接口获得stream输入执行 streamworkspace需要调用方在 Device 侧通过aclrtMalloc按workspaceSize大小申请若workspaceSize为 0可以传nullptr。从 Roll 的 Tiling 实现看roll_tiling.cpp、L78-L84WORKSPACE_SIZE恒为 0即本算子在当前实现下不需要额外 workspace——这从侧面说明 Roll 的全部中间结果均在片上UB/寄存器完成无需片外中转缓冲。三、返回值与参数校验aclnnRollGetWorkspaceSize返回aclnnStatus。第一段接口负责完成全部参数校验出现以下场景时返回对应错误码ACLNN_ERR_PARAM_NULLPTRx、shifts、out、workspaceSize、executor为空dims允许为空。ACLNN_ERR_PARAM_INVALID涵盖以下全部情形输入或输出 dtype 不在支持范围内输入与输出 dtype 不一致输入或输出格式不是ND具体到实现是检测到私有格式见下文输入与输出 shape 不一致rank 大于 80 维输入时shifts长度不为 1 或dims非空dims为空但shifts长度不为 1dims非空但shifts与dims长度不一致dims元素越界超出[-rank, rank)。这些校验规则可以在 op_api/aclnn_roll.cpp 中找到一一对应的实现CheckDtypeValiddtype 支持列表校验 输入输出 dtype 一致性校验CheckFormatValid通过op::IsPrivateFormat拒绝私有存储格式等价于“仅支持 ND”CheckShapeValidshape 一致性校验 MAX_SUPPORT_DIMS_NUMS 8的维度上限校验CheckArraySize0 维输入、dims空/非空三种分支下的shifts/dims长度约束CheckDimsRangedims元素越界校验。此外还有两个值得注意的实现细节空 Tensor 快速返回若x-IsEmpty()第一段接口直接令*workspaceSize 0并返回ACLNN_SUCCESS不构建计算图aclnn_roll.cpp。非连续 Tensor 处理aclnnRollGetWorkspaceSize内部会检查输入视图是否“稠密连续”HasDenseViewLayout。若输入非连续会先插入l0op::Contiguous算子整理为连续视图再执行 Roll若输出不能直接写入非连续输出则先生成结果再通过l0op::ViewCopy回写到目标输出aclnn_roll.cpp。这正对应文档“约束说明”中“非连续输入会先整理为连续视图后执行、非连续输出会在算子结果生成后做回写”两条约束。四、约束说明使用aclnnRoll时需遵守以下约束格式仅支持ND。数据类型仅支持bool、uint8、int8、bfloat16、float16、float32、int32、uint32、complex64九种。这一列表在 op_api/aclnn_roll.cpp 的 DTYPE_SUPPORT_LIST 与 roll_def.cpp 的算子定义 中均有完全一致的定义二者共同约束了入口与图编译两个层面。0 维输入shifts长度必须为 1且dims必须为空。非连续输入会先整理为连续视图后执行内部自动插入 Contiguous。非连续输出会在算子结果生成后做回写内部自动插入 ViewCopy。另外补充两点从源码确认的限制维度上限 8 维含 0 维标量超出返回ACLNN_ERR_PARAM_INVALIDKernel 侧同样通过ROLL_MAX_DIM_NUM做兜底校验roll_tiling.cpp。输入输出元素总数必须一致Tiling 阶段会校验inputElements outputElementsroll_tiling.cpp不满足则 Tiling 失败。五、调用示例完整可运行的示例位于 experimental/math/roll/examples/test_aclnn_roll.cppUT 用例位于 tests/ut/op_api/test_aclnn_roll.cpp。下面以该示例为主线分步讲解一个基于complex64的 2×3 二维 Roll 调用。5.1 环境初始化int Init(int32_t deviceId, aclrtStream* stream) { auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; }依次调用aclInit→aclrtSetDevice→aclrtCreateStream完成 ACL 运行时初始化。5.2 创建输入输出 Tensortemplate typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * static_castint64_t(sizeof(T)); auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); ... ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); ... *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; }创建 Tensor 的要点先在 Device 侧aclrtMalloc并aclrtMemcpyH2D搬入数据按连续行主序row-major计算出strides用aclCreateTensor创建aclTensor格式固定为ACL_FORMAT_ND。示例选用ACL_COMPLEX64类型并在 Host 侧自定义了struct Complex64 { float real; float imag; }结构体表示复数8 字节其输入数据为std::vectorint64_t xShape {2, 3}; std::vectorComplex64 xHostData { {0.0F, 10.0F}, {1.0F, 11.0F}, {2.0F, 12.0F}, {3.0F, 13.0F}, {4.0F, 14.0F}, {5.0F, 15.0F}, };5.3 创建 shifts / dims 属性并查询 workspacestd::vectorint64_t shiftsData {1}; std::vectorint64_t dimsData {1}; aclIntArray* shifts aclCreateIntArray(shiftsData.data(), shiftsData.size()); aclIntArray* dims aclCreateIntArray(dimsData.data(), dimsData.size()); uint64_t workspaceSize 0; aclOpExecutor* executor nullptr; ret aclnnRollGetWorkspaceSize(x, shifts, dims, y, workspaceSize, executor);本例对 shape 为{2, 3}的输入在最后一维dims1上循环位移 1 个位置即每一行内部向右滚动一格。5.4 申请 workspace 并执行void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); ... } ret aclnnRoll(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnRoll failed. ERROR: %d\n, ret); return ret); ret aclrtSynchronizeStream(stream);执行后必须aclrtSynchronizeStream同步再通过aclrtMemcpyD2H把结果拷回 Host 侧进行验证。5.5 结果验证与资源释放const std::vectorComplex64 expected { {2.0F, 12.0F}, {0.0F, 10.0F}, {1.0F, 11.0F}, {5.0F, 15.0F}, {3.0F, 13.0F}, {4.0F, 14.0F}, }; const bool resultOk std::memcmp(resultData.data(), expected.data(), expected.size() * sizeof(Complex64)) 0; LOG_PRINT(complex64 bit-exact check: %s\n, resultOk ? PASS : FAIL);预期的结果完全符合“最后一维正向位移 1”的语义第一行[0,1,2] → [2,0,1]第二行[3,4,5] → [5,3,4]且示例采用逐位精确比对bit-exact。最后依次释放aclIntArray、aclTensor、Device 内存与 stream并调用aclrtResetDevice、aclFinalize收尾。六、编译部署与运行参照 experimental/math/roll/README.md 中“编译部署”一节的说明在克隆仓库后按以下命令构建并安装包含 Roll 算子的开发套件包cd ${git_clone_path}/ops-math bash build.sh --pkg --experimental --socascend910b --opsroll ./build_out/cann-ops-vendor_name-linux.arch.run其中--experimental指定构建 experimental 目录下的实验算子--socascend910b指定目标 SoC 平台Roll 同时配置了 ascend910b / ascend910_93 / ascend950--opsroll只构建 roll 这一个算子加快构建速度。安装完成后即可编译运行 examples/test_aclnn_roll.cpp 验证算子行为。此外仓库还提供了对应平台与框架的 UT 用例tests/ut下的 op_api、op_host、op_kernel 三级测试可用于回归验证。七、源码级实现原理7.1 算子定义OpDefroll_def.cpp 通过OP_ADD(Roll)注册算子输入x与输出y均声明为必选dtype 与 format 列表与 aclnn 层校验完全一致ND 格式、九种类型属性shifts为必选ListIntdims为可选ListInt默认值为空列表{}——这解释了为什么dims可以省略AICore 配置覆盖ascend910b、ascend910_93、ascend950。7.2 形状推导InferShaperoll_infershape.cpp 的实现非常简洁*outputShape *inputShape即输出 shape 完全等于输入 shape符合“循环位移不改变形状”的语义。7.3 Tiling 策略roll_tiling.cpp 中RollTiling是核心整体流程为维度与元素数校验rank 不超过 8输入输出元素总数一致属性解析与归一化dims为空时按一维处理shapes[0] totalNumdims非空时做负索引归一化、重复维度位移合并、PositiveMod取模统计出非零位移的“活动维度”activeDimCount、activeDim当只有一个活动维度时记录innerSize/dimSize/outerSize/activeShift供 Kernel 直接使用多核切分根据平台 AIV 核数Ops::Base::GetAivCoreNum、数据类型字节数与 GM 带宽对齐GM_BANDWIDTH_ALIGN_BYTES 512、UB 容量UB_BYTES 64 * 1024计算blockDim与每个核负责的元素区间perCoreElements、lastCoreElements。其中包含大量针对不同 dtypebfloat16/uint8/float16 等、不同 shape 特征的边界对齐与单核回退优化分支结果输出通过context-SetBlockDim设置核数、SetTilingKey(GET_TPL_TILING_KEY(ROLL_TPL_SCH_MODE_0))指定调度模式并把RollTilingData写入 tiling buffer。从 Tiling 数据结构的字段totalNum、dimNum、shapes[]、strides[]、shifts[]、activeDim、perCoreElements等可以推断Kernel 端正是依据这些切分信息把总元素按核分块每个核只负责自己区间内涉及循环位移的部分避免重复搬运整块数据。7.4 Kernel 实现op_kernel/roll.cpp 中值得注意的类型特化#if defined(ORIG_DTYPE_X) ORIG_DTYPE_X DT_COMPLEX64 using RollDataType uint64_t; #elif defined(ORIG_DTYPE_X) ORIG_DTYPE_X DT_BOOL using RollDataType uint8_t; #else using RollDataType DTYPE_X; #endifcomplex64 以 64 位无符号整数搬运复数本质上是“实部 虚部”的 8 字节对位移不关心内部字节含义因此直接按uint64_t整体搬运保证复数对的完整性这也与示例中“bit-exact 比对”的验证方式一致bool 以 uint8_t 处理规避布尔类型在寄存器上的特化问题按单字节搬运Kernel 入口声明为KERNEL_TYPE_AIV_ONLY即由 AIV 向量核执行经TPipe流水与RollKernel::RollT::Process()完成处理。7.5 aclnn 入口的完整调用链结合 op_api/roll.cpp 与 op_api/aclnn_roll.cpp一次调用的完整链路为aclnnRollGetWorkspaceSize └─ CheckParams空指针/dtype/format/shape/数组长度/dims 越界校验 └─ 空 Tensor 快速返回 └─ 非连续输入 → 插入 Contiguous └─ l0op::Roll → RollAiCore → ADD_TO_LAUNCHER_LIST_AICORE(Roll, ...) └─ 非连续输出 → 插入 ViewCopy 回写 └─ 返回 workspaceSize 与 executor aclnnRoll └─ CommonOpExecutorRun(workspace, workspaceSize, executor, stream)其中 roll.cpp 还提供了一个便利重载未显式传入输出 Tensor 时会通过executor-AllocTensor(x-GetViewShape(), x-GetDataType(), x-GetViewFormat())自动分配一个与输入同 shape、同 dtype、同格式的输出 Tensor。八、总结aclnnRoll是 CANN ops-math 仓库中 Roll 算子的标准 aclnn 两级接口支持 Atlas A2/A3/A5 训练系列产品覆盖 0-8 维、九种数据类型、ND 格式的循环位移计算。本文从接口语义、参数表、错误码、约束条件、完整调用示例到 OpDef / InferShape / Tiling / Kernel 的源码实现完成了从“会调用”到“懂原理”的全链路梳理。实践时建议严格遵循dims为空/非空两种场景下的shifts长度约束善用负索引与负位移量利用第一段接口返回的workspaceSize动态申请 workspace参考 test_aclnn_roll.cpp 的 bit-exact 比对方式验证结果需要了解算子适配与平台细节时可继续阅读 roll_def.cpp、roll_tiling.cpp 与 roll.cpp 等实现文件。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考