CANN ops-math 算子 Pdist基于 Ascend NPU 的 p-范数成对距离计算详解【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN ops-math 仓库中experimental/math/pdist目录下的 Pdist 算子展开系统讲解其在 NPU 上实现输入二维 tensor 各行之间 p-范数成对距离计算的数学原理、算子原型、两段式 ACLNN 调用方式以及从 Host 侧算子定义/Tiling 到 Kernel 侧向量计算实现的完整源码脉络。读完本文你将掌握在 Atlas A2/A3 系列产品上通过aclnnPdist与aclnnPdistForward接口调用该算子并能理解其参数校验、内存预算、多核切分与确定性实现背后的工程细节。算子功能与数学原理Pdist 算子计算输入二维 tensor 各行之间的 p-范数成对距离功能上等价于 PyTorch 的torch.nn.functional.pdist。设输入 x 的形状为 (N, M)则输出为一维 tensor形状为 (N*(N-1)/2,)其中第 i 行与第 j 行i j之间的距离按上三角行优先顺序依次排列。计算公式按 p 的取值分为三种情况见 README.md 与接口文档 aclnnPdist.md0 p ∞dist(i, j) (Σ_k |x_ik − x_jk|^p)^(1/p)即闵可夫斯基距离Minkowski distance。当 p1 时退化为曼哈顿距离p2 时退化为欧氏距离p 0dist(i, j) Σ_k (x_ik ≠ x_jk)即汉明距离Hamming distance统计两行对应位置不相等的元素个数p ∞dist(i, j) max_k |x_ik − x_jk|即切比雪夫距离Chebyshev distance取两行逐元素差绝对值的最大值。从 Kernel 源码 pdist.h 可以印证这一分支设计tiling 阶段会依据 p 的取值生成 tilingKeyp0 时为 00p∞ 时为 1p∞ 时为 2Kernel 侧ComputeChunkDistance按 tilingKey 分派到三条计算路径——汉明距离路径使用CompareScalar与 0 比较判不等SelectReduceSum闵可夫斯基距离路径使用AbsMulp2或Ln/Muls/Exp其他 pReduceSum切比雪夫距离路径使用AbsReduceMax。最后ApplyInvP对求和结果施加 1/p 次幂p1、p2 分别有快速分支其余 p 走Ln/Muls/Exp。产品支持情况Pdist 算子当前支持的产品如下产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√对应到算子定义中的硬件配置见 pdist_def.cppAICore().AddConfig(ascend910b)与AddConfig(ascend910_93)分别对应 A2 与 A3 的 AICore 配置项说明该算子是在这两个平台的昇腾 AI Core 上编译执行的。算子原型设计Pdist 算子的 IR 原型定义如下见 pdist_def.cpp参数名类别描述数据类型数据格式x输入张量二维输入张量形状 (N, M)N ≥ 2M ≥ 1FLOAT16、FLOAT32NDp属性距离参数p ≥ 0含 inf可选默认 2.0FLOAT-y输出张量一维输出张量形状 (N*(N-1)/2,)与 x 一致ND在 Atlas A2/A3 系列产品上数据类型支持 FLOAT16、FLOAT32。源码中this-Attr(p).AttrType(OPTIONAL).Float(2.0f)明确了 p 是可选属性且默认值为 2.0Tiling 阶段解析属性时同样以 2.0f 作为缺省值见 pdist_tiling.cpp。约束说明使用 Pdist 算子需满足以下约束输入 x 必须为二维 tensor且 dim(0) ≥ 2即 N ≥ 2输出 y 的数据类型必须与 x 一致p 的取值范围为 [0, ∞]即 p ≥ 0且支持正无穷N 上限为 65535PDIST_MAX_SUPPORTED_ROWS见 pdist_constants.h。computeNum 使用 uint64_t 计算N 过大时输出规模 N*(N-1)/2 会超出实际内存限制因此 Tiling 阶段在 pdist_tiling.cpp 中对 rows 超限直接报错。这些约束在多层均有对应校验Host 侧 InferShape 检查输入必须是 2 维且 N ≥ 2见 pdist_infershape.cppTiling 阶段在ParseInputParams中校验 rows ≥ 2、cols ≥ 1 及 rows ≤ 65535ACLNN 接口层在CheckParamsLogic中校验维度、输出 shape 与 p ≥ 0含 NaN 拒绝见 aclnn_pdist.cpp。两段式调用接口每个算子采用两段式接口必须先调用 GetWorkspaceSize 接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用执行接口完成计算。aclnnPdistp 以 float 传入aclnnStatus aclnnPdistGetWorkspaceSize( const aclTensor *self, double p, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnPdist( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnPdistGetWorkspaceSize的参数说明参数名输入/输出描述数据类型数据格式维度(shape)selfconst aclTensor*输入输入二维 tensor对应公式中 xFLOAT、FLOAT16ND2 维 (N, M)N≥2M≥1pdouble输入距离参数对应公式中 pp≥0含 inf---outaclTensor*输出输出一维 tensor对应公式中 distFLOAT、FLOAT16ND1 维 (N*(N-1)/2,)workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小---executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程---第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、out 存在空指针ACLNN_ERR_PARAM_INVALID161002self 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002self 不是二维 tensor 或 N 2ACLNN_ERR_PARAM_INVALID161002p 0ACLNN_ERR_PARAM_INVALID161002out 的 shape 与 N*(N-1)/2 不匹配aclnnPdist的参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 StreamaclnnPdistForwardp 以 aclScalar 传入Forward 接口功能与 aclnnPdist 完全一致唯一区别是 p 参数以aclScalar指针传入通过aclCreateScalar创建便于上层框架以统一的标量对象传递aclnnStatus aclnnPdistForwardGetWorkspaceSize( const aclTensor *self, const aclScalar *p, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnPdistForward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)该接口的第一段校验与 aclnnPdist 相同只是空指针检查范围扩展到 self、p、out 三者。两个接口均为默认确定性实现同一输入多次调用结果一致。调用示例以 4×3 的 float32 输入矩阵为例N4 时输出长度为 N*(N-1)/2 6。完整可运行代码见 test_aclnn_pdist.cpp 与 test_aclnn_pdist_forward.cpp核心调用流程如下省略 Init、CreateAclTensor 等通用辅助函数#include iostream #include vector #include acl/acl.h #include aclnn_pdist.h int main() { int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 构造输入: 4x3 矩阵 std::vectorint64_t inputShape {4, 3}; std::vectorfloat inputData {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; aclTensor* inputTensor nullptr; void* inputDeviceAddr nullptr; CreateAclTensor(inputData, inputShape, inputDeviceAddr, ACL_FLOAT, inputTensor); // 构造输出: N*(N-1)/2 6 std::vectorint64_t outputShape {6}; std::vectorfloat outputData(6, 0); aclTensor* outputTensor nullptr; void* outputDeviceAddr nullptr; CreateAclTensor(outputData, outputShape, outputDeviceAddr, ACL_FLOAT, outputTensor); // 调用第一段接口 float p 2.0f; uint64_t workspaceSize 0; aclOpExecutor* executor; aclnnPdistGetWorkspaceSize(inputTensor, p, outputTensor, workspaceSize, executor); void* workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 调用第二段接口 aclnnPdist(workspaceAddr, workspaceSize, executor, stream); aclrtSynchronizeStream(stream); // 释放资源 aclDestroyTensor(inputTensor); aclDestroyTensor(outputTensor); aclrtFree(inputDeviceAddr); aclrtFree(outputDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }使用aclnnPdistForward时仅需将 p 的构造方式替换为 aclScalarfloat pValue 2.0f; aclScalar* pScalar aclCreateScalar(pValue, ACL_FLOAT); aclnnPdistForwardGetWorkspaceSize(inputTensor, pScalar, outputTensor, workspaceSize, executor); // ... aclDestroyScalar(pScalar); // 释放 aclScalar示例中的CreateAclTensor辅助函数展示了标准张量构造方式先在 Device 侧aclrtMalloc并aclrtMemcpy拷贝数据再通过aclCreateTensor以 ND 格式、连续 strides 创建 aclTensor 描述PrintOutResult通过aclrtMemcpyDEVICE_TO_HOST回读结果。源码实现深度解析Host 侧算子定义、InferShape 与 Tiling算子定义pdist_def.cpp通过OpDef注册 Pdist声明输入 x、输出 y均支持 DT_FLOAT/DT_FLOAT16ND 格式AutoContiguous以及可选属性 p默认 2.0f并挂接 ascend910b、ascend910_93 两个 AICore 编译配置。InferShapepdist_infershape.cpp校验输入为 2 维且 N ≥ 2然后直接推导输出一维 shape 为N * (N - 1) / 2这是整个算子的核心几何事实——N 行两两组合的成对数量。Tilingpdist_tiling.cpp完成四个关键决策Workspace 申请GetWorkspaceSize为每个核申请sizeof(float)的 workspace见第 37-43 行用于 Kernel 内跨 chunk 累积距离值UB 内存预算ComputeUbBudget根据平台 UB 大小、reduce 临时 bufferGetReduceSumMaxMinTmpSize/GetReduceMaxMaxMinTmpSize计算、求和 tensor 保留区等预留字节数反推出每轮循环可处理的列数ubTensorEachLoop并按 8 元素对齐FP16 输入还需额外预留 fp16 转换队列多核切分ComputeCoreSplit将 computeNum N*(N-1)/2 个距离任务按 8 个一组分成 block均匀分配到各核剩余不足一个 block 的部分由尾核lastNumsBlocks、lastNumsNoneFullBlock承接并通过context-SetBlockDim(usedCores)设置 block 维度TilingKey 生成FillTilingData依据 p 的取值0 / 常规 / inf生成 0/1/2 三种 tilingKey并通过ASCENDC_TPL_SEL_PARAM依据输入数据类型完成 Kernel 模板实例化选择。Tiling 数据通过PdistTilingData结构体pdist_tiling_data.h在 Host 与 Kernel 之间传递包含 rows、cols、pValue、computeNum、ubTensorEachLoop、coreNumVar、tilingKey、reduceBufSize 及核间 block 分配信息。Kernel 侧NPU 向量计算Kernel 入口pdist.cpp通过GET_TILING_DATA_WITH_STRUCT解析 tiling 数据实例化NsPdist::PdistD_T_X并执行InitProcess。Process按核 ID 定位本核负责的输出区间含尾块处理对每个 block 内的距离任务逐条执行ProcessBlock。每条距离的完整计算链路为行号反解GetIJFromIndex根据输出线性下标 idx 反解出行对 (rowI, rowJ)。由于输出按下三角展开的计数规律组织源码采用整数平方根牛顿迭代 累加和校正的方式避免了逐行遍历的开销分块加载与差分LoadDiff按列分块chunk 大小不超过 ubTensorEachLoop从 Global Memory 加载两行数据。FP16 输入先经DataCopyPad拷入 fp16 队列再Cast为 float 参与计算随后Sub得到逐元素差chunk 距离计算ComputeChunkDistance按 tilingKey 分派到汉明/闵可夫斯基/切比雪夫三条向量计算路径详见前文数学原理部分返回该 chunk 的局部累加值chunk 间归并常规范数与汉明距离对 chunk 结果做累加切比雪夫距离取各 chunk 最大值ProcessBlock中accum的两种更新方式施加 1/p 次幂并写回ApplyInvP仅对常规范数分支生效并对结果做0 才开方/幂运算否则置 0的保护处理避免对 0 求 log随后WriteOutput将结果按 8 元素块写回 Global MemoryFP16 输出时先CastCAST_RINT 舍入再拷贝。值得留意的是输入按列分块循环while (remaining 0)的设计使得该算子能够处理任意列宽 M而 UB 占用始终保持恒定这正是ubTensorEachLoop由 UB 预算反推的意义所在。ACLNN 接口层参数校验与图构建接口层实现见 aclnn_pdist.cpp。第一段接口aclnnPdistGetWorkspaceSize依次完成CheckNotNull空指针检查对应错误码 161001CheckDtypeValid校验 self/out 的数据类型均在 {DT_FLOAT16, DT_FLOAT} 支持列表内且两者一致CheckParamsLogic校验维度self 必须 2 维、out 必须 1 维、输出 shape 与 N*(N-1)/2 的匹配关系以及 p ≥ 0 且非 NaN对应错误码 161002构建计算图当 N ≤ 1 时直接返回空 workspace输出 0 维空 tensor当 M 0第二维为 0时走 L0Fill算子填充 0常规情况则先l0op::Contiguous保证输入连续再执行l0op::Pdist最后以l0op::ViewCopy将结果写入用户 out通过uniqueExecutor-GetWorkspaceSize()返回 workspace 大小并将执行器ReleaseTo(executor)交给用户。第二段接口aclnnPdist则直接调用框架能力CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成异步执行。目录结构与测试保障Pdist 算子的完整目录结构如下见 README.mdpdist/ ├── README.md ├── op_host/ # Host 侧算子定义、Tiling、InferShape │ ├── pdist_def.cpp │ ├── pdist_infershape.cpp │ └── pdist_tiling.cpp ├── op_kernel/ # Kernel 侧NPU 计算逻辑 │ ├── pdist.cpp │ ├── pdist.h │ ├── pdist_constants.h │ ├── pdist_tiling_data.h │ └── pdist_tiling_key.h ├── op_api/ # ACLNN 接口层 │ ├── aclnn_pdist.h │ ├── aclnn_pdist.cpp │ ├── aclnn_pdist_forward.h │ └── aclnn_pdist_forward.cpp ├── docs/ # 接口文档 │ ├── aclnnPdist.md │ └── aclnnPdistForward.md ├── examples/ # 调用示例 │ ├── test_aclnn_pdist.cpp │ └── test_aclnn_pdist_forward.cpp └── tests/ ├── ut/ # 单元测试 │ ├── op_host/ # Tiling InferShape UT (14 cases) │ ├── op_api/ # ACLNN 接口 UT (9 cases) │ └── op_kernel/ # Kernel CPU 模拟器 UT ├── st/ # 系统测试 (97 cases, Real NPU) │ ├── test_aclnn_pdist.cpp │ └── run.sh └── reports/ └── iter3-acceptance-report.md仓库中实际存在的测试代码进一步印证了该算子具备完善的测试覆盖Host 侧 UTtest_pdist_infershape.cpp 验证 InferShape 的输出 shape 推导与非法输入拒绝逻辑test_pdist_tiling.cpp 验证 Tiling 参数计算接口层 UTtest_aclnn_pdist.cpp 覆盖 ACLNN 接口的参数校验与调用路径Kernel 层 UTtest_pdist.cpp 在 CPU 模拟器上运行 Kernel 逻辑配套 gen_data.py 与 compare_data.py 生成测试数据并与期望结果比对系统测试test_aclnn_pdist.cpp 与 run.sh 在真实 NPU 上执行 97 个用例。贡献说明Pdist 算子由个人开发者 xiaoy2459 于 2026/6/3 贡献至开源仓贡献内容为Pdist 算子适配开源仓详情见 README.md 中的贡献说明表格。该算子的落地也展示了 CANN ops-math 开源社区中接口文档docs/ 示例examples/ 分层测试tests/ 源码op_host/op_kernel/op_api/标准算子交付模板的完整实践。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
CANN ops-math 算子 Pdist：基于 Ascend NPU 的 p-范数成对距离计算详解
CANN ops-math 算子 Pdist基于 Ascend NPU 的 p-范数成对距离计算详解【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math导读本文围绕 CANN ops-math 仓库中experimental/math/pdist目录下的 Pdist 算子展开系统讲解其在 NPU 上实现输入二维 tensor 各行之间 p-范数成对距离计算的数学原理、算子原型、两段式 ACLNN 调用方式以及从 Host 侧算子定义/Tiling 到 Kernel 侧向量计算实现的完整源码脉络。读完本文你将掌握在 Atlas A2/A3 系列产品上通过aclnnPdist与aclnnPdistForward接口调用该算子并能理解其参数校验、内存预算、多核切分与确定性实现背后的工程细节。算子功能与数学原理Pdist 算子计算输入二维 tensor 各行之间的 p-范数成对距离功能上等价于 PyTorch 的torch.nn.functional.pdist。设输入 x 的形状为 (N, M)则输出为一维 tensor形状为 (N*(N-1)/2,)其中第 i 行与第 j 行i j之间的距离按上三角行优先顺序依次排列。计算公式按 p 的取值分为三种情况见 README.md 与接口文档 aclnnPdist.md0 p ∞dist(i, j) (Σ_k |x_ik − x_jk|^p)^(1/p)即闵可夫斯基距离Minkowski distance。当 p1 时退化为曼哈顿距离p2 时退化为欧氏距离p 0dist(i, j) Σ_k (x_ik ≠ x_jk)即汉明距离Hamming distance统计两行对应位置不相等的元素个数p ∞dist(i, j) max_k |x_ik − x_jk|即切比雪夫距离Chebyshev distance取两行逐元素差绝对值的最大值。从 Kernel 源码 pdist.h 可以印证这一分支设计tiling 阶段会依据 p 的取值生成 tilingKeyp0 时为 00p∞ 时为 1p∞ 时为 2Kernel 侧ComputeChunkDistance按 tilingKey 分派到三条计算路径——汉明距离路径使用CompareScalar与 0 比较判不等SelectReduceSum闵可夫斯基距离路径使用AbsMulp2或Ln/Muls/Exp其他 pReduceSum切比雪夫距离路径使用AbsReduceMax。最后ApplyInvP对求和结果施加 1/p 次幂p1、p2 分别有快速分支其余 p 走Ln/Muls/Exp。产品支持情况Pdist 算子当前支持的产品如下产品是否支持Atlas A2 训练系列产品/Atlas A2 推理系列产品√Atlas A3 训练系列产品/Atlas A3 推理系列产品√对应到算子定义中的硬件配置见 pdist_def.cppAICore().AddConfig(ascend910b)与AddConfig(ascend910_93)分别对应 A2 与 A3 的 AICore 配置项说明该算子是在这两个平台的昇腾 AI Core 上编译执行的。算子原型设计Pdist 算子的 IR 原型定义如下见 pdist_def.cpp参数名类别描述数据类型数据格式x输入张量二维输入张量形状 (N, M)N ≥ 2M ≥ 1FLOAT16、FLOAT32NDp属性距离参数p ≥ 0含 inf可选默认 2.0FLOAT-y输出张量一维输出张量形状 (N*(N-1)/2,)与 x 一致ND在 Atlas A2/A3 系列产品上数据类型支持 FLOAT16、FLOAT32。源码中this-Attr(p).AttrType(OPTIONAL).Float(2.0f)明确了 p 是可选属性且默认值为 2.0Tiling 阶段解析属性时同样以 2.0f 作为缺省值见 pdist_tiling.cpp。约束说明使用 Pdist 算子需满足以下约束输入 x 必须为二维 tensor且 dim(0) ≥ 2即 N ≥ 2输出 y 的数据类型必须与 x 一致p 的取值范围为 [0, ∞]即 p ≥ 0且支持正无穷N 上限为 65535PDIST_MAX_SUPPORTED_ROWS见 pdist_constants.h。computeNum 使用 uint64_t 计算N 过大时输出规模 N*(N-1)/2 会超出实际内存限制因此 Tiling 阶段在 pdist_tiling.cpp 中对 rows 超限直接报错。这些约束在多层均有对应校验Host 侧 InferShape 检查输入必须是 2 维且 N ≥ 2见 pdist_infershape.cppTiling 阶段在ParseInputParams中校验 rows ≥ 2、cols ≥ 1 及 rows ≤ 65535ACLNN 接口层在CheckParamsLogic中校验维度、输出 shape 与 p ≥ 0含 NaN 拒绝见 aclnn_pdist.cpp。两段式调用接口每个算子采用两段式接口必须先调用 GetWorkspaceSize 接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器再调用执行接口完成计算。aclnnPdistp 以 float 传入aclnnStatus aclnnPdistGetWorkspaceSize( const aclTensor *self, double p, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnPdist( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)aclnnPdistGetWorkspaceSize的参数说明参数名输入/输出描述数据类型数据格式维度(shape)selfconst aclTensor*输入输入二维 tensor对应公式中 xFLOAT、FLOAT16ND2 维 (N, M)N≥2M≥1pdouble输入距离参数对应公式中 pp≥0含 inf---outaclTensor*输出输出一维 tensor对应公式中 distFLOAT、FLOAT16ND1 维 (N*(N-1)/2,)workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小---executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程---第一段接口完成入参校验出现以下场景时报错返回值错误码描述ACLNN_ERR_PARAM_NULLPTR161001self、out 存在空指针ACLNN_ERR_PARAM_INVALID161002self 的数据类型不在支持范围之内ACLNN_ERR_PARAM_INVALID161002self 不是二维 tensor 或 N 2ACLNN_ERR_PARAM_INVALID161002p 0ACLNN_ERR_PARAM_INVALID161002out 的 shape 与 N*(N-1)/2 不匹配aclnnPdist的参数说明参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 StreamaclnnPdistForwardp 以 aclScalar 传入Forward 接口功能与 aclnnPdist 完全一致唯一区别是 p 参数以aclScalar指针传入通过aclCreateScalar创建便于上层框架以统一的标量对象传递aclnnStatus aclnnPdistForwardGetWorkspaceSize( const aclTensor *self, const aclScalar *p, aclTensor *out, uint64_t *workspaceSize, aclOpExecutor **executor)aclnnStatus aclnnPdistForward( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)该接口的第一段校验与 aclnnPdist 相同只是空指针检查范围扩展到 self、p、out 三者。两个接口均为默认确定性实现同一输入多次调用结果一致。调用示例以 4×3 的 float32 输入矩阵为例N4 时输出长度为 N*(N-1)/2 6。完整可运行代码见 test_aclnn_pdist.cpp 与 test_aclnn_pdist_forward.cpp核心调用流程如下省略 Init、CreateAclTensor 等通用辅助函数#include iostream #include vector #include acl/acl.h #include aclnn_pdist.h int main() { int32_t deviceId 0; aclrtStream stream; Init(deviceId, stream); // 构造输入: 4x3 矩阵 std::vectorint64_t inputShape {4, 3}; std::vectorfloat inputData {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; aclTensor* inputTensor nullptr; void* inputDeviceAddr nullptr; CreateAclTensor(inputData, inputShape, inputDeviceAddr, ACL_FLOAT, inputTensor); // 构造输出: N*(N-1)/2 6 std::vectorint64_t outputShape {6}; std::vectorfloat outputData(6, 0); aclTensor* outputTensor nullptr; void* outputDeviceAddr nullptr; CreateAclTensor(outputData, outputShape, outputDeviceAddr, ACL_FLOAT, outputTensor); // 调用第一段接口 float p 2.0f; uint64_t workspaceSize 0; aclOpExecutor* executor; aclnnPdistGetWorkspaceSize(inputTensor, p, outputTensor, workspaceSize, executor); void* workspaceAddr nullptr; if (workspaceSize 0) { aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); } // 调用第二段接口 aclnnPdist(workspaceAddr, workspaceSize, executor, stream); aclrtSynchronizeStream(stream); // 释放资源 aclDestroyTensor(inputTensor); aclDestroyTensor(outputTensor); aclrtFree(inputDeviceAddr); aclrtFree(outputDeviceAddr); if (workspaceSize 0) aclrtFree(workspaceAddr); aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }使用aclnnPdistForward时仅需将 p 的构造方式替换为 aclScalarfloat pValue 2.0f; aclScalar* pScalar aclCreateScalar(pValue, ACL_FLOAT); aclnnPdistForwardGetWorkspaceSize(inputTensor, pScalar, outputTensor, workspaceSize, executor); // ... aclDestroyScalar(pScalar); // 释放 aclScalar示例中的CreateAclTensor辅助函数展示了标准张量构造方式先在 Device 侧aclrtMalloc并aclrtMemcpy拷贝数据再通过aclCreateTensor以 ND 格式、连续 strides 创建 aclTensor 描述PrintOutResult通过aclrtMemcpyDEVICE_TO_HOST回读结果。源码实现深度解析Host 侧算子定义、InferShape 与 Tiling算子定义pdist_def.cpp通过OpDef注册 Pdist声明输入 x、输出 y均支持 DT_FLOAT/DT_FLOAT16ND 格式AutoContiguous以及可选属性 p默认 2.0f并挂接 ascend910b、ascend910_93 两个 AICore 编译配置。InferShapepdist_infershape.cpp校验输入为 2 维且 N ≥ 2然后直接推导输出一维 shape 为N * (N - 1) / 2这是整个算子的核心几何事实——N 行两两组合的成对数量。Tilingpdist_tiling.cpp完成四个关键决策Workspace 申请GetWorkspaceSize为每个核申请sizeof(float)的 workspace见第 37-43 行用于 Kernel 内跨 chunk 累积距离值UB 内存预算ComputeUbBudget根据平台 UB 大小、reduce 临时 bufferGetReduceSumMaxMinTmpSize/GetReduceMaxMaxMinTmpSize计算、求和 tensor 保留区等预留字节数反推出每轮循环可处理的列数ubTensorEachLoop并按 8 元素对齐FP16 输入还需额外预留 fp16 转换队列多核切分ComputeCoreSplit将 computeNum N*(N-1)/2 个距离任务按 8 个一组分成 block均匀分配到各核剩余不足一个 block 的部分由尾核lastNumsBlocks、lastNumsNoneFullBlock承接并通过context-SetBlockDim(usedCores)设置 block 维度TilingKey 生成FillTilingData依据 p 的取值0 / 常规 / inf生成 0/1/2 三种 tilingKey并通过ASCENDC_TPL_SEL_PARAM依据输入数据类型完成 Kernel 模板实例化选择。Tiling 数据通过PdistTilingData结构体pdist_tiling_data.h在 Host 与 Kernel 之间传递包含 rows、cols、pValue、computeNum、ubTensorEachLoop、coreNumVar、tilingKey、reduceBufSize 及核间 block 分配信息。Kernel 侧NPU 向量计算Kernel 入口pdist.cpp通过GET_TILING_DATA_WITH_STRUCT解析 tiling 数据实例化NsPdist::PdistD_T_X并执行InitProcess。Process按核 ID 定位本核负责的输出区间含尾块处理对每个 block 内的距离任务逐条执行ProcessBlock。每条距离的完整计算链路为行号反解GetIJFromIndex根据输出线性下标 idx 反解出行对 (rowI, rowJ)。由于输出按下三角展开的计数规律组织源码采用整数平方根牛顿迭代 累加和校正的方式避免了逐行遍历的开销分块加载与差分LoadDiff按列分块chunk 大小不超过 ubTensorEachLoop从 Global Memory 加载两行数据。FP16 输入先经DataCopyPad拷入 fp16 队列再Cast为 float 参与计算随后Sub得到逐元素差chunk 距离计算ComputeChunkDistance按 tilingKey 分派到汉明/闵可夫斯基/切比雪夫三条向量计算路径详见前文数学原理部分返回该 chunk 的局部累加值chunk 间归并常规范数与汉明距离对 chunk 结果做累加切比雪夫距离取各 chunk 最大值ProcessBlock中accum的两种更新方式施加 1/p 次幂并写回ApplyInvP仅对常规范数分支生效并对结果做0 才开方/幂运算否则置 0的保护处理避免对 0 求 log随后WriteOutput将结果按 8 元素块写回 Global MemoryFP16 输出时先CastCAST_RINT 舍入再拷贝。值得留意的是输入按列分块循环while (remaining 0)的设计使得该算子能够处理任意列宽 M而 UB 占用始终保持恒定这正是ubTensorEachLoop由 UB 预算反推的意义所在。ACLNN 接口层参数校验与图构建接口层实现见 aclnn_pdist.cpp。第一段接口aclnnPdistGetWorkspaceSize依次完成CheckNotNull空指针检查对应错误码 161001CheckDtypeValid校验 self/out 的数据类型均在 {DT_FLOAT16, DT_FLOAT} 支持列表内且两者一致CheckParamsLogic校验维度self 必须 2 维、out 必须 1 维、输出 shape 与 N*(N-1)/2 的匹配关系以及 p ≥ 0 且非 NaN对应错误码 161002构建计算图当 N ≤ 1 时直接返回空 workspace输出 0 维空 tensor当 M 0第二维为 0时走 L0Fill算子填充 0常规情况则先l0op::Contiguous保证输入连续再执行l0op::Pdist最后以l0op::ViewCopy将结果写入用户 out通过uniqueExecutor-GetWorkspaceSize()返回 workspace 大小并将执行器ReleaseTo(executor)交给用户。第二段接口aclnnPdist则直接调用框架能力CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成异步执行。目录结构与测试保障Pdist 算子的完整目录结构如下见 README.mdpdist/ ├── README.md ├── op_host/ # Host 侧算子定义、Tiling、InferShape │ ├── pdist_def.cpp │ ├── pdist_infershape.cpp │ └── pdist_tiling.cpp ├── op_kernel/ # Kernel 侧NPU 计算逻辑 │ ├── pdist.cpp │ ├── pdist.h │ ├── pdist_constants.h │ ├── pdist_tiling_data.h │ └── pdist_tiling_key.h ├── op_api/ # ACLNN 接口层 │ ├── aclnn_pdist.h │ ├── aclnn_pdist.cpp │ ├── aclnn_pdist_forward.h │ └── aclnn_pdist_forward.cpp ├── docs/ # 接口文档 │ ├── aclnnPdist.md │ └── aclnnPdistForward.md ├── examples/ # 调用示例 │ ├── test_aclnn_pdist.cpp │ └── test_aclnn_pdist_forward.cpp └── tests/ ├── ut/ # 单元测试 │ ├── op_host/ # Tiling InferShape UT (14 cases) │ ├── op_api/ # ACLNN 接口 UT (9 cases) │ └── op_kernel/ # Kernel CPU 模拟器 UT ├── st/ # 系统测试 (97 cases, Real NPU) │ ├── test_aclnn_pdist.cpp │ └── run.sh └── reports/ └── iter3-acceptance-report.md仓库中实际存在的测试代码进一步印证了该算子具备完善的测试覆盖Host 侧 UTtest_pdist_infershape.cpp 验证 InferShape 的输出 shape 推导与非法输入拒绝逻辑test_pdist_tiling.cpp 验证 Tiling 参数计算接口层 UTtest_aclnn_pdist.cpp 覆盖 ACLNN 接口的参数校验与调用路径Kernel 层 UTtest_pdist.cpp 在 CPU 模拟器上运行 Kernel 逻辑配套 gen_data.py 与 compare_data.py 生成测试数据并与期望结果比对系统测试test_aclnn_pdist.cpp 与 run.sh 在真实 NPU 上执行 97 个用例。贡献说明Pdist 算子由个人开发者 xiaoy2459 于 2026/6/3 贡献至开源仓贡献内容为Pdist 算子适配开源仓详情见 README.md 中的贡献说明表格。该算子的落地也展示了 CANN ops-math 开源社区中接口文档docs/ 示例examples/ 分层测试tests/ 源码op_host/op_kernel/op_api/标准算子交付模板的完整实践。【免费下载链接】ops-math本项目是CANN提供的数学类基础计算算子库实现网络在NPU上加速计算。项目地址: https://gitcode.com/cann/ops-math创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考