broadcast

Applicability

Product

Supported

Atlas A3 training products/Atlas A3 inference products

Atlas A2 training products/Atlas A2 inference products

Atlas 200I/500 A2 inference products

Atlas inference products

Atlas training products

For the Atlas inference products, only the Atlas 300I Duo inference card is supported.

Description

Broadcasts the data of the root rank in the communicator to other ranks.

Prototype

1
def broadcast(tensor, root_rank, fusion=2,fusion_id=0, group="hccl_world_group")

Parameters

Parameter

Input/Output

Description

tensor

Input

TensorFlow tensor type, which is a list.

For the Atlas A3 training products/Atlas A3 inference products, the supported data types are int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64, and bfp16.

For the Atlas A2 training products/Atlas A2 inference products, the supported data types are int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64, and bfp16.

For the Atlas training products, the supported data types are int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, and float64.

For the Atlas 300I Duo inference card, the supported data types are int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, and float64.

root_rank

Input

Int type.

Root rank ID in the group.

group

Input

A string containing a maximum of 128 bytes, including the end character.

Group name, which can be a user-defined value or hccl_world_group.

fusion

Input

Int type.

Broadcast operator fusion flag. The value can be one of the following:

  • 0: The Broadcast operator is not fused with other Broadcast operators during network compilation.
  • 2: Broadcast operators with the same fusion_id are fused during network compilation.

fusion_id

Input

Int type.

Broadcast operator fusion ID.

When fusion is set to 2, Broadcast operators with the same fusion_id are fused during network compilation.

Returns

The result tensor

Restrictions

  • The caller rank must be within the range defined by the group argument passed to this API call. Otherwise, the API call fails.
  • If the input and output of two Broadcast operators depend on each other, they cannot be fused. Otherwise, a graph loop may occur.

    As shown in the following figure, the input and output dependency exists between broadcast2 and broadcast1. Therefore, the broadcast1 and broadcast2 operators cannot be fused. That is, when the broadcast API is called, the fusion parameter must be set to 0.

Example

1
2
3
4
5
from npu_bridge.hccl import hccl_ops
tensor = tf.random_uniform((1, 3), minval=1, maxval=10, dtype=tf.float32)
inputs = [tensor]
root = 0
result = hccl_ops.broadcast(inputs, root)