CANN/pypto分布式共享内存信号API

pypto.distributed.shmem_signal【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Atlas A3 推理系列产品√Atlas A2 推理系列产品√功能说明根据 offsets 指定的索引位置将信号值 signal 写入 target_pe 对应的 shared memory tensor 的部分视图从而通知 target_pe。函数原型shmem_signal( src: ShmemTensor, src_pe: Union[int, SymbolicScalar], signal: int, shape: list[int] None, offsets: list[Union[int, SymbolicScalar]] None, *, target_pe: Union[int, SymbolicScalar], sig_op: AtomicType AtomicType.SET, pred: list[Tensor] None, ) - Tensor参数说明参数名输入/输出说明src输入触发信号的 shared memory tensor。src_pe输入shared memory tensor 所属的 pe0 pe n_pes。支持的数据类型为 int 或 SymbolicScalar 类型。signal输入发送到 src 中的信号值。支持的数据类型为int类型。shape输入需要写入信号的 shared memory tensor 的视图大小。参数类型为 list[int] 类型。offsets输入需要写入信号的 shared memory tensor 的视图的偏移量。支持 int 或 SymbolicScalar 类型的列表。offsets 的维度应与 src 的维度一致且每个维度的偏移量值应小于 src 对应维度的大小。target_pe输入接收信号的 pe。如果 target_pe -1则广播信号给所有 pe。支持 int 或 SymbolicScalar 类型的列表。sig_op输入数据传输时应用的原子操作类型。支持的数据类型为: AtomicType.SETAtomicType.ADD。默认为 AtomicType.SET 类型。pred输入用于控制操作执行的依赖关系张量列表。对数据类型无要求。不支持空 Tensor。返回值说明返回一个输出 Tensor用于表示操作完成的依赖关系。约束说明shmem_signal 和 shmem_wait_until 必须配合使用且设置 TileShape 时切块大小保持一致。调用示例TileShape 设置示例说明调用 shmem_signal 前应通过 set_vec_tile_shapes 设置 TileShape TileShape 维度应和参数 shape 保持一致。示例1参数 shape 为 [m, n]TileShape设置为 [m1, n1]则 m1n1 分别用于切分 mn 轴。pypto.set_vec_tile_shapes(4, 8)接口调用示例示例 1将信号值 2 写入 pe 1 的 shared memory tensor 的全部视图中并与该视图原本的值进行累加操作从而通知 pe 1。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[64, 128]) pypto.set_vec_tile_shapes(32, 64) signal_out pypto.distributed.shmem_signal( srcshmem_tensor, src_pe1, signal2, target_pe1, sig_oppypto.AtomicType.ADD, predpredToken, )示例 2将信号值 2 写入 pe 1 的 shared memory tensor 的部分视图中从而通知 pe 1。该部分视图的 shape 为 [64, 64]offset 为 [0, 0] 并与该视图原本的值进行累加操作。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[64, 128]) pypto.set_vec_tile_shapes(32, 64) signal_out pypto.distributed.shmem_signal( srcshmem_tensor, src_pe1, signal2, shape[64, 64], offsets[0, 0], target_pe1, sig_oppypto.AtomicType.ADD, predpredToken, )示例 3将信号值 4 写入 pe 3 的 shared memory tensor 的部分视图中从而通知 pe 5。该部分视图的 shape 为 [64, 64]offset 为 [0, 1] 并覆盖该视图原本的值。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[64, 128]) pypto.set_vec_tile_shapes(32, 64) signal_out pypto.distributed.shmem_signal( srcshmem_tensor, src_pe3, signal4, shape[64, 64], offsets[0, 1], target_pe5, sig_oppypto.AtomicType.SET, predpredToken, )示例 4将信号值 4 写入 pe 3 的 shared memory tensor 的部分视图中从而通知所有 pe。该部分视图的 shape 为 [64, 64]offset 为 [0, 1] 并覆盖该视图原本的值。shmem_tensor pypto.distributed.create_shmem_tensor(group_nametp, n_pes8, dtypepypto.DT_FP16, shape[64, 128]) pypto.set_vec_tile_shapes(32, 64) signal_out pypto.distributed.shmem_signal( srcshmem_tensor, sec_pe3, signal4, shape[64, 64], offsets[0, 1], target_pe-1, predpredToken, )【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考