CANN ops-math 算子 aclnnMaxDim 完全指南两段式接口调用 ArgMaxWithValue 获取张量最大值与索引【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读aclnnMaxDim是 CANN ops-math 数学算子库中ArgMaxWithValue算子对外暴露的 aclnn 计算接口用于在 NPU昇腾设备上返回 Tensor 指定维度的最大值及其索引位置语义上对应 PyTorch 的torch.max(input, dim, keepdim)。本文以 aclnnMaxDim.md 为主体结合 op_api、op_host、op_kernel 等目录下的仓库源码与测试用例完整讲解其产品支持情况、两段式接口的调用方式、每个参数的约束与取值、返回值错误码以及从 aclnn 入口到 kernel 执行的底层实现链路帮助你直接在昇腾设备上编写可编译、可运行的 C 调用样例。产品支持情况aclnnMaxDim由 ArgMaxWithValue 算子承载在不同昇腾硬件平台上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持不同平台之间还存在数据类型支持差异详见下文参数说明例如Atlas 推理系列产品/Atlas 训练系列产品不支持 BFLOAT16 与 INT32Atlas A2/A3系列不支持 INT32而Ascend 950PR/Ascend 950DT全量支持。这一差异在源码中体现为平台相关的数据类型校验列表见 aclnn_max_dim.cpp 中的DTYPE_SUPPORT_910_LIST、DTYPE_SUPPORT_GE910B_LIST、DTYPE_SUPPORT_REGBASE_LIST。功能说明aclnnMaxDim返回 Tensor 指定维度的最大值及其索引位置最大值保存到out中最大值的索引保存到indices中如果keepdim为false则不保留对应的轴如果keepdim为true则保留指定轴的维度值为 1。索引的含义是在指定维度上第一个出现最大值的位置。若同一维度上存在多个相同的最大值返回第一个出现的索引。两段式接口Two-Phase API架构CANN 算子库为每个算子提供两段式接口aclnnMaxDim也不例外。完整机制参见 two_phase_api.md。两段式调用流程如下第一段aclnnMaxDimGetWorkspaceSize完成入参校验空指针、数据类型、数据格式、shape、dim 取值范围并根据算子计算流程计算出执行所需的 workspace 大小同时生成一个包含完整算子计算流程的 op 执行器aclOpExecutor。第二段aclnnMaxDim传入第一段申请到的 workspace 内存与 executor在指定 Stream 上真正执行异步计算。这种设计把计算资源预算与实际执行解耦调用者可以根据第一段返回的workspaceSize精确申请 Device 侧内存避免盲目预留导致的内存浪费。函数原型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)两个接口的声明位于头文件 aclnn_max_dim.h其中以domain aclnn_math标识其归属于数学计算域并以 mermaid 图给出了 API 内部的计算路径aclnnMaxDimGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待计算最大值的输入张量-FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32ND1-8√dimint64_t输入指定的计算维度取值范围在 [-self.dim(), self.dim())----keepdimbool输入是否保留 reduce 轴的维度false不保留对应轴true保留指定轴且维度值为 1----outaclTensor*输出最大值输出张量数据类型与 self 一致keepdimfalse时输出维度为 self 维度减 1keepdimtrue时输出维度等于 self 维度FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32ND-√indicesaclTensor*输出最大值索引输出张量-INT32、INT64ND-√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台相关的数据类型差异Atlas 推理系列产品、Atlas 训练系列产品self和out不支持 BFLOAT16 和 INT32 数据类型indices不支持 BFLOAT16 数据类型。Atlas A2 训练/推理系列产品、Atlas A3 训练/推理系列产品self和out不支持 INT32 数据类型。对照 README 的说明README.md可以进一步确认各平台完整清单Atlas 推理/训练系列产品self/out 支持 FLOAT、FLOAT16、INT64、BOOLindices 支持 BOOL、FLOAT32、FLOAT16、INT8、INT16、UINT16、UINT8、INT32、INT64、UINT32、UINT64、DOUBLE、COMPLEX64、COMPLEX128。Atlas A2/A3 系列self/out 支持 FLOAT、FLOAT16、INT64、BOOL、BFLOAT16indices 在上述基础上额外支持 BFLOAT16。Ascend 950PR/Ascend 950DTself/out 支持 FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32。注indices的输出索引数据类型仅支持 INT32、INT64 两种具体实现中会根据平台选择内部索引类型后再 Cast 到用户指定的类型详见下文源码剖析。返回值与错误码两段接口均返回aclnnStatus状态码具体取值参见 aclnn_return_code.md。第一段接口完成入参校验出现如下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的self、out或indices是空指针时ACLNN_ERR_PARAM_INVALID161002self、out或indices的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self、out或indices的数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002dim超出输入self的维度范围这些校验逻辑在 aclnn_max_dim.cpp 中均有对应实现CheckNotNull检查空指针、CheckDtypeValid按平台选择数据类型白名单并检查self/out类型一致、CheckShape检查最大维度不超过MAX_DIM_LEN 8、CheckDim校验dim落在[-shapeSize, shapeSize-1]区间标量 shape 时区间为[-1, 0]。aclnnMaxDim 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnMaxDimGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream约束说明确定性计算aclnnMaxDim默认确定性实现。即相同的输入与运行环境多次执行结果一致便于调试与精度比对相关概念可参考 determinism_compute.md。输入self与输出out支持非连续 Tensor数据格式为ND输入维度为 1-8 维非连续 Tensor 说明参见 non_contiguous_tensor.md数据格式说明参见 data_format.md。dim支持负数索引从尾部数维度取值范围[-self.dim(), self.dim())。调用示例以下示例代码完整演示了通过 aclnn 两段式接口调用ArgMaxWithValue算子的流程具体编译和执行过程请参考 compile_and_run_sample.md。仓库内同款示例还可见 examples/test_aclnn_arg_max_with_value.cpp使用 RAII 智能指针管理资源的版本。#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::vectorfloat indicesHostData {1, 1}; 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资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(indicesDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例运行结果分析输入self的 shape 为{4, 2}数据为{0, 1, 2, 3, 4, 5, 6, 7}dim 0、keepDim false。沿第 0 维行方向对 4 个 2 元素行逐列求最大值得到out {6, 7}对应索引indices {3, 3}输出 shape 为{2}正是输入维度减 1。若将keepDim置为true则indicesShape与outShape均应为{1, 2}见示例注释。调用要点第 1、4、6、7 步是 acl 资源管理的固定写法aclInit/aclrtSetDevice/aclrtCreateStream初始化aclrtSynchronizeStream同步等待异步任务最后释放 Tensor、Device 内存、Stream 并aclFinalize。workspace 仅在workspaceSize 0时才需要aclrtMalloc申请示例中对out/indices的 host 初始值并不影响计算最终结果以 Device 侧计算为准。两段式接口必须成对使用先GetWorkspaceSize拿到 workspace 大小与 executor再执行第二段跳过第一段直接调用第二段会导致参数不完整。源码级原理剖析从 aclnn 入口到 NPU Kernel1. API 入口的完整计算管线aclnn_max_dim.cpp 中aclnnMaxDimGetWorkspaceSize的实现揭示了内部计算管线入参校验按CheckParams依次执行空指针、数据类型、shape、dim 校验对应文档中的 161001/161002 错误码。空输入短路若self-IsEmpty()直接返回workspaceSize 0不构造计算图。连续性转换l0op::Contiguous将非连续输入转换为连续 Tensor。数据类型适配为兼容不同平台对 BOOL以及非 910 平台的 INT64先执行l0op::Cast到 FLOAT再送入 ArgMaxWithValue 计算。核心计算l0op::ArgMaxWithValue一次调用同时产出最大值与索引两个输出索引内部先按平台选择RegBase 平台用用户指定的indices数据类型其余平台固定 INT32。形状校验CheckShapeAndScalarSame校验内部计算结果的 shape 与用户传入的out/indices一致。结果回填对两个输出分别做l0op::Cast转换到用户指定数据类型再通过l0op::ViewCopy将结果写入用户提供的out/indicesTensor。资源预算*workspaceSize uniqueExecutor-GetWorkspaceSize()汇总整条管线的 workspace 需求。aclnnMaxDim第二段接口则直接委托CommonOpExecutorRun(workspace, workspaceSize, executor, stream)执行异步计算aclnn_max_dim.cpp。2. l0op 层算子如何注册到 AICoreargmax_with_value.cpp 实现了l0op::ArgMaxWithValue通过OP_TYPE_REGISTER(ArgMaxWithValue)注册算子类型并调用INFER_SHAPE(ArgMaxWithValue, ...)触发 shape 推导特别处理了一维输入的情况当self与indices均为 1 维时输出被重分配为标量 shape{}对应整条维度求最大值的语义keepdimtrue时会将输出 view shape 中 reduce 轴对应的维度显式置为 1支持负数 dim 的换算最终通过ADD_TO_LAUNCHER_LIST_AICORE(ArgMaxWithValue, OP_INPUT(self), OP_OUTPUT(indices, out), OP_ATTR(dim, keepdim))将算子下发到 AICore 执行器。3. Host 侧OpDef 注册与 InferShape算子定义 arg_max_with_value_def.cpp 中注册了输入x支持 FLOAT16、BF16、FLOAT、INT64、INT32格式 ND输出indice与values其中indice支持 INT32/INT64属性dimension必选INT、keep_dims可选默认 false、indice_dtype可选默认 3即 INT32AICore 配置同时面向ascend950与ascend350两个平台开启DynamicRankSupportFlag与DynamicShapeSupportFlag即支持动态 shape/动态 rank。Shape 推导逻辑位于 arg_common_base_infershape.cppkeep_dimstrue时调用ReduceDimsWithKeepDims将 reduce 轴置 1keep_dimsfalse时调用ReduceDimsWithoutKeepDims直接删除该轴indice_shape与value_shape始终保持一致。这解释了文档中out 维度在 keepdimfalse 时为 self 维度减 1、keepdimtrue 时等于 self 维度的规则来源。4. Kernel 侧arch35 下的多种归约策略op_kernel/arch35 目录下按不同的数据排布与归约方式实现了多种 kernel 头文件包括arg_max_with_value_base.h公共基类arg_max_with_value_ar.h/arg_max_with_value_ara.h/arg_max_with_value_ra.h分别对应按行归约Row Reduce、沿轴归约Axis Reduce等不同访问模式arg_max_with_value_ara_cut_a_and_next_a.h、arg_max_with_value_ara_gather.h、arg_max_with_value_ar_gather.h面向大数据切分与 gather 场景的优化实现arg_max_with_value_group_reduce.h分组归约arg_max_with_value_copy.h结果拷贝。入口 arg_max_with_value_apt.cpp 根据运行时 shape 与 dim 选择合适的分发/归约路径Tiling 策略则在 op_host/arch35/arg_max_with_value_tiling.cpp 与arg_max_with_value_tiling_arch35.cpp中完成ascend350/ascend950各自的 tiling 配置arg_max_with_value_simplified_key.ini与arg_max_with_value_binary.json位于 op_host/config 目录。测试验证UT 与 ST 覆盖情况仓库为aclnnMaxDim/ArgMaxWithValue提供了完整的单测与场景测试可作为理解行为与自行验证的参考InferShape 单元测试test_arg_max_with_value_infershape.cpp 覆盖了{4,2}输入在dim0时keepdimsfalse输出{2}、keepdimstrue输出{1,2}以及{2,3,4}在dim1、keepdimsfalse输出{2,4}等场景与文档中的 shape 规则一一对应。Tiling 单元测试tests/ut/op_host/arch35/test_arg_max_with_value_tiling.cpp。场景测试STatk_aclnnMaxDim.json 以torch.max为基准对齐name: torch.max覆盖 fp32 输入从 2 维到 5 维如[912,16,64]、[4,7,16,576,224]、[10,1,4,288,112]等、不同dim/keepdim组合以及 ND/NCL/NCHW/NCDHW 格式精度标准为high_precision配套的执行脚本为 executor_aclnnMaxDim.py。OP API 单元测试test_aclnn_arg_max_with_value.cpp 直接以两段式接口驱动计算是上文调用示例的测试级等价实现。总结aclnnMaxDim是 CANN ops-math 中实现按维度求最大值及其索引的标准 aclnn 接口核心要点可归纳为两段式调用先aclnnMaxDimGetWorkspaceSize校验参数并获取 workspace 大小与 executor再aclnnMaxDim在指定 Stream 上执行参数约束清晰输入支持 FLOAT/FLOAT16/INT64/BOOL/BFLOAT16/INT32视平台而定、ND 格式、1-8 维、可非连续dim支持负数且范围[-self.dim(), self.dim())indices仅 INT32/INT64平台差异明确Atlas 910 系、A2/A3 系与 Ascend 950 系在数据类型支持上存在差异编码前应确认目标平台确定性计算默认确定性实现便于精度对齐与问题复现源码可循从 aclnn_max_dim.cpp 的 Contiguous → Cast → ArgMaxWithValue → Cast → ViewCopy 管线到 arg_common_base_infershape.cpp 的 shape 推导再到 arch35 kernel 的多策略归约整条链路在仓库中均有完整实现与测试支撑可直接作为二次开发或算子移植的参考样板。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 CANN ops-math 算子 aclnnMaxDim 完全指南：两段式接口调用 ArgMaxWithValue 获取张量最大值与索引
CANN ops-math 算子 aclnnMaxDim 完全指南两段式接口调用 ArgMaxWithValue 获取张量最大值与索引【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读aclnnMaxDim是 CANN ops-math 数学算子库中ArgMaxWithValue算子对外暴露的 aclnn 计算接口用于在 NPU昇腾设备上返回 Tensor 指定维度的最大值及其索引位置语义上对应 PyTorch 的torch.max(input, dim, keepdim)。本文以 aclnnMaxDim.md 为主体结合 op_api、op_host、op_kernel 等目录下的仓库源码与测试用例完整讲解其产品支持情况、两段式接口的调用方式、每个参数的约束与取值、返回值错误码以及从 aclnn 入口到 kernel 执行的底层实现链路帮助你直接在昇腾设备上编写可编译、可运行的 C 调用样例。产品支持情况aclnnMaxDim由 ArgMaxWithValue 算子承载在不同昇腾硬件平台上的支持情况如下产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I/500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品支持不同平台之间还存在数据类型支持差异详见下文参数说明例如Atlas 推理系列产品/Atlas 训练系列产品不支持 BFLOAT16 与 INT32Atlas A2/A3系列不支持 INT32而Ascend 950PR/Ascend 950DT全量支持。这一差异在源码中体现为平台相关的数据类型校验列表见 aclnn_max_dim.cpp 中的DTYPE_SUPPORT_910_LIST、DTYPE_SUPPORT_GE910B_LIST、DTYPE_SUPPORT_REGBASE_LIST。功能说明aclnnMaxDim返回 Tensor 指定维度的最大值及其索引位置最大值保存到out中最大值的索引保存到indices中如果keepdim为false则不保留对应的轴如果keepdim为true则保留指定轴的维度值为 1。索引的含义是在指定维度上第一个出现最大值的位置。若同一维度上存在多个相同的最大值返回第一个出现的索引。两段式接口Two-Phase API架构CANN 算子库为每个算子提供两段式接口aclnnMaxDim也不例外。完整机制参见 two_phase_api.md。两段式调用流程如下第一段aclnnMaxDimGetWorkspaceSize完成入参校验空指针、数据类型、数据格式、shape、dim 取值范围并根据算子计算流程计算出执行所需的 workspace 大小同时生成一个包含完整算子计算流程的 op 执行器aclOpExecutor。第二段aclnnMaxDim传入第一段申请到的 workspace 内存与 executor在指定 Stream 上真正执行异步计算。这种设计把计算资源预算与实际执行解耦调用者可以根据第一段返回的workspaceSize精确申请 Device 侧内存避免盲目预留导致的内存浪费。函数原型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)两个接口的声明位于头文件 aclnn_max_dim.h其中以domain aclnn_math标识其归属于数学计算域并以 mermaid 图给出了 API 内部的计算路径aclnnMaxDimGetWorkspaceSize 参数说明参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorselfaclTensor*输入待计算最大值的输入张量-FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32ND1-8√dimint64_t输入指定的计算维度取值范围在 [-self.dim(), self.dim())----keepdimbool输入是否保留 reduce 轴的维度false不保留对应轴true保留指定轴且维度值为 1----outaclTensor*输出最大值输出张量数据类型与 self 一致keepdimfalse时输出维度为 self 维度减 1keepdimtrue时输出维度等于 self 维度FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32ND-√indicesaclTensor*输出最大值索引输出张量-INT32、INT64ND-√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----平台相关的数据类型差异Atlas 推理系列产品、Atlas 训练系列产品self和out不支持 BFLOAT16 和 INT32 数据类型indices不支持 BFLOAT16 数据类型。Atlas A2 训练/推理系列产品、Atlas A3 训练/推理系列产品self和out不支持 INT32 数据类型。对照 README 的说明README.md可以进一步确认各平台完整清单Atlas 推理/训练系列产品self/out 支持 FLOAT、FLOAT16、INT64、BOOLindices 支持 BOOL、FLOAT32、FLOAT16、INT8、INT16、UINT16、UINT8、INT32、INT64、UINT32、UINT64、DOUBLE、COMPLEX64、COMPLEX128。Atlas A2/A3 系列self/out 支持 FLOAT、FLOAT16、INT64、BOOL、BFLOAT16indices 在上述基础上额外支持 BFLOAT16。Ascend 950PR/Ascend 950DTself/out 支持 FLOAT、FLOAT16、INT64、BOOL、BFLOAT16、INT32。注indices的输出索引数据类型仅支持 INT32、INT64 两种具体实现中会根据平台选择内部索引类型后再 Cast 到用户指定的类型详见下文源码剖析。返回值与错误码两段接口均返回aclnnStatus状态码具体取值参见 aclnn_return_code.md。第一段接口完成入参校验出现如下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的self、out或indices是空指针时ACLNN_ERR_PARAM_INVALID161002self、out或indices的数据类型不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002self、out或indices的数据格式不在支持的范围之内ACLNN_ERR_PARAM_INVALID161002dim超出输入self的维度范围这些校验逻辑在 aclnn_max_dim.cpp 中均有对应实现CheckNotNull检查空指针、CheckDtypeValid按平台选择数据类型白名单并检查self/out类型一致、CheckShape检查最大维度不超过MAX_DIM_LEN 8、CheckDim校验dim落在[-shapeSize, shapeSize-1]区间标量 shape 时区间为[-1, 0]。aclnnMaxDim 参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnMaxDimGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream约束说明确定性计算aclnnMaxDim默认确定性实现。即相同的输入与运行环境多次执行结果一致便于调试与精度比对相关概念可参考 determinism_compute.md。输入self与输出out支持非连续 Tensor数据格式为ND输入维度为 1-8 维非连续 Tensor 说明参见 non_contiguous_tensor.md数据格式说明参见 data_format.md。dim支持负数索引从尾部数维度取值范围[-self.dim(), self.dim())。调用示例以下示例代码完整演示了通过 aclnn 两段式接口调用ArgMaxWithValue算子的流程具体编译和执行过程请参考 compile_and_run_sample.md。仓库内同款示例还可见 examples/test_aclnn_arg_max_with_value.cpp使用 RAII 智能指针管理资源的版本。#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::vectorfloat indicesHostData {1, 1}; 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资源需要根据具体API的接口定义修改 aclrtFree(selfDeviceAddr); aclrtFree(indicesDeviceAddr); aclrtFree(outDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }示例运行结果分析输入self的 shape 为{4, 2}数据为{0, 1, 2, 3, 4, 5, 6, 7}dim 0、keepDim false。沿第 0 维行方向对 4 个 2 元素行逐列求最大值得到out {6, 7}对应索引indices {3, 3}输出 shape 为{2}正是输入维度减 1。若将keepDim置为true则indicesShape与outShape均应为{1, 2}见示例注释。调用要点第 1、4、6、7 步是 acl 资源管理的固定写法aclInit/aclrtSetDevice/aclrtCreateStream初始化aclrtSynchronizeStream同步等待异步任务最后释放 Tensor、Device 内存、Stream 并aclFinalize。workspace 仅在workspaceSize 0时才需要aclrtMalloc申请示例中对out/indices的 host 初始值并不影响计算最终结果以 Device 侧计算为准。两段式接口必须成对使用先GetWorkspaceSize拿到 workspace 大小与 executor再执行第二段跳过第一段直接调用第二段会导致参数不完整。源码级原理剖析从 aclnn 入口到 NPU Kernel1. API 入口的完整计算管线aclnn_max_dim.cpp 中aclnnMaxDimGetWorkspaceSize的实现揭示了内部计算管线入参校验按CheckParams依次执行空指针、数据类型、shape、dim 校验对应文档中的 161001/161002 错误码。空输入短路若self-IsEmpty()直接返回workspaceSize 0不构造计算图。连续性转换l0op::Contiguous将非连续输入转换为连续 Tensor。数据类型适配为兼容不同平台对 BOOL以及非 910 平台的 INT64先执行l0op::Cast到 FLOAT再送入 ArgMaxWithValue 计算。核心计算l0op::ArgMaxWithValue一次调用同时产出最大值与索引两个输出索引内部先按平台选择RegBase 平台用用户指定的indices数据类型其余平台固定 INT32。形状校验CheckShapeAndScalarSame校验内部计算结果的 shape 与用户传入的out/indices一致。结果回填对两个输出分别做l0op::Cast转换到用户指定数据类型再通过l0op::ViewCopy将结果写入用户提供的out/indicesTensor。资源预算*workspaceSize uniqueExecutor-GetWorkspaceSize()汇总整条管线的 workspace 需求。aclnnMaxDim第二段接口则直接委托CommonOpExecutorRun(workspace, workspaceSize, executor, stream)执行异步计算aclnn_max_dim.cpp。2. l0op 层算子如何注册到 AICoreargmax_with_value.cpp 实现了l0op::ArgMaxWithValue通过OP_TYPE_REGISTER(ArgMaxWithValue)注册算子类型并调用INFER_SHAPE(ArgMaxWithValue, ...)触发 shape 推导特别处理了一维输入的情况当self与indices均为 1 维时输出被重分配为标量 shape{}对应整条维度求最大值的语义keepdimtrue时会将输出 view shape 中 reduce 轴对应的维度显式置为 1支持负数 dim 的换算最终通过ADD_TO_LAUNCHER_LIST_AICORE(ArgMaxWithValue, OP_INPUT(self), OP_OUTPUT(indices, out), OP_ATTR(dim, keepdim))将算子下发到 AICore 执行器。3. Host 侧OpDef 注册与 InferShape算子定义 arg_max_with_value_def.cpp 中注册了输入x支持 FLOAT16、BF16、FLOAT、INT64、INT32格式 ND输出indice与values其中indice支持 INT32/INT64属性dimension必选INT、keep_dims可选默认 false、indice_dtype可选默认 3即 INT32AICore 配置同时面向ascend950与ascend350两个平台开启DynamicRankSupportFlag与DynamicShapeSupportFlag即支持动态 shape/动态 rank。Shape 推导逻辑位于 arg_common_base_infershape.cppkeep_dimstrue时调用ReduceDimsWithKeepDims将 reduce 轴置 1keep_dimsfalse时调用ReduceDimsWithoutKeepDims直接删除该轴indice_shape与value_shape始终保持一致。这解释了文档中out 维度在 keepdimfalse 时为 self 维度减 1、keepdimtrue 时等于 self 维度的规则来源。4. Kernel 侧arch35 下的多种归约策略op_kernel/arch35 目录下按不同的数据排布与归约方式实现了多种 kernel 头文件包括arg_max_with_value_base.h公共基类arg_max_with_value_ar.h/arg_max_with_value_ara.h/arg_max_with_value_ra.h分别对应按行归约Row Reduce、沿轴归约Axis Reduce等不同访问模式arg_max_with_value_ara_cut_a_and_next_a.h、arg_max_with_value_ara_gather.h、arg_max_with_value_ar_gather.h面向大数据切分与 gather 场景的优化实现arg_max_with_value_group_reduce.h分组归约arg_max_with_value_copy.h结果拷贝。入口 arg_max_with_value_apt.cpp 根据运行时 shape 与 dim 选择合适的分发/归约路径Tiling 策略则在 op_host/arch35/arg_max_with_value_tiling.cpp 与arg_max_with_value_tiling_arch35.cpp中完成ascend350/ascend950各自的 tiling 配置arg_max_with_value_simplified_key.ini与arg_max_with_value_binary.json位于 op_host/config 目录。测试验证UT 与 ST 覆盖情况仓库为aclnnMaxDim/ArgMaxWithValue提供了完整的单测与场景测试可作为理解行为与自行验证的参考InferShape 单元测试test_arg_max_with_value_infershape.cpp 覆盖了{4,2}输入在dim0时keepdimsfalse输出{2}、keepdimstrue输出{1,2}以及{2,3,4}在dim1、keepdimsfalse输出{2,4}等场景与文档中的 shape 规则一一对应。Tiling 单元测试tests/ut/op_host/arch35/test_arg_max_with_value_tiling.cpp。场景测试STatk_aclnnMaxDim.json 以torch.max为基准对齐name: torch.max覆盖 fp32 输入从 2 维到 5 维如[912,16,64]、[4,7,16,576,224]、[10,1,4,288,112]等、不同dim/keepdim组合以及 ND/NCL/NCHW/NCDHW 格式精度标准为high_precision配套的执行脚本为 executor_aclnnMaxDim.py。OP API 单元测试test_aclnn_arg_max_with_value.cpp 直接以两段式接口驱动计算是上文调用示例的测试级等价实现。总结aclnnMaxDim是 CANN ops-math 中实现按维度求最大值及其索引的标准 aclnn 接口核心要点可归纳为两段式调用先aclnnMaxDimGetWorkspaceSize校验参数并获取 workspace 大小与 executor再aclnnMaxDim在指定 Stream 上执行参数约束清晰输入支持 FLOAT/FLOAT16/INT64/BOOL/BFLOAT16/INT32视平台而定、ND 格式、1-8 维、可非连续dim支持负数且范围[-self.dim(), self.dim())indices仅 INT32/INT64平台差异明确Atlas 910 系、A2/A3 系与 Ascend 950 系在数据类型支持上存在差异编码前应确认目标平台确定性计算默认确定性实现便于精度对齐与问题复现源码可循从 aclnn_max_dim.cpp 的 Contiguous → Cast → ArgMaxWithValue → Cast → ViewCopy 管线到 arg_common_base_infershape.cpp 的 shape 推导再到 arch35 kernel 的多策略归约整条链路在仓库中均有完整实现与测试支撑可直接作为二次开发或算子移植的参考样板。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考