pypto.distributed.shmem_wait_until【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto产品支持情况产品是否支持Atlas A3 推理系列产品√Atlas A2 推理系列产品√功能说明根据 offsets 指定的索引位置从 src_pe 对应的 shared memory tensor 的部分视图中等待直到该视图的值达到目标数值 cmp_value。当条件满足时当前 pe 会接收到信号。函数原型shmem_wait_until( src: ShmemTensor, src_pe: Union[int, SymbolicScalar], cmp_value: int 0, shape: list[int] None, offsets: list[Union[int, SymbolicScalar]] None, *, cmp: OpType OpType.EQ, clear_signal: bool False, pred: list[Tensor] None, ) - Tensor参数说明参数名输入/输出说明src输入触发信号的 shared memory tensor。src_pe输入shared memory tensor 所属的 pe。支持的数据类型为int 或 SymbolicScalar。0 src_pe n_pes。cmp_value输入要等待的目标数值。支持的数据类型为 int 类型。shape输入需要等待信号的 shared memory tensor 的视图大小。参数类型为 list[int] 类型。offsets输入需要等待信号的 shared memory tensor 的视图的偏移量。支持 int 或 SymbolicScalar 类型的列表。offsets 的维度应与 src 的维度一致且每个维度的偏移量值应小于 src 对应维度的大小。cmp输入用于条件判断的比较操作类型。目前仅支持 EQ等于类型。clear_signal输入是否在等待完成后重置信号true/false。支持的数据类型为: bool类型。默认为 false。pred输入用于控制操作执行的依赖关系张量列表。对数据类型无要求。不支持空 Tensor。返回值说明返回一个输出Tensor用于表示操作完成的依赖关系。约束说明shmem_signal 和 shmem_wait_until必须配合使用且设置 TileShape 时切块大小保持一致。调用示例TileShape 设置示例说明调用 shmem_wait_until 前应通过set_vec_tile_shapes设置TileShape。TileShape 维度应和参数 shape 保持一致。示例 1参数 shape 为 [m, n]TileShape设置为 [m1, n1]则 m1n1 分别用于切分 mn 轴。pypto.set_vec_tile_shapes(4, 8)接口调用示例示例 1当前 pe 1 在给定的 pe 1 的 shared memory tensor 全部视图上等待直到该视图的值达到目标值 cmp_value 4。一旦条件满足当前 pe 收到信号。等待完成后不重置该视图的值。注意shmem_signal 和 shmem_wait_until必须配合使用且设置的切块大小保持一致。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, ) wait_until_out pypto.distributed.shmem_wait_until( srcshmem_tensor, src_pe1, cmp_value4, clear_signalFalse, pred[signal_out], )示例 2当前 pe 1 在给定的 pe 1 的 shared memory tensor 部分视图上等待直到该视图的值达到目标值 cmp_value 4。一旦条件满足当前 pe 收到信号。等待完成后不重置该视图的值。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, ) wait_until_out pypto.distributed.shmem_wait_until( srcshmem_tensor, src_pe1, cmp_value4, shape[64, 64], offsets[0, 0], cmppypto.OpType.EQ, clear_signalFalse, pred[signal_out], )示例 3当前 pe 5 在给定的 pe 3 的 shared memory tensor 部分视图上等待直到该视图的值达到目标值 cmp_value 4。一旦条件满足当前 pe 收到信号。等待完成后该视图的值重置为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_pe3, signal4, shape[64, 64], offsets[0, 1], target_pe5, sig_oppypto.AtomicType.SET, predpredToken, ) wait_until_out pypto.distributed.shmem_wait_until( srcshmem_tensor, src_pe3, cmp_value4, shape[64, 64], offstes[0, 1], clear_signalTrue, pred[signal_out], )【免费下载链接】pyptoPyPTO发音: pai p-t-oParallel Tensor/Tile Operation编程范式。项目地址: https://gitcode.com/cann/pypto创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考