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.md）：

• 0 < p < ∞：dist(i, j) = (Σ_k |x_ik − x_jk|^p)^(1/p)，即闵可夫斯基距离（Minkowski distance）。当 p=1 时退化为曼哈顿距离，p=2 时退化为欧氏距离；
• p = 0：dist(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 的取值生成 tilingKey（p=0 时为 0，0<p<∞ 时为 1，p=∞ 时为 2），Kernel 侧ComputeChunkDistance按 tilingKey 分派到三条计算路径——汉明距离路径使用CompareScalar（与 0 比较判不等）+Select+ReduceSum；闵可夫斯基距离路径使用Abs+Mul（p=2）或Ln/Muls/Exp（其他 p）+ReduceSum；切比雪夫距离路径使用Abs+ReduceMax。最后ApplyInvP对求和结果施加 1/p 次幂（p=1、p=2 分别有快速分支，其余 p 走Ln/Muls/Exp）。

产品支持情况

Pdist 算子当前支持的产品如下：

对应到算子定义中的硬件配置（见 pdist_def.cpp），AICore().AddConfig("ascend910b")与AddConfig("ascend910_93")分别对应 A2 与 A3 的 AICore 配置项，说明该算子是在这两个平台的昇腾 AI Core 上编译执行的。

算子原型设计

Pdist 算子的 IR 原型定义如下（见 pdist_def.cpp）：

在 Atlas A2/A3 系列产品上，数据类型支持 FLOAT16、FLOAT32。源码中this->Attr("p").AttrType(OPTIONAL).Float(2.0f)明确了 p 是可选属性且默认值为 2.0；Tiling 阶段解析属性时同样以 2.0f 作为缺省值（见 pdist_tiling.cpp）。

约束说明

使用 Pdist 算子需满足以下约束：

• 输入 x 必须为二维 tensor，且 dim(0) ≥ 2（即 N ≥ 2）；
• 输出 y 的数据类型必须与 x 一致；
• p 的取值范围为 [0, +∞]，即 p ≥ 0，且支持正无穷；
• N 上限为 65535（PDIST_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.cpp）；Tiling 阶段在ParseInputParams中校验 rows ≥ 2、cols ≥ 1 及 rows ≤ 65535；ACLNN 接口层在CheckParamsLogic中校验维度、输出 shape 与 p ≥ 0（含 NaN 拒绝，见 aclnn_pdist.cpp）。

两段式调用接口

每个算子采用两段式接口：必须先调用 GetWorkspaceSize 接口获取计算所需 workspace 大小以及包含了算子计算流程的执行器，再调用执行接口完成计算。

aclnnPdist（p 以 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的参数说明：

第一段接口完成入参校验，出现以下场景时报错：

aclnnPdist的参数说明：

aclnnPdistForward（p 以 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 输入矩阵为例，N=4 时输出长度为 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::vector<int64_t> inputShape = {4, 3}; std::vector<float> 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::vector<int64_t> outputShape = {6}; std::vector<float> 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 的构造方式替换为 aclScalar：

float 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通过aclrtMemcpy（DEVICE_TO_HOST）回读结果。

源码实现深度解析

Host 侧：算子定义、InferShape 与 Tiling

算子定义（pdist_def.cpp）：通过OpDef注册 Pdist，声明输入 x、输出 y（均支持 DT_FLOAT/DT_FLOAT16，ND 格式，AutoContiguous）以及可选属性 p（默认 2.0f），并挂接 ascend910b、ascend910_93 两个 AICore 编译配置。

InferShape（pdist_infershape.cpp）：校验输入为 2 维且 N ≥ 2，然后直接推导输出一维 shape 为N * (N - 1) / 2，这是整个算子的核心几何事实——N 行两两组合的成对数量。

Tiling（pdist_tiling.cpp）完成四个关键决策：

1. Workspace 申请：GetWorkspaceSize为每个核申请sizeof(float)的 workspace（见第 37-43 行），用于 Kernel 内跨 chunk 累积距离值；
2. UB 内存预算：ComputeUbBudget根据平台 UB 大小、reduce 临时 buffer（GetReduceSumMaxMinTmpSize/GetReduceMaxMaxMinTmpSize计算）、求和 tensor 保留区等预留字节数，反推出每轮循环可处理的列数ubTensorEachLoop，并按 8 元素对齐（FP16 输入还需额外预留 fp16 转换队列）；
3. 多核切分：ComputeCoreSplit将 computeNum = N*(N-1)/2 个距离任务按 8 个一组分成 block，均匀分配到各核，剩余不足一个 block 的部分由尾核（lastNumsBlocks、lastNumsNoneFullBlock）承接，并通过context->SetBlockDim(usedCores)设置 block 维度；
4. 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::Pdist<D_T_X>并执行Init+Process。

Process按核 ID 定位本核负责的输出区间（含尾块处理），对每个 block 内的距离任务逐条执行ProcessBlock。每条距离的完整计算链路为：

1. 行号反解：GetIJFromIndex根据输出线性下标 idx 反解出行对 (rowI, rowJ)。由于输出按下三角展开的计数规律组织，源码采用"整数平方根（牛顿迭代）+ 累加和校正"的方式，避免了逐行遍历的开销；
2. 分块加载与差分：LoadDiff按列分块（chunk 大小不超过 ubTensorEachLoop）从 Global Memory 加载两行数据。FP16 输入先经DataCopyPad拷入 fp16 队列，再Cast为 float 参与计算；随后Sub得到逐元素差；
3. chunk 距离计算：ComputeChunkDistance按 tilingKey 分派到汉明/闵可夫斯基/切比雪夫三条向量计算路径（详见前文数学原理部分），返回该 chunk 的局部累加值；
4. chunk 间归并：常规范数与汉明距离对 chunk 结果做累加；切比雪夫距离取各 chunk 最大值（ProcessBlock中accum的两种更新方式）；
5. 施加 1/p 次幂并写回：ApplyInvP仅对常规范数分支生效，并对结果做">0 才开方/幂运算，否则置 0"的保护处理（避免对 0 求 log），随后WriteOutput将结果按 8 元素块写回 Global Memory，FP16 输出时先Cast（CAST_RINT 舍入）再拷贝。

值得留意的是，输入按列分块循环（while (remaining > 0)）的设计使得该算子能够处理任意列宽 M，而 UB 占用始终保持恒定，这正是ubTensorEachLoop由 UB 预算反推的意义所在。

ACLNN 接口层：参数校验与图构建

接口层实现见 aclnn_pdist.cpp。第一段接口aclnnPdistGetWorkspaceSize依次完成：

1. CheckNotNull空指针检查（对应错误码 161001）；
2. CheckDtypeValid校验 self/out 的数据类型均在 {DT_FLOAT16, DT_FLOAT} 支持列表内且两者一致；
3. CheckParamsLogic校验维度（self 必须 2 维、out 必须 1 维）、输出 shape 与 N*(N-1)/2 的匹配关系，以及 p ≥ 0 且非 NaN（对应错误码 161002）；
4. 构建计算图：当 N ≤ 1 时直接返回空 workspace（输出 0 维空 tensor）；当 M == 0（第二维为 0）时走 L0Fill算子填充 0；常规情况则先l0op::Contiguous保证输入连续，再执行l0op::Pdist，最后以l0op::ViewCopy将结果写入用户 out；
5. 通过uniqueExecutor->GetWorkspaceSize()返回 workspace 大小，并将执行器ReleaseTo(executor)交给用户。

第二段接口aclnnPdist则直接调用框架能力CommonOpExecutorRun(workspace, workspaceSize, executor, stream)完成异步执行。

目录结构与测试保障

Pdist 算子的完整目录结构如下（见 README.md）：

pdist/ ├── 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 侧 UT：test_pdist_infershape.cpp 验证 InferShape 的输出 shape 推导与非法输入拒绝逻辑，test_pdist_tiling.cpp 验证 Tiling 参数计算；
• 接口层 UT：test_aclnn_pdist.cpp 覆盖 ACLNN 接口的参数校验与调用路径；
• Kernel 层 UT：test_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