三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

CANN/tensorflow HCCL发送API

CANN/tensorflow HCCL发送API

send

【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow

功能说明

提供group内点对点通信数据的send功能。

函数原型

def send(tensor, sr_tag, dest_rank, group="hccl_world_group")

参数说明

参数名输入/输出描述
tensor输入TensorFlow的tensor类型。
针对Ascend 950PR/Ascend 950DT,支持数据类型:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64、bfp16。
针对Atlas A3 训练系列产品/Atlas A3 推理系列产品,支持数据类型:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64、bfp16。
针对Atlas A2 训练系列产品/Atlas A2 推理系列产品,支持数据类型:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64、bfp16。
针对Atlas 训练系列产品,支持数据类型:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64。
针对Atlas 300I Duo 推理卡,支持数据类型:int8、uint8、int16、uint16、int32、uint32、int64、uint64、float16、float32、float64。
sr_tag输入消息标签,相同sr_tag的send/recv对可以收发数据,int类型。
dest_rank输入数据的目标节点,该rank是group中的rank id,int类型。
group输入String类型,最大长度为128字节,含结束符。
group名称,可以为用户自定义group或者"hccl_world_group"。

返回值

进行send操作之后发送到对端的tensor。

约束说明

  • 调用该接口的rank必须在当前接口入参group定义的范围内,不在此范围内的rank调用该接口会失败。
  • send和receive必须配对使用,即调用send接口后,需要等到与之配对的receive接口接收数据后,才可以进行下一个接口调用。

调用示例

from npu_bridge.hccl import hccl_ops tensor = tf.random_uniform((1, 3), minval=1, maxval=10, dtype=tf.float32) sr_tag = 0 dest_rank = 1 result = hccl_ops.send(tensor, sr_tag, dest_rank)

【免费下载链接】tensorflowAscend TensorFlow Adapter项目地址: https://gitcode.com/cann/tensorflow

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

← 返回列表