算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载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]]即每行内部循环滚动行间顺序不变。算子原型与支持范围原型信息项目说明算子类型OpTypeRoll算子输入xtensor支持 bool、uint8、int8、bfloat16、float16、float32、int32、uint32、complex64格式 ND算子输出ytensordtype 与x相同格式 ND属性shiftslistInt整型列表必选dimslistInt整型列表可选默认空核函数名roll该原型在 roll_def.cpp 中通过算子定义注册实现输入x与输出y均声明了 9 种支持的DataTypeDT_BOOL/DT_UINT8/DT_INT8/DT_BF16/DT_FLOAT16/DT_FLOAT/DT_INT32/DT_UINT32/DT_COMPLEX64且格式限定为ND属性shifts为REQUIRED的ListIntdims为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.cppdims为空时shifts长度必须为 1dims非空时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 --socascend910b --opsroll ./build_out/cann-ops-vendor_name-linux.arch.run命令说明--pkg生成可安装的算子软件包--experimental开启实验性算子目录Roll 位于experimental/math/roll的构建--socascend910b指定目标 SoC本算子同样支持ascend910_93、ascend950--opsroll仅构建 Roll 算子缩短编译时间第二个命令将生成的.run包安装到 CANN 运行环境。算子目录的构建接入方式见 roll/CMakeLists.txt 中的add_all_modules_sources(OPTYPE roll ACLNNTYPE aclnn_exclude)。aclnnRoll 接口详解Roll 算子通过标准的aclnn 两段式接口调用函数原型定义在 aclnn_roll.haclnnStatus 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 参数参数名输入/输出描述数据类型数据格式维度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**--返回值与参数校验第一段接口aclnnRollGetWorkspaceSize完成全部参数校验实现见 aclnn_roll.cpp 中的CheckParams及若干子检查函数返回aclnnStatus出错场景包括ACLNN_ERR_PARAM_NULLPTRx、shifts、out、workspaceSize、executor为空ACLNN_ERR_PARAM_INVALID输入或输出 dtype 不在支持范围内CheckDtypeValid比对DTYPE_SUPPORT_LIST输入与输出 dtype 不一致输入或输出为私有存储格式非 NDCheckFormatValid使用IsPrivateFormat判定输入与输出 shape 不一致CheckShapeValidrank 大于 8x-GetViewShape().GetDimNum() MAX_SUPPORT_DIMS_NUMS0 维输入时shifts长度不为 1 或dims非空dims为空但shifts长度不为 1dims非空但shifts与dims长度不一致dims元素越界dim -tensorDim || dim tensorDim见CheckDimsRange。特别地当输入张量为空x-IsEmpty()时接口直接返回ACLNN_SUCCESS且workspaceSize 0不会下发 Kernel 任务。测试用例 test_aclnn_roll.cpp 覆盖了上述校验路径例如case_invalid_dtypeACL_DOUBLE 返回ACLNN_ERR_PARAM_INVALID、case_invalid_dims_rangedims[2]越界。aclnnRoll 参数参数名输入/输出描述workspace输入Device 侧 workspace 地址workspaceSize输入Device 侧 workspace 大小executor输入执行器stream输入执行 stream第二段接口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::vectorint64_t shiftsData {1}; std::vectorint64_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 的调用链路为aclnnRollGetWorkspaceSize完成参数校验后根据输入输出张量的布局决定执行路径aclnn_roll.cpp若输入非连续!HasDenseViewLayout(x)先调用l0op::Contiguous将其整理为连续视图若输出可以直接写入CanWriteOutDirectly密集布局且 storage shape 与 view shape 一致则直接以out为目标执行否则先计算到临时张量再通过l0op::ViewCopy回写到非连续输出L0 层l0op::Rollroll.cpp通过ADD_TO_LAUNCHER_LIST_AICORE(Roll, OP_INPUT(x), OP_OUTPUT(rollOut), OP_ATTR(shifts, dims))下发 AICore 任务Host 侧 Tilingroll_tiling.cpp计算分核参数并填充RollTilingDataKernel 端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 恒为 0FillWorkspace写入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_rollfloat、[2,3]、shifts[1], dims[1]断言totalNum6、dimNum2、activeDim1、dimSize3、innerSize1、activeShift1flatten_roll_when_dims_emptyfloat16、[2,3,4]、dims为空断言dimNum1、totalNum24。Kernel 端实现多路径数据搬移Kernel 核函数入口位于 roll.cpp模板类型做了特殊映射complex64 按 8 字节的uint64_t处理、bool 按uint8_t处理其余类型按原 dtype任务类型固定为KERNEL_TYPE_AIV_ONLY纯向量核执行。核心类RollKernel::RollTroll.h在Init中根据GetBlockIdx()计算本核负责的起始下标startIndex_与元素数elementCount_并在 UB 上申请inQueue_/outQueue_双缓冲。Process按位移形态选择执行路径包括无位移CopyIdentity直接按源对齐拷贝展平一维 rollCopyFlattenRoll/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 侧归一化合并。调用方式采用aclnnRollGetWorkspaceSizeaclnnRoll两段式接口完整可运行样例见 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【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 中 Roll 算子的 aclnnRoll 接口使用指南与实现原理CANN ops math 中 Roll 算子的 aclnnRoll 接口使用指南与实现原理 本文以 CANN ops math 开源仓库中 experimen算子库人工智能CANNCANN ops-math 算子详解KLDivV2 的接口、实现原理与 aclnn 调用实战CANN ops math 算子详解KLDivV2 的接口、实现原理与 aclnn 调用实战 KLDivV2 是 CANN ops math 算子库中用于计算算子库人工智能CANNCANN ops-math 算子实战aclnnCircularPad3dBackward 两段式接口原理与调用指南CANN ops math 算子实战aclnnCircularPad3dBackward 两段式接口原理与调用指南 本文围绕 CANN ops math 开源算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考