CANN/pypto 分布式共享内存存储

pypto.distributed.shmem_store【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Atlas A3 推理系列产品√Atlas A2 推理系列产品√功能说明以 offsets 指定的 shared memory tensor 索引位置为基准将输入的 Tensor 赋值到 shared memory tensor 的对应区域。函数原型shmem_store( src: Tensor, offsets: list[Union[int, SymbolicScalar]], dst: ShmemTensor, dst_pe: Union[int, SymbolicScalar], *, put_op: AtomicType AtomicType.SET, pred: list[Tensor] None, ) - Tensor参数说明参数名输入/输出说明src输入源操作数。支持的数据类型为DT_INT32DT_FP16DT_FP32DT_BF16。不支持空 TensorShape 支持 2 - 4 维Shape Size 不大于 2147483647即 INT32_MAX。支持的数据格式为 ND。offsets输入dst 的偏移量。支持 int 或 SymbolicScalar 类型的列表。offsets 的维度应与 dst 的维度一致且每个维度的偏移量值应小于 dst 对应维度的大小。dst输入目的操作数一个 shared memory tensor其形状为src.shape。dst_pe输入shared memory tensor 所属的 pe。支持的数据类型为 int 或 SymbolicScalar 类型。0 pe n_pes。put_op输入数据传输时应用的原子操作类型。支持的数据类型为: AtomicType.SETAtomicType.ADD。默认为 AtomicType.SET 类型。pred输入用于控制操作执行的依赖关系张量列表。对数据类型无要求。不支持空 Tensor。返回值说明返回输出 Tensor用于表示操作完成的依赖关系。约束说明pred 不能包含 src即 src 不可出现在 pred 中。src 的 dtype 必须和 dst 的 dtype 一致。调用示例TileShape 设置示例说明调用该接口前应通过 set_vec_tile_shapes 设置 TileShape。TileShape 维度应和 src 一致。示例 1输入的 shape 为 [m, n]TileShape 设置为 [m1, n1]则 m1n1 分别用于切分 mn 轴。pypto.set_vec_tile_shapes(4, 8)接口调用示例示例 1先创建一个 shared memory tensor。将输入数据赋值到 pe 2 的 shared memory tensor 的指定区域并与该视图原本的数据进行累加操作。注意shared memory tensor 的 dtype 和 输入数据的 dtype 必须一致。input_tensor pypto.tensor([16, 64], pypto.DT_BF16, input_tensor) shmem_shape input_tensor.shape shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shapeshmem_shape) pypto.set_vec_tile_shapes(16, 64) store_out pypto.experimental.shmem_store( srcinput_tensor, offsets[0, 0], dstshmem_tensor, dst_pe2, put_oppypto.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_nametp, n_pes8, dtypepypto.DT_FP16, shapeshmem_shape) pypto.set_vec_tile_shapes(16, 64) store_out pypto.experimental.shmem_store( srcinput_tensor, offsets[0, 0], dstshmem_tensor, dst_pe3, put_oppypto.AtomicType.SET, )【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考