拓冰网站定制 
PyPTO 的 pl.sum 归约算子详解沿指定维度对 Tile 求和的行列归约编程指南【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto_pro.language.sum在 Python 代码中通常写作pl.sum是 PyPTO 编程框架中面向 Tile 数据的归约算子用于沿指定维度对源 Tile 的全部元素求和是行求和、列求和以及 softmax、flash attention 等算子内部归约阶段的常用基础能力。本文基于 sum 官方文档结合仓库内函数声明、IR 指令分发与单元测试完整讲解该算子的函数原型、参数语义、数据类型约束、FP16 精度注意事项并提供 dim0 与 dim1 两个可直接运行的调用示例帮助读者在 Ascend 950 系列硬件上正确编写行/列归约 Kernel。产品支持情况pl.sum的可用性取决于目标 NPU 平台具体支持矩阵如下平台支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持也就是说该算子仅在 Ascend 950 系列950PR/950DT上可用A2/A3 系列产品无法直接调用。仓库源码同样印证了这一点在 block_ops.py 中归约指令的 dtype 校验注释明确说明TROWSUM/TCOLSUM 在 Ascend 950 上按维度有不同的数据类型支持范围说明pl.sum在 IR 层最终会分发到 TROWSUM/TCOLSUM 这类硬件归约指令而这两条指令并未在 A2/A3 上开放。功能说明pl.sum沿指定维度对源 Tile 求和语义与常见的矩阵归约一致dim0沿行方向归约对每一行求和输出形状为[M, 1]dim1沿列方向归约对每一列求和输出形状为[1, N]。当源 Tile 的 shape 为[M, N]时dim0 的 out shape 为[M, 1]dim1 的 out shape 为[1, N]。归约后维度被压扁成大小为 1 的维便于后续通过 expand 类算子如 row_expand / col_expand恢复广播形状。在 PyPTO 的 API 设计中sum 归属于统一归约接口族B8. Reductions / expands。_api.py中的声明指出dim0沿最后一个轴行方向归约dim1沿第一个轴列方向归约且tmp临时 Tile 的 dtype 必须与src一致、shape 与src完全一致1:1 工作区而非归约后的尺寸、memory 必须位于 UBVec存储空间见 API 声明注释。函数原型pypto_pro.language.sum( out: Tile, src: Tile, tmp: Tile, *, dim: int 0, ) - None与大多数 PyPTO 算子一样sum采用输出在前、源操作数在后的目的操作数风格所有 Tile 参数均为显式传入dim为仅关键字参数keyword-only默认值为 0。函数声明位于 python/pypto_pro/language/_api.py。参数说明参数输入/输出说明out输出目的操作数Tile 类型存放归约结果支持的数据类型详见约束说明。如果 src 的 shape 为 [M, N]dim0 时 shape 为 [M, 1]dim1 时 shape 为 [1, N]。src输入源操作数Tile 类型支持的数据类型详见约束说明。tmp输入临时存储Tile 类型兼容性参数。dim0 归约到 TROWSUM 时不读写该参数dim1 归约到 TCOLSUM默认非 binary 路径不使用该参数。dim输入归约维度0 表示沿行方向做归约1 表示沿列方向做归约。关于tmp参数需要特别说明从硬件指令角度看TROWSUM/TCOLSUM 在大部分路径下并不真正读写临时存储因此tmp被定位为兼容性参数——即为了保持 API 形状统一而保留的形参。仓库中统一归约接口的声明也指出 tmp 是hardware required workspaceAPI 声明但实际是否读写取决于 dim 取值与最终下降到的指令路径用户只需按规范分配一块与 src 同 shape、同 dtype、位于 Vec 存储空间的 Tile 传入即可。dim的取值范围必须严格为 0 或 1。IR 层在分发时会校验若 dim 不是 0 或 1会抛出InvalidShape异常报错信息为{op}: dim must be 0 (row) or 1 (column), got {dim}见 block_ops.py 的 _create_dim_op 实现。约束说明数据类型约束pl.sum的 src 与 out 支持的数据类型随 dim 取值不同而不同具体如下参数dim0行归约dim1列归约srcDT_INT16、DT_INT32、DT_INT64、DT_UINT64、DT_FP16、DT_FP32DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32out与 src 保持一致与 src 保持一致可以看出列归约dim1支持的数据类型范围明显更宽额外支持 DT_INT8、DT_UINT8、DT_UINT16、DT_UINT32 以及 DT_BF16而行归约dim0则不支持整型小类型和 BF16。这一差异与底层指令集有关仓库源码中_ROW_SUM_DTYPES与_COL_SUM_DTYPES分别定义了行/列归约的合法 dtype 集合block_ops.py其中_ROW_SUM_DTYPES (FP16, FP32, INT32, INT64, UINT64, INT16)与文档中 dim0 的表格完全一致源码注释还特别提到TROWSUM/TCOLSUM 没有 int8/uint8以及TROWMAX 缺少 bf16 而 TCOLMAX 支持block_ops.py说明不同归约指令在 950 硬件上的指令级数据类型支持存在差异框架在 IR 构建阶段就会对 src 的 dtype 做合法性校验并报错拦截。此外out 的 dtype 必须与 src 保持一致IR 层的_check_dtype_match会校验二者一致见 block_ops.py 的 row_sum/col_sum 构建函数。FP16 精度说明FP16 归约会受到三方面因素影响输入量化FP16 本身只有 10 位有效尾数大数值输入在量化表示时已经产生误差有限精度累加中间累加结果以 FP16 精度保存累加项越多舍入误差累积越明显输出舍入最终结果写回 FP16 输出时再次舍入。当输入规模较大或数值较大时设备计算结果可能与高精度参考结果存在差异此外归约指令采用的累加顺序也可能影响最终结果不同硬件/不同指令的累加树结构不同。因此在对精度敏感的 FP16 场景中建议结合实际测试评估误差必要时改用 FP32 累加后再降精度输出。返回值说明该函数无返回值- None归约结果直接写入out参数对应的 Tile 中。调用示例dim 0行求和以下示例将一个[64, 128]的 FP32 Tensor 沿行方向归约输出[64, 1]的结果import pypto_pro.language as pl pl.jit(auto_mutexTrue) def row_sum_kernel(a: pl.Tensor[[64, 128], pl.DT_FP32], out: pl.Tensor[[64, 1], pl.DT_FP32]): tt pl.TileType(shape[64, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tt_out pl.TileType(shape[64, 1], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec, layoutpl.DN) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[1]) tile_out pl.make_tile_group(typett_out, addrs0x10000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.sum(cur_out, cur_a, cur_tmp) pl.store(out, cur_out, [0, 0])要点拆解tt定义源 Tile 与 tmp 工作区 Tile 的类型shape[64, 128]、FP32、位于 VecUB存储空间tt_out定义输出 Tile 类型shape[64, 1]、FP32、Vec 存储空间并显式指定layoutpl.DN行归约输出的列向量通常以 DN 布局承载三个 Tile 通过make_tile_group分配在互不重叠的地址0x0000 / 0x8000 / 0x10000并用mutex_ids声明互斥auto_mutexTrue下框架会据此做访存同步整个计算序列放在pl.section_vector()向量节内先pl.load把全局 Tensor 搬入 UB再执行pl.sumdim 缺省为 0最后pl.store写回。设备实测输出节选如下输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据out[[1.008000e03], [5.104000e03], [9.200000e03], [1.329600e04], ...]可以看到每一行求和后收敛为一个标量值输出 shape 为[64, 1]。dim 1列求和将[64, 128]的 FP32 Tensor 沿列方向归约输出[1, 128]的结果pl.jit(auto_mutexTrue) def col_sum_kernel(a: pl.Tensor[[64, 128], pl.DT_FP32], out: pl.Tensor[[1, 128], pl.DT_FP32]): tt pl.TileType(shape[64, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tt_out pl.TileType(shape[1, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[1]) tile_out pl.make_tile_group(typett_out, addrs0x10000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.sum(cur_out, cur_a, cur_tmp, dim1) pl.store(out, cur_out, [0, 0])与 dim0 示例的差异仅在两点输出 Tensor/Tile 的 shape 改为[1, 128]且pl.sum调用显式传入dim1。列归约输出不需要指定 DN 布局。设备实测输出节选如下输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据out[[6.400000e04 6.401600e04 6.403200e04 6.404800e04 6.406400e04 6.408000e04 6.409600e04 6.411200e04 ...]]每一列 64 个元素求和后得到一个输出元素最终构成一行 128 个元素的结果。FP16 输入下的精度差异以下为sum(..., dim1)使用 FP16 输入时的设备实测结果输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据z[[6.390400e04 6.390400e04 6.390400e04 6.416000e04 6.416000e04 6.416000e04 6.419200e04 6.419200e04 ...]]与 FP32 列求和输出6.400000e04起、数值连续递增的结果相比FP16 输出出现了明显的数值跳跃如6.390400e04 → 6.416000e04与相邻输出相同的现象这正是 FP16 有限精度累加与输出舍入的直观体现。在精度敏感场景下建议优先使用 FP32 归约。与算子实现的关联从 pl.sum 到硬件指令了解pl.sum的底层实现有助于理解上文中的 dtype 约束与 tmp 参数语义。整个调用链可概括为Python API 层pl.sum声明于 language/_api.py经_api_decl注册为pypto_pro.language的公开算子前端解析层算子名sum登记在 parser/_op_pipeline.py 的操作数角色表中签名[W, R, None]表示 out 为写W、src 为读R、tmp 为编译期附加参数None该表还用于跨核同步等编译决策IR 构建层_create_dim_op根据dim值将pl.sum分发为row_sum或col_sum两个 IR 算子block_ops.py并在分发前按_DIM_REDUCE_DTYPES表校验 src 的 dtype 是否符合对应维度的 ISA 支持范围block_ops.pyIR 算子构建row_sum/col_sum分别通过_ir_row_sum/_ir_col_sum创建block_ir_op(row_sum)/block_ir_op(col_sum)调用同时再次校验 dtype 与 out/src dtype 一致性block_ops.py硬件指令最终row_sum/col_sum下降为 Ascend 950 的 TROWSUM/TCOLSUM 向量归约指令源码注释中的TROWSUM/TCOLSUM即为指令名见 block_ops.py。仓库中还提供了覆盖pl.sum的端到端测试例如 test_reduce_expand_fp32.py同时覆盖 dim0 与 dim1与 test_reduce_expand_fp16.py以及 flash attention 前向算子中通过pl.sum(..., dim1)对 QK 得分矩阵做列归约的典型用法见 test_fa_bnsd_dn.py可作为实际工程中行/列归约写法的参考。常见问题与使用建议dim 只能取 0 或 1Tile 是二维结构硬件归约指令只有行/列两种形态传其他值会在编译期抛出InvalidShape异常不会静默降级out 的 shape 与布局dim0 输出为[M, 1]示例中显式指定layoutpl.DNdim1 输出为[1, N]无需 DN 布局。归约后如需广播回原 shape可与 row_expand/col_expand 类算子配合使用tmp 的分配规范tmp 的 dtype 与 shape 需与 src 完全一致1:1 工作区且必须位于 Vec 存储空间虽然大部分路径不真正读写 tmp仍建议按规范分配以保证兼容性FP16 精度大输入规模或大数值下结果可能与高精度参考存在差异且累加顺序会影响结果精度敏感场景建议用 FP32 归约平台限制仅在 Ascend 950PR/950DT 上支持A2/A3 系列产品不可用跨平台开发时需做能力判断或降级方案。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyPTO 的 pl.sum 归约算子详解：沿指定维度对 Tile 求和的行列归约编程指南
PyPTO 的 pl.sum 归约算子详解沿指定维度对 Tile 求和的行列归约编程指南【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto导读pypto_pro.language.sum在 Python 代码中通常写作pl.sum是 PyPTO 编程框架中面向 Tile 数据的归约算子用于沿指定维度对源 Tile 的全部元素求和是行求和、列求和以及 softmax、flash attention 等算子内部归约阶段的常用基础能力。本文基于 sum 官方文档结合仓库内函数声明、IR 指令分发与单元测试完整讲解该算子的函数原型、参数语义、数据类型约束、FP16 精度注意事项并提供 dim0 与 dim1 两个可直接运行的调用示例帮助读者在 Ascend 950 系列硬件上正确编写行/列归约 Kernel。产品支持情况pl.sum的可用性取决于目标 NPU 平台具体支持矩阵如下平台支持情况Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品不支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品不支持也就是说该算子仅在 Ascend 950 系列950PR/950DT上可用A2/A3 系列产品无法直接调用。仓库源码同样印证了这一点在 block_ops.py 中归约指令的 dtype 校验注释明确说明TROWSUM/TCOLSUM 在 Ascend 950 上按维度有不同的数据类型支持范围说明pl.sum在 IR 层最终会分发到 TROWSUM/TCOLSUM 这类硬件归约指令而这两条指令并未在 A2/A3 上开放。功能说明pl.sum沿指定维度对源 Tile 求和语义与常见的矩阵归约一致dim0沿行方向归约对每一行求和输出形状为[M, 1]dim1沿列方向归约对每一列求和输出形状为[1, N]。当源 Tile 的 shape 为[M, N]时dim0 的 out shape 为[M, 1]dim1 的 out shape 为[1, N]。归约后维度被压扁成大小为 1 的维便于后续通过 expand 类算子如 row_expand / col_expand恢复广播形状。在 PyPTO 的 API 设计中sum 归属于统一归约接口族B8. Reductions / expands。_api.py中的声明指出dim0沿最后一个轴行方向归约dim1沿第一个轴列方向归约且tmp临时 Tile 的 dtype 必须与src一致、shape 与src完全一致1:1 工作区而非归约后的尺寸、memory 必须位于 UBVec存储空间见 API 声明注释。函数原型pypto_pro.language.sum( out: Tile, src: Tile, tmp: Tile, *, dim: int 0, ) - None与大多数 PyPTO 算子一样sum采用输出在前、源操作数在后的目的操作数风格所有 Tile 参数均为显式传入dim为仅关键字参数keyword-only默认值为 0。函数声明位于 python/pypto_pro/language/_api.py。参数说明参数输入/输出说明out输出目的操作数Tile 类型存放归约结果支持的数据类型详见约束说明。如果 src 的 shape 为 [M, N]dim0 时 shape 为 [M, 1]dim1 时 shape 为 [1, N]。src输入源操作数Tile 类型支持的数据类型详见约束说明。tmp输入临时存储Tile 类型兼容性参数。dim0 归约到 TROWSUM 时不读写该参数dim1 归约到 TCOLSUM默认非 binary 路径不使用该参数。dim输入归约维度0 表示沿行方向做归约1 表示沿列方向做归约。关于tmp参数需要特别说明从硬件指令角度看TROWSUM/TCOLSUM 在大部分路径下并不真正读写临时存储因此tmp被定位为兼容性参数——即为了保持 API 形状统一而保留的形参。仓库中统一归约接口的声明也指出 tmp 是hardware required workspaceAPI 声明但实际是否读写取决于 dim 取值与最终下降到的指令路径用户只需按规范分配一块与 src 同 shape、同 dtype、位于 Vec 存储空间的 Tile 传入即可。dim的取值范围必须严格为 0 或 1。IR 层在分发时会校验若 dim 不是 0 或 1会抛出InvalidShape异常报错信息为{op}: dim must be 0 (row) or 1 (column), got {dim}见 block_ops.py 的 _create_dim_op 实现。约束说明数据类型约束pl.sum的 src 与 out 支持的数据类型随 dim 取值不同而不同具体如下参数dim0行归约dim1列归约srcDT_INT16、DT_INT32、DT_INT64、DT_UINT64、DT_FP16、DT_FP32DT_INT8、DT_UINT8、DT_INT16、DT_UINT16、DT_INT32、DT_UINT32、DT_INT64、DT_UINT64、DT_FP16、DT_BF16、DT_FP32out与 src 保持一致与 src 保持一致可以看出列归约dim1支持的数据类型范围明显更宽额外支持 DT_INT8、DT_UINT8、DT_UINT16、DT_UINT32 以及 DT_BF16而行归约dim0则不支持整型小类型和 BF16。这一差异与底层指令集有关仓库源码中_ROW_SUM_DTYPES与_COL_SUM_DTYPES分别定义了行/列归约的合法 dtype 集合block_ops.py其中_ROW_SUM_DTYPES (FP16, FP32, INT32, INT64, UINT64, INT16)与文档中 dim0 的表格完全一致源码注释还特别提到TROWSUM/TCOLSUM 没有 int8/uint8以及TROWMAX 缺少 bf16 而 TCOLMAX 支持block_ops.py说明不同归约指令在 950 硬件上的指令级数据类型支持存在差异框架在 IR 构建阶段就会对 src 的 dtype 做合法性校验并报错拦截。此外out 的 dtype 必须与 src 保持一致IR 层的_check_dtype_match会校验二者一致见 block_ops.py 的 row_sum/col_sum 构建函数。FP16 精度说明FP16 归约会受到三方面因素影响输入量化FP16 本身只有 10 位有效尾数大数值输入在量化表示时已经产生误差有限精度累加中间累加结果以 FP16 精度保存累加项越多舍入误差累积越明显输出舍入最终结果写回 FP16 输出时再次舍入。当输入规模较大或数值较大时设备计算结果可能与高精度参考结果存在差异此外归约指令采用的累加顺序也可能影响最终结果不同硬件/不同指令的累加树结构不同。因此在对精度敏感的 FP16 场景中建议结合实际测试评估误差必要时改用 FP32 累加后再降精度输出。返回值说明该函数无返回值- None归约结果直接写入out参数对应的 Tile 中。调用示例dim 0行求和以下示例将一个[64, 128]的 FP32 Tensor 沿行方向归约输出[64, 1]的结果import pypto_pro.language as pl pl.jit(auto_mutexTrue) def row_sum_kernel(a: pl.Tensor[[64, 128], pl.DT_FP32], out: pl.Tensor[[64, 1], pl.DT_FP32]): tt pl.TileType(shape[64, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tt_out pl.TileType(shape[64, 1], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec, layoutpl.DN) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[1]) tile_out pl.make_tile_group(typett_out, addrs0x10000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.sum(cur_out, cur_a, cur_tmp) pl.store(out, cur_out, [0, 0])要点拆解tt定义源 Tile 与 tmp 工作区 Tile 的类型shape[64, 128]、FP32、位于 VecUB存储空间tt_out定义输出 Tile 类型shape[64, 1]、FP32、Vec 存储空间并显式指定layoutpl.DN行归约输出的列向量通常以 DN 布局承载三个 Tile 通过make_tile_group分配在互不重叠的地址0x0000 / 0x8000 / 0x10000并用mutex_ids声明互斥auto_mutexTrue下框架会据此做访存同步整个计算序列放在pl.section_vector()向量节内先pl.load把全局 Tensor 搬入 UB再执行pl.sumdim 缺省为 0最后pl.store写回。设备实测输出节选如下输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据out[[1.008000e03], [5.104000e03], [9.200000e03], [1.329600e04], ...]可以看到每一行求和后收敛为一个标量值输出 shape 为[64, 1]。dim 1列求和将[64, 128]的 FP32 Tensor 沿列方向归约输出[1, 128]的结果pl.jit(auto_mutexTrue) def col_sum_kernel(a: pl.Tensor[[64, 128], pl.DT_FP32], out: pl.Tensor[[1, 128], pl.DT_FP32]): tt pl.TileType(shape[64, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tt_out pl.TileType(shape[1, 128], dtypepl.DT_FP32, target_memorypl.MemorySpace.Vec) tile_a pl.make_tile_group(typett, addrs0x0000, mutex_ids[0]) tile_tmp pl.make_tile_group(typett, addrs0x8000, mutex_ids[1]) tile_out pl.make_tile_group(typett_out, addrs0x10000, mutex_ids[2]) with pl.section_vector(): cur_a tile_a.current() cur_tmp tile_tmp.current() cur_out tile_out.current() pl.load(cur_a, a, [0, 0]) pl.sum(cur_out, cur_a, cur_tmp, dim1) pl.store(out, cur_out, [0, 0])与 dim0 示例的差异仅在两点输出 Tensor/Tile 的 shape 改为[1, 128]且pl.sum调用显式传入dim1。列归约输出不需要指定 DN 布局。设备实测输出节选如下输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据out[[6.400000e04 6.401600e04 6.403200e04 6.404800e04 6.406400e04 6.408000e04 6.409600e04 6.411200e04 ...]]每一列 64 个元素求和后得到一个输出元素最终构成一行 128 个元素的结果。FP16 输入下的精度差异以下为sum(..., dim1)使用 FP16 输入时的设备实测结果输入数据a[[-8 -7.75 -7.5 -7.25 -7 -6.75 -6.5 -6.25 ...], [24 24.25 24.5 24.75 25 25.25 25.5 25.75 ...], [56 56.25 56.5 56.75 57 57.25 57.5 57.75 ...], [88 88.25 88.5 88.75 89 89.25 89.5 89.75 ...], ...] 输出数据z[[6.390400e04 6.390400e04 6.390400e04 6.416000e04 6.416000e04 6.416000e04 6.419200e04 6.419200e04 ...]]与 FP32 列求和输出6.400000e04起、数值连续递增的结果相比FP16 输出出现了明显的数值跳跃如6.390400e04 → 6.416000e04与相邻输出相同的现象这正是 FP16 有限精度累加与输出舍入的直观体现。在精度敏感场景下建议优先使用 FP32 归约。与算子实现的关联从 pl.sum 到硬件指令了解pl.sum的底层实现有助于理解上文中的 dtype 约束与 tmp 参数语义。整个调用链可概括为Python API 层pl.sum声明于 language/_api.py经_api_decl注册为pypto_pro.language的公开算子前端解析层算子名sum登记在 parser/_op_pipeline.py 的操作数角色表中签名[W, R, None]表示 out 为写W、src 为读R、tmp 为编译期附加参数None该表还用于跨核同步等编译决策IR 构建层_create_dim_op根据dim值将pl.sum分发为row_sum或col_sum两个 IR 算子block_ops.py并在分发前按_DIM_REDUCE_DTYPES表校验 src 的 dtype 是否符合对应维度的 ISA 支持范围block_ops.pyIR 算子构建row_sum/col_sum分别通过_ir_row_sum/_ir_col_sum创建block_ir_op(row_sum)/block_ir_op(col_sum)调用同时再次校验 dtype 与 out/src dtype 一致性block_ops.py硬件指令最终row_sum/col_sum下降为 Ascend 950 的 TROWSUM/TCOLSUM 向量归约指令源码注释中的TROWSUM/TCOLSUM即为指令名见 block_ops.py。仓库中还提供了覆盖pl.sum的端到端测试例如 test_reduce_expand_fp32.py同时覆盖 dim0 与 dim1与 test_reduce_expand_fp16.py以及 flash attention 前向算子中通过pl.sum(..., dim1)对 QK 得分矩阵做列归约的典型用法见 test_fa_bnsd_dn.py可作为实际工程中行/列归约写法的参考。常见问题与使用建议dim 只能取 0 或 1Tile 是二维结构硬件归约指令只有行/列两种形态传其他值会在编译期抛出InvalidShape异常不会静默降级out 的 shape 与布局dim0 输出为[M, 1]示例中显式指定layoutpl.DNdim1 输出为[1, N]无需 DN 布局。归约后如需广播回原 shape可与 row_expand/col_expand 类算子配合使用tmp 的分配规范tmp 的 dtype 与 shape 需与 src 完全一致1:1 工作区且必须位于 Vec 存储空间虽然大部分路径不真正读写 tmp仍建议按规范分配以保证兼容性FP16 精度大输入规模或大数值下结果可能与高精度参考存在差异且累加顺序会影响结果精度敏感场景建议用 FP32 归约平台限制仅在 Ascend 950PR/950DT 上支持A2/A3 系列产品不可用跨平台开发时需做能力判断或降级方案。【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考