企业官网建设流程全解析

pypto.atomic_add

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto

产品支持情况

• Ascend 950PR：支持
• Atlas A3 训练系列产品/Atlas A3 推理系列产品：支持
• Atlas A2 训练系列产品/Atlas A2 推理系列产品：支持

功能说明

以offsets指定的dst索引位置为基准，将输入Tensor src以原子（Atomic）方式累加到输出Tensor dst的对应区域。以2维为例，计算公式如下：

$$ dst\left[ offsets\left[0\right] : offsets\left[0\right] + src.shape\left[0\right],\ offsets\left[1\right] : offsets\left[1\right] + src.shape\left[1\right] \right]\ += src $$

其他维度以此类推。

函数原型

atomic_add(src: Tensor, offsets: List[Union[int, SymbolicScalar]], dst: Tensor) -> None

参数说明

返回值说明

无返回值，结果直接写入dst（inplace操作）。

约束说明

1. 输出Tensor dst的valid shape需由用户在调用atomic_add前确保正确，该接口不会自动推导。

2. offsets的长度需等于src/dst的维度数，即 $len(offsets) == src.dim == dst.dim$，否则会产生编译报错。

3. 在调用atomic_add前，需保证dst对应区域的数据已有效，否则会导致未定义行为。

4. 为保证写入区域不越界，对任意维度 $i$ 需满足 $offsets[i] + src.shape[i] \le dst.shape[i]$，否则会导致未定义行为。

5. 当多个核并发对dst的重叠区域执行atomic_add时，各次累加的执行顺序不确定。

调用示例

x = pypto.tensor([2, 2], pypto.DT_FP32) out = pypto.tensor([4, 4], pypto.DT_FP32) pypto.atomic_add(x, [0, 0], out)

结果示例如下：

输入数据x: [[1, 1], [1, 1]] 输出数据out（运算前）: [[1, 1, 0, 0], [1, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]] 输出数据out（运算后）: [[2, 2, 0, 0], [2, 2, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]

【免费下载链接】pyptoPyPTO（发音: pai p-t-o）：Parallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto