CANN ops-math Range 算子技术指南：等差数列生成算子的接口调用与源码实现剖析

算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载Range 算子是 CANN ops-math 数学算子库中用于生成等差数列的基础算子它接收 start、end、step 三个标量输入在 NPU 上高效产出区间[start, end)内均匀分布的浮点序列。本文以仓库中 experimental/math/range/README.md 为骨架结合该算子目录下的算子原型、Host 侧定义、形状推导、Tiling 与 Kernel 源码以及完整的 aclnn 调用样例系统讲解 Range 算子的功能语义、参数规格、动态 Shape 推导机制、多核切分策略与 Ascend C 内核实现帮助读者在 Atlas A2 系列产品上快速完成 Range 算子的工程集成与二次开发。功能说明Range 算子完成如下数学映射给定标量start、end、step生成区间[start, end)范围内的等差数列序列元素为start, start step, start 2*step, ..., 小于 end 的最大值序列总长度由ceil(|end - start| / |step|)决定输出为一维1DTensor。该语义与 TensorFlow 的RANGE算子保持一致这一点在算子原型 range_proto.h 的注释中有明确说明Compatible with the TensorFlow operator RANGE因此对于从 TensorFlow 图迁移到 CANN 场景的开发者具有天然的语义兼容性。在 op_graph/range_proto.h 中算子通过REG_OP(Range)注册了完整的输入输出原型REG_OP(Range) .INPUT(start, TensorType({ge::DT_FLOAT})) .INPUT(end, TensorType({ge::DT_FLOAT})) .INPUT(step, TensorType({ge::DT_FLOAT})) .OUTPUT(z, TensorType({ge::DT_FLOAT})) .OP_END_FACTORY_REG(Range)即三个输入start、end、step与一个输出z全部为float32类型。产品支持情况原文档明确列出 Range 算子的支持范围产品是否支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品√这一产品约束在 Host 侧算子定义 range_def.cpp 中得到印证算子通过AICore().AddConfig(ascend910b, aicoreConfig)仅为ascend910b对应 Atlas A2 系列配置了 AI Core 实现。也就是说Range 算子的 Kernel 编译产物面向 Atlas A2 系列昇腾处理器生成使用前请确认目标设备属于该产品线。参数说明Range 算子的输入输出参数如下表所示参数名输入/输出/属性描述数据类型数据格式start输入等差数列的起始值包含在序列中float32NDend输入等差数列的结束边界不包含在序列中float32NDstep输入等差数列的步长float32NDout输出生成的等差数列结果1D Tensorfloat32ND结合 range_def.cpp 的注册信息可以进一步提炼出以下实现层面的关键特征值依赖ValueDepend三个输入start、end、step均声明为ValueDepend(REQUIRED)且ParamType(REQUIRED)表示算子执行必须读取这三个输入的标量数值本身而非仅依赖 shape这是动态 Shape 推导和动态编译的前提。格式约束三个输入均限定为FORMAT_ND并调用AutoContiguous()确保数据以连续内存布局传入输出z同样为FORMAT_ND的float32数据。动态能力DynamicCompileStaticFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)表明算子支持动态 Rank 与动态 Shape实际输出长度在运行时根据输入标量值确定PrecisionReduceFlag(true)允许在精度允许范围内做精度规约优化。使用要点step不能为 0。虽然原文档约束说明章节标记为无但从形状推导与 Tiling 实现来看stepAbs 0时除法将无意义相关代码会直接返回失败详见下文因此实际工程使用中应保证step ! 0。动态 Shape 推导机制Range 算子的输出长度不是静态 shape 能表达的它依赖输入标量值因此采用了输入值依赖 动态 Shape的推导链路。形状推导实现在 range_infershape.cpp 中通过InputsDataDependency({0, 1, 2})声明三个输入均为数据依赖项range_infershape.cpp在InferShape中读取三个标量值计算diff |end - start|与stepAbs |step|若stepAbs 0直接返回GRAPH_FAILED输出长度按output_length ceil(diff / stepAbs)计算并将输出 shape 设为维度数为 1、第 0 维长度为output_length的向量range_infershape.cpp。由此可知输出长度公式为N ceil(|end - start| / |step|)以示例程序中的参数start 2.0, end 29.0, step 3.0为例输出长度为ceil(|29 - 2| / 3) 9即生成序列2, 5, 8, 11, 14, 17, 20, 23, 26。值得注意的是长度计算对end - start与step均取绝对值因此该公式同时覆盖了 step 为负递减序列的情况只要start、end、step符号关系自洽即可。算子内部实现原理Range 算子的执行路径分为 Host 侧 Tiling 与设备侧 Ascend C Kernel 两部分二者通过RangeTilingData结构定义见 range_tiling_data.h传递切分信息。Host 侧 Tiling多核均衡切分Tiling 实现在 range_tiling.cpp 中核心流程如下读取平台信息通过platform_ascendc::PlatformAscendC获取设备 AI Core 数量coreNum与 UB 内存大小range_tiling.cpp读取输入标量值TilingInputsDataDependency({0, 1, 2})声明 Tiling 阶段同样依赖三个输入值运行时从输入 Tensor 中读出startVal、endVal、stepValrange_tiling.cpp计算输出总长度并对齐totalLength ceil(|end - start| / |step|)并按 32B 对齐基准alignNum 32 / 4 8即每个对齐单元容纳 8 个 float 元素向上取整得到totalLengthAlignedrange_tiling.cpp多核分配在ComputeTiling中优先做均匀切分isEvenCore 1当总块数不能被核数整除时将前若干核分配前段较长块、其余核分配尾段较短块并把前段/尾段的块数、长度、Tile 数、末 Tile 长度分别写入former*与tail*字段range_tiling.cpp单核内部切 TileTilingParamsCalc根据 UB 容量默认UB_BLOCK_NUM 100块、双缓冲BUFFER_NUM 2计算每核的tileNum、tileLength与lastTileLength并显式设置SetBlockDim(coreNum)range_tiling.cpp。设备侧 Kernel按核索引生成序列Kernel 入口 range.cpp 注册了带REGISTER_TILING_DEFAULT/GET_TILING_DATA_WITH_STRUCT的__global__ __aicore__函数将 tiling 数据解析后调用NsRange::Rangefloat, float, float完成计算。核心类实现在 range.h 中设计要点包括始终以 float 计算注释明确说明不论输入是什么类型kernel 侧始终使用 float 进行计算输入标量通过__gm__ float*直接读取range.h按核定位输出区通过AscendC::GetBlockIdx()获取当前核 ID结合 Tiling 阶段算好的blockLength、formerNum等字段计算当前核的全局偏移blockOffset并据此绑定 GlobalTensorzGmrange.h序列生成每个 Tile 固定处理 8 个 float 元素TILE_ELEM_NUM 32B / 4B 8对第globalIdx个元素直接计算value start step * globalIdx超出有效长度origSize的位置补 0随后通过 VECOUT 队列EnQue/DeQue配合双缓冲与DataCopy落盘range.h边界处理当输出长度恰为 1 时走独立分支仅写入start一个有效元素range.h。调用说明原文档给出的调用方式为aclnn 接口通过aclnnRange两段式接口完成算子调用。仓库中提供了可直接编译运行的完整样例 test_aclnn_range.cpp下面以它为例说明完整调用流程。1. 初始化 ACL 运行环境auto ret aclInit(nullptr); // ACL 全局初始化 ret aclrtSetDevice(deviceId); // 设置 device ret aclrtCreateStream(stream); // 创建 stream对应样例 test_aclnn_range.cpp2. 构造输入与输出 Tensor三个输入为形状{1}的 float32 标量 Tensor同时还需用aclCreateFloatArray构造对应的标量数组句柄供第一段接口传入float startValue 2.0; float endValue 29.0; float stepValue 3.0; // 以 start 为例创建 shape{1} 的 ACL_FLOAT Tensor 及 FloatArray ret CreateAclTensor(selfStartHostData, selfStartShape, selfStartDeviceAddr, aclDataType::ACL_FLOAT, selfStart); const std::vectorfloat startVec {startValue}; aclFloatArray* startArray aclCreateFloatArray(startVec.data(), startVec.size());对应样例 test_aclnn_range.cppend、step构造方式相同输出长度按公式预计算并创建一维输出 Tensorint64_t outputLength static_castint64_t(abs(ceil((endValue - startValue) / stepValue))); std::vectorint64_t outShape {outputLength}; ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out);对应样例 test_aclnn_range.cpp3. 两段式接口调用先调用第一段接口查询 workspace 大小再按需申请 workspace 内存随后调用第二段接口在指定 stream 上异步执行// 第一段查询 workspace 大小并创建 executor ret aclnnRangeGetWorkspaceSize(startArray, endArray, stepArray, out, workspaceSize, executor); // 按需申请 workspace 内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 第二段下发执行 ret aclnnRange(workspaceAddr, workspaceSize, executor, stream); // 同步等待任务完成 ret aclrtSynchronizeStream(stream);对应样例 test_aclnn_range.cpp4. 结果回拷与资源释放执行完成后通过aclrtMemcpy将 device 侧结果拷贝回 host 并打印随后依次释放aclTensor、aclFloatArray、device 内存、stream最后aclrtResetDevice与aclFinalize收尾对应样例 test_aclnn_range.cpp。样例运行输出形如Output data length: 9 mean result[0] is: 2.000000 mean result[1] is: 5.000000 ...调用注意事项三个输入参数以aclFloatArray而非 Tensor形式传给aclnnRangeGetWorkspaceSize这是 Range 算子 aclnn 接口与一般张量算子不同的签名特征每个aclTensor创建后必须用aclDestroyTensor释放aclFloatArray用aclDestroyFloatArray释放避免句柄泄漏样例头文件包含aclnn_range.h编译时需链接 CANN 运行时与算子库相关依赖可参考仓库中其他算子的构建方式。总结Range 算子以极简的三标量输入模型在 Atlas A2 系列产品上实现了区间等差数列的高效生成。从仓库源码可以看到一条完整的技术链路REG_OP原型声明兼容 TensorFlow RANGE→ Host 侧值依赖的动态 Shape 推导输出长度ceil(|end-start|/|step|)→ 基于 AI Core 数目的多核均衡 Tiling 与 32B 对齐 → 双缓冲驱动的 Ascend C Kernel 逐元素生成。对于需要在 NPU 上构造索引序列、时间轴或采样点的开发者既可以直接使用aclnnRange两段式接口完成集成参考 test_aclnn_range.cpp也可以对照 range_tiling.cpp 与 range.h 深入理解其并行策略作为自研标量驱动型算子的参考范本。赞分享算子库人工智能CANN【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-math点击查看免费下载相关推荐CANN ops-math 算子开发指南Arange 等差序列算子的 aclnn 接口实现与两段式调用实战CANN ops math 算子开发指南Arange 等差序列算子的 aclnn 接口实现与两段式调用实战 导读 本文以 CANN ops math 开源仓库算子库人工智能CANNCANN ops-math 算子 aclnnTensorMove 接口指南纯数据搬运算子的两段式调用与源码实现解析CANN ops math 算子 aclnnTensorMove 接口指南纯数据搬运算子的两段式调用与源码实现解析 本文是 CANN ops math 仓库中算子库人工智能CANNCANN ops-math Greater 算子实战aclnnGtScalar / aclnnGtTensor 两段式接口调用与 NPU 源码实现剖析CANN ops math Greater 算子实战aclnnGtScalar / aclnnGtTensor 两段式接口调用与 NPU 源码实现剖析 导读算子库人工智能CANN创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考