CANN/cann-competitions：Pdist算子Ascend C开发设计

余印榕

698人浏览 · 2026-05-19 07:05:24

余印榕 · 2026-05-19 07:05:24 发布

需求背景（required）

【免费下载链接】cann-competitions 本仓库用于 CANN 开源社区各类竞赛、开源课题、社区任务等课题发布、开发者作品提交和展示。项目地址: https://gitcode.com/cann/cann-competitions

需求来源

昇腾社区任务：aclnnPdist算子Ascend C开发

背景介绍

Pdist算子实现优化

基于Pdist算子历史TBE版本使用Ascend C编程语言进行优化，并修复p=inf场景存在的精度问题。

Pdist算子（TBE）实现路径和相关API路径

Pdist算子实现路径为：/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/impl/dynamic/

Pdist算子原型路径：/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_proto/inc/

Pdist算子信息库路径：/usr/local/Ascend/ascend-toolkit/latest/opp/built-in/op_impl/ai_core/tbe/config/ascend910b

Pdist算子TBE实现现状分析

通过对Pdist算子TBE版本的功能分析，当前支持的能力如下：

参数	参数含义	数据类型	支持数据类型	约束	形状
x	输入tensor	tensor	float16, float32	N>=2	(N, M)
p	距离参数	scalar	float	p>=0	标量
output	输出tensor	tensor	float16, float32	无	(N*(N-1)/2,)

计算公式：$output[k] = \left( \sum_{m=0}^{M-1} |x[r][m] - x[c][m]|^p \right)^{1/p}$，其中 $k$ 为上三角索引，$r < c$ 为行对索引。

TBE版本有两个计算分支：p=0时标量逐元素判断非零后求和；p>0时统一走 Sub → Abs → Ln → Muls(p) → Exp → ReduceSum → Ln → Muls(1/p) → Exp 通用路径。p=inf时该通用路径会产生NaN，存在精度问题。

TBE实现流程图

mermaid

Pdist算子功能分析

Pdist算子功能：计算输入矩阵各行之间的p范数成对距离。

输入：x（二维tensor，形状为(N, M)），p（距离参数，标量）

输出：output（一维tensor，形状为(N*(N-1)/2,)）

支持数据类型：float16、float32

支持的p值场景：

p = 0：汉明距离（非零元素计数）
0 < p < inf：闵可夫斯基距离（包含p=1曼哈顿距离、p=2欧几里得距离等）
p = inf：切比雪夫距离（取最大绝对差值）

需求分析（required）

需求描述

使用Ascend C编程语言实现Pdist算子，支持float16、float32数据类型，支持所有合法p值场景（p>=0，包括inf），并修复原TBE实现中p=inf场景的精度问题。

需求拆解

支持float16、float32数据类型
支持p=0、0<p<inf、p=inf等各种距离计算场景
修复p=inf场景精度问题
实现算子泛化功能，满足任意合法(N, M)形状输入
性能不低于TBE版本的95%

详细设计（required）

算子分析

数学公式

$$ \text{dist}(i, j) = \begin{cases} \left( \sum_{k=0}^{M-1} |x_{ik} - x_{jk}|^p \right)^{1/p} & 0 < p < \infty \ \sum_{k=0}^{M-1} \mathbb{1}(x_{ik} \neq x_{jk}) & p = 0 \ \max_{k=0}^{M-1} |x_{ik} - x_{jk}| & p = \infty \end{cases} $$

输出为一维tensor，长度为 $N(N-1)/2$，按上三角行优先顺序排列。

支持数据类型

float16、float32

支持形状

输入：(N, M)，N >= 2，M >= 1

输出：(N*(N-1)/2,)

算子实现

实现方案

3.2.1 host侧设计：

tiling策略：

严格参照TBE实现逻辑。host侧获取输入shape(N, M)和距离参数p，将输出视为一维向量（长度为N*(N-1)/2），按输出索引均分到各核心。每核每次处理8个行对（data_each_block = 32B / 4B = 8个float32），根据UB空间大小计算M维度的分块参数。

1. 分核策略：

按输出索引均分到各核心（与TBE一致）：

compute_num = N * (N - 1) / 2
num_block_each_core = compute_num / core_num_var / data_each_block
last_nums = compute_num % (data_each_block * core_num_var)
last_nums_blocks = last_nums / data_each_block
last_nums_none_full_block = last_nums % data_each_block

各核心按core_id分配对应的输出索引段，尾块由前若干核心或单独核心处理。

2. 数据分块和内存优化策略：

UB空间划分（与TBE一致）：

src1_ub             : ub_tensor_each_loop * 4B    // x[i]数据（float32）
src2_ub             : ub_tensor_each_loop * 4B    // x[j]数据（float32）
work_tensor         : 256 * 4B                    // Reduce工作空间
temp_sum_tensor     : 8 * 4B                      // 临时求和
dst_sum_tensor      : 8 * 4B                      // 8个行对累加结果
src_temp_fp16       : ub_tensor_each_loop * 2B    // fp16临时缓冲（仅fp16时分配）
dst_sum_fp16        : 8 * 2B                      // fp16输出缓冲（仅fp16时分配）

M维度分块：当M > ub_tensor_each_loop时，对M维度循环处理，每次处理num_each_loop个元素，M维度分块循环中使用thread_num=2做简单双缓冲（与TBE一致）。

3. tilingKey规划策略：

在TBE两分支（p=0、p>0）的基础上，新增p=inf分支以修复精度问题：

tilingKey	条件	说明
0	p == 0	与TBE一致
1	0 < p < inf	与TBE一致
2	p == inf	新增分支，修复精度

host侧根据p值设置tilingKey，传递到kernel侧走对应分支。

3.2.2 kernel侧设计：

进行Init和Process两个阶段。

Init阶段：从tiling_gm搬入tiling参数（rows, cols, p, compute_num, num_each_core, num_each_loop, ub_tensor_each_loop, core_num_var），分配UB缓冲区，计算分块参数。float16输入额外分配fp16临时缓冲用于精度转换。
Process阶段包括三层循环结构（与TBE一致）：
- 外层：遍历本核分配的行对块，每块8个行对
- 中层：遍历块内8个行对k=0~7，对每个行对通过标量sqrt反算行号(i, j)，然后对M维度分块循环
- 内层：CopyIn搬入x[i]和x[j]的M分块 → fp16转fp32 → 按tilingKey计算 → 累加部分和
索引反算（与TBE一致）：从输出索引k反算(i, j)：$i = \lfloor (N-0.5) - \sqrt{(N-0.5)^2 - 2k} \rfloor$，$j = k - Ni + i(i+1)/2 + i + 1$，含边界修正。
三个tilingKey的计算逻辑：
- tilingKey=0（p=0，与TBE一致）：Sub → 标量逐元素判断非零置1 → ReduceSum → 累加
- tilingKey=1（0<p<inf，与TBE一致）：Sub → Abs → Ln → Muls(p) → Exp → ReduceSum → 累加，块内8个行对完成后做最终还原 Ln → Muls(1/p) → Exp
- tilingKey=2（p=inf，新增修复）：Sub → Abs → ReduceMax → 更新最大值，无需最终还原
CopyOut（与TBE一致）：块内8个行对计算完成后，若为float16则将dst_sum_tensor从fp32转回fp16，然后DataCopy搬出到GM。

Ascend C实现流程图

mermaid

支持硬件

支持的芯片版本	涉及勾选
Atlas A2 训练系列产品	√
Atlas A3 系列产品	√

算子约束限制

输入tensor必须为二维，形状为(N, M)，N >= 2，M >= 1
p值必须 >= 0（包括inf）
输出tensor形状为(N*(N-1)/2,)
输入输出数据类型必须一致

可维可测分析

精度标准/性能标准

验收标准	描述	标准来源
精度标准	满足AscendOpTest工具默认阈值，修复p=inf场景精度问题	任务书要求
性能标准	所有核参与计算场景下，不低于TBE版本的95%	任务书要求