华为CANN PyPTO分布式共享内存写入API

pypto.distributed.shmem_put

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

产品支持情况

功能说明

以 offsets 指定的 shared memory tensor 索引位置为基准，将输入的 Tensor 赋值到 shared memory tensor 的对应区域。

函数原型

shmem_put( src: Tensor, offsets: list[Union[int, SymbolicScalar]], dst: ShmemTensor, dst_pe: Union[int, SymbolicScalar], *, put_op: AtomicType = AtomicType.SET, pred: list[Tensor] = None, ) -> Tensor

参数说明

返回值说明

返回输出 Tensor：用于表示操作完成的依赖关系。

约束说明

无

调用示例

TileShape 设置示例

说明：调用该接口前，应通过 set_vec_tile_shapes 设置 TileShape。TileShape 维度应和 src 一致。

• 示例 1：src 的 shape 为 [m, n]，TileShape 设置为 [m1, n1]，则 m1，n1 分别用于切分 m，n 轴。

  pypto.set_vec_tile_shapes(4, 8)

接口调用示例

• 示例 1：先创建一个 shared memory tensor。将输入数据赋值到 pe = 1 的 shared memory tensor 的指定区域，并与该视图原本的数据进行累加操作。

  input_tensor = pypto.tensor([16, 64], pypto.DT_BF16, "input_tensor") shmem_shape = input_tensor.shape shmem_tensor = pypto.distributed.create_shmem_tensor(group_name="tp", n_pes=8, dtype=pypto.DT_FP32, shape=shmem_shape) pypto.set_vec_tile_shapes(16, 64) put_out = pypto.distributed.shmem_put( src=input_tensor , offsets=[0, 0], dst=shmem_tensor, dst_pe=1, put_op=pypto.AtomicType.ADD, )

• 示例 2：先创建一个 shared memory tensor。将输入数据赋值到 pe = 3 的 shared memory tensor 的指定区域，并覆盖该视图原本的数据。

  input_tensor = pypto.tensor([16, 64], pypto.DT_BF16, "input_tensor") shmem_shape = input_tensor.shape shmem_tensor = pypto.distributed.create_shmem_tensor(group_name="tp", n_pes=8, dtype=pypto.DT_FP32, shape=shmem_shape) pypto.set_vec_tile_shapes(16, 64) put_out = pypto.distributed.shmem_put( src=input_tensor, offsets=[0, 0], dst=shmem_tensor, dst_pe=3, put_op=pypto.AtomicType.SET, )

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