CANN ops-math aclnnMaxDim 算子详解指定维度最大值与索引的两段式调用实战【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读aclnnMaxDim是 CANN ops-math 数学算子库experimental/math/arg_max_with_value中一个典型的规约类算子接口它返回输入张量self在指定维度上的最大值及其首次出现的下标语义上与torch.max(dim)对齐。本文以 aclnnMaxDim.md 为主体完整梳理产品支持情况、函数原型、参数约束、返回码与调用示例并结合仓库内 op_api、op_host、op_kernel 与测试源码解释两段式接口背后的执行器构建、计算图拼装Contiguous → ArgMaxWithValue → Cast → ViewCopy与核函数调度细节。读完本文你将能够独立编写、编译并运行一段aclnnMaxDim的 host 侧调用代码并理解其底层实现机制。产品支持情况aclnnMaxDim当前仅支持 Atlas A2/A3 系列产品不支持 Atlas 200I/500 A2 推理产品。以 op_api/aclnn_max_dim.cpp 中IsA2OrA3()的源码为准该接口通过获取当前平台SocVersion判断是否等于ASCEND910BAtlas A2或ASCEND910_93Atlas A3不满足时直接返回ACLNN_ERR_PARAM_INVALID并打印 aclnnMaxDim is only supported on Atlas A2/A3 products.。产品是否支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品√Atlas A2 训练系列产品 / Atlas A2 推理系列产品√Atlas 200I/500 A2 推理产品×在算子定义层op_host/arg_max_with_value_def.cppArgMaxWithValue的 AICore 配置也仅注册了ascend910b与ascend910_93两个平台与文档保持一致。功能说明与计算公式接口功能返回self中指定维度的最大值及其索引位置。最大值保存到out中最大值的索引保存到indices中。若keepdim为false则不保留对应的轴若为true则保留指定轴的维度值为 1。语义细节当出现多个相等最大值时取首次出现的下标支持 NaN 传播语义与torch.max(dim)一致这两点在 README.md 中有明确说明。计算公式如下最大值计算$$ \text{out}i \max{\text{dim}}(\text{self}) $$索引计算$$ \text{indices}i \arg\max{\text{dim}}(\text{self}) $$从算子注册角度看aclnnMaxDim的底层 IR 即为ArgMaxWithValue算子见 op_api/aclnn_max_dim.h 中的 mermaid 计算路径注释该算子由 op_host/arg_max_with_value_def.cpp 完成定义输入x支持FLOAT16 / BF16 / FLOAT / INT16四种数据类型输出indice固定为INT32输出values与输入同类型数据格式均为 ND。函数原型两段式接口每个 CANN aclnn 算子都采用两段式接口设计aclnnMaxDim也不例外必须先调用aclnnMaxDimGetWorkspaceSize获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用aclnnMaxDim执行计算。aclnnStatus aclnnMaxDimGetWorkspaceSize( const aclTensor *self, int64_t dim, bool keepdim, aclTensor *out, aclTensor *indices, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnMaxDim( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)从实现看op_api/aclnn_max_dim.cpp第一段接口完成参数校验空指针、数据类型、shape 维度数、dim 范围构建执行器CREATE_EXECUTOR()拼装计算图见下文计算路径小节通过uniqueExecutor-GetWorkspaceSize()返回 workspace 大小并将执行器ReleaseTo(executor)释放给调用方。第二段接口则非常轻量调用CommonOpExecutorRun(workspace, workspaceSize, executor, stream)即可它负责真正在指定 stream 上执行此前封装好的算子计算流程。aclnnMaxDimGetWorkspaceSize 参数说明第一段接口的参数如下表所示参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待计算的目标张量。数据类型与 out 一致。FLOAT、FLOAT16、BFLOAT16、INT16ND1-8维√dimint64_t输入指定的维度。取值范围在 [-self.dim(), self.dim())INT64---keepdimbool输入reduce 轴的维度是否保留。-BOOL---outaclTensor*输出存储最大值结果。数据类型和 self 一致。keepdim 为 false 时输出维度为 self 维度减 1keepdim 为 true 时输出维度等于 self 维度。FLOAT、FLOAT16、BFLOAT16、INT16ND根据 keepdim 决定√indicesaclTensor*输出存储最大值的索引。-INT32ND与 out 相同√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小。-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程。-----补充说明规约轴长度必须大于 0空规约轴属于未定义输入。非连续输入是受支持的一等公民第一段接口内部会先对self做l0op::Contiguous归一化再送入后续计算测试用例 examples/test_aclnn_arg_max_with_value.cpp 中专门有一组{2, 3}视图挂载在{2, 4}storage 上的非连续 fp32 用例。dim支持负索引源码中dimMin -1 * shapeSizeaclnn_max_dim.cpp对 0 维标量输入源码做了特殊处理shapeSize 0时合法范围为[-1, 0]。返回值与错误码aclnnStatus为返回状态码具体参见 aclnn 返回码。第一段接口完成入参校验出现如下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 self、out 和 indices 是空指针时。ACLNN_ERR_PARAM_INVALID161002self、out、indices 数据类型不在支持的范围内时。ACLNN_ERR_PARAM_INVALID161002dim 超出输入 self 的维度范围时。这些校验在源码CheckParams中按顺序执行aclnn_max_dim.cpp先CheckNotNull空指针 → 161001再CheckDtypeValid平台与数据类型 → 161002、CheckShape维度数上限 8 → 161002、CheckDimdim 范围 → 161002。其中CheckDtypeValid还额外要求out与self数据类型一致、indices必须是INT32OP_CHECK_DTYPE_NOT_SAME/OP_CHECK_DTYPE_NOT_MATCH。另外需要注意当self为空张量self-IsEmpty()时第一段接口直接返回workspaceSize 0和空执行器不会继续构建计算图aclnn_max_dim.cpp。aclnnMaxDim 参数说明第二段接口的参数如下参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址。workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口 aclnnMaxDimGetWorkspaceSize 获取。executor输入op 执行器包含了算子计算流程。stream输入指定执行任务的 Stream。返回值aclnnStatus返回状态码具体参见 aclnn 返回码。约束说明确定性计算aclnnMaxDim默认确定性实现。这意味着在相同输入与环境下多次运行最大值索引结果可复现不会因核间并行导致索引不确定性。计算路径从 aclnn 接口到 AICore 核函数aclnnMaxDim虽只是一个对外 API但其背后是一条完整的计算链。头文件中的 mermaid 注释op_api/aclnn_max_dim.h给出了基本路径self ──(Contiguous)── ArgMaxWithValue ──(Cast)── ViewCopy ── Out/Indices ▲ └──── dim对应 aclnn_max_dim.cpp 的实现每个步骤的职责如下Contiguousl0op::Contiguous(self, executor)将可能非连续的输入张量统一为连续布局后续规约按连续内存处理ArgMaxWithValue核心规约算子返回(indices, values)二元组。在 op_api/argmax_with_value.cpp 中先通过INFER_SHAPE推导输出 shape再通过ADD_TO_LAUNCHER_LIST_AICORE注册 AICore 核函数调度若输入为 1 维且输出为 1 维会将输出重分配为 0 维标量 tensorkeepdimtrue时还会把归约轴维度显式置 1Cast对中间结果argmax_indices、argmax_out分别做类型转换转换为用户指定的indicesINT32与out数据类型ViewCopy将转换后的结果写回用户传入的out/indicestensor。在 shape 推导侧op_host/arg_max_with_value_infershape.cppInferShapeArgMaxWithValue通过ReduceDimsWithKeepDims/ReduceDimsWithoutKeepDims计算输出 shape并令indiceShape valueShape索引与最大值同 shape0 维标量输入时输出保持 0 维。核函数与 tiling 策略在算子内核侧ArgMaxWithValue采用编译期分派的调度策略op_kernel/arg_max_with_value.cpphost 侧SetTilingKey选定schModeCOPY / LAST / NLAST 三种模式编译期通过模板参数展开为独立的二进制运行期零分支。三种模式分别对应不同的数据布局规约轴在最后一维LAST、非最后一维NLAST以及拷贝型场景。Host 侧 tiling 逻辑op_host/arg_max_with_value_tiling.cpp会将规约轴展开为firstDim x axisSize x lastDim按三种 pattern 选择对应调度并用 32B 对齐的不均等输出切分ComputeOutputSplit把输出分到至多coreNum个核上保证不会出现两个核共享同一 32B 输出块的情况所有 UB tile 尺寸均从实际 shape 与 dtype 推导无魔法常数PIECE_AXIS4096、WORKSET_CAP8192、UB_MARGIN16384仅为保护性上限。调用示例以下示例代码摘自 aclnnMaxDim.md与仓库 examples/test_aclnn_arg_max_with_value.cpp 同源仅供参考具体编译和执行过程请参考编译与运行样例。它完整演示了初始化 → 构造 tensor → 两段式调用 → 同步取结果 → 释放资源的标准流程#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_max_dim.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } 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; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape {4, 2}; // 如果keepDim的值为true则indicesShape和outShape的shape为{1, 2} std::vectorint64_t indicesShape {2}; std::vectorint64_t outShape {2}; void* selfDeviceAddr nullptr; void* indicesDeviceAddr nullptr; void* outDeviceAddr nullptr; aclTensor* self nullptr; aclTensor* indices nullptr; aclTensor* out nullptr; std::vectorfloat selfHostData {0, 1, 2, 3, 4, 5, 6, 7}; std::vectorint32_t indicesHostData {0, 0}; std::vectorfloat outHostData {1, 1}; int64_t dim 0; bool keepDim false; // 创建self aclTensor ret CreateAclTensor(selfHostData, selfShape, selfDeviceAddr, aclDataType::ACL_FLOAT, self); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建indices aclTensor ret CreateAclTensor(indicesHostData, indicesShape, indicesDeviceAddr, aclDataType::ACL_INT32, indices); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建out aclTensor ret CreateAclTensor(outHostData, outShape, outDeviceAddr, aclDataType::ACL_FLOAT, out); CHECK_RET(ret ACL_SUCCESS, return ret); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnMaxDim第一段接口 ret aclnnMaxDimGetWorkspaceSize(self, dim, keepDim, out, indices, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMaxDimGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnMaxDim第二段接口 ret aclnnMaxDim(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnMaxDim failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果拷贝至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(outShape); std::vectorfloat resultData(size, 0); ret aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outDeviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(result[%ld] is: %f\n, i, resultData[i]); } // 6. 释放aclTensor和aclScalar需要根据具体API的接口定义修改 aclDestroyTensor(self); aclDestroyTensor(indices); aclDestroyTensor(out); // 7. 释放device 资源 aclrtFree(selfDeviceAddr); aclrtFree(indicesDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例要点解读两段式调用是硬性要求workspaceSize与executor由第一段接口产出第二段接口必须原样传入workspaceSize可能为 0例如空输入张量场景此时无需申请 workspace 内存示例中if (workspaceSize 0)的判断正是为此设计。shape 由 keepDim 决定示例输入selfShape {4, 2}、dim 0、keepDim false因此outShape indicesShape {2}若keepDim true两者应为{1, 2}对应源码中inputShape.SetDim(normalizedDim, 1)的逻辑。结果语义dim 0表示沿第 0 维长度 4求最大值self [[0,1],[2,3],[4,5],[6,7]]得到out [6, 7]、indices [3, 3]。若存在并列最大值返回首次出现的下标。测试与验证仓库为该算子提供了多层次的验证手段example 自检用例examples/test_aclnn_arg_max_with_value.cpp 内置 5 组用例覆盖 fp32 并列最大值ties、fp16 负轴 keepdim、bf16 沿第 0 维、int16 全类型、以及非连续 fp32 输入每组用例同时对values做字节级比较、对indices做逐元素比较并打印 PASS/FAIL 与最终 ALL PASS。ST 测试tests/st/aclnnMaxDim/executor_aclnnMaxDim.py 通过 ATK 框架将aclnn_max_dim与 CPU 参考实现torch.max(input, dim, keepdim)对拍function_aclnn_max_dim在 cpu 设备上以 torch 计算作为 golden用例定义位于 tests/st/aclnnMaxDim/atk_aclnnMaxDim.json。UT 测试tests/ut/op_host/test_arg_max_with_value_infershape.cpp 与 tests/ut/op_host/test_arg_max_with_value_tiling.cpp 分别针对 shape 推导与 tiling 计算做 host 侧单元验证。常见问题与注意事项平台不支持在非 A2/A3 产品如 Atlas 200I/500 A2上调用会返回ACLNN_ERR_PARAM_INVALID161002请先确认SocVersion属于ASCEND910B/ASCEND910_93。错误码 161001空指针self、out、indices任一为空指针即触发三个 tensor 必须全部先经aclCreateTensor创建成功。dim 越界dim取值范围是[-self.dim(), self.dim())含负索引越界返回 161002对 0 维标量输入合法范围被特殊放宽为[-1, 0]。规约轴长度为 0 属于未定义输入请确保规约轴至少包含一个元素。空输入张量当self-IsEmpty()时第一段接口直接返回workspaceSize 0此时不需要也无法分配 workspace但仍需正常走完第二段接口的调用流程。类型一致性out必须与self同数据类型indices固定为INT32文档表格中列出的数据类型即为全部支持范围FLOAT、FLOAT16、BFLOAT16、INT16传入其他类型会触发 161002。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考